Update modeling_quiet.py

modeling_quiet.py

@@ -1169,7 +1169,7 @@ def nonzero_mean(x, axis=None):
 def loss_mean(x):
     return x.sum() / (x != 0).sum()
 
-class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
+class QuietForCausalLM(QuietPreTrainedModel, QuietGenerationMixin):
     _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config):
@@ -1324,91 +1324,91 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         elif isinstance(module, nn.Embedding):
             nn.init.xavier_uniform_(module.weight)
             
-    @torch.no_grad()
-    def generate(
-        self,
-        input_ids: torch.LongTensor,
-        attention_mask: Optional[torch.LongTensor] = None,
-        **generate_kwargs,
-    ) -> torch.LongTensor:
-        n_ahead = 8
-        n_ahead_talk = 4
-        merged_talk_heads = True
-
-        if attention_mask is None:
-            attention_mask = torch.ones_like(input_ids)
-
-        generate_kwargs.update({
-            "max_thoughts": n_ahead + n_ahead_talk + 1,
-            "merged_talk_heads": merged_talk_heads,
-            "merged_lm_and_talk_heads": False,
-            "merged_lm_and_think_heads": True,
-            "use_concat_talk_head": True,
-            "use_shallow_think": True,
-            "use_shallow_talk": False,
-            "use_complex_think_head": False,
-            "use_complex_talk_head": True,
-            "use_weighted_talk_head": True,
-        })
+    # @torch.no_grad()
+    # def generate(
+    #     self,
+    #     input_ids: torch.LongTensor,
+    #     attention_mask: Optional[torch.LongTensor] = None,
+    #     **generate_kwargs,
+    # ) -> torch.LongTensor:
+    #     n_ahead = 8
+    #     n_ahead_talk = 4
+    #     merged_talk_heads = True
+
+    #     if attention_mask is None:
+    #         attention_mask = torch.ones_like(input_ids)
+
+    #     generate_kwargs.update({
+    #         "max_thoughts": n_ahead + n_ahead_talk + 1,
+    #         "merged_talk_heads": merged_talk_heads,
+    #         "merged_lm_and_talk_heads": False,
+    #         "merged_lm_and_think_heads": True,
+    #         "use_concat_talk_head": True,
+    #         "use_shallow_think": True,
+    #         "use_shallow_talk": False,
+    #         "use_complex_think_head": False,
+    #         "use_complex_talk_head": True,
+    #         "use_weighted_talk_head": True,
+    #     })
         
-        # Validate stopping criteria
-        stopping_criteria = generate_kwargs.pop("stopping_criteria", None)
-        if stopping_criteria is not None:
-            stopping_criteria = validate_stopping_criteria(
-                stopping_criteria,
-                self.config,
-            )
-            stopping_criteria = StoppingCriteriaList(stopping_criteria)
-        else:
-            stopping_criteria = StoppingCriteriaList()
-
-        streamer = generate_kwargs.pop("streamer", None)
-        temp = generate_kwargs.pop("temperature", 1.0)
-        max_length = generate_kwargs.pop("max_length", 20)
-
-        finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=input_ids.device)
-
-        for cur_token_idx in range(max_length):
-            # Sample the next token
-            new_ids = self(
-                input_ids[~finished_generating],
-                attention_mask=attention_mask[~finished_generating]
-            )['logits']
+    #     # Validate stopping criteria
+    #     stopping_criteria = generate_kwargs.pop("stopping_criteria", None)
+    #     if stopping_criteria is not None:
+    #         stopping_criteria = validate_stopping_criteria(
+    #             stopping_criteria,
+    #             self.config,
+    #         )
+    #         stopping_criteria = StoppingCriteriaList(stopping_criteria)
+    #     else:
+    #         stopping_criteria = StoppingCriteriaList()
+
+    #     streamer = generate_kwargs.pop("streamer", None)
+    #     temp = generate_kwargs.pop("temperature", 1.0)
+    #     max_length = generate_kwargs.pop("max_length", 20)
+
+    #     finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=input_ids.device)
+
+    #     for cur_token_idx in range(max_length):
+    #         # Sample the next token
+    #         new_ids = self(
+    #             input_ids[~finished_generating],
+    #             attention_mask=attention_mask[~finished_generating]
+    #         )['logits']
             
