Update modeling_quiet.py

modeling_quiet.py

@@ -1425,15 +1425,29 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         logits = self.lm_head(mixed_hidden_states)
         return logits
     
-    def prepare_inputs_for_generation(self, input_ids, **kwargs):
-        return {"input_ids": input_ids}
+    
+    def generate(self, *args, **kwargs):
+        # Save the original input_ids and attention_mask
+        original_input_ids = kwargs.pop("input_ids", None)
+        original_attention_mask = kwargs.pop("attention_mask", None)
+
+        # Call the infer method to get the logits
+        logits = self.infer(
+            input_ids=original_input_ids,
+            attention_mask=original_attention_mask,
+            position_ids=kwargs.pop("position_ids", None),
+            past_key_values=kwargs.pop("past_key_values", None),
+            inputs_embeds=kwargs.pop("inputs_embeds", None),
+            use_cache=kwargs.pop("use_cache", None),
+            output_attentions=kwargs.pop("output_attentions", None),
+            output_hidden_states=kwargs.pop("output_hidden_states", None),
+            return_dict=kwargs.pop("return_dict", None),
+        )
 
-    def _prepare_encoder_decoder_kwargs_for_generation(self, input_ids, model_kwargs):
-        return model_kwargs
+        # Generate output using the logits
+        output_ids = torch.argmax(logits, dim=-1)
 
-    @torch.no_grad()
-    def generate(self, **kwargs):
-        return self.infer(**kwargs)
+        return output_ids
 
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)