Update modeling_quiet.py

modeling_quiet.py

@@ -1429,7 +1429,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 		self,
 		input_ids: torch.LongTensor,
 		attention_mask: Optional[torch.Tensor] = None,
-		max_length: Optional[int] = None,
+		max_new_tokens: Optional[int] = None,
 		temperature: float = 1.1,
 		**kwargs,
 	):
@@ -1438,7 +1438,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 			attention_mask = torch.ones_like(input_ids)
 
 		from .generate import generate
-		return generate(self, input_ids=input_ids, attention_mask=attention_mask, max_length=max_length, temperature=temperature, **kwargs)
+		return generate(self, input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=max_new_tokens, temperature=temperature, **kwargs)
 
 	@add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
 	@replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)