Update modeling_quiet.py

modeling_quiet.py

@@ -1110,34 +1110,31 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         # Apply the language model head to get the final logits
         logits = self.lm_head(mixed_hidden_states)
         return logits
-
-    @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,
-    ):	
+    
+    def generate_with_callback(self, input_ids: torch.LongTensor = torch.LongTensor(), attention_mask: Optional[torch.Tensor] = None, max_new_tokens: Optional[int] = None, temperature: float = 1.1, callback=None, **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
         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
+
+        if callback is not None:
+            callback(generated_text)
+
+        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
+
+        return self.generate_with_callback(input_ids, attention_mask, max_new_tokens, temperature, callback=callback, **kwargs)
 
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)