modelling.py

# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/12_modelling.ipynb.

# %% auto 0
__all__ = ['VTDEConfig', 'VTDEModel', 'SamVisionPreTrainedModel', 'SamVisionModel']

# %% ../notebooks/12_modelling.ipynb 1
from transformers.models.clip.modeling_clip import CLIPOutput, clip_loss
from typing import Optional, Tuple, Union
from transformers import PreTrainedModel, VisionTextDualEncoderModel
import torch
from transformers import VisionTextDualEncoderConfig

class VTDEConfig(VisionTextDualEncoderConfig):
    model_type = "vtde"
    
    def __init__(self, projection_dim=512, logit_scale_init_value=2.6592, 
    text_pooling_mode='mean',
    vision_pooling_mode='max',
    **kwargs):
        """
        pooling_mode in ['mean', 'max', 'cls']
        https://arxiv.org/pdf/2210.09996.pdf
        https://github.com/kahnchana/clippy/blob/3c102c29c32f7c66c6e52e09b795fe9c061bbb03/src/open_clip/hf_model.py#L56
        also
        https://arxiv.org/pdf/2301.07836.pdf
        """
        self.text_pooling_mode = text_pooling_mode
        self.vision_pooling_mode = vision_pooling_mode
        super().__init__(projection_dim, logit_scale_init_value, **kwargs)

VTDEConfig.register_for_auto_class()


class VTDEModel(VisionTextDualEncoderModel):
    config_class = VTDEConfig
    base_model_prefix = "vtde"

    def __init__(
        self,
        config: Optional[VTDEConfig] = None,
        vision_model: Optional[PreTrainedModel] = None,
        text_model: Optional[PreTrainedModel] = None,
    ):
        # You can customize the constructor if needed
        super().__init__(config, vision_model, text_model)
        self.text_pooling_mode = config.text_pooling_mode
        self.vision_pooling_mode = config.vision_pooling_mode

    def get_text_features(
        self,
        input_ids=None,
        attention_mask=None,
        position_ids=None,
        token_type_ids=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        text_outputs = self.text_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        if self.text_pooling_mode == 'cls':
            pooled_output = text_outputs[1]
        elif self.text_pooling_mode == 'mean':
            pooled_output = torch.mean(text_outputs[0], dim=1)
        elif self.text_pooling_mode == 'max':
            pooled_output = torch.max(text_outputs[0], dim=1)[0]
        elif self.text_pooling_mode == 'norm':
            """we select the patch with the largest norm"""
            last_hidden_states = text_outputs[0]
            patch_norms = torch.norm(last_hidden_states[:, 1:, :], dim=-1)
            max_norm_idx = torch.argmax(patch_norms, dim=1)
            pooled_output = last_hidden_states[:, max_norm_idx, :][:, 0, :]
        else:
            "We want to raise the name of the pooling mode"
            raise NotImplementedError

        text_features = self.text_projection(pooled_output)

        return text_features

    def get_image_features(
        self,
        pixel_values=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        vision_outputs = self.vision_model(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        if self.vision_pooling_mode == 'cls':
            pooled_output = vision_outputs[1]
        elif self.vision_pooling_mode == 'mean':
            pooled_output = torch.mean(vision_outputs[0], dim=1)
        elif self.vision_pooling_mode == 'max':
            pooled_output = torch.max(vision_outputs[0], dim=1)[0]
        elif self.vision_pooling_mode == 'norm':
            """we select the patch with the largest norm"""
            last_hidden_states = vision_outputs[0]
            patch_norms = torch.norm(last_hidden_states[:, 1:, :], dim=-1)
            max_norm_idx = torch.argmax(patch_norms, dim=1)
            pooled_output = last_hidden_states[:, max_norm_idx, :][:, 0, :]
        else:
            raise NotImplementedError

        image_features = self.visual_projection(pooled_output)

        return image_features

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        return_loss: Optional[bool] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], CLIPOutput]:

        return_dict = return_dict if return_dict is not None else self.config.return_dict

        image_embeds = self.get_image_features(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        text_embeds = self.get_text_features(
            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,
        )

        # normalized features
        image_embeds = image_embeds / image_embeds.norm(dim=-1, keepdim=True)
        text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)

        # cosine similarity as logits
        logit_scale = self.logit_scale.exp()
        logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
        logits_per_image = logits_per_text.T

        loss = None
        if return_loss:
            loss = clip_loss(logits_per_text)

        if not return_dict:
            output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_embeds, image_embeds)
            return ((loss,) + output) if loss is not None else output

        return CLIPOutput(
            loss=loss,
            logits_per_image=logits_per_image,
            logits_per_text=logits_per_text,
            text_embeds=text_embeds,
            image_embeds=image_embeds,
            text_model_output=text_embeds,
            vision_model_output=image_embeds,
        )
    

VTDEModel.register_for_auto_class("AutoModel")
VTDEModel.register_for_auto_class("AutoModelForZeroShotImageClassification")

# %% ../notebooks/12_modelling.ipynb 2
# we want to create a vision-text encoder model for SAM
from transformers import PreTrainedModel
from transformers.models.sam.modeling_sam import SamPositionalEmbedding, SamVisionEncoder, SamVisionEncoderOutput
from transformers.models.sam.configuration_sam import SamVisionConfig
from torch import nn

class SamVisionPreTrainedModel(PreTrainedModel):
    config_class = SamVisionConfig
    base_model_prefix = "sam_vision_encoder"
    main_input_name = "pixel_values"

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

class SamVisionModel(SamVisionPreTrainedModel):

    def __init__(self, config):
        super().__init__(config)
        self.shared_image_embedding = SamPositionalEmbedding(config)
        self.vision_encoder = SamVisionEncoder(config)

    def forward(
        self,
        pixel_values=None,
        attention_mask=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ) -> SamVisionEncoderOutput:
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        image_embeddings = self.shared_image_embedding(pixel_values)
        vision_encoder_outputs = self.vision_encoder(
            image_embeddings,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        return vision_encoder_outputs
    
SamVisionModel.register_for_auto_class("AutoModel")
# SamVisionConfig.register_for_auto_class()