Update modeling_quiet.py

modeling_quiet.py

@@ -1100,32 +1100,54 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
     @torch.no_grad()
     def generate(
         self,
-        input_ids: torch.LongTensor = torch.LongTensor(),
-        attention_mask: Optional[torch.Tensor] = None,
-        max_new_tokens: Optional[int] = None,
-        temperature: float = 1.1,
-        **kwargs,
+        input_ids=None,
+        attention_mask=None,
+        max_new_tokens=None,
+        min_length=None,
+        do_sample=None,
+        early_stopping=None,
+        num_beams=None,
+        temperature=1.0,
+        top_k=None,
+        top_p=None,
+        repetition_penalty=None,
+        bad_words_ids=None,
+        bos_token_id=None,
+        pad_token_id=None,
+        eos_token_id=None,
+        length_penalty=None,
+        no_repeat_ngram_size=None,
+        num_return_sequences=None,
+        decoder_start_token_id=None,
+        use_cache=None,
+        num_beam_groups=None,
+        diversity_penalty=None,
+        prefix_allowed_tokens_fn=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        output_scores=None,
+        return_dict_in_generate=None,
+        forced_bos_token_id=None,
+        forced_eos_token_id=None,
+        remove_invalid_values=None,
+        synced_gpus=None,
+        **model_kwargs,
     ):
-        if isinstance(input_ids, str):
-            input_ids = self.tokenizer(input_ids, return_tensors="pt").input_ids
-
-        if attention_mask is None:
-            # Create a default attention mask if not provided
-            attention_mask = torch.ones_like(input_ids)
-
-        from .generate import generate
-
-        output = generate(
+        # Prepare the generation process with customized settings
+        model_inputs = self.prepare_inputs_for_generation(
+            input_ids, past_key_values=None, attention_mask=attention_mask, use_cache=use_cache, **model_kwargs
+        )
+        
+        # Call the external custom generation function, ensuring it's integrated properly
+        return custom_generate(
             self,
-            input_ids,
+            input_ids=input_ids,
             attention_mask=attention_mask,
             max_new_tokens=max_new_tokens,
             temperature=temperature,
-            **kwargs,
+            **model_kwargs
         )
 
-        return output.sequences
-
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
     def forward(