Update modeling_quiet.py

modeling_quiet.py

@@ -1111,21 +1111,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         logits = self.lm_head(mixed_hidden_states)
         return logits
     
-    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:
-            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)
-
-        if callback is not None:
-            callback(generated_text)
-
-        return generated_text
-
     @torch.no_grad()
     def generate(
         self,
@@ -1143,16 +1128,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
             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
+        return generate(self, 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)
@@ -2084,5 +2060,4 @@ class QuietForSequenceClassification(QuietPreTrainedModel):
             past_key_values=transformer_outputs.past_key_values,
             hidden_states=transformer_outputs.hidden_states,
             attentions=transformer_outputs.attentions,
-        )
-            
+        )