[Fix] Aria and LLama Vision and OpenAI compatible models (#641)

* Enhance Qwen model with additional parameters and improved visual handling

- Added `system_prompt`, `interleave_visuals`, and `max_length` parameters to Qwen2_VL class.
- Simplified device assignment logic for single process scenarios.
- Improved visual processing by refining how visuals are handled and ensuring proper mapping to contexts.
- Enhanced message construction to support interleaving of visuals and text based on placeholders.
- Set default generation parameters and refined handling of generated outputs to ensure proper trimming and formatting.

* Refactor VisualPuzzles utility functions for improved readability and consistency

- Updated string formatting to use double quotes for consistency.
- Enhanced the `VisualPuzzles_doc_to_text` and `parse_response` functions for better clarity and structure.
- Simplified conditional checks and improved whitespace handling in response parsing.
- Ensured consistent handling of options and answers throughout the utility functions.

* Update model imports and parameters for Aria, LlamaVision, and OpenAICompatible

- Replaced AutoModelForCausalLM and AutoProcessor with AriaForConditionalGeneration and AriaProcessor in the Aria model.
- Updated the pretrained model string in LlamaVision to "meta-llama/Llama-3.2-11B-Vision-Instruct".
- Enhanced OpenAICompatible to support AzureOpenAI and modified API key handling for better flexibility.
- Adjusted timeout parameter and refined token handling in OpenAICompatible for improved functionality.

Author: Luodian

lmms_eval/models/aria.py

@@ -10,7 +10,7 @@
 from decord import VideoReader, cpu
 from PIL import Image
 from tqdm import tqdm
-from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers import AriaForConditionalGeneration, AriaProcessor
 
 from lmms_eval import utils
 from lmms_eval.api.instance import Instance
@@ -19,6 +19,8 @@
 
 warnings.filterwarnings("ignore")
 
+import re
+
 from loguru import logger as eval_logger
 
 DEFAULT_IMAGE_TOKEN = "<image>"
@@ -72,10 +74,10 @@ def __init__(
             dtype = getattr(torch, dtype)
 
         self.max_frames_num = max_frames_num
-        self._model = AutoModelForCausalLM.from_pretrained(pretrained, revision=revision, device_map=self.device_map, torch_dtype=torch.bfloat16, trust_remote_code=True, attn_implementation=attn_implementation)
+        self._model = AriaForConditionalGeneration.from_pretrained(pretrained, revision=revision, device_map=self.device_map, torch_dtype=torch.bfloat16, trust_remote_code=True, attn_implementation=attn_implementation)
 
         self.pretrained = pretrained
-        self._image_processor = AutoProcessor.from_pretrained(pretrained, revision=revision, trust_remote_code=True)
+        self._image_processor = AriaProcessor.from_pretrained(pretrained, revision=revision, trust_remote_code=True)
         self._tokenizer = self._image_processor.tokenizer
 
         self._config = self._model.config

lmms_eval/models/llama_vision.py

@@ -27,7 +27,7 @@
 class LlamaVision(lmms):
     def __init__(
         self,
-        pretrained: str = "meta-llama/Llama-3.2-11B-Vision",
+        pretrained: str = "meta-llama/Llama-3.2-11B-Vision-Instruct",
         revision: str = "main",
         device: str = "cuda",
         dtype: Optional[Union[str, torch.dtype]] = "auto",

lmms_eval/models/openai_compatible.py

@@ -21,7 +21,7 @@
 
 from dotenv import find_dotenv, load_dotenv
 from loguru import logger as eval_logger
-from openai import OpenAI
+from openai import AzureOpenAI, OpenAI
 from PIL import Image
 
 load_dotenv(verbose=True)
@@ -32,11 +32,12 @@ class OpenAICompatible(lmms):
     def __init__(
         self,
         model_version: str = "grok-2-latest",
-        timeout: int = 120,
+        timeout: int = 10,
         max_retries: int = 5,
         max_size_in_mb: int = 20,
         continual_mode: bool = False,
         response_persistent_folder: str = None,
+        azure_openai: bool = False,
         **kwargs,
     ) -> None:
         super().__init__()
@@ -61,7 +62,11 @@ def __init__(
                 self.response_cache = {}
                 self.cache_mode = "start"
 
-        self.client = OpenAI(api_key=os.getenv("OPENAI_COMPATIBLE_API_KEY"), base_url=os.getenv("OPENAI_COMPATIBLE_API_URL"))
+        self.client = (
+            OpenAI(api_key=os.getenv("OPENAI_API_KEY"), base_url=os.getenv("OPENAI_API_BASE"))
+            if not azure_openai
+            else AzureOpenAI(api_key=os.getenv("AZURE_OPENAI_API_KEY"), azure_endpoint=os.getenv("AZURE_OPENAI_API_BASE"), api_version=os.getenv("AZURE_OPENAI_API_VERSION"))
+        )
 
         accelerator = Accelerator()
         # assert self.batch_size_per_gpu == 1, "Llava currently does not support batched generation. See https://github.com/haotian-liu/LLaVA/issues/754. HF Llava also has this issue."
@@ -186,14 +191,17 @@ def generate_until(self, requests) -> List[str]:
             if "num_beams" not in gen_kwargs:
                 gen_kwargs["num_beams"] = 1
 
-            payload["max_completion_tokens"] = gen_kwargs["max_new_tokens"]
+            # payload["max_completion_tokens"] = gen_kwargs["max_new_tokens"]
+            payload["max_tokens"] = gen_kwargs["max_new_tokens"]
             payload["temperature"] = gen_kwargs["temperature"]
 
             if "o1" in self.model_version or "o3" in self.model_version:
                 # del payload["max_output_tokens"]
                 del payload["temperature"]
                 payload["reasoning_effort"] = "medium"
                 payload["response_format"] = {"type": "text"}
+                payload.pop("max_tokens")
+                payload["max_completion_tokens"] = gen_kwargs["max_tokens"]
 
             for attempt in range(self.max_retries):
                 try:

lmms_eval/models/qwen2_5_vl.py

@@ -1,4 +1,5 @@
 import base64
+import re
 from io import BytesIO
 from typing import List, Optional, Tuple, Union
 
@@ -48,6 +49,9 @@ def __init__(
         use_custom_video_loader: Optional[bool] = False,
         fps: Optional[float] = None,  # Only applicable if use_custom_video_loader is True
         max_image_size: Optional[int] = None,  # Only applicable if use_custom_video_loader is True
+        system_prompt: Optional[str] = "You are a helpful assistant.",
+        interleave_visuals: Optional[bool] = False,
+        reasoning_prompt: Optional[str] = None,
         **kwargs,
     ) -> None:
         super().__init__()
@@ -66,12 +70,9 @@ def __init__(
         if accelerator.num_processes > 1:
             self._device = torch.device(f"cuda:{accelerator.local_process_index}")
             self.device_map = f"cuda:{accelerator.local_process_index}"
-        elif accelerator.num_processes == 1 and device_map == "auto":
-            self._device = torch.device(device)
-            self.device_map = device_map
         else:
-            self._device = torch.device(f"cuda:{accelerator.local_process_index}")
-            self.device_map = f"cuda:{accelerator.local_process_index}"
+            self._device = torch.device(device)
+            self.device_map = device_map if device_map else device
 
         if use_flash_attention_2:
             self._model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
@@ -85,10 +86,18 @@ def __init__(
         self.max_pixels = max_pixels
         self.min_pixels = min_pixels
         self.max_num_frames = max_num_frames
-        self.processor = AutoProcessor.from_pretrained(pretrained, max_pixels=max_pixels, min_pixels=min_pixels, padding_side="left")
-        self._tokenizer = AutoTokenizer.from_pretrained(pretrained, padding_side="left")
+
+        if reasoning_prompt:
+            self.reasoning_prompt = reasoning_prompt.replace("\\n", "\n")
+        else:
+            self.reasoning_prompt = None
+        self.processor = AutoProcessor.from_pretrained(pretrained, max_pixels=max_pixels, min_pixels=min_pixels)
+        self._tokenizer = AutoTokenizer.from_pretrained(pretrained)
+        self.system_prompt = system_prompt
+        self.interleave_visuals = interleave_visuals
 
         self._config = self.model.config
+        self._max_length = kwargs.get("max_length", 2048)
         self.batch_size_per_gpu = int(batch_size)
         self.use_cache = use_cache
 
@@ -184,8 +193,11 @@ def _collate(x):
             contexts, all_gen_kwargs, doc_to_visual, doc_id, task, split = zip(*chunk)
             task = task[0]
             split = split[0]
-            visuals = [doc_to_visual[0](self.task_dict[task][split][ids]) for ids in doc_id]
-            visuals = self.flatten(visuals)
+            visual_list = [doc_to_visual[0](self.task_dict[task][split][ids]) for ids in doc_id]
+            if None in visual_list:
+                visual_list = []
+            else:
+                visual_list = self.flatten(visual_list)
 
             gen_kwargs = all_gen_kwargs[0]
 
@@ -200,112 +212,116 @@ def _collate(x):
                 elif not isinstance(until, list):
                     raise ValueError(f"Expected `gen_kwargs['until']` to be of type Union[str,list] but got {type(until)}")
 
-            # if isinstance(contexts, tuple):
-            #     contexts = list(contexts)
-
-            # for i in range(len(contexts)):
-            #     for j in range(32):
-            #         if f"<image {j}>" in contexts[i]:
-            #             contexts[i] = contexts[i].replace(f"<image {j}>", "<image>")
-            #         if f"\\<image {j}\\>" in contexts[i]:
-            #             contexts[i] = contexts[i].replace(f"\\<image {j}\\>", "<image>")
-            # if "<image>" in contexts[i]:
-            #     contexts[i] = contexts[i].replace("<image>", "")
-            # print(contexts[i])
-
-            # for i in range(len(contexts)):
-            #     if "<image>" in contexts[i]:
-            #         contexts[i] = contexts[i].replace("<image>", "")
-
-            messages = []
-            processed_visuals = []
+            if isinstance(contexts, tuple):
+                contexts = list(contexts)
+
+            for i in range(len(contexts)):
+                if "<image>" in contexts[i]:
+                    contexts[i] = contexts[i].replace("<image>", "")
+
+            batched_messages = []
             for i, context in enumerate(contexts):
-                # context += "\nPlease think step by step."
-                # if "<image>" in context:
-                #     context = context.replace("<image>", "")
+                if "<image>" in context:
+                    context = context.replace("<image>", "")
 
-                message = [{"role": "system", "content": "You are a helpful assistant."}]
+                message = [{"role": "system", "content": self.system_prompt}]
+                if self.reasoning_prompt:
+                    context = context.strip() + self.reasoning_prompt
+                    contexts[i] = context
 
-                if len(visuals) > 0:
-                    visual = visuals[i] if i < len(visuals) else None
+                processed_visuals = []
+                for visual in visual_list:
                     if isinstance(visual, str) and visual.endswith((".mp4", ".avi", ".mov")):  # Video file
-                        if self.use_custom_video_loader:
-                            visual = read_video_pyav_base64(visual, num_frm=self.max_num_frames, fps=self.fps, img_format="JPEG", max_image_size=self.max_image_size)
-                            image_contents = list(map(lambda x: f"data:image/jpeg;base64,{x}", visual))
-                            message.append({"role": "user", "content": [{"type": "video", "video": image_contents}, {"type": "text", "text": context}]})
-                        else:
-                            vr = decord.VideoReader(visual)
-                            first_frame = vr[0].asnumpy()
-                            height, width = first_frame.shape[:2]
-                            # max_pixels = height * width
-                            message.append({"role": "user", "content": [{"type": "video", "video": visual, "max_pixels": 360 * 420}, {"type": "text", "text": context}]})
-                    elif isinstance(visual, Image.Image):  # Single image
+                        vr = decord.VideoReader(visual)
+                        first_frame = vr[0].asnumpy()
+                        height, width = first_frame.shape[:2]
+                        # max_pixels = height * width
+                        processed_visuals.append({"type": "video", "video": visual, "max_pixels": self.max_pixels, "min_pixels": self.min_pixels})
+                    elif isinstance(visual, Image.Image):  # Handle both single and multiple images
                         base64_image = visual.convert("RGB")
                         buffer = BytesIO()
                         base64_image.save(buffer, format="JPEG")
                         base64_bytes = base64.b64encode(buffer.getvalue())
                         base64_string = base64_bytes.decode("utf-8")
-                        message.append({"role": "user", "content": [{"type": "image", "image": f"data:image/jpeg;base64,{base64_string}"}, {"type": "text", "text": context}]})
-                    elif isinstance(visual, (list, tuple)) and all(isinstance(v, Image.Image) for v in visual):  # Multiple images
-                        image_content = []
-                        for v in visual:
-                            base64_image = v.convert("RGB")
-                            buffer = BytesIO()
-                            base64_image.save(buffer, format="JPEG")
-                            base64_bytes = base64.b64encode(buffer.getvalue())
-                            base64_string = base64_bytes.decode("utf-8")
-                            image_content.append({"type": "image", "image": f"data:image/jpeg;base64,{base64_string}"})
-                        message.append({"role": "user", "content": image_content + [{"type": "text", "text": context}]})
-                    else:
-                        message.append({"role": "user", "content": [{"type": "text", "text": context}]})
-                else:
-                    message.append({"role": "user", "content": [{"type": "text", "text": context}]})
-
-                messages.append(message)
-            # print("message")
-
-            text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-            image_inputs, video_inputs = process_vision_info(messages)
-            inputs = self.processor(
-                text=text,
-                images=image_inputs,
-                videos=video_inputs,
-                # fps=self.fps,
-                padding=True,
-                return_tensors="pt",
-            )
+                        processed_visuals.append({"type": "image", "image": f"data:image/jpeg;base64,{base64_string}", "max_pixels": self.max_pixels, "min_pixels": self.min_pixels})
+
+                if self.interleave_visuals is False:
+                    message.append(
+                        {
+                            "role": "user",
+                            "content": processed_visuals + [{"type": "text", "text": context}],
+                        }
+                    )
+                else:  # currently support find <image x> in the context
+                    image_placeholders = re.findall(r"<image \d+>", context)
+                    content_parts = []
+                    text_parts = re.split(r"<image \d+>", context)
+                    if text_parts[0]:
+                        content_parts.append({"type": "text", "text": text_parts[0]})
+
+                    for i, placeholder in enumerate(image_placeholders):
+                        img_idx = int(re.search(r"<image (\d+)>", placeholder).group(1)) - 1
+                        image_idx = min(img_idx, len(processed_visuals) - 1) if processed_visuals else 0
+                        if processed_visuals and image_idx < len(processed_visuals):
+                            content_parts.append(processed_visuals[image_idx])
+                        if i + 1 < len(text_parts) and text_parts[i + 1]:
+                            content_parts.append({"type": "text", "text": text_parts[i + 1]})
+
+                    message.append(
+                        {
+                            "role": "user",
+                            "content": content_parts,
+                        }
+                    )
+
+                batched_messages.append(message)
+
+            texts = [self.processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True) for msg in batched_messages]
+            image_inputs, video_inputs = process_vision_info(batched_messages)
+            if video_inputs is not None:
+                total_frames = video_inputs[0].shape[0]
+                indices = np.linspace(0, total_frames - 1, self.max_num_frames, dtype=int)
+                # Append the last frame index if not already included
+                if total_frames - 1 not in indices:
+                    indices = np.append(indices, total_frames - 1)
+                video_inputs[0] = video_inputs[0][indices]
+            inputs = self.processor(text=texts, images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt")
 
             if self.device_map == "auto":
                 inputs = inputs.to("cuda")
             else:
                 inputs = inputs.to(self.device)
 
-            if "max_new_tokens" not in gen_kwargs:
-                gen_kwargs["max_new_tokens"] = 4096
-            if "temperature" not in gen_kwargs:
-                gen_kwargs["temperature"] = 0
-            if "top_p" not in gen_kwargs:
-                gen_kwargs["top_p"] = None
-            if "num_beams" not in gen_kwargs:
-                gen_kwargs["num_beams"] = 1
+            # Set default generation kwargs
+            default_gen_kwargs = {
+                "max_new_tokens": 128,
+                "temperature": 0.0,  # Set to 0 for greedy default
+                "top_p": None,
+                "num_beams": 1,
+            }
+            # Update with provided kwargs
+            current_gen_kwargs = {**default_gen_kwargs, **gen_kwargs}
 
             pad_token_id = self.tokenizer.pad_token_id
 
             cont = self.model.generate(
                 **inputs,
                 eos_token_id=self.tokenizer.eos_token_id,
                 pad_token_id=pad_token_id,
-                do_sample=True if gen_kwargs["temperature"] > 0 else False,
-                temperature=gen_kwargs["temperature"],
-                top_p=gen_kwargs["top_p"],
-                num_beams=gen_kwargs["num_beams"],
-                max_new_tokens=gen_kwargs["max_new_tokens"],
+                do_sample=True if current_gen_kwargs["temperature"] > 0 else False,
+                temperature=current_gen_kwargs["temperature"],
+                top_p=current_gen_kwargs["top_p"],
+                num_beams=current_gen_kwargs["num_beams"],
+                max_new_tokens=current_gen_kwargs["max_new_tokens"],
                 use_cache=self.use_cache,
             )
 
             generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, cont)]
             answers = self.processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)
             for i, ans in enumerate(answers):
+                for term in until:
+                    if len(term) > 0:
+                        ans = ans.split(term)[0]
                 answers[i] = ans
 
             for ans, context in zip(answers, contexts):