Update modeling_quiet.py

modeling_quiet.py

@@ -1111,54 +1111,54 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
             logits = self.lm_head(mixed_hidden_states)
             return logits
     
-    @torch.no_grad()
-    def generate(
-        self,
-        input_ids: torch.LongTensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        max_new_tokens: Optional[int] = None,
-        temperature: float = 1.0,
-        **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)
-
-        batch_size, seq_len = input_ids.shape
-        max_length = seq_len + max_new_tokens if max_new_tokens is not None else self.config.max_length
-
-        position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device)
-        position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
-
-        past_key_values = None
-        hidden_states = None
-        all_hidden_states = ()
-
-        for _ in range(max_length - seq_len):
-            logits = self.infer(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_values=past_key_values,
-                inputs_embeds=hidden_states,
-                use_cache=True,
-                output_attentions=False,
-                output_hidden_states=False,
-                return_dict=False,
-            )
-
-            next_token_logits = logits[:, -1, :] / temperature
-            next_token_id = torch.argmax(next_token_logits, dim=-1)
-
-            input_ids = torch.cat([input_ids, next_token_id.unsqueeze(-1)], dim=-1)
-            attention_mask = torch.cat([attention_mask, torch.ones((batch_size, 1), device=attention_mask.device)], dim=-1)
-            position_ids = torch.cat([position_ids, (position_ids[:, -1] + 1).unsqueeze(-1)], dim=-1)
-
-            all_hidden_states = all_hidden_states + (hidden_states,)
-
-        return input_ids, all_hidden_states
+    # @torch.no_grad()
+    # def generate(
+    #     self,
+    #     input_ids: torch.LongTensor,
+    #     attention_mask: Optional[torch.Tensor] = None,
+    #     max_new_tokens: Optional[int] = None,
+    #     temperature: float = 1.0,
+    #     **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)
+
+    #     batch_size, seq_len = input_ids.shape
+    #     max_length = seq_len + max_new_tokens if max_new_tokens is not None else self.config.max_length
+
+    #     position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device)
+    #     position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
+
+    #     past_key_values = None
+    #     hidden_states = None
+    #     all_hidden_states = ()
+
+    #     for _ in range(max_length - seq_len):
+    #         logits = self.infer(
+    #             input_ids=input_ids,
+    #             attention_mask=attention_mask,
+    #             position_ids=position_ids,
+    #             past_key_values=past_key_values,
+    #             inputs_embeds=hidden_states,
+    #             use_cache=True,
+    #             output_attentions=False,
+    #             output_hidden_states=False,
+    #             return_dict=False,
+    #         )
+
+    #         next_token_logits = logits[:, -1, :] / temperature
+    #         next_token_id = torch.argmax(next_token_logits, dim=-1)
+
+    #         input_ids = torch.cat([input_ids, next_token_id.unsqueeze(-1)], dim=-1)
+    #         attention_mask = torch.cat([attention_mask, torch.ones((batch_size, 1), device=attention_mask.device)], dim=-1)
+    #         position_ids = torch.cat([position_ids, (position_ids[:, -1] + 1).unsqueeze(-1)], dim=-1)
+
+    #         all_hidden_states = all_hidden_states + (hidden_states,)
+
+    #     return input_ids, all_hidden_states
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
     def forward(