Update generate.py

generate.py

@@ -80,7 +80,7 @@ def custom_generate(
                 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)
@@ -103,8 +103,10 @@ def custom_generate(
             if streamer is not None:
                 streamer.put(new_ids_sampled)
 
-        return generated_token_ids
+        # Convert generated token IDs to text
+        generated_text = self.tokenizer.decode(generated_token_ids[0], skip_special_tokens=False)
 
+        return generated_token_ids, generated_text
 
 def generate(
     self,
@@ -158,8 +160,8 @@ def generate(
 ):
 
     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
@@ -191,9 +193,9 @@ def generate(
     if attention_mask is not None:
         attention_mask = attention_mask.to(self.device)
 
-    generated_token_ids = custom_generate(
+    generated_token_ids, generated_text = 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,
@@ -228,6 +230,4 @@ def generate(
         **model_kwargs,
     )
 
-
-    generated_text = self.tokenizer.decode(generated_token_ids[0], skip_special_tokens=False)
-    return generated_token_ids, generated_text
+    return generated_token_ids, generated_text