Update modeling_quiet.py

modeling_quiet.py

@@ -37,7 +37,7 @@ import transformers
 
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
-from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask # _prepare_4d_causal_attention_mask_for_sdpa
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import (
@@ -1110,7 +1110,126 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         # Apply the language model head to get the final logits
         logits = self.lm_head(mixed_hidden_states)
         return logits
+    
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+    def _generate_no_beam_search(
+        self,
+        input_ids,
+        cur_len,
+        max_length,
+        min_length,
+        do_sample,
+        temperature,
+        top_k,
+        top_p,
+        repetition_penalty,
+        no_repeat_ngram_size,
+        bad_words_ids,
+        pad_token_id,
+        eos_token_id,
+        batch_size,
+        attention_mask,
+        use_cache,
+        model_kwargs,
+    ):
+        if input_ids is None or input_ids.nelement() == 0:
+            input_ids = torch.LongTensor([[self.tokenizer.bos_token_id]]).to(self.device)
+            attention_mask = torch.ones_like(input_ids).to(self.device)
+
+        device = input_ids.device
+        with torch.no_grad():
+            batch_size = input_ids.shape[0]
+            finished_generating = torch.zeros(batch_size, dtype=torch.bool, device=device)
+            generated_token_ids = torch.full((batch_size, max_length - cur_len), self.tokenizer.pad_token_id, dtype=torch.long, device=device)
+
+            for cur_token_idx in range(max_length - cur_len):
+                new_ids = self(
+                    input_ids[~finished_generating],
+                    attention_mask=attention_mask[~finished_generating] if attention_mask is not None else None,
+                    **model_kwargs
+                )['logits']
+
+                new_ids[:, :, self.tokenizer.vocab_size:] = -float("inf")
+
+                for list_idx, answer_idx in enumerate((~finished_generating).nonzero(as_tuple=True)[0]):
+                    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] / temperature, dim=-1), 1)
+
+                    if last_token_idx + 1 >= len(base_answer_ids):
+                        new_padding = torch.full((batch_size, 1), self.tokenizer.pad_token_id, dtype=torch.long,
+                                                device=device)
+                        input_ids = torch.cat([input_ids, new_padding], dim=-1)
+                        if attention_mask is not None:
+                            attention_mask = torch.cat([attention_mask, torch.zeros_like(new_padding)], dim=-1)
+
+                    if attention_mask is not None:
+                        attention_mask[answer_idx, last_token_idx + 1] = 1
+                    input_ids[answer_idx, last_token_idx + 1] = new_ids_sampled
+                    generated_token_ids[answer_idx, cur_token_idx] = 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.convert_tokens_to_ids("</s>"):
+                        finished_generating[answer_idx] = 1
+                
+                if finished_generating.all():
+                    break
+
+        return generated_token_ids
+
+    @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:
+            attention_mask = torch.ones_like(input_ids)
+
+        max_length = max_new_tokens + input_ids.shape[1] if max_new_tokens is not None else None
+
+        # Set model attributes
+        self.max_thoughts = kwargs.get('n_ahead', 4) + kwargs.get('n_ahead_talk', 4) + 1
+        self.merged_talk_heads = kwargs.get('merged_talk_heads', True)
+        self.merged_lm_and_talk_heads = kwargs.get('merged_lm_and_talk_heads', False)
+        self.merged_lm_and_think_heads = kwargs.get('merged_lm_and_think_heads', True)
+        self.use_concat_talk_head = kwargs.get('use_concat_talk_head', True)
+        self.use_shallow_think = kwargs.get('use_shallow_think', True)
+        self.use_shallow_talk = kwargs.get('use_shallow_talk', False)
+        self.use_complex_think_head = kwargs.get('use_complex_think_head', False)
+        self.use_complex_talk_head = kwargs.get('use_complex_talk_head', True)
+        self.use_weighted_talk_head = kwargs.get('use_weighted_talk_head', True)
+
+        # Set model properties
+        self.use_end_thought_token = True
+        self.use_start_thought_token = True
+        self.n_ahead = kwargs.get('n_ahead', 4)
+        self.n_passes = 1
+        self.eval_mode = True
+        self.first_run = False
+        self.rm_initialized = True
+        self.original_mode = False
 
+        return super().generate(
+            input_ids,
+            attention_mask=attention_mask,
+            max_length=max_length,
+            temperature=temperature,
+            **kwargs,
+        )
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
     def forward(
@@ -1137,7 +1256,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         top_p: Optional[float] = None,
         min_p: Optional[float] = None,
         top_k: Optional[int] = None,
-        cache_position: Optional[bool] = None,
         repetition_penalty: Optional[float] = None,
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
@@ -1412,17 +1530,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                         past_key_values_length,
                         sliding_window=self.config.sliding_window,
                     )
-            if attention_mask is not None:
-                if attention_mask.dim() == 2:
-                    # Expand the attention mask to have dimensions (batch_size, 1, 1, seq_length)
-                    attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
-                elif attention_mask.dim() == 3:
-                    # Expand the attention mask to have dimensions (batch_size, 1, seq_length, seq_length)
-                    attention_mask = attention_mask.unsqueeze(1)
-                else:
-                    raise ValueError(
-                        f"Attention mask should have 2 or 3 dimensions, but got {attention_mask.dim()} dimensions."
-                    )
+
             outputs = self.model(
                 # input_ids=input_ids,
                 attention_mask=attention_mask,
@@ -1861,43 +1969,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
         )
-        
-
-    from .generate import custom_generate
-
-    def prepare_inputs_for_generation(self, input_ids, **kwargs):
-        return {"input_ids": input_ids, **kwargs}
-
-    def _generate_no_beam_search(
-        self,
-        input_ids,
-        cur_len,
-        max_length,
-        min_length,
-        do_sample,
-        temperature,
-        top_k,
-        top_p,
-        repetition_penalty,
-        no_repeat_ngram_size,
-        bad_words_ids,
-        pad_token_id,
-        eos_token_id,
-        batch_size,
-        attention_mask,
-        use_cache,
-        model_kwargs,
-    ):
-        generated_token_ids = custom_generate(
-            self,
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            max_new_tokens=max_length - cur_len,
-            temperature=temperature,
-            **model_kwargs,
-        )
-
-        return generated_token_ids
 
     @staticmethod
     def _reorder_cache(past_key_values, beam_idx):