From 3e392ba070f121d0c7660a2df6dc8cd74da8ab9e Mon Sep 17 00:00:00 2001
From: zhangzef <814499375@qq.com>
Date: Sun, 16 Feb 2025 23:07:46 +0800
Subject: [PATCH] support qwen2.5-vl in vllm_grpo_trainer

---
 .../src/open_r1/trainer/vllm_grpo_trainer.py    | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/r1-v/src/open_r1/trainer/vllm_grpo_trainer.py b/src/r1-v/src/open_r1/trainer/vllm_grpo_trainer.py
index 154c4d28..9c1bb31c 100644
--- a/src/r1-v/src/open_r1/trainer/vllm_grpo_trainer.py
+++ b/src/r1-v/src/open_r1/trainer/vllm_grpo_trainer.py
@@ -36,6 +36,7 @@
     PreTrainedModel,
     PreTrainedTokenizerBase,
     Qwen2VLForConditionalGeneration,
+    Qwen2_5_VLForConditionalGeneration,
     Trainer,
     TrainerCallback,
     is_wandb_available,
@@ -175,6 +176,10 @@ def __init__(
                 model = Qwen2VLForConditionalGeneration.from_pretrained(
                     model, **model_init_kwargs
                 )
+            elif "Qwen2.5-VL" in model_id:
+                model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                    model, **model_init_kwargs
+                )
             elif "Aria" in model_id:
                 model_init_kwargs.pop("use_cache")
                 model = AriaForConditionalGeneration.from_pretrained(
@@ -199,6 +204,10 @@ def __init__(
                 self.ref_model = Qwen2VLForConditionalGeneration.from_pretrained(
                     model_id, **model_init_kwargs
                 )
+            elif "Qwen2.5-VL" in model_id:
+                self.ref_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                    model_id, **model_init_kwargs
+                )
             elif "Aria" in model_id:
                 self.ref_model = AriaForConditionalGeneration.from_pretrained(
                     model_id, **model_init_kwargs
@@ -217,12 +226,12 @@ def __init__(
 
         # Processing class
         if processing_class is None:
-            if "Qwen2-VL" in model_id or "Aria" in model_id:
+            if "Qwen2-VL" in model_id or "Aria" in model_id or "Qwen2.5-VL" in model_id:
                 processing_class = AutoProcessor.from_pretrained(model_id)
                 pad_token_id = processing_class.tokenizer.pad_token_id
                 processing_class.pad_token_id = pad_token_id
                 processing_class.eos_token_id = processing_class.tokenizer.eos_token_id
-                if "Qwen2-VL" in model_id:
+                if "Qwen2-VL" in model_id or "Qwen2.5-VL" in model_id:
                     processing_class.image_processor.max_pixels = max_pixels
                     processing_class.image_processor.min_pixels = min_pixels
             else:
@@ -304,12 +313,12 @@ def data_collator(features):  # No data collation is needed in GRPO
         # rewrite the processing AutoTokenizer -> AutoProcessor
         model_id = model if isinstance(model, str) else model.config._name_or_path
         if processing_class is None:
-            if "Qwen2-VL" in model_id or "Aria" in model_id:
+            if "Qwen2-VL" in model_id or "Aria" in model_id or "Qwen2.5-VL" in model_id:
                 processing_class = AutoProcessor.from_pretrained(model_id)
                 pad_token_id = processing_class.tokenizer.pad_token_id
                 processing_class.pad_token_id = pad_token_id
                 processing_class.eos_token_id = processing_class.tokenizer.eos_token_id
-                if "Qwen2-VL" in model_id:
+                if "Qwen2-VL" in model_id or "Qwen2.5-VL" in model_id:
                     processing_class.image_processor.max_pixels = max_pixels
                     processing_class.image_processor.min_pixels = min_pixels
             else: