long_clip_hf.py

"""
LongCLIP: Unlocking the Long-Text Capability of CLIP

This module provides HuggingFace Transformers-compatible implementations of LongCLIP,
which extends CLIP's text encoder to support 248 tokens (vs 77 in original CLIP).

Repository: https://github.com/beichenzbc/Long-CLIP
Paper: https://arxiv.org/abs/2403.15378
"""

import logging
from typing import Any, Dict, List, Optional, Union

import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPTextConfig, CLIPVisionConfig
from transformers import CLIPTextModel, CLIPVisionModel, CLIPModel
from transformers import CLIPImageProcessor, CLIPTokenizer
from transformers.configuration_utils import PretrainedConfig
from transformers.models.clip.modeling_clip import CLIPTextTransformer
from transformers.processing_utils import ProcessorMixin

logger = logging.getLogger(__name__)


# ================== Configuration Classes ==================


class LongCLIPTextConfig(CLIPTextConfig):
    """
    Configuration class for LongCLIP text model.

    Extends CLIPTextConfig to support 248 token context length
    and custom positional embedding interpolation.

    Args:
        max_position_embeddings (int, optional): Maximum sequence length. Defaults to 248.
        use_position_interpolation (bool, optional): Whether to use position interpolation.
            Defaults to True.
        interpolation_keep_length (int, optional): Number of positions to keep from
            original embeddings before interpolation. Defaults to 20.
        **kwargs: Additional arguments passed to CLIPTextConfig.
    """

    model_type = "longclip_text_model"

    def __init__(
        self,
        max_position_embeddings: int = 248,
        use_position_interpolation: bool = True,
        interpolation_keep_length: int = 20,
        **kwargs,
    ):
        super().__init__(max_position_embeddings=max_position_embeddings, **kwargs)

        self.use_position_interpolation = use_position_interpolation
        self.interpolation_keep_length = interpolation_keep_length


class LongCLIPVisionConfig(CLIPVisionConfig):
    """
    Configuration class for LongCLIP vision model.

    This is identical to the standard CLIPVisionConfig as LongCLIP
    does not modify the vision encoder.

    Args:
        **kwargs: Arguments passed to CLIPVisionConfig.
    """

    model_type = "longclip_vision_model"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)


class LongCLIPConfig(CLIPConfig):
    """
    Configuration class for LongCLIP model.

    Combines LongCLIPTextConfig and LongCLIPVisionConfig to create
    a complete LongCLIP model configuration.

    Args:
        text_config (Dict[str, Any] or LongCLIPTextConfig, optional):
            Configuration for the text model. If None, uses default LongCLIPTextConfig.
        vision_config (Dict[str, Any] or LongCLIPVisionConfig, optional):
            Configuration for the vision model. If None, uses default LongCLIPVisionConfig.
        projection_dim (int, optional): Dimensionality of text and vision projection layers.
            Defaults to 512.
        **kwargs: Additional arguments passed to CLIPConfig.
    """

    model_type = "longclip"
    is_composition = True

    def __init__(
        self,
        text_config: Dict[str, Any] | None = None,
        vision_config: Dict[str, Any] | None = None,
        projection_dim: int = 512,
        **kwargs,
    ):
        # Initialize text config
        if text_config is None:
            text_config = {}
            logger.info(
                "text_config is None. Initializing the LongCLIPTextConfig with default values."
            )

        if vision_config is None:
            vision_config = {}
            logger.info(
                "vision_config is None. Initializing the LongCLIPVisionConfig with default values."
            )

        # Create config objects if they're dictionaries
        if isinstance(text_config, dict):
            text_config = LongCLIPTextConfig(**text_config)

        if isinstance(vision_config, dict):
            vision_config = LongCLIPVisionConfig(**vision_config)

        # Call parent init with config dicts
        super().__init__(
            text_config=text_config.to_dict(),
            vision_config=vision_config.to_dict(),
            projection_dim=projection_dim,
            **kwargs,
        )

        # Store as config objects for easier access
        self.text_config = text_config
        self.vision_config = vision_config

    @classmethod
    def from_text_vision_configs(
        cls,
        text_config: LongCLIPTextConfig,
        vision_config: LongCLIPVisionConfig,
        **kwargs,
    ):
        """
        Instantiate a LongCLIPConfig from text and vision configs.

        Args:
            text_config (LongCLIPTextConfig): Text model configuration.
            vision_config (LongCLIPVisionConfig): Vision model configuration.
            **kwargs: Additional keyword arguments.

        Returns:
            LongCLIPConfig: Configuration object.
        """
        return cls(
            text_config=text_config.to_dict(),
            vision_config=vision_config.to_dict(),
            **kwargs,
        )

    def to_dict(self) -> Dict[str, Any]:
        """
        Serializes this instance to a Python dictionary.

        Returns:
            Dict[str, Any]: Dictionary of all attributes.
        """
        output = super().to_dict()
        # Ensure text_config and vision_config are properly serialized
        if hasattr(self, "text_config") and isinstance(
            self.text_config, PretrainedConfig
        ):
            output["text_config"] = self.text_config.to_dict()
        if hasattr(self, "vision_config") and isinstance(
            self.vision_config, PretrainedConfig
        ):
            output["vision_config"] = self.vision_config.to_dict()
        return output


