Integrate to AutoModel (#10)

- Try to integrate AutoModel (24370b049274f4d06c01d479db9cae41ade39e56)
- Update gme_inference.py (02d3017aa1d162ad0ee332345548a1b28351e07a)
- Update README.md (2944f212de0c1ccbe1bc993d660f4a3919b532ee)
- Update README.md (5fa49f4b1ef97eacaf12285d8aca835728b89f68)
- lint (e0e7250f5772018b48d567e195e682d7f4a52586)
- Create modeling_gme_qwen2vl.py (fd9d1a41f746fa74148ccffaf660f7378a0a1026)
- Update config.json (e93daf4018038593fb92d37d3fd1c2fe61f88d12)
- revert gme_inference.py (facd32fbf0276e7365b453eb1400abdf5c0fee5e)
- Update README.md (9e37c1412199889941439551345ead1317bc7e46)
- revert README.md (4e639a8908b3d2d559abe9feaebe99cd19c5806e)


Co-authored-by: Solomatin Roman <[email protected]>

README.md

@@ -3879,4 +3879,4 @@ If you find our paper or models helpful, please consider cite:
       primaryClass={cs.CL},
       url={http://arxiv.org/abs/2412.16855}, 
 }
-```
+```

config.json

@@ -1,8 +1,10 @@
 {
-  "_name_or_path": "gme-Qwen2-VL-2B-Instruct",
-  "architectures": [
-    "Qwen2VLForConditionalGeneration"
-  ],
+  "_name_or_path": "Alibaba-NLP/gme-Qwen2-VL-2B-Instruct",
+  "architectures": ["GmeQwen2VLForVision2Seq"],
+  "auto_map": {
+    "AutoModel": "modeling_gme_qwen2vl.GmeQwen2VLForVision2Seq",
+    "AutoConfig": "modeling_gme_qwen2vl.GmeQwen2VLConfig"
+  },
   "attention_dropout": 0.0,
   "bos_token_id": 151643,
   "eos_token_id": 151645,
@@ -13,17 +15,13 @@
   "intermediate_size": 8960,
   "max_position_embeddings": 32768,
   "max_window_layers": 28,
-  "model_type": "qwen2_vl",
+  "model_type": "gme_qwen2_vl",
   "num_attention_heads": 12,
   "num_hidden_layers": 28,
   "num_key_value_heads": 2,
-  "rms_norm_eps": 1e-06,
+  "rms_norm_eps": 1e-6,
   "rope_scaling": {
-    "mrope_section": [
-      16,
-      24,
-      24
-    ],
+    "mrope_section": [16, 24, 24],
     "type": "mrope"
   },
   "rope_theta": 1000000.0,

modeling_gme_qwen2vl.py

@@ -0,0 +1,314 @@
+from __future__ import annotations
+
+import base64
+import logging
+import math
+import os
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Union
+
+import requests
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+from transformers import (
+    AutoProcessor,
+    PreTrainedModel,
+    Qwen2VLConfig,
+    Qwen2VLForConditionalGeneration,
+)
+import os
+
+
+class GmeQwen2VLConfig(Qwen2VLConfig):
+    def __init__(
+        self,
+        min_image_tokens: int = 256,
+        max_image_tokens: int = 1280,
+        max_length: int = 1800,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.min_image_tokens = min_image_tokens
+        self.max_image_tokens = max_image_tokens
+        self.max_length = max_length
+
+
+class GmeQwen2VLForVision2Seq(PreTrainedModel):
+    config_class = GmeQwen2VLConfig
+    base_model_prefix: str = "base"
+
+    def __init__(self, config: GmeQwen2VLConfig, **kwargs: Any) -> None:
+        super().__init__(config)
+        self.base = Qwen2VLForConditionalGeneration.from_pretrained(config._name_or_path)
+        self.base.tie_weights()  # It's important to produce same outputs.
+
+        min_pixels: int = config.min_image_tokens * 28 * 28
+        max_pixels: int = config.max_image_tokens * 28 * 28
+        self.processor = AutoProcessor.from_pretrained(
+            config._name_or_path, min_pixels=min_pixels, max_pixels=max_pixels, **kwargs
+        )
+        self.max_length: int = config.max_length
+        self.normalize: bool = True
+        self.processor.tokenizer.padding_side = "right"
+        self.default_instruction: str = "You are a helpful assistant."
+        self.sep: str = " "
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        # pixel_values_videos: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        # video_grid_thw: Optional[torch.LongTensor] = None,
+        pooling_mask: Optional[torch.LongTensor] = None,
+        **kwargs
+    ) -> torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.base.model.embed_tokens(input_ids)
+            if pixel_values is not None:
+                pixel_values = pixel_values.type(self.base.visual.get_dtype())
+                image_embeds = self.base.visual(pixel_values, grid_thw=image_grid_thw).to(inputs_embeds.device)
+                image_mask = input_ids == self.base.config.image_token_id
+                inputs_embeds[image_mask] = image_embeds
+            # if pixel_values_videos is not None:
+            #     pixel_values_videos = pixel_values_videos.type(self.base.visual.get_dtype())
+            #     video_embeds = self.base.visual(pixel_values_videos, grid_thw=video_grid_thw).to(inputs_embeds.device)
+            #     video_mask = input_ids == self.base.config.video_token_id
+            #     inputs_embeds[video_mask] = video_embeds
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(inputs_embeds.device)
+
+        outputs = self.base.model(
+            input_ids=None,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+        )
+
+        pooling_mask = attention_mask if pooling_mask is None else pooling_mask
+        left_padding = (pooling_mask[:, -1].sum() == pooling_mask.shape[0])  # TODO
+        if left_padding:
+            embeddings = outputs.last_hidden_state[:, -1]
+        else:
+            sequence_lengths = pooling_mask.sum(dim=1) - 1
+            batch_size = outputs.last_hidden_state.shape[0]
+            embeddings = outputs.last_hidden_state[torch.arange(
+                batch_size, device=outputs.last_hidden_state.device
+            ), sequence_lengths]
+        if self.normalize:
+            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+        return embeddings.contiguous()
+
+    def embed(self, texts: list[str], images: list[Image.Image], is_query=True, instruction=None, **kwargs):
+        self.eval()
+        # Inputs must be batched
+        input_texts, input_images = list(), list()
+        for t, i in zip(texts, images):
+            if not is_query or instruction is None:
+                instruction = self.default_instruction
+            input_str = ''
+            if i is None:
+                input_images = None  # All examples in the same batch are consistent
+            else:
+                input_str += '<|vision_start|><|image_pad|><|vision_end|>'
+                i = fetch_image(i)
+                input_images.append(i)
+            if t is not None:
+                input_str += t
+            msg = f'<|im_start|>system\n{instruction}<|im_end|>\n<|im_start|>user\n{input_str}<|im_end|>\n<|im_start|>assistant\n<|endoftext|>'
+            input_texts.append(msg)
+
+        inputs = self.processor(
+            text=input_texts,
+            images=input_images,
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors='pt'
+        )
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}  # TODO
+        with torch.inference_mode():
+            embeddings = self.forward(**inputs)
+        return embeddings
+
+    def encode(self, sentences: list[str], *, prompt_name=None, **kwargs):
+        return self.get_fused_embeddings(texts=sentences, prompt_name=prompt_name, **kwargs)
+
+    def encode_queries(self, queries: List[str], **kwargs):
+        embeddings = self.encode(queries, **kwargs)
+        return embeddings
+
+    def encode_corpus(self, corpus: List[Dict[str, str]], **kwargs):
+        if type(corpus) is dict:
+            sentences = [
+                (corpus["title"][i] + self.sep + corpus["text"][i]).strip()
+                if "title" in corpus
+                else corpus["text"][i].strip()
+                for i in range(len(corpus["text"]))
+            ]
+        else:
+            sentences = [
+                (doc["title"] + self.sep + doc["text"]).strip() if "title" in doc else doc["text"].strip()
+                for doc in corpus
+            ]
+        embeddings = self.encode(sentences, is_query=False, **kwargs)
+        return embeddings
+
+    def get_image_embeddings(self, images: list[Image.Image] | DataLoader, **kwargs):
+        return self.get_fused_embeddings(images=images, **kwargs)
+
+    def get_text_embeddings(self, texts: list[str], **kwargs):
+        return self.get_fused_embeddings(texts=texts, **kwargs)
+
+    def get_fused_embeddings(self, texts: list[str] = None, images: list[Image.Image] | DataLoader = None, **kwargs):
+        if isinstance(images, DataLoader):
+            image_loader = images
+            batch_size = image_loader.batch_size
+            image_loader.dataset.transform = None
+        else:
+            batch_size = kwargs.pop('batch_size', 32)
+            if images is None:
+                image_loader = None
+            else:
+                image_loader = DataLoader(
+                    images,
+                    batch_size=batch_size,
+                    shuffle=False,
+                    collate_fn=custom_collate_fn,
+                    num_workers=min(math.floor(os.cpu_count() / 2), 8),
+                )
+
+        if texts is None:
+            assert image_loader is not None
+            n_batch = len(image_loader)
+        else:
+            n_batch = len(texts) // batch_size + int(len(texts) % batch_size > 0)
+            image_loader = image_loader or [None] * n_batch
+
+        all_embeddings = list()
+        none_batch = [None] * batch_size
+        show_progress_bar = kwargs.pop('show_progress_bar', False)
+        pbar = tqdm(total=n_batch, disable=not show_progress_bar, mininterval=1, miniters=10, desc='encode')
+        for n, img_batch in zip(range(0, n_batch * batch_size, batch_size), image_loader):
+            text_batch = none_batch if texts is None else texts[n: n+batch_size]
+            img_batch = none_batch if img_batch is None else img_batch
+            embeddings = self.embed(texts=text_batch, images=img_batch, **kwargs)
+            pbar.update(1)
+            all_embeddings.append(embeddings.cpu())
+        pbar.close()
+        all_embeddings = torch.cat(all_embeddings, dim=0)
+        return all_embeddings
+
+
+def custom_collate_fn(batch):
+    return batch
+
+
+### Copied from qwen_vl_utils.vision_process.py
+import base64
+from io import BytesIO
+import requests
+
+IMAGE_FACTOR = 28
+MIN_PIXELS = 4 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(
+    height: int, width: int, factor: int = IMAGE_FACTOR, min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+
+    if max(h_bar, w_bar) / min(h_bar, w_bar) > MAX_RATIO:
+        logging.warning(
+            f"Absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(h_bar, w_bar) / min(h_bar, w_bar)}"
+        )
+        if h_bar > w_bar:
+            h_bar = w_bar * MAX_RATIO
+        else:
+            w_bar = h_bar * MAX_RATIO
+    return h_bar, w_bar
+
+
+def fetch_image(image: str | Image.Image, size_factor: int = IMAGE_FACTOR) -> Image.Image:
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        image_obj = Image.open(requests.get(image, stream=True).raw)
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
+    image = image_obj.convert("RGB")
+    ## resize
+    # if "resized_height" in ele and "resized_width" in ele:
+    #     resized_height, resized_width = smart_resize(
+    #         ele["resized_height"],
+    #         ele["resized_width"],
+    #         factor=size_factor,
+    #     )
+    # else:
+    width, height = image.size
+    # min_pixels = ele.get("min_pixels", MIN_PIXELS)
+    # max_pixels = ele.get("max_pixels", MAX_PIXELS)
+    resized_height, resized_width = smart_resize(
+        height,
+        width,
+        factor=size_factor,
+        min_pixels=MIN_PIXELS,
+        max_pixels=MAX_PIXELS,
+    )
+    image = image.resize((resized_width, resized_height))
+
+    return image
+###
+