-            # Mask out the start and end thought tokens so we don't accidentally sample them
-            new_ids[:, :, self.tokenizer.vocab_size:] = -float("inf")
+    #         # Mask out the start and end thought tokens so we don't accidentally sample them
+    #         new_ids[:, :, self.tokenizer.vocab_size:] = -float("inf")
             
-            for list_idx, answer_idx in enumerate((~finished_generating).nonzero(as_tuple=True)[0]):
-                # Find the index of the last token that is not padding
-                base_answer_ids = input_ids[answer_idx]
-                new_answer_ids = new_ids[list_idx]
-                last_token_idx = (base_answer_ids != self.tokenizer.pad_token_id).nonzero(as_tuple=True)[0].max()
-
-                new_ids_sampled = torch.multinomial(
-                    torch.nn.functional.softmax(new_answer_ids[last_token_idx] / temp, dim=-1), 1)
-
-                # Assign the new id to the last token
-                if last_token_idx + 1 >= len(base_answer_ids):
-                    # Add padding everywhere
-                    new_padding = torch.full((len(input_ids), 1), self.tokenizer.pad_token_id, dtype=torch.long,
-                                            device=input_ids.device)
-                    input_ids = torch.cat([input_ids, new_padding], dim=-1)
-                    attention_mask = torch.cat([attention_mask, torch.zeros_like(new_padding)], dim=-1)
+    #         for list_idx, answer_idx in enumerate((~finished_generating).nonzero(as_tuple=True)[0]):
+    #             # Find the index of the last token that is not padding
+    #             base_answer_ids = input_ids[answer_idx]
+    #             new_answer_ids = new_ids[list_idx]
+    #             last_token_idx = (base_answer_ids != self.tokenizer.pad_token_id).nonzero(as_tuple=True)[0].max()
+
+    #             new_ids_sampled = torch.multinomial(
+    #                 torch.nn.functional.softmax(new_answer_ids[last_token_idx] / temp, dim=-1), 1)
+
+    #             # Assign the new id to the last token
+    #             if last_token_idx + 1 >= len(base_answer_ids):
+    #                 # Add padding everywhere
+    #                 new_padding = torch.full((len(input_ids), 1), self.tokenizer.pad_token_id, dtype=torch.long,
+    #                                         device=input_ids.device)
+    #                 input_ids = torch.cat([input_ids, new_padding], dim=-1)
+    #                 attention_mask = torch.cat([attention_mask, torch.zeros_like(new_padding)], dim=-1)
                 
-                attention_mask[answer_idx, last_token_idx + 1] = 1
-                input_ids[answer_idx, last_token_idx + 1] = new_ids_sampled
+    #             attention_mask[answer_idx, last_token_idx + 1] = 1
+    #             input_ids[answer_idx, last_token_idx + 1] = new_ids_sampled
                 
-                if new_ids_sampled == self.tokenizer.eos_token_id or new_ids_sampled == self.tokenizer.bos_token_id or new_ids_sampled == self.tokenizer.pad_token_id:
-                    finished_generating[answer_idx] = 1
+    #             if new_ids_sampled == self.tokenizer.eos_token_id or new_ids_sampled == self.tokenizer.bos_token_id or new_ids_sampled == self.tokenizer.pad_token_id:
+    #                 finished_generating[answer_idx] = 1
             
-            if finished_generating.all():
-                break
+    #         if finished_generating.all():
+    #             break
 
-        if streamer is not None:
-            streamer.put(input_ids)
-            streamer.end()
+    #     if streamer is not None:
+    #         streamer.put(input_ids)
+    #         streamer.end()
 
-        return input_ids
+    #     return input_ids
 
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
@@ -1421,31 +1421,41 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         inputs_embeds: Optional[torch.FloatTensor] = None,
         labels: Optional[torch.LongTensor] = None,
         use_cache: Optional[bool] = None,
