eval/lmms_eval_wrapper.py

"""
LMMS-Eval wrapper for nanoVLM model.
This allows using lmms-eval for intermediate evaluation during training.
"""

import torch
from typing import List, Tuple, Optional, Union
from PIL import Image
import numpy as np
import torch.distributed as dist

from tqdm import tqdm

from lmms_eval import utils
from lmms_eval.api.model import lmms
from lmms_eval.api.instance import Instance

from models.vision_language_model import VisionLanguageModel
from data.processors import get_tokenizer, get_image_processor, get_image_string
from data.collators import VQACollator


class NanoVLMWrapper(lmms):
    """Wrapper to make nanoVLM compatible with lmms-eval framework."""
    
    def __init__(
        self,
        model: str | VisionLanguageModel = "lusxvr/nanoVLM-450M",
        device: str = "cuda",
        batch_size: int = 32,
        **kwargs
    ):
        super().__init__()
        if isinstance(model, str):
            self.model = VisionLanguageModel.from_pretrained(model).to(device)
        else:
            self.model = model.to(device)
        self.device = device
        self.batch_size = batch_size
        
        if dist.is_available() and dist.is_initialized():
            self._rank = dist.get_rank()
            self._world_size = dist.get_world_size()
        else:
            # Fallback for non-distributed execution
            self._rank = 0
            self._world_size = 1
        
        # Get tokenizer and image processor from model config if not provided
        self.tokenizer = get_tokenizer(self.model.cfg.lm_tokenizer, self.model.cfg.vlm_extra_tokens, self.model.cfg.lm_chat_template)
        resize_to_max_side_len = False
        if hasattr(self.model.cfg, "resize_to_max_side_len"):
            resize_to_max_side_len = self.model.cfg.resize_to_max_side_len
        print(f"Resize to max side len: {resize_to_max_side_len}")
        self.image_processor = get_image_processor(self.model.cfg.max_img_size, self.model.cfg.vit_img_size, resize_to_max_side_len)
            
    def _prepare_visual_input(self, visual_list: List[Image.Image]) -> Optional[torch.Tensor]:
        """Convert visual inputs to model format."""
        if not visual_list or visual_list[0] is None: # Still check if the list is empty or contains None
            return None, None
            
        images = []
        splitted_image_ratios = []
        for visual in visual_list:
            image = None
            if isinstance(visual, Image.Image):
                image = visual
            elif isinstance(visual, str): # Keep path loading for convenience
                image = Image.open(visual).convert("RGB")
            elif isinstance(visual, np.ndarray): # Keep numpy array loading for convenience
                image = Image.fromarray(visual)
            else:
                # If it's not an Image, a path string, or a numpy array, it's an error
                raise ValueError(f"Unsupported visual type: {type(visual)}. Expected PIL Image, path string, or numpy array.")
            
            # Process image
            processed_images, splitted_image_ratio = self.image_processor(image)
            if not hasattr(self.tokenizer, "global_image_token") and splitted_image_ratio[0]*splitted_image_ratio[1] == len(processed_images) - 1:
                # If the tokenizer doesn't have a global image token, but the processor generated it, remove it
                processed_images = processed_images[1:]

            images.append(processed_images)
            splitted_image_ratios.append(splitted_image_ratio)
        
        if images:
            return images, splitted_image_ratios
        return None, None
        
    def loglikelihood(self, requests: List[Instance]) -> List[Tuple[float, bool]]:
        raise NotImplementedError("Loglikelihood is not implemented for nanoVLM")

    def flatten(self, input):
        new_list = []
        for sublist in input:
            if sublist is None:
                new_list.append(None)
            else:
                for i in sublist:
                    new_list.append(i)
        return new_list
    
    def get_benchmark_formatting(self, task_name: str) -> dict:
        """Get benchmark-specific formatting rules."""
        benchmark_formats = {
            ("ai2d", "mmstar", "seedbench", "scienceqa"): { #   
                "text_replacements": {
                    "\nOptions:": "\nChoices:",
                    "\nA. ": "\nChoices:\nA. ",
                    "Please select the correct answer from the options above.": "Answer with the letter.",
                    "Answer with the option's letter from the given choices directly": "Answer with the letter directly",
                },
                "assistant_prefix": "Answer:",
                "user_prefix": "",
                "user_suffix": ""
            },
            ("docvqa_val", "docvqa_test"): {
                "text_replacements": {},
                "assistant_prefix": "",
                "user_prefix": "Give a short and terse answer to the following question. "
                                + "Do not paraphrase or reformat the text you see in the image. Do not include any full stops. "
                                + "Just give the answer without additional explanation. Question: ",
                "user_suffix": ""
            },
            "chartvqa": {
                "text_replacements": {},
                "assistant_prefix": "",
                "user_prefix": "For the question below, follow the following instructions:\n"
                                + "-The answer should contain as few words as possible.\n"
                                + "-Don't paraphrase or reformat the text you see in the image.\n"
                                + "-Answer a binary question with Yes or No.\n"
                                + "-When asked to give a numerical value, provide a number like 2 instead of Two.\n"
                                + "-If the final answer has two or more items, provide it in the list format like [1, 2].\n"
                                + "-When asked to give a ratio, give out the decimal value like 0.25 instead of 1:4.\n"
                                + "-When asked to give a percentage, give out the whole value like 17 instead of decimal like 0.17%.\n"
                                + "-Don't include any units in the answer.\n"
                                + "-Do not include any full stops at the end of the answer.\n"
                                + "-Try to include the full label from the graph when asked about an entity.\n"
                                + "Question: ",
                "user_suffix": ""
            },
            ("textvqa_val", "textvqa_test"): {
                "text_replacements": {},
                "assistant_prefix": "",
                "user_prefix": "Answer the following question about the image using as few words as possible. "
                                + "Follow these additional instructions:\n"
                                + "-Always answer a binary question with Yes or No.\n"
                                + "-When asked what time it is, reply with the time seen in the image.\n"
                                + "-Do not put any full stops at the end of the answer.\n"
                                + "-Do not put quotation marks around the answer.\n"
                                + "-An answer with one or two words is favorable.\n"
                                + "-Do not apply common sense knowledge. The answer can be found in the image.\n"
                                + "Question: ",
                "user_suffix": ""
            },
            ("mmmu_val", "mmmu_test"): {
                "text_replacements": {
                    "Question:": "",
                    "Answer with the option's letter from the given choices directly.": "Answer with the letter directly.",
                    "\nA. ": "\nChoices:\nA. "
                },
                "assistant_prefix": "Answer:",
                "user_prefix": "",
                "user_suffix": ""
            },
            ("infovqa_val", "mme", "ocrbench"): {
                "text_replacements": {},
                "assistant_prefix": "",
                "user_prefix": "",
                "user_suffix": "\nGive a very brief answer."
            }
        }
        
        # Check individual task names first
        if task_name in benchmark_formats:
            return benchmark_formats[task_name]
        
        # Check if task is in any list/tuple keys
        for key, formatting in benchmark_formats.items():
            if isinstance(key, (list, tuple)) and task_name in key:
                return formatting
        
        # Default formatting
        return {"text_replacements": {}, "assistant_prefix": "", "user_prefix": "", "user_suffix": ""}
    
    def apply_benchmark_formatting(self, context_str: str, prompt: str, task_name: str) -> tuple[str, str]:
        """Apply benchmark-specific formatting to context and prompt."""
        formatting = self.get_benchmark_formatting(task_name)
        
        # Add user prefix to context
        if formatting["user_prefix"]:
            context_str = formatting["user_prefix"] + context_str
        
        # Apply text replacements to context
        for old_text, new_text in formatting["text_replacements"].items():
            context_str = context_str.replace(old_text, new_text)
        
        # Add user suffix to context
        if formatting["user_suffix"]:
            context_str = context_str + formatting["user_suffix"]
        
        # Add assistant prefix to prompt
        if formatting["assistant_prefix"]:
            prompt = prompt + formatting["assistant_prefix"]
        
        return context_str, prompt
    
    def generate_until(self, requests: List[Instance]) -> List[str]:
        res = []

        def _collate(x):
            # the negative sign on len(toks) sorts descending - this has a few advantages:
            # - time estimates will always be over not underestimates, which is more useful for planning
            # - to know the size of a batch when going through the list, you know the first one is always the batch
            #   padded context length. this is useful to simplify the batching logic and more importantly to make
            #   automatic adaptive batches much much easier to implement
            # - any OOMs will happen right away rather than near the end
            toks = self.tokenizer.encode(x[0])
            return -len(toks), x[0]

        pbar = tqdm(total=len(requests), disable=(self.rank != 0), desc="Model Responding")
        # we group requests by their generation_kwargs,
        # so that we don't try to execute e.g. greedy sampling and temp=0.8 sampling
        # in the same batch.
        re_ords = utils.Collator([reg.args for reg in requests], _collate, grouping=True)
        chunks = re_ords.get_batched(n=self.batch_size, batch_fn=None)
        for chunk in chunks:
            try:
                contexts, all_gen_kwargs, doc_to_visual, doc_id, task, split = zip(*chunk)
                visuals = [dtv(self.task_dict[t][s][i]) for dtv, i, t, s in zip(doc_to_visual, doc_id, task, split)]
                images, splitted_image_ratio = self._prepare_visual_input(self.flatten(visuals))
            except Exception as e:
                print(f"Error preparing visual input: {e}")
                if len(contexts) > 0:
                    pbar.update(len(contexts))
                    generated_texts = [""] * len(contexts)
                    res.extend(generated_texts)
                continue

            messages = []
            splitted_image_idx = 0
            for i in range(len(contexts)):
                current_context_str = contexts[i]
                
                # Apply benchmark-specific text replacements
                current_context_str, _ = self.apply_benchmark_formatting(current_context_str, "", task[i])
                
                if visuals[i] is None:
                    image_count = 0
                else:
                    image_count = len(visuals[i])
                image_string = ""
                for _ in range(image_count):
                    image_string += get_image_string(self.tokenizer, [splitted_image_ratio[splitted_image_idx]], self.model.cfg.mp_image_token_length)
                    splitted_image_idx += 1

                prompt_content = image_string + current_context_str
                
                # Format text_data as a list of message dictionaries
                messages_for_item = [{"role": "user", "content": prompt_content}]
                messages.append(messages_for_item)
                
                # # Process images; _prepare_visual_input returns a stacked tensor or None
                # processed_images_tensor = self._prepare_visual_input(current_visuals_list) if current_visuals_list else None
                # images.append(processed_images_tensor)
                
            prompts = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            pr = False
            if pr:
                print(task[0])
                print("Original Prompt")
                print(prompts[0])

            # Apply benchmark-specific assistant prefixes
            for i in range(len(prompts)):
                _, prompts[i] = self.apply_benchmark_formatting("", prompts[i], task[i])

            if pr:
                print("Formatted Prompt")
                print(prompts[0])

            inputs = self.tokenizer(
                prompts,
                return_tensors="pt",
                padding="longest",
                padding_side="left",
                truncation=True,
                max_length=self.max_length
            )

            input_ids = inputs["input_ids"].to(self.device)
            attention_mask = inputs["attention_mask"].to(self.device)
            # images = images.to(self.device)

            # Extract generation parameters for the batch
            # We use the gen_kwargs from the first item in the chunk, assuming they are uniform for the batch.
            # lmms-eval groups requests by gen_kwargs, so this assumption should hold.
            current_gen_kwargs = all_gen_kwargs[0] if all_gen_kwargs else {}
            max_new_tokens = current_gen_kwargs.get("max_new_tokens", 50)
            temperature = current_gen_kwargs.get("temperature", 0.0) # Default to greedy
            top_p = current_gen_kwargs.get("top_p", 1.0)
            # Check if greedy generation is explicitly requested or implied by temperature 0
            greedy = current_gen_kwargs.get("do_sample", False) is False or temperature == 0.0
            # Pass None for temperature/top_p if greedy, as some HF models expect this
            gen_temperature = temperature if not greedy else None
            gen_top_p = top_p if not greedy else None
            
            # Generate
            generated_ids_batch = self.model.generate(
                input_ids,
                images,
                attention_mask,
                max_new_tokens=max_new_tokens,
                greedy=greedy,
                temperature=gen_temperature,
                top_p=gen_top_p,
            )

            # Decode generated sequences
            # generated_ids_batch from model.generate usually contains only the generated tokens (excluding prompt)
            generated_texts = self.tokenizer.batch_decode(
                generated_ids_batch,
                skip_special_tokens=True
            )
            if pr:
                print(generated_texts[0])
            res.extend(generated_texts)
            pbar.update(len(contexts))

        pbar.close()

        # print(res)
        # re_ords.get_original() will sort the results back to the original order of requests
        return re_ords.get_original(res)

    def generate_until_multi_round(self, requests: List[Instance]) -> List[str]:
        raise NotImplementedError("Multi Round Generation is not implemented for nanoVLM")
    
    @property
    def max_length(self):
        """Return the maximum sequence length."""
        return self.model.cfg.lm_max_position_embeddings 
    
    @property
    def batch_size_per_gpu(self):
        """Return the batch size."""
        return self.batch_size