Update modeling_quiet.py

modeling_quiet.py

@@ -1169,60 +1169,7 @@ def nonzero_mean(x, axis=None):
 def loss_mean(x):
     return x.sum() / (x != 0).sum()
 
-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
-
-class QuietForCausalLM(QuietPreTrainedModel, QuietGenerationMixin):
+class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
     _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config):
@@ -1377,91 +1324,46 @@ class QuietForCausalLM(QuietPreTrainedModel, QuietGenerationMixin):
         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, attention_mask=None, streamer=None, **kwargs):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
         
-    #     # 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")
-            
-    #         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
-
-    #     if streamer is not None:
-    #         streamer.put(input_ids)
-    #         streamer.end()
-
-    #     return input_ids
+        max_length = kwargs.get("max_length", 20)
+        temp = kwargs.get("temperature", 1.0)
+
+        with torch.no_grad():
+            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
+        return input_ids, attention_mask
 
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
@@ -1474,41 +1376,31 @@ class QuietForCausalLM(QuietPreTrainedModel, QuietGenerationMixin):
         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,
-        **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,
-        )
+    ) -> 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."
+        ```"""
+
         if not self.training:
             n_ahead_talk_to_restore = self.n_ahead_talk
             n_passes_to_restore = self.n_passes
@@ -2280,7 +2172,6 @@ class QuietForCausalLM(QuietPreTrainedModel, QuietGenerationMixin):
     """,
     QUIET_START_DOCSTRING,
 )
-
 # Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with Llama->Quiet, LLAMA->QUIET
 class QuietForSequenceClassification(QuietPreTrainedModel):
     def __init__(self, config):