Update generate.py

generate.py

@@ -101,11 +101,10 @@ def generate(
     max_length=None,
     min_length=None,
     do_sample=None,
-	n_ahead=12,
-    n_ahead_talk=4,
     early_stopping=None,
     num_beams=None,
-    temperature=0.9,
+    temperature=1.1,
+    streamer=None,
     top_k=None,
     top_p=None,
     repetition_penalty=None,
@@ -129,9 +128,21 @@ def generate(
     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=torch.bfloat16,
     **model_kwargs,
 ):
-
     # Set model attributes
     self.max_thoughts = n_ahead + n_ahead_talk + 1
     self.merged_talk_heads = merged_talk_heads
@@ -154,7 +165,6 @@ def generate(
     self.rm_initialized = True
     self.original_mode = False
 
-	
     streamer = TextStreamer(self.tokenizer, skip_prompt=False, skip_special_tokens=True)
 
     generated_token_ids, attention_mask = custom_generate(