Update modeling_quiet.py

modeling_quiet.py

@@ -2136,6 +2136,8 @@ class QuietForCausalLM(QuietPreTrainedModel):
             attentions=outputs.attentions,
         )
         
+    
+        
     def compute_complexity_scores(self, input_ids, attention_mask):
         # Compute complexity scores based on input sequence characteristics
         # Example: Normalize sequence lengths and consider the presence of rare tokens
@@ -2228,6 +2230,81 @@ class QuietForCausalLM(QuietPreTrainedModel):
                 tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
             )
         return reordered_past
+    
+    @torch.no_grad()
+    def generate(self, input_ids, attention_mask=None, **generate_kwargs):
+        batch_size, seq_len = input_ids.shape
+
+        max_length = generate_kwargs.get("max_length", self.config.max_length)
+        min_length = generate_kwargs.get("min_length", self.config.min_length)
+        do_sample = generate_kwargs.get("do_sample", self.config.do_sample)
+        early_stopping = generate_kwargs.get("early_stopping", self.config.early_stopping)
+        num_beams = generate_kwargs.get("num_beams", self.config.num_beams)
+        temperature = generate_kwargs.get("temperature", self.config.temperature)
+        top_k = generate_kwargs.get("top_k", self.config.top_k)
+        top_p = generate_kwargs.get("top_p", self.config.top_p)
+        repetition_penalty = generate_kwargs.get("repetition_penalty", self.config.repetition_penalty)
+        pad_token_id = generate_kwargs.get("pad_token_id", self.config.pad_token_id)
+        eos_token_id = generate_kwargs.get("eos_token_id", self.config.eos_token_id)
+
+        # Prepend the start thought token to the input sequence
+        start_thought_token_id = self.tokenizer.convert_tokens_to_ids("<|startthought|>")
+        input_ids = torch.cat([input_ids, torch.tensor([[start_thought_token_id]], device=input_ids.device)], dim=-1)
+        if attention_mask is not None:
+            attention_mask = torch.cat([attention_mask, torch.ones((batch_size, 1), device=attention_mask.device)], dim=-1)
+        
+        thought_embeds = self.model.embed_tokens(input_ids)
+        
+        past_key_values = None
+        unfinished_sequences = input_ids.new(batch_size).fill_(1)
+        sequence_lengths = input_ids.new(batch_size).fill_(max_length)
+
+        while True:
+            model_inputs = self.prepare_inputs_for_generation(input_ids, past_key_values=past_key_values, attention_mask=attention_mask, use_cache=True)
+            thought_outputs = self.model(**model_inputs, output_hidden_states=True, return_dict=True)
+            next_thought_embeds = self.prepare_thought_embeds(thought_outputs.hidden_states[-1], temperature=temperature)
+
+            thought_embeds = torch.cat([thought_embeds, next_thought_embeds], dim=1)
+            
+            lm_logits = self.lm_head(thought_outputs.last_hidden_state) 
+            lm_logits = lm_logits[:, -1, :] / temperature
+
+            if do_sample:
+                next_tokens = torch.multinomial(F.softmax(lm_logits, dim=-1), num_samples=1).squeeze(1)
+            else:
+                next_tokens = torch.argmax(lm_logits, dim=-1)
+
+            # Update input_ids, attention_mask and past_key_values
+            input_ids = torch.cat([input_ids, next_tokens.unsqueeze(-1)], dim=-1)
+            if attention_mask is not None:
+                attention_mask = torch.cat([attention_mask, attention_mask.new_ones((batch_size, 1))], dim=-1)
+            past_key_values = thought_outputs.past_key_values
+
+            # Check if generation is complete
+            if eos_token_id is not None:
+                unfinished_sequences = unfinished_sequences & (next_tokens != eos_token_id)
+            unfinished_sequences = unfinished_sequences & (input_ids.shape[-1] < max_length)
+            if unfinished_sequences.max() == 0:
+                break
+            elif input_ids.shape[-1] >= max_length:
+                input_ids[:, 0] = eos_token_id
+                break
+
+        return input_ids
+        
+    def prepare_thought_embeds(self, hidden_states, temperature=1.0):
+        if self.use_start_thought_token:
+            start_embed = self.start_embedding[0].unsqueeze(0) * temperature
+        else:
+            start_embed = hidden_states[:, 0, :]
+            
+        if self.use_end_thought_token:  
+            end_embed = self.end_embedding[0].unsqueeze(0) * temperature
+            thought_embeds = torch.cat([start_embed, hidden_states[:, 1:-1, :], end_embed], dim=1)
+        else:
+            thought_embeds = torch.cat([start_embed, hidden_states[:, 1:, :]], dim=1)
+            
+        return thought_embeds
 
 
 @add_start_docstrings(