Update modeling_quiet.py

modeling_quiet.py

@@ -1665,7 +1665,37 @@ class QuietForCausalLM(QuietPreTrainedModel):
                 else:
                     with torch.set_grad_enabled(not self.train_only_thinking_embedding):
                         inputs_embeds = self.model.embed_tokens(input_ids)
-            
+                        
+                def _update_inputs_for_thought_tokens(
+                    self, input_ids, attention_mask, contains_start, contains_end
+                ):
+                    batch_size = input_ids.size(0)
+                    seq_len = input_ids.size(1)
+
+                    if contains_start:
+                        start_token_ids = torch.tensor(
+                            [[self.start_token_id]] * batch_size, device=input_ids.device
+                        )
+                        input_ids = torch.cat([input_ids, start_token_ids], dim=1)
+                        if attention_mask is not None:
+                            start_attention_mask = torch.ones(
+                                (batch_size, 1), device=attention_mask.device
+                            )
+                            attention_mask = torch.cat([attention_mask, start_attention_mask], dim=1)
+
+                    if contains_end:
+                        end_token_ids = torch.tensor(
+                            [[self.end_token_id]] * batch_size, device=input_ids.device
+                        )
+                        input_ids = torch.cat([input_ids, end_token_ids], dim=1)
+                        if attention_mask is not None:
+                            end_attention_mask = torch.ones(
+                                (batch_size, 1), device=attention_mask.device
+                            )
+                            attention_mask = torch.cat([attention_mask, end_attention_mask], dim=1)
+
+                    return input_ids, attention_mask
+                            
             if self.n_ahead != 1 or self.n_ahead_talk != 1 or self.comparison_mode:
                 if attention_mask is None:
                     base_attention_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=0).to(input_ids.device)