Update custom_st.py

custom_st.py

@@ -59,73 +59,163 @@ class MultiModalTransformer(BaseTransformer):
                 )
                 image_mask = features["input_ids"] == self.auto_model.config.image_token_id
                 features["inputs_embeds"][image_mask] = image_embeds
-                features.pop("pixel_values")
-                features.pop("image_grid_thw")
-        features.pop("input_ids")       
+                # features.pop("pixel_values")
+                # features.pop("image_grid_thw")
+        # features.pop("input_ids")
+        inputs = {k: v for k, v in features.items() if k in 'position_ids,attention_mask,inputs_embeds'}
         outputs = self.auto_model.model(
-            **features,
+            **inputs,
             return_dict=True,
             output_hidden_states=True,
             # **kwargs
         )
-        pooling_mask = features["attention_mask"] if features.get("pooling_mask", None) is None else features["pooling_mask"]
-        left_padding = (pooling_mask[:, -1].sum() == pooling_mask.shape[0])  # TODO
-        if left_padding:
-            embeddings = outputs.last_hidden_state
-        else:
-            sequence_lengths = pooling_mask.sum(dim=1) - 1
-            embeddings = outputs.last_hidden_state[torch.arange(
-                outputs.last_hidden_state.shape[0], device=outputs.last_hidden_state.device
-            ), sequence_lengths]
-        features.update({"token_embeddings": embeddings})
+        # pooling_mask = features["attention_mask"] if features.get("pooling_mask", None) is None else features["pooling_mask"]
+        # left_padding = (pooling_mask[:, -1].sum() == pooling_mask.shape[0])  # TODO
+        # if left_padding:
+        #     embeddings = outputs.last_hidden_state
+        # else:
+        #     sequence_lengths = pooling_mask.sum(dim=1) - 1
+        #     embeddings = outputs.last_hidden_state[torch.arange(
+        #         outputs.last_hidden_state.shape[0], device=outputs.last_hidden_state.device
+        #     ), sequence_lengths]
+        features.update({"token_embeddings": outputs.last_hidden_state})
         return features 
 
-    def tokenize(self, texts: List[List[Dict[str, Image.Image]]] | List[str]) -> Dict[str, torch.Tensor]:
-        split_token = "<|im_end|>\n"
-        def process_text_item(item):
+    def tokenize(self, texts: List[List[Dict[str, Any]]] | List[str]) -> Dict[str, torch.Tensor]:
+        default_instruction = 'You are a helpful assistant.'
+
+        all_texts, all_images = list(), list()
+        for item in texts:
             if isinstance(item, str):
-                return item, None
-
-            text, img = "", None
-            if "image" in item:
-                text += "<|vision_start|><|image_pad|><|vision_end|>"
-                img = item["image"]
-                if isinstance(img, bytes):
-                    img = Image.open(BytesIO(img)).convert("RGB")
-                elif isinstance(img, str):
-                    img = Image.open(img).convert("RGB")
-                elif not isinstance(img, Image):
-                    raise ValueError(f"Unknown image type {type(img)}")
-            if "text" in item:
-                text += item["text"].lstrip()
-            if split_token in text:
-                instruction, text = text.split(split_token, 1)
-                text = f'{instruction}{split_token}<|im_start|>user\n{input_str}<|im_end|>\n<|im_start|>assistant\n<|endoftext|>'
+                txt, img, inst = item, None, default_instruction
+            elif isinstance(item, dict):
+                txt = item.get('text', None)
+                img = item.get('image', None)
+                inst = item.get('prompt', default_instruction)
             else:
-                text = f"<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n<|endoftext|>"
-            return text, img
+                raise RuntimeError(f'Input format not supported! {item=}')
 
-        all_texts, all_images = [], []
-        for item in texts:
-            text, images = process_text_item(item)
-            all_texts.append(text)
-            all_images.append(images)
-        
-        if all_images != [None] * len(all_images):
-            inputs = self.processor(
-                text=all_texts,
-                images=all_images,
-                padding="longest",
-                truncation=True,
-                max_length=self.max_seq_length,
-                return_tensors="pt"
-            )
-        else:
-            inputs = self.processor(
-                text=all_texts,
-                padding="longest",
-                truncation=True,
-                max_length=self.max_seq_length,
-                return_tensors="pt"
-            )
+            input_str = ''
+            if img is None:
+                all_images = None  # All examples in the same batch are consistent
+                # or will have ValueError: Could not make a flat list of images from xxxx
+            else:
+                input_str += '<|vision_start|><|image_pad|><|vision_end|>'
+                img = fetch_image(img)
+                all_images.append(img)
+            if txt is not None:
+                input_str += txt
+            msg = f'<|im_start|>system\n{inst}<|im_end|>\n<|im_start|>user\n{input_str}<|im_end|>\n<|im_start|>assistant\n<|endoftext|>'
+            all_texts.append(msg)
+
+        inputs = self.processor(
+            text=all_texts,
+            images=all_images,
+            padding="longest",
+            truncation=True,
+            max_length=self.max_seq_length,
+            return_tensors='pt'
+        )
         return inputs
+
+
+### 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
+###