Update generate.py

generate.py

@@ -129,6 +129,30 @@ def generate(
     synced_gpus=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
+
+	
     streamer = TextStreamer(self.tokenizer, skip_prompt=False, skip_special_tokens=True)
 
     generated_token_ids, attention_mask = custom_generate(