# ================== Model Classes ==================


class LongCLIPTextEmbeddings(nn.Module):
    """
    Text embeddings for LongCLIP with custom positional embedding mechanism.

    This module implements the dual positional embedding approach used in LongCLIP:
    - The first 20 positions use the original CLIP positional embeddings (mask1)
    - The remaining positions (21-248) use interpolated embeddings (mask2)
    - position_embedding: Fixed base embeddings
    - position_embedding_res: Trainable residual embeddings

    Args:
        config (LongCLIPTextConfig): Configuration for text embeddings.
    """

    def __init__(self, config: LongCLIPTextConfig):
        super().__init__()
        self.config = config
        embed_dim = config.hidden_size

        # Token embeddings
        self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)

        # Dual positional embeddings (LongCLIP approach)
        # position_embedding: Base embeddings (typically loaded from checkpoint)
        self.position_embedding = nn.Embedding(
            config.max_position_embeddings, embed_dim
        )

        # position_embedding_res: Trainable residual embeddings
        self.position_embedding_res = nn.Parameter(
            torch.zeros(config.max_position_embeddings, embed_dim)
        )

        # Create masks for applying embeddings
        # mask1: Use original embeddings for first interpolation_keep_length positions
        # mask2: Use interpolated embeddings for remaining positions
        self.register_buffer(
            "mask1", self._create_mask(config, use_first=True), persistent=False
        )
        self.register_buffer(
            "mask2", self._create_mask(config, use_first=False), persistent=False
        )

        # Store position IDs for efficiency
        self.register_buffer(
            "position_ids",
            torch.arange(config.max_position_embeddings).expand((1, -1)),
            persistent=False,
        )

    def _create_mask(self, config: LongCLIPTextConfig, use_first: bool) -> torch.Tensor:
        """
        Create mask for positional embeddings.

        Args:
            config: Configuration object.
            use_first: If True, mask first `interpolation_keep_length` positions.
                      If False, mask remaining positions.

        Returns:
            Mask tensor of shape [max_position_embeddings, 1].
        """
        mask = torch.zeros(config.max_position_embeddings, 1)
        if use_first:
            # mask1: First interpolation_keep_length positions
            mask[: config.interpolation_keep_length] = 1.0
        else:
            # mask2: Remaining positions
            mask[config.interpolation_keep_length :] = 1.0
        return mask

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
    ) -> torch.Tensor:
        """
        Forward pass for text embeddings.

        Args:
            input_ids: Token IDs of shape [batch_size, seq_length].
            position_ids: Position IDs of shape [batch_size, seq_length].
            inputs_embeds: Pre-computed token embeddings.

        Returns:
            Embeddings of shape [batch_size, seq_length, hidden_size].
        """
        seq_length = (
            input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
        )

        if position_ids is None:
            position_ids = self.position_ids[:, :seq_length]

        # Get token embeddings
        if inputs_embeds is None:
            inputs_embeds = self.token_embedding(input_ids)

        # Get positional embeddings
        position_embeddings = self.position_embedding(position_ids)

        # Add residual positional embeddings (for positions > interpolation_keep_length)
        # Expand position_embedding_res for batch dimension
        position_embeddings_res = self.position_embedding_res.unsqueeze(0).expand(
            position_ids.shape[0], -1, -1
        )[:, :seq_length, :]

        # Apply masks: mask1 for first 20, mask2 for rest
        # Broadcasting: [seq_length, 1] * [batch, seq_length, hidden_size]
        mask1 = self.mask1[:seq_length].transpose(0, 1)  # [1, seq_length]
        mask2 = self.mask2[:seq_length].transpose(0, 1)  # [1, seq_length]

        # Combine embeddings with masking
        embeddings = (
            inputs_embeds
            + position_embeddings * mask1.unsqueeze(-1)
            + position_embeddings_res * mask2.unsqueeze(-1)
        )

        return embeddings


