Update modeling_quiet.py

modeling_quiet.py

@@ -1280,8 +1280,8 @@ class QuietForCausalLM(QuietPreTrainedModel):
         # Generate the continuation
         continuation_length = self.n_ahead - 2
         new_key_values = past_key_values
-        
         start_time = time.time()
+
         for continuation_idx in range(continuation_length):
             outputs = self.model(
                 input_ids=input_ids if continuation_idx == 0 else next_token_id.unsqueeze(-1).to(input_ids.device),
@@ -1295,9 +1295,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
                 return_dict=return_dict,
             )
             new_key_values = outputs.past_key_values
-
             hidden_states = outputs[0]
-
             logits = self.lm_head(hidden_states)
             logits = logits[:, -1, :]  # Only consider the last token
 
@@ -1336,9 +1334,9 @@ class QuietForCausalLM(QuietPreTrainedModel):
         )
         hidden_states_before = outputs_before[0][:, -1:, :]
 
-        # two new tokens: last continuation token and end thought token
+        # Get the hidden states after the thought
         outputs_after = self.model(
-            input_ids=torch.cat([next_token_id.unsqueeze(-1).to(input_ids.device), torch.tensor(end_thought_token_id).unsqueeze(-1).unsqueeze(-1).to(input_ids.device)], dim=-1),
+            input_ids=input_ids,
             attention_mask=attention_mask,
             position_ids=position_ids,
             past_key_values=new_key_values,
@@ -1358,6 +1356,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
 
         # Apply the language model head to get the final logits
         logits = self.lm_head(mixed_hidden_states)
+
         return logits
 
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)