Update modeling_quiet.py

modeling_quiet.py

@@ -1127,14 +1127,32 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         return generated_text
 
     @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):
-        return self.generate_with_callback(input_ids, attention_mask, max_new_tokens, temperature, callback=None, **kwargs)
-    
-    def generate_with_streaming(self, input_ids: torch.LongTensor = torch.LongTensor(), attention_mask: Optional[torch.Tensor] = None, max_new_tokens: Optional[int] = None, temperature: float = 1.1, **kwargs):
-        def callback(generated_text):
-            yield generated_text
+    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,
+    ):	
+        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)
 
-        return self.generate_with_callback(input_ids, attention_mask, max_new_tokens, temperature, callback=callback, **kwargs)
+        from .generate import generate
+        generated_token_ids, generated_text = generate(self, input_ids, attention_mask=attention_mask, max_new_tokens=max_new_tokens, temperature=temperature, **kwargs)
+        
+        # Convert the generated token IDs to a tensor
+        generated_token_ids = torch.tensor(generated_token_ids)
+        
+        # Return the generated text if it's a string, otherwise return the token IDs
+        if isinstance(generated_text, str):
+            return generated_text
+        else:
+            return generated_token_ids
 
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
@@ -2066,4 +2084,5 @@ class QuietForSequenceClassification(QuietPreTrainedModel):
             past_key_values=transformer_outputs.past_key_values,
             hidden_states=transformer_outputs.hidden_states,
             attentions=transformer_outputs.attentions,
-        )
+        )
+