[Feat] LLaMA-3.2-Vision (#314)

* Fix extra calling in qwen_vl_api, and use tempfile to replace storing tmp files

* Add cache mode for qwen_vl_api

* Add llama vision for video

* Pass gen_kwargs in llama_vision

* Add llama vision instruct

lmms_eval/models/__init__.py

@@ -48,6 +48,7 @@
     "xcomposer2d5": "XComposer2D5",
     "oryx": "Oryx",
     "videochat2": "VideoChat2",
+    "llama_vision": "LlamaVision",
 }
 
 

lmms_eval/models/llama_vision.py

@@ -0,0 +1,234 @@
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import PIL
+import torch
+from accelerate import Accelerator, DistributedType
+from accelerate.state import AcceleratorState
+from decord import VideoReader, cpu
+from torchvision.transforms.functional import to_pil_image
+from tqdm import tqdm
+from transformers import AutoConfig, AutoProcessor, MllamaForConditionalGeneration
+
+from lmms_eval import utils
+from lmms_eval.api.instance import Instance
+from lmms_eval.api.model import lmms
+from lmms_eval.api.registry import register_model
+
+warnings.filterwarnings("ignore")
+
+from loguru import logger as eval_logger
+
+DEFAULT_IMAGE_TOKEN = "<|image|>"
+
+
+@register_model("llama_vision")
+class LlamaVision(lmms):
+    """
+    Llava Model for Hugging Face Transformers: https://huggingface.co/docs/transformers/v4.39.3/en/model_doc/llava
+
+    Adapted from the InstructBLIP model in lmms_eval/models/instructblip.py
+
+    Example usage:
+
+    accelerate launch --num_processes=8 --main_process_port 12345 -m lmms_eval \
+        --model llava_hf \
+        --model_args pretrained=llava-hf/llava-1.5-7b-hf \
+        --tasks seedbench \
+        --batch_size 1 \
+        --output_path ./logs/ \
+        --log_samples
+    """
+
+    def __init__(
+        self,
+        pretrained: str = "meta-llama/Llama-3.2-11B-Vision",
+        revision: str = "main",
+        device: str = "cuda",
+        dtype: Optional[Union[str, torch.dtype]] = "auto",
+        batch_size: int = 1,
+        trust_remote_code: Optional[bool] = False,
+        attn_implementation: Optional[str] = None,
+        device_map: str = "",
+        max_frames_num: Optional[int] = 32,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+        # Do not use kwargs for now
+        assert kwargs == {}, f"Unexpected kwargs: {kwargs}"
+
+        accelerator = Accelerator()
+        if accelerator.num_processes > 1 and device_map == "":
+            self._device = torch.device(f"cuda:{accelerator.local_process_index}")
+            self.device_map = f"cuda:{accelerator.local_process_index}"
+        else:
+            self._device = torch.device(device)
+            self.device_map = device_map
+        if isinstance(dtype, str) and dtype != "auto":
+            dtype = getattr(torch, dtype)
+
+        self.max_frames_num = max_frames_num
+        self._model = MllamaForConditionalGeneration.from_pretrained(pretrained, revision=revision, torch_dtype=dtype, device_map=self.device_map, trust_remote_code=trust_remote_code, attn_implementation=attn_implementation)
+        self.model.eval()
+        self.processor = AutoProcessor.from_pretrained(pretrained)
+        if accelerator.num_processes > 1 and device_map == "":
+            assert accelerator.distributed_type in [DistributedType.FSDP, DistributedType.MULTI_GPU, DistributedType.DEEPSPEED], "Unsupported distributed type provided. Only DDP and FSDP are supported."
+            # If you want to use DistributedType.DEEPSPEED, you have to run accelerate config before using the model
+            # Also, you have to select zero stage 0 (equivalent to DDP) in order to make the prepare model works
+            # I tried to set different parameters in the kwargs to let default zero 2 stage works, but it didn't work.
+            if accelerator.distributed_type == DistributedType.DEEPSPEED:
+                kwargs = {
+                    "train_micro_batch_size_per_gpu": self.batch_size_per_gpu,
+                    "train_batch_size": self.batch_size_per_gpu * accelerator.num_processes,
+                }
+                AcceleratorState().deepspeed_plugin.deepspeed_config_process(must_match=True, **kwargs)
+                eval_logger.info("Detected that you are using DistributedType.DEEPSPEED. Make sure you run `accelerate config` and set zero stage to 0")
+            if accelerator.distributed_type == DistributedType.FSDP or accelerator.distributed_type == DistributedType.DEEPSPEED:
+                self._model = accelerator.prepare(self.model)
+            else:
+                self._model = accelerator.prepare_model(self.model, evaluation_mode=True)
+            self.accelerator = accelerator
+            if self.accelerator.is_local_main_process:
+                eval_logger.info(f"Using {accelerator.num_processes} devices with data parallelism")
+            self._rank = self.accelerator.local_process_index
+            self._world_size = self.accelerator.num_processes
+        elif accelerator.num_processes == 1 and device_map == "auto":
+            eval_logger.info(f"Using {accelerator.num_processes} devices with pipeline parallelism")
+            self._rank = 0
+            self._word_size = 1
+        else:
+            eval_logger.info(f"Using single device: {self._device}")
+            self.model.to(self._device)
+            self._rank = 0
+            self._word_size = 1
+        self.accelerator = accelerator
+
+    @property
+    def config(self):
+        # return the associated transformers.AutoConfig for the given pretrained model.
+        return self._config
+
+    @property
+    def tokenizer(self):
+        return self._tokenizer
+
+    @property
+    def model(self):
+        # returns the model, unwrapping it if using Accelerate
+        if hasattr(self, "accelerator"):
+            return self.accelerator.unwrap_model(self._model)
+        else:
+            return self._model
+
+    @property
+    def eot_token_id(self):
+        # we use EOT because end of *text* is more accurate for what we're doing than end of *sentence*
+        return self.tokenizer.eos_token_id
+
+    @property
+    def max_length(self):
+        return self._max_length
+
+    @property
+    def batch_size(self):
+        return self.batch_size_per_gpu
+
+    @property
+    def device(self):
+        return self._device
+
+    @property
+    def rank(self):
+        return self._rank
+
+    @property
+    def world_size(self):
+        return self._world_size
+
+    def tok_encode(self, string: str, left_truncate_len=None, add_special_tokens=None) -> List[int]:
+        """ """
+        add_special_tokens = False if add_special_tokens is None else add_special_tokens
+        encoding = self.tokenizer.encode(string, add_special_tokens=add_special_tokens)
+        # left-truncate the encoded context to be at most `left_truncate_len` tokens long
+        if left_truncate_len:
+            encoding = encoding[-left_truncate_len:]
+        return encoding
+
+    def tok_decode(self, tokens):
+        return self.tokenizer.decode(tokens)
+
+    def loglikelihood(self, requests: List[Instance]) -> List[Tuple[float, bool]]:
+        assert False, "Not implemented"
+
+    def flatten(self, input):
+        new_list = []
+        for i in input:
+            for j in i:
+                new_list.append(j)
+        return new_list
+
+    def load_video(self, video_path, max_frames_num):
+        if type(video_path) == str:
+            vr = VideoReader(video_path, ctx=cpu(0))
+        else:
+            vr = VideoReader(video_path[0], ctx=cpu(0))
+        total_frame_num = len(vr)
+        uniform_sampled_frames = np.linspace(0, total_frame_num - 1, max_frames_num, dtype=int)
+        frame_idx = uniform_sampled_frames.tolist()
+        spare_frames = vr.get_batch(frame_idx).asnumpy()
+        return spare_frames  # (frames, height, width, channels)
+
+    def generate_until(self, requests: List[Instance]) -> List[str]:
+        res = []
+
+        pbar = tqdm(total=len(requests), disable=(self.rank != 0), desc="Model Responding")
+
+        for contexts, gen_kwargs, doc_to_visual, doc_id, task, split in [reg.args for reg in requests]:
+            visuals = [doc_to_visual(self.task_dict[task][split][doc_id])]
+            visuals = self.flatten(visuals)
+
+            messages = [{"role": "user", "content": []}]
+            images = []
+
+            for visual in visuals:
+                if isinstance(visual, str):
+                    frames = self.load_video(visual, self.max_frames_num)
+                    frames = torch.from_numpy(frames).permute(0, 3, 1, 2)
+                    images.extend([to_pil_image(frame) for frame in frames])
+                elif isinstance(visual, PIL.Image.Image):
+                    images.append(visual)
+
+            for _ in range(len(images)):
+                messages[-1]["content"].append({"type": "image"})
+            messages[-1]["content"].append({"type": "text", "content": contexts})
+            prompt = self.processor.apply_chat_template(messages, add_generation_prompt=True)
+            inputs = self.processor(images, prompt, return_tensors="pt").to(self.model.device)
+
+            if "max_new_tokens" not in gen_kwargs:
+                gen_kwargs["max_new_tokens"] = 1024
+            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
+            if "do_sample" not in gen_kwargs:
+                gen_kwargs["do_sample"] = False
+
+            with torch.no_grad():
+                output = self.model.generate(
+                    **inputs,
+                    max_new_tokens=gen_kwargs["max_new_tokens"],
+                    temperature=gen_kwargs["temperature"],
+                    do_sample=gen_kwargs["do_sample"],
+                )
+                output = output[:, inputs["input_ids"].shape[-1] :]
+                res.append(self.processor.decode(output[0]))
+
+            pbar.update(1)
+        pbar.close()
+        return res
+
+    def generate_until_multi_round(self, requests) -> List[str]:
+        raise NotImplementedError("TODO: Implement multi-round generation for LLaVAHF")

lmms_eval/models/qwen_vl_api.py

@@ -1,5 +1,7 @@
 import base64
+import json
 import os
+import tempfile
 import time
 from copy import deepcopy
 from io import BytesIO
@@ -33,14 +35,32 @@ def __init__(
         image_token: str = "<image>",  # Use to separate interleaved image and text
         system_prompt: str = "",  # Whether you want some special system prompt here
         tmp_folder: str = "./tmp",  # Due to qwen's api restriction,
+        continual_mode: bool = False,
+        response_persistent_folder: str = None,
         **kwargs,
     ) -> None:
         super().__init__()
+        self.continual_mode = continual_mode
 
         self.model_version = model_version
         self.image_token = image_token
         self.system_prompt = system_prompt
         self.tmp_folder = tmp_folder
+        if self.continual_mode:
+            if response_persistent_folder is None:
+                raise ValueError("Continual mode requires a persistent path for the response. Please provide a valid path.")
+
+            os.makedirs(response_persistent_folder, exist_ok=True)
+            self.response_persistent_folder = response_persistent_folder
+            self.response_persistent_file = os.path.join(self.response_persistent_folder, f"{self.model_version}_response.json")
+
+            if os.path.exists(self.response_persistent_file):
+                with open(self.response_persistent_file, "r") as f:
+                    self.response_cache = json.load(f)
+                self.cache_mode = "resume"
+            else:
+                self.response_cache = {}
+                self.cache_mode = "start"
 
     @property
     def rank(self):
@@ -50,20 +70,32 @@ def rank(self):
     def world_size(self):
         return self._world_size
 
+    def save_image_to_temp_file(self, image):
+        temp_file = tempfile.NamedTemporaryFile(suffix=".jpeg", delete=True)
+        image.save(temp_file.name)
+        return temp_file
+
     def generate_until(self, requests) -> List[str]:
         res = []
         pbar = tqdm(total=len(requests), disable=(self.rank != 0), desc="Model Responding")
-        os.makedirs(self.tmp_folder, exist_ok=True)
 
         for contexts, gen_kwargs, doc_to_visual, doc_id, task, split in [reg.args for reg in requests]:
+            if self.continual_mode is True and self.cache_mode == "resume":
+                doc_uuid = f"{task}___{split}___{doc_id}"
+                if doc_uuid in self.response_cache:
+                    response_text = self.response_cache[doc_uuid]
+                    if response_text:
+                        res.append(response_text)
+                        pbar.update(1)
+                        continue
             # encode, pad, and truncate contexts for this batch
             visuals = [doc_to_visual(self.task_dict[task][split][doc_id])]
             visuals = self.flatten(visuals)
             imgs = []
 
             for idx, visual in enumerate(visuals):
-                visual.save(os.path.join(self.tmp_folder, f"tmp_{idx}_{self.rank}_{self.world_size}.jpg"))
-                imgs.append(os.path.join(self.tmp_folder, f"tmp_{idx}_{self.rank}_{self.world_size}.jpg"))
+                temp_file = self.save_image_to_temp_file(visual)
+                imgs.append(temp_file.name)
 
             messages = [{"role": "user", "content": []}]
 
@@ -90,7 +122,8 @@ def generate_until(self, requests) -> List[str]:
 
             for attempt in range(5):
                 try:
-                    response_data = dashscope.MultiModalConversation.call(model=self.model_version, messages=messages, api_key=API_KEY, max_length=gen_kwargs["max_new_tokens"])
+                    response_data = dashscope.MultiModalConversation.call(model=self.model_version, messages=messages, api_key=API_KEY, max_length=gen_kwargs["max_new_tokens"], temperature=gen_kwargs["temperature"])
+                    break
                 except Exception as e:
                     eval_logger.info(f"Attempt {attempt + 1} failed with error: {str(e)}")
                     if attempt < 5 - 1:  # If we have retries left, sleep and then continue to next attempt
@@ -106,6 +139,12 @@ def generate_until(self, requests) -> List[str]:
                 eval_logger.error(f"Error {e} happens when parsing input.")
                 eval_logger.error(f"{response_data}")
                 res.append("")
+
+            if self.continual_mode is True:  # Cache the response
+                doc_uuid = f"{task}___{split}___{doc_id}"
+                self.response_cache[doc_uuid] = res[-1]
+                with open(self.response_persistent_file, "w") as f:
+                    json.dump(self.response_cache, f)
             pbar.update(1)
 
         pbar.close()