Update modeling_quiet.py

modeling_quiet.py

@@ -1261,6 +1261,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        max_length: Optional[int] = None,  # Add the max_length argument
     ):
         batch_size, seq_len = input_ids.shape
 
@@ -1278,11 +1279,12 @@ class QuietForCausalLM(QuietPreTrainedModel):
             attention_mask = torch.cat([attention_mask, torch.ones((batch_size, 1)).to(attention_mask.device)], dim=-1)
 
         # Generate the continuation
-        continuation_length = self.n_ahead - 2
+        if max_length is None:
+            max_length = self.n_ahead - 2  # Use the default value if max_length is not provided
         new_key_values = past_key_values
         start_time = time.time()
 
-        for continuation_idx in range(continuation_length):
+        for continuation_idx in range(max_length):
             outputs = self.model(
                 input_ids=input_ids if continuation_idx == 0 else next_token_id.unsqueeze(-1).to(input_ids.device),
                 attention_mask=attention_mask,
@@ -1358,7 +1360,6 @@ class QuietForCausalLM(QuietPreTrainedModel):
         logits = self.lm_head(mixed_hidden_states)
 
         return logits
-
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
     def forward(