Update modeling_ovis2_5.py

modeling_ovis2_5.py

@@ -894,7 +894,54 @@ class Ovis2_5(OvisPreTrainedModel):
             pixel_values=kwargs.pop('pixel_values', None),
             grid_thws=kwargs.pop('grid_thws', None)
         )
-        return self.llm.generate(inputs=None, inputs_embeds=inputs_embeds, attention_mask=attention_mask, **kwargs)
+        enable_thinking = kwargs.pop('enable_thinking', False)
+        enable_thinking_budget = kwargs.pop('enable_thinking_budget', False)
+        thinking_budget = kwargs.pop('thinking_budget', 1024)
+        
+        if enable_thinking and enable_thinking_budget:
+            actual_max_new_tokens = kwargs['max_new_tokens']
+            kwargs['max_new_tokens'] = thinking_budget
+            generated_ids = self.llm.generate(inputs=None, inputs_embeds=inputs_embeds, attention_mask=attention_mask, **kwargs)
+            output_ids = generated_ids
+            output_ids_list = generated_ids[0]
+            
+            # check if the generation has already finished (151645 is <|im_end|>)
+            if 151645 not in output_ids_list:
+                # check if the thinking process has finished (151668 is </think>)
+                # and prepare the second model input
+                if 151668 not in output_ids_list:
+                    print("thinking budget is reached")
+                    early_stopping_text = "\n\nConsidering the limited time by the user, I have to give the solution based on the thinking directly now.\n</think>\n\n"
+                    early_stopping_ids = self.text_tokenizer(early_stopping_text, return_tensors="pt", return_attention_mask=False).input_ids.to(inputs.device)
+                    input_ids_appendent = torch.cat([output_ids, early_stopping_ids], dim=-1)
+                    kwargs['streamer'].put(early_stopping_ids) if 'streamer' in kwargs else None
+                else:
+                    input_ids_appendent = output_ids
+                
+
+                # second generation
+                new_inputs = torch.cat([inputs, input_ids_appendent], dim=-1)
+                attention_mask = torch.ne(new_inputs, self.text_tokenizer.pad_token_id).to(device=inputs.device)
+                inputs_embeds_appendent = self.merge_multimodal(
+                    input_ids=input_ids_appendent,
+                    pixel_values=None,
+                    grid_thws=None
+                )
+                new_inputs_embeds = torch.cat([inputs_embeds, inputs_embeds_appendent], dim=-2)
+                
+                kwargs['max_new_tokens'] = inputs_embeds.size(-2) + actual_max_new_tokens - new_inputs_embeds.size(-2)
+                generated_ids2 = self.llm.generate(inputs=None, inputs_embeds=new_inputs_embeds, attention_mask=attention_mask, **kwargs)
+                kwargs['streamer'].manual_end() if 'streamer' in kwargs else None
+                return torch.cat([input_ids_appendent, generated_ids2], dim=-1)
+            
+            else:
+                kwargs['streamer'].manual_end() if 'streamer' in kwargs else None
+                return generated_ids
+            
+        else:
+            generated_ids = self.llm.generate(inputs=None, inputs_embeds=inputs_embeds, attention_mask=attention_mask, **kwargs)
+            kwargs['streamer'].manual_end() if 'streamer' in kwargs else None
+            return generated_ids
 
 
 AutoConfig.register('siglip2_navit', Siglip2NavitConfig)