Update modeling_quiet.py

modeling_quiet.py

@@ -869,7 +869,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         self.n_tokens_print = 1
         self.gradient_accumulation_steps = 1
         self.training_steps = 0
-        self.tokenizer = None #AutoTokenizer.from_pretrained("Crystalcareai/Quiet-Star-Custom")
+        self.tokenizer = AutoTokenizer.from_pretrained("Crystalcareai/Quiet-Star-Custom")
         self.start_token_id = None
         self.end_token_id = None
         self.rm_initialized = False
@@ -1111,24 +1111,24 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         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,
-    # ):	
-    #     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
-    #     return generate(self, input_ids, attention_mask=attention_mask, max_new_tokens=max_new_tokens, temperature=temperature, **kwargs)
+    @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,
+    ):	
+        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
+        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)