-        # output_router_logits: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, CausalLMOutputWithPast]:
-        r"""
-        Args:
-            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
-                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
-                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-        Returns:
-        Example:
-        ```python
-        >>> from transformers import AutoTokenizer, QuietForCausalLM
-        >>> model = QuietForCausalLM.from_pretrained("quietai/Quiet-7B-v0.1")
-        >>> tokenizer = AutoTokenizer.from_pretrained("quietai/Quiet-7B-v0.1")
-        >>> prompt = "Hey, are you conscious? Can you talk to me?"
-        >>> inputs = tokenizer(prompt, return_tensors="pt")
-        >>> # Generate
-        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
-        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
-        ```"""
+        **kwargs,
+    ):
+        n_ahead = 8
+        n_ahead_talk = 4
+        merged_talk_heads = True
 
+        kwargs.update({
+            "max_thoughts": n_ahead + n_ahead_talk + 1,
+            "merged_talk_heads": merged_talk_heads,
+            "merged_lm_and_talk_heads": False,
+            "merged_lm_and_think_heads": True,
+            "use_concat_talk_head": True,
+            "use_shallow_think": True,
+            "use_shallow_talk": False,
+            "use_complex_think_head": False,
+            "use_complex_talk_head": True,
+            "use_weighted_talk_head": True,
+        })
+
+        return super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            **kwargs,
+        )
         if not self.training:
             n_ahead_talk_to_restore = self.n_ahead_talk
             n_passes_to_restore = self.n_passes
@@ -2217,6 +2227,59 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
     """,
     QUIET_START_DOCSTRING,
 )
+
+class QuietGenerationMixin(GenerationMixin):
+    def generate(self, input_ids, attention_mask=None, **generate_kwargs):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        max_length = generate_kwargs.get("max_length", 20)
+        temp = generate_kwargs.get("temperature", 1.0)
+
+        finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=input_ids.device)
+
+        for cur_token_idx in range(max_length):
+            # Sample the next token
+            new_ids = self(
+                input_ids[~finished_generating],
+                attention_mask=attention_mask[~finished_generating]
+            )['logits']
+            
+            # Mask out the start and end thought tokens so we don't accidentally sample them
+            new_ids[:, :, self.tokenizer.vocab_size:] = -float("inf")
+            
+            for list_idx, answer_idx in enumerate((~finished_generating).nonzero(as_tuple=True)[0]):
+                # Find the index of the last token that is not padding
+                base_answer_ids = input_ids[answer_idx]
+                new_answer_ids = new_ids[list_idx]
+                last_token_idx = (base_answer_ids != self.tokenizer.pad_token_id).nonzero(as_tuple=True)[0].max()
+
+                new_ids_sampled = torch.multinomial(
+                    torch.nn.functional.softmax(new_answer_ids[last_token_idx] / temp, dim=-1), 1)
+
+                # Assign the new id to the last token
+                if last_token_idx + 1 >= len(base_answer_ids):
+                    # Add padding everywhere
+                    new_padding = torch.full((len(input_ids), 1), self.tokenizer.pad_token_id, dtype=torch.long,
+                                             device=input_ids.device)
+                    input_ids = torch.cat([input_ids, new_padding], dim=-1)
+                    attention_mask = torch.cat([attention_mask, torch.zeros_like(new_padding)], dim=-1)
+                
+                attention_mask[answer_idx, last_token_idx + 1] = 1
+                input_ids[answer_idx, last_token_idx + 1] = new_ids_sampled
+                
+                if new_ids_sampled == self.tokenizer.eos_token_id or new_ids_sampled == self.tokenizer.bos_token_id or new_ids_sampled == self.tokenizer.pad_token_id:
+                    finished_generating[answer_idx] = 1
+            
+            if finished_generating.all():
+                break
+
+        streamer = generate_kwargs.get("streamer")
+        if streamer is not None:
+            streamer.put(input_ids)
+            streamer.end()
+
+        return input_ids
 # Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with Llama->Quiet, LLAMA->QUIET
 class QuietForSequenceClassification(QuietPreTrainedModel):
     def __init__(self, config):