Update modeling_quiet.py

modeling_quiet.py

@@ -928,30 +928,6 @@ class QuietModel(QuietPreTrainedModel):
     def set_input_embeddings(self, value):
         self.embed_tokens = value
 
-    def _generate_thoughts(self, hidden_states, max_length):
-        batch_size = hidden_states.size(0)
-        thought_ids = torch.zeros((batch_size, self.config.num_thoughts, max_length), dtype=torch.long, device=hidden_states.device)
-        thought_embeddings = []
-
-        for i in range(self.config.num_thoughts):
-            thought_input_ids = torch.zeros((batch_size, 1), dtype=torch.long, device=hidden_states.device)
-            thought_outputs = self.model.generate(
-                input_ids=thought_input_ids,
-                max_length=max_length,
-                do_sample=True,
-                top_k=50,
-                top_p=0.95,
-                pad_token_id=self.config.pad_token_id,
-                eos_token_id=self.config.eos_token_id,
-            )
-            thought_ids[:, i, :] = thought_outputs
-            thought_embeddings.append(self.model.get_input_embeddings()(thought_outputs))
-
-        thought_embeddings = torch.stack(thought_embeddings, dim=1)
-        return thought_ids, thought_embeddings
-
-
-
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     def forward(
         self,
@@ -1116,14 +1092,34 @@ class QuietForCausalLM(QuietPreTrainedModel):
             nn.ReLU(),
             nn.Linear(config.hidden_size, 1),
         )
-        
         self.max_thoughts = config.max_thoughts
         self.thought_length = config.thought_length
         self.use_policy_loss = True
         self.remove_negative_rewards = True
-        
         self.post_init()
 
+    def _generate_thoughts(self, hidden_states, max_length):
+        batch_size = hidden_states.size(0)
+        thought_ids = torch.zeros((batch_size, self.config.num_thoughts, max_length), dtype=torch.long, device=hidden_states.device)
+        thought_embeddings = []
+
+        for i in range(self.config.num_thoughts):
+            thought_input_ids = torch.zeros((batch_size, 1), dtype=torch.long, device=hidden_states.device)
+            thought_outputs = self.model.generate(
+                input_ids=thought_input_ids,
+                max_length=max_length,
+                do_sample=True,
+                top_k=50,
+                top_p=0.95,
+                pad_token_id=self.config.pad_token_id,
+                eos_token_id=self.config.eos_token_id,
+            )
+            thought_ids[:, i, :] = thought_outputs
+            thought_embeddings.append(self.model.get_input_embeddings()(thought_outputs))
+
+        thought_embeddings = torch.stack(thought_embeddings, dim=1)
+        return thought_ids, thought_embeddings
+
     def calculate_policy_loss(self, thoughts, rewards):
         thought_log_probs = []
         for thought in thoughts: