Update generate.py

generate.py

@@ -1,12 +1,7 @@
 import torch
-from transformers.generation.utils import (
-    GenerationMixin,
-    validate_stopping_criteria,
-    StoppingCriteriaList,
-)
+from transformers.generation.utils import GenerationMixin, validate_stopping_criteria, StoppingCriteriaList
 from transformers import TextStreamer
 
-
 def custom_generate(
     self,
     input_ids,
@@ -72,13 +67,14 @@ 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)
@@ -94,15 +90,20 @@ def custom_generate(
                 # Check if the end token is generated
                 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)
 
-        return generated_token_ids
+        # Check if dynamic_temperature argument is present
+        if 'dynamic_temperature' in kwargs and kwargs['dynamic_temperature'] is not None:
+            # Convert generated token IDs to strings and return them
+            generated_text = self.tokenizer.batch_decode(generated_token_ids, skip_special_tokens=True)
+            return generated_text
 
+    return generated_token_ids
 
 def generate(
     self,
@@ -153,10 +154,9 @@ def generate(
     torch_dtype=torch.bfloat16,
     **model_kwargs,
 ):
-
     if max_new_tokens is None:
-        max_new_tokens = 128 
-    
+        max_new_tokens = 128
+
     # Set model attributes
     self.max_thoughts = n_ahead + n_ahead_talk + 1
     self.merged_talk_heads = merged_talk_heads
@@ -190,7 +190,7 @@ def 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,
@@ -225,4 +225,6 @@ def generate(
         **model_kwargs,
     )
 
-    return generated_token_ids
+    # Convert generated token IDs to strings and return them
+    generated_text = self.tokenizer.batch_decode(generated_token_ids, skip_special_tokens=True)
+    return generated_text