class LongCLIPTextTransformer(CLIPTextTransformer):
    """
    Text transformer for LongCLIP.

    This extends CLIPTextTransformer to use LongCLIPTextEmbeddings
    with custom positional embedding mechanism.

    Args:
        config (LongCLIPTextConfig): Configuration for text transformer.
    """

    def __init__(self, config: LongCLIPTextConfig):
        super().__init__(config)
        # Replace embeddings with LongCLIP version
        self.embeddings = LongCLIPTextEmbeddings(config)


class LongCLIPTextModel(CLIPTextModel):
    """
    LongCLIP text model compatible with HuggingFace Transformers.

    This model extends CLIPTextModel to support 248 token context length
    with custom positional embedding interpolation.

    Args:
        config (LongCLIPTextConfig): Configuration for the text model.
    """

    config_class = LongCLIPTextConfig

    def __init__(self, config: LongCLIPTextConfig):
        super().__init__(config)
        # Replace text_model with LongCLIP version
        self.text_model = LongCLIPTextTransformer(config)
        # Initialize weights
        self.post_init()

    def get_input_embeddings(self) -> nn.Module:
        """Get token embedding layer."""
        return self.text_model.embeddings.token_embedding

    def set_input_embeddings(self, value: nn.Module):
        """Set token embedding layer."""
        self.text_model.embeddings.token_embedding = value


class LongCLIPVisionModel(CLIPVisionModel):
    """
    LongCLIP vision model.

    This is identical to CLIPVisionModel as LongCLIP does not modify
    the vision encoder. Provided for API consistency.

    Args:
        config (LongCLIPVisionConfig): Configuration for the vision model.
    """

    config_class = LongCLIPVisionConfig


