Update generate.py

generate.py

@@ -6,6 +6,7 @@ from transformers.generation.utils import (
 )
 from transformers import TextStreamer
 
+
 def custom_generate(
     self,
     input_ids,
@@ -42,17 +43,12 @@ def custom_generate(
     synced_gpus=None,
     **kwargs,
 ):
-    if input_ids is None or input_ids.nelement() == 0:
-        # If input_ids is None or an empty tensor, create a default input tensor
-        input_ids = torch.LongTensor([[self.tokenizer.bos_token_id]]).to(self.device)
-        attention_mask = torch.ones_like(input_ids).to(self.device)
-
     device = input_ids.device
     with torch.no_grad():
-        batch_size = input_ids.shape[0]
-        finished_generating = torch.zeros(batch_size, dtype=torch.bool, device=device)
-        generated_token_ids = torch.full((batch_size, max_new_tokens), self.tokenizer.pad_token_id, dtype=torch.long, device=device)
+        finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=device)
 
+        if max_new_tokens is None:
+            max_new_tokens = 50  # Default value if not specified
         for cur_token_idx in range(max_new_tokens):
             # Sample the next token
             new_ids = self(
@@ -76,7 +72,7 @@ def custom_generate(
                 # Assign the new id to the last token
                 if last_token_idx + 1 >= len(base_answer_ids):
                     # Add padding everywhere
-                    new_padding = torch.full((batch_size, 1), self.tokenizer.pad_token_id, dtype=torch.long,
+                    new_padding = torch.full((len(input_ids), 1), self.tokenizer.pad_token_id, dtype=torch.long,
                                             device=device)
                     input_ids = torch.cat([input_ids, new_padding], dim=-1)
                     if attention_mask is not None:
@@ -85,7 +81,6 @@ def custom_generate(
                 if attention_mask is not None:
                     attention_mask[answer_idx, last_token_idx + 1] = 1
                 input_ids[answer_idx, last_token_idx + 1] = new_ids_sampled
-                generated_token_ids[answer_idx, cur_token_idx] = new_ids_sampled
 
                 if new_ids_sampled == self.tokenizer.eos_token_id or new_ids_sampled == self.tokenizer.bos_token_id or new_ids_sampled == self.tokenizer.pad_token_id:
                     finished_generating[answer_idx] = 1
@@ -100,7 +95,8 @@ def custom_generate(
             if streamer is not None:
                 streamer.put(new_ids_sampled)
 
-        return generated_token_ids
+        generated_token_ids = input_ids.tolist()
+        return generated_token_ids, attention_mask
 
 
 def generate(
@@ -137,7 +133,7 @@ def generate(
     forced_eos_token_id=None,
     remove_invalid_values=None,
     synced_gpus=None,
-    n_ahead=8,
+    n_ahead=12,
     n_ahead_talk=4,
     merged_talk_heads=True,
     merged_lm_and_talk_heads=False,
@@ -152,10 +148,6 @@ def generate(
     torch_dtype=torch.bfloat16,
     **model_kwargs,
 ):
-
-    if max_new_tokens is None:
-        max_new_tokens = 128 
-    
     # Set model attributes
     self.max_thoughts = n_ahead + n_ahead_talk + 1
     self.merged_talk_heads = merged_talk_heads
@@ -178,18 +170,9 @@ def generate(
     self.rm_initialized = True
     self.original_mode = False
 
-    # Check if the input is a string (for compatibility with text-generation-webui)
-    if isinstance(input_ids, str):
-        input_ids = self.tokenizer.encode(input_ids, return_tensors='pt')
-
-    # Move input_ids and attention_mask to the same device as the model
-    input_ids = input_ids.to(self.device)
-    if attention_mask is not None:
-        attention_mask = attention_mask.to(self.device)
-
-    generated_token_ids = custom_generate(
+    generated_token_ids, attention_mask = custom_generate(
         self,
-        input_ids=input_ids, 
+        input_ids=input_ids,
         attention_mask=attention_mask,
         max_new_tokens=max_new_tokens,
         min_length=min_length,
@@ -224,4 +207,4 @@ def generate(
         **model_kwargs,
     )
 
-    return generated_token_ids
+    return generated_token_ids, attention_mask