Upload modeling_longclip.py with huggingface_hub

modeling_longclip.py

@@ -0,0 +1,400 @@
+"""
+LongCLIP model implementation compatible with HuggingFace Transformers.
+
+This module provides transformers-compatible implementations of LongCLIP models.
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from transformers import CLIPTextModel, CLIPVisionModel, CLIPModel
+from transformers.models.clip.modeling_clip import (
+    CLIPTextTransformer,
+)
+
+from .configuration_longclip import (
+    LongCLIPConfig,
+    LongCLIPTextConfig,
+    LongCLIPVisionConfig,
+)
+
+
+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.
+
+    Example:
+        ```python
+        >>> from long_clip_hf import LongCLIPTextConfig, LongCLIPTextModel
+        >>> from transformers import CLIPTokenizer
+        >>>
+        >>> # Initialize model
+        >>> config = LongCLIPTextConfig()
+        >>> model = LongCLIPTextModel(config)
+        >>>
+        >>> # Tokenize text
+        >>> tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+        >>> inputs = tokenizer(
+        ...     ["a photo of a cat"],
+        ...     return_tensors="pt",
+        ...     padding="max_length",
+        ...     max_length=248,
+        ...     truncation=True,
+        ... )
+        >>>
+        >>> # Get text features
+        >>> outputs = model(**inputs)
+        >>> text_features = outputs.pooler_output
+        ```
+    """
+
+    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.
+
+    Example:
+        ```python
+        >>> from long_clip_hf import LongCLIPVisionConfig, LongCLIPVisionModel
+        >>> from transformers import CLIPImageProcessor
+        >>> from PIL import Image
+        >>>
+        >>> # Initialize model
+        >>> config = LongCLIPVisionConfig()
+        >>> model = LongCLIPVisionModel(config)
+        >>>
+        >>> # Process image
+        >>> processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
+        >>> image = Image.open("path/to/image.jpg")
+        >>> inputs = processor(images=image, return_tensors="pt")
+        >>>
+        >>> # Get image features
+        >>> outputs = model(**inputs)
+        >>> image_features = outputs.pooler_output
+        ```
+    """
+
+    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.
+
+    Example:
+        ```python
+        >>> from long_clip_hf import LongCLIPConfig, LongCLIPModel
+        >>> from transformers import CLIPTokenizer, CLIPImageProcessor
+        >>> from PIL import Image
+        >>>
+        >>> # Initialize model
+        >>> config = LongCLIPConfig()
+        >>> model = LongCLIPModel(config)
+        >>>
+        >>> # Prepare inputs
+        >>> tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+        >>> processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
+        >>>
+        >>> text = "a photo of a cat"
+        >>> image = Image.open("path/to/image.jpg")
+        >>>
+        >>> text_inputs = tokenizer(
+        ...     [text],
+        ...     return_tensors="pt",
+        ...     padding="max_length",
+        ...     max_length=248,
+        ...     truncation=True,
+        ... )
+        >>> image_inputs = processor(images=image, return_tensors="pt")
+        >>>
+        >>> # Get features
+        >>> outputs = model(
+        ...     input_ids=text_inputs["input_ids"],
+        ...     pixel_values=image_inputs["pixel_values"],
+        ... )
+        >>>
+        >>> # Compute similarity
+        >>> logits_per_image = outputs.logits_per_image
+        >>> probs = logits_per_image.softmax(dim=1)
+        ```
+    """
+
+    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