class LongCLIPModel(CLIPModel):
    """
    LongCLIP model combining text and vision encoders.

    This model extends CLIPModel to use LongCLIPTextModel with 248 token
    context length while keeping the standard vision encoder.

    Args:
        config (LongCLIPConfig): Configuration for the complete model.
    """

    config_class = LongCLIPConfig

    def __init__(self, config: LongCLIPConfig):
        super().__init__(config)

        # Replace text model with LongCLIP version
        if not isinstance(config.text_config, LongCLIPTextConfig):
            text_config = LongCLIPTextConfig(**config.text_config)
        else:
            text_config = config.text_config

        self.text_model = LongCLIPTextModel(text_config)

        # Vision model stays the same (standard CLIP)
        if not isinstance(config.vision_config, LongCLIPVisionConfig):
            vision_config = LongCLIPVisionConfig(**config.vision_config)
        else:
            vision_config = config.vision_config

        self.vision_model = LongCLIPVisionModel(vision_config)

        # Initialize weights
        self.post_init()

    def get_text_features(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> torch.FloatTensor:
        """
        Get text features from the text encoder.

        Args:
            input_ids: Token IDs.
            attention_mask: Attention mask.
            position_ids: Position IDs.
            output_attentions: Whether to output attention weights.
            output_hidden_states: Whether to output hidden states.
            return_dict: Whether to return a ModelOutput object.

        Returns:
            Text features of shape [batch_size, projection_dim].
        """
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        text_outputs = self.text_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        pooled_output = (
            text_outputs[1] if not return_dict else text_outputs.pooler_output
        )
        text_features = self.text_projection(pooled_output)

        return text_features

    def get_image_features(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> torch.FloatTensor:
        """
        Get image features from the vision encoder.

        Args:
            pixel_values: Pixel values.
            output_attentions: Whether to output attention weights.
            output_hidden_states: Whether to output hidden states.
            return_dict: Whether to return a ModelOutput object.

        Returns:
            Image features of shape [batch_size, projection_dim].
        """
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        vision_outputs = self.vision_model(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        pooled_output = (
            vision_outputs[1] if not return_dict else vision_outputs.pooler_output
        )
        image_features = self.visual_projection(pooled_output)

        return image_features


# ================== Processor Class ==================


class LongCLIPProcessor(ProcessorMixin):
    """
    Processor for LongCLIP that combines image and text preprocessing.

    This processor wraps CLIPImageProcessor and CLIPTokenizer to provide
    a unified interface for preprocessing inputs for LongCLIP models.

    Args:
        image_processor (CLIPImageProcessor): Image processor for preprocessing images.
        tokenizer (CLIPTokenizer): Tokenizer for preprocessing text.

    Attributes:
        image_processor_class (str): Name of the image processor class.
        tokenizer_class (str): Name of the tokenizer class.
    """

    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "CLIPImageProcessor"
    tokenizer_class = "CLIPTokenizer"

    def __init__(
        self,
        image_processor: Optional[CLIPImageProcessor] = None,
        tokenizer: Optional[CLIPTokenizer] = None,
        **kwargs,
    ):
        if image_processor is None:
            raise ValueError("You need to specify an `image_processor`.")
        if tokenizer is None:
            raise ValueError("You need to specify a `tokenizer`.")

        super().__init__(image_processor, tokenizer)

    def __call__(
        self,
        text: Union[str, List[str], None] = None,
        images=None,
        return_tensors: Optional[str] = "pt",
        padding: Union[bool, str] = True,
        max_length: Optional[int] = 248,
        truncation: Optional[bool] = True,
        **kwargs,
    ):
        """
        Preprocess text and images for LongCLIP model.

        Args:
            text (str, List[str], optional): Text or list of texts to process.
            images: Image or list of images to process. Can be PIL Image, numpy array, or tensor.
            return_tensors (str, optional): Type of tensors to return ('pt' for PyTorch).
            padding (bool or str, optional): Padding strategy. Defaults to True.
            max_length (int, optional): Maximum sequence length. Defaults to 248 for LongCLIP.
            truncation (bool, optional): Whether to truncate sequences. Defaults to True.
            **kwargs: Additional keyword arguments.

        Returns:
            BatchEncoding: Dictionary containing processed inputs with keys:
                - input_ids: Tokenized text (if text provided)
                - attention_mask: Attention mask for text (if text provided)
                - pixel_values: Processed images (if images provided)
        """
        # Process text
        if text is not None:
            text_inputs = self.tokenizer(
                text,
                return_tensors=return_tensors,
                padding=padding,
                max_length=max_length,
                truncation=truncation,
                **kwargs,
            )
        else:
            text_inputs = {}

        # Process images
        if images is not None:
            image_inputs = self.image_processor(
                images,
                return_tensors=return_tensors,
            )
        else:
            image_inputs = {}

        # Combine inputs
        return {**text_inputs, **image_inputs}

    def batch_decode(self, *args, **kwargs):
        """
        Decode token IDs back to text.

        This method is forwarded to the tokenizer's batch_decode method.
        """
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        """
        Decode token IDs back to text.

        This method is forwarded to the tokenizer's decode method.
        """
        return self.tokenizer.decode(*args, **kwargs)

    @property
    def model_input_names(self):
        """
        Get the names of model inputs.

        Returns:
            List[str]: List of input names.
        """
        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))


# Register configuration for auto classes
from transformers import AutoConfig, AutoModel

AutoConfig.register("longclip", LongCLIPConfig)
AutoModel.register(LongCLIPConfig, LongCLIPModel)