Update modeling_quiet.py

modeling_quiet.py

@@ -929,32 +929,32 @@ class QuietModel(QuietPreTrainedModel):
         self.embed_tokens = value
 
     def _generate_thoughts(self, hidden_states, max_length):
-    thought_ids = []
-    thought_embeddings = []
-    
-    for _ in range(self.config.max_thoughts):
-        thought_id = torch.LongTensor([[self.config.start_token_id]]).to(hidden_states.device)
-        thought_embedding = self.embed_tokens(thought_id)
+        thought_ids = []
+        thought_embeddings = []
         
-        for _ in range(max_length):
-            outputs = self.forward(
-                inputs_embeds=thought_embedding,
-                attention_mask=None,
-                use_cache=True,
-            )
-            logits = outputs.logits[:, -1, :]
-            next_token_id = torch.argmax(logits, dim=-1)
+        for _ in range(self.config.max_thoughts):
+            thought_id = torch.LongTensor([[self.config.start_token_id]]).to(hidden_states.device)
+            thought_embedding = self.embed_tokens(thought_id)
             
-            if next_token_id == self.config.end_token_id:
-                break
+            for _ in range(max_length):
+                outputs = self.forward(
+                    inputs_embeds=thought_embedding,
+                    attention_mask=None,
+                    use_cache=True,
+                )
+                logits = outputs.logits[:, -1, :]
+                next_token_id = torch.argmax(logits, dim=-1)
                 
-            thought_id = torch.cat([thought_id, next_token_id.unsqueeze(0)], dim=-1)  
-            thought_embedding = torch.cat([thought_embedding, self.embed_tokens(next_token_id.unsqueeze(0))], dim=1)
+                if next_token_id == self.config.end_token_id:
+                    break
+                    
+                thought_id = torch.cat([thought_id, next_token_id.unsqueeze(0)], dim=-1)  
+                thought_embedding = torch.cat([thought_embedding, self.embed_tokens(next_token_id.unsqueeze(0))], dim=1)
 
-        thought_ids.append(thought_id.squeeze(0))
-        thought_embeddings.append(thought_embedding.squeeze(0))
+            thought_ids.append(thought_id.squeeze(0))
+            thought_embeddings.append(thought_embedding.squeeze(0))
 
-    return thought_ids, thought_embeddings
+        return thought_ids, thought_embeddings
 
 
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)