Update generate.py

generate.py

@@ -47,7 +47,7 @@ def custom_generate(
     with torch.no_grad():
         batch_size = input_ids.shape[0]
         finished_generating = torch.zeros(batch_size, dtype=torch.bool, device=device)
-        generated_text = [''] * batch_size
+        generated_token_ids = torch.full((batch_size, max_new_tokens), self.tokenizer.pad_token_id, dtype=torch.long, device=device)
 
         for cur_token_idx in range(max_new_tokens):
             # Sample the next token
@@ -67,14 +67,13 @@ def custom_generate(
                 last_token_idx = (base_answer_ids != self.tokenizer.pad_token_id).nonzero(as_tuple=True)[0].max()
 
                 new_ids_sampled = torch.multinomial(
-                    torch.nn.functional.softmax(new_answer_ids[last_token_idx] / temperature, dim=-1), 1
-                )
+                    torch.nn.functional.softmax(new_answer_ids[last_token_idx] / temperature, dim=-1), 1)
 
                 # 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,
-                                         device=device)
+                                            device=device)
                     input_ids = torch.cat([input_ids, new_padding], dim=-1)
                     if attention_mask is not None:
                         attention_mask = torch.cat([attention_mask, torch.zeros_like(new_padding)], dim=-1)
@@ -82,27 +81,23 @@ 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_id = new_ids_sampled.item()
-                generated_text[answer_idx] += self.tokenizer.decode([generated_token_id])
+                generated_token_ids[answer_idx, cur_token_idx] = new_ids_sampled
 
-                if generated_token_id == self.tokenizer.eos_token_id or generated_token_id == self.tokenizer.bos_token_id or generated_token_id == self.tokenizer.pad_token_id:
+                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
 
                 # Check if the end token is generated
-                if generated_token_id == self.tokenizer.convert_tokens_to_ids("</s>"):
+                if new_ids_sampled == self.tokenizer.convert_tokens_to_ids("</s>"):
                     finished_generating[answer_idx] = 1
-
+            
             if finished_generating.all():
                 break
 
             if streamer is not None:
                 streamer.put(new_ids_sampled)
 
-        # Check if dynamic_temperature argument is present
-        if 'dynamic_temperature' in kwargs and kwargs['dynamic_temperature'] is not None:
-            return generated_text
+        return generated_token_ids
 
-    return generated_text
 
 def generate(
     self,
@@ -153,9 +148,10 @@ def generate(
     torch_dtype=torch.bfloat16,
     **model_kwargs,
 ):
-    if max_new_tokens is None:
-        max_new_tokens = 128
 
+    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
@@ -187,9 +183,9 @@ def generate(
     if attention_mask is not None:
         attention_mask = attention_mask.to(self.device)
 
-    generated_text = custom_generate(
+    generated_token_ids = 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 +220,4 @@ def generate(
         **model_kwargs,
     )
 
-    return generated_text
+    return generated_token_ids