spec2_processor.py

# Copyright 2025 SVECTOR AI and The Spec-2 Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Dict, List, Optional, Union

import numpy as np
import PIL.Image
import torch
from transformers import ProcessorMixin
from transformers.image_utils import PILImageResampling, is_vision_available

from .image_processor import Spec2ImageProcessor
from .tokenizer import Spec2Tokenizer


if is_vision_available():
    from PIL import Image


class Spec2Processor(ProcessorMixin):
    """
    Constructs a Spec2 processor which combines a Spec2 image processor and a Spec2 tokenizer into a single processor.

    The processor can be used to prepare inputs for the model by processing text and images appropriately.

    Args:
        image_processor (`Spec2ImageProcessor`):
            An instance of `Spec2ImageProcessor`.
        tokenizer (`Spec2Tokenizer`):
            An instance of `Spec2Tokenizer`.
    """

    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "Spec2ImageProcessor"
    tokenizer_class = "Spec2Tokenizer"

    def __init__(self, image_processor, tokenizer):
        if not is_vision_available():
            raise ImportError("Vision libraries are not available. Make sure to install PIL and Pillow: pip install Pillow")

        self.image_processor = image_processor
        self.tokenizer = tokenizer
        self.image_token = self.tokenizer.image_token
        self.image_token_id = self.tokenizer.image_token_id

    def __call__(
        self,
        text=None,
        images=None,
        return_tensors=None,
        **kwargs
    ):
        """
        Process text and images for the model.

        Args:
            text (`str`, `List[str]`, `List[List[str]]`):
                The text to be processed. Can be a string, a list of strings or a list of lists of strings.
            images (`PIL.Image.Image`, `List[PIL.Image.Image]`, `torch.Tensor`, `List[torch.Tensor]`):
                The images to be processed. Can be a PIL Image, a list of PIL Images, a tensor or a list of tensors.
            return_tensors (`str`, optional):
                The type of tensors to return. Can be one of: 'pt' (PyTorch), 'tf' (TensorFlow), 'np' (NumPy).

        Returns:
            A dictionary containing the processed inputs with keys like 'input_ids', 'attention_mask', 'pixel_values', etc.
        """
        encoding = {}

        # Process text inputs
        if text is not None:
            text_inputs = self._process_text(text, **kwargs)
            encoding.update(text_inputs)

        # Process image inputs
        if images is not None:
            image_features = self._process_images(images, **kwargs)
            encoding.update(image_features)

        # Handle multimodal case - if we have both text and images
        if text is not None and images is not None:
            encoding = self._merge_text_and_image_features(encoding, **kwargs)

        # Convert to tensors if requested
        if return_tensors is not None:
            encoding = self._convert_to_tensors(encoding, return_tensors=return_tensors)

        return encoding

    def _process_text(self, text, **kwargs):
        """Process text inputs."""
        if isinstance(text, str):
            # Check if text already contains image token
            if self.image_token not in text:
                # For single text with images, we add the image token at the end
                text = f"{text} {self.image_token}"
        elif isinstance(text, list):
            # For a list of texts, add image token if not already present
            if all(isinstance(t, str) for t in text):
                text = [f"{t} {self.image_token}" if self.image_token not in t else t for t in text]

        # Tokenize text
        text_encoding = self.tokenizer(text, return_tensors=None, **kwargs)
        return text_encoding

    def _process_images(self, images, **kwargs):
        """Process image inputs."""
        # Convert single image to list
        if not isinstance(images, list):
            images = [images]

        # Process images with image processor
        image_features = self.image_processor(images, return_tensors=None, **kwargs)
        return image_features

    def _merge_text_and_image_features(self, encoding, **kwargs):
        """Merge text and image features for multimodal inputs."""
        # This function handles the specific logic for merging text tokens with image embeddings
        # For Spec-2, we maintain the tokens order and ensure image token is properly placed
        
        input_ids = encoding.get("input_ids", [])
        pixel_values = encoding.get("pixel_values", [])
        
        if isinstance(input_ids[0], list):  # batch case
            batch_size = len(input_ids)
            merged_encoding = {
                "input_ids": input_ids,
                "pixel_values": pixel_values,
                "image_token_indices": []
            }
            
            # For each item in the batch, find the position of the image token
            for i, ids in enumerate(input_ids):
                image_token_indices = [j for j, id_val in enumerate(ids) if id_val == self.image_token_id]
                if image_token_indices:
                    merged_encoding["image_token_indices"].append(image_token_indices[0])
                else:
                    # If no image token found, append it at the end
                    input_ids[i].append(self.image_token_id)
                    merged_encoding["image_token_indices"].append(len(input_ids[i]) - 1)
            
            # Update attention masks if present
            if "attention_mask" in encoding:
                merged_encoding["attention_mask"] = encoding["attention_mask"]
                
        else:  # single item case
            image_token_indices = [i for i, id_val in enumerate(input_ids) if id_val == self.image_token_id]
            if image_token_indices:
                image_token_index = image_token_indices[0]
            else:
                # If no image token found, append it at the end
                input_ids.append(self.image_token_id)
                image_token_index = len(input_ids) - 1
                
            merged_encoding = {
                "input_ids": input_ids,
                "pixel_values": pixel_values[0] if pixel_values else None,
                "image_token_index": image_token_index
            }
            
            # Update attention mask if present
            if "attention_mask" in encoding:
                merged_encoding["attention_mask"] = encoding["attention_mask"]
        
        return merged_encoding

    def _convert_to_tensors(self, encoding, return_tensors="pt"):
        """Convert processed features to tensors."""
        # Convert all features to tensors of the requested type
        for key, value in encoding.items():
            if key in ["pixel_values", "input_ids", "attention_mask"]:
                if return_tensors == "pt":
                    if isinstance(value, list) and all(isinstance(v, list) for v in value):
                        # For batched inputs
                        encoding[key] = torch.tensor(value)
                    elif isinstance(value, list):
                        # For single inputs
                        encoding[key] = torch.tensor([value])
                    elif isinstance(value, np.ndarray):
                        encoding[key] = torch.tensor(value)
                elif return_tensors == "np":
                    if isinstance(value, list):
                        encoding[key] = np.array(value)
                    elif isinstance(value, torch.Tensor):
                        encoding[key] = value.numpy()
                # Add other tensor types (tf, etc.) as needed
        
        return encoding

    @property
    def model_input_names(self):
        tokenizer_input_names = self.tokenizer.model_input_names
        image_processor_input_names = self.image_processor.model_input_names
        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))