Update generate.py

generate.py

@@ -1,12 +1,10 @@
 import torch
-from transformers.utils import logging
 from transformers.generation.utils import (
     GenerationMixin,
     validate_stopping_criteria,
     StoppingCriteriaList,
 )
-
-logger = logging.get_logger(__name__)
+from transformers import TextStreamer
 
 
 def custom_generate(
@@ -45,8 +43,9 @@ def custom_generate(
     synced_gpus=None,
     **kwargs,
 ):
+    device = input_ids.device
     with torch.no_grad():
-        finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=input_ids.device)
+        finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=device)
 
         while not finished_generating.all() and input_ids.shape[1] < max_length:
             # Sample the next token
@@ -72,7 +71,7 @@ def custom_generate(
                 if last_token_idx + 1 >= len(base_answer_ids):
                     # Add padding everywhere
                     new_padding = torch.full((len(input_ids), 1), self.tokenizer.pad_token_id, dtype=torch.long,
-                                            device=input_ids.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)
@@ -92,8 +91,7 @@ def custom_generate(
                 streamer.put(new_ids_sampled)
 
         generated_token_ids = input_ids.tolist()
-    
-        return generated_token_ids
+        return generated_token_ids, attention_mask
 
 
 def generate(
@@ -105,8 +103,7 @@ def generate(
     do_sample=None,
     early_stopping=None,
     num_beams=None,
-    temperature= 0.9,
-    streamer=None,
+    temperature=1.1,
     top_k=None,
     top_p=None,
     repetition_penalty=None,
@@ -126,49 +123,15 @@ def generate(
     output_hidden_states=None,
     output_scores=None,
     return_dict_in_generate=None,
-    forced_bos_token_id= True,
-    forced_eos_token_id= True,
+    forced_bos_token_id=None,
+    forced_eos_token_id=None,
     remove_invalid_values=None,
     synced_gpus=None,
-    n_ahead=12,
-    n_ahead_talk=4,
-    merged_talk_heads=True,
-    merged_lm_and_talk_heads=False,
-    merged_lm_and_think_heads=True,
-    use_concat_talk_head=True,
-    use_shallow_think=True,
-    use_shallow_talk=False,
-    use_complex_think_head=False,
-    use_complex_talk_head=True,
-    use_weighted_talk_head=True,
-    trust_remote_code=True,
-    torch_dtype=None,
     **model_kwargs,
 ):
-    # Set model attributes
-    self.max_thoughts = n_ahead + n_ahead_talk + 1
-    self.merged_talk_heads = merged_talk_heads
-    self.merged_lm_and_talk_heads = merged_lm_and_talk_heads
-    self.merged_lm_and_think_heads = merged_lm_and_think_heads
-    self.use_concat_talk_head = use_concat_talk_head
-    self.use_shallow_think = use_shallow_think
-    self.use_shallow_talk = use_shallow_talk
-    self.use_complex_think_head = use_complex_think_head
-    self.use_complex_talk_head = use_complex_talk_head
-    self.use_weighted_talk_head = use_weighted_talk_head
-
-    # Set model properties
-    self.use_end_thought_token = True
-    self.use_start_thought_token = True
-    self.n_ahead = n_ahead
-    self.n_passes = 1
-    self.eval_mode = True
-    self.first_run = False
-    self.rm_initialized = True
-    self.original_mode = False
-
-    # Generate using the custom generate function
-    generated_token_ids = custom_generate(
+    streamer = TextStreamer(self.tokenizer, skip_prompt=False, skip_special_tokens=True)
+
+    generated_token_ids, attention_mask = custom_generate(
         self,
         input_ids=input_ids,
         attention_mask=attention_mask,
@@ -205,4 +168,4 @@ def generate(
         **model_kwargs,
     )
 
-    return generated_token_ids
+    return generated_token_ids, attention_mask