Upload nllb_mrl.py

nllb_mrl.py

@@ -0,0 +1,150 @@
+from typing import List, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import PyTorchModelHubMixin
+from open_clip import create_model_and_transforms, get_tokenizer
+from PIL import Image
+from transformers import PretrainedConfig, PreTrainedModel
+
+
+class MatryoshkaNllbClipConfig(PretrainedConfig):
+    def __init__(
+        self,
+        clip_model_name: str = "",
+        clip_model_version: str = "",
+        target_resolution: int = -1,
+        mrl_resolutions: List[int] = [],
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.clip_model_name = clip_model_name
+        self.clip_model_version = clip_model_version
+        self.target_resolution = target_resolution
+        self.mrl_resolutions = mrl_resolutions
+
+
+class MatryoshkaLayer(nn.Module):
+    def __init__(self, resolutions: List[int], target_resolution: int = 768):
+        super().__init__()
+        self.resolutions = resolutions
+        self.layers = nn.ModuleDict()
+        for resolution in resolutions:
+            self.layers[str(resolution)] = nn.Linear(target_resolution, resolution)
+
+    def forward(self, x, resolution: Union[int, None] = None):
+        if resolution is not None:
+            if resolution not in self.resolutions:
+                raise ValueError(f"Resolution {resolution} not in {self.resolutions}")
+            return self.layers[str(resolution)](x)
+        outputs = []
+        for resolution in self.resolutions:
+            outputs.append(self.layers[str(resolution)](x))
+        return outputs
+
+
+class MatryoshkaNllbClip(PreTrainedModel):
+    def __init__(self, config: MatryoshkaNllbClipConfig, device):
+        super().__init__(config)
+        if isinstance(device, str):
+            device = torch.device(device)
+        self.config = config
+        self.model, _, self.transform = create_model_and_transforms(
+            config.clip_model_name, config.clip_model_version, output_dict=True
+        )
+        self._device = device
+        self.model.to(device)
+        self.matryoshka_layer = MatryoshkaLayer(
+            config.mrl_resolutions, config.target_resolution
+        )
+        self.matryoshka_layer.to(device)
+        self.tokenizer = get_tokenizer(config.clip_model_name)
+
+    def forward(self, image_inputs, input_ids, resolution: Union[int, None] = None):
+        image_inputs = image_inputs.to(self._device)
+        input_ids = input_ids.to(self._device)
+        outputs = self.model(
+            image=image_inputs,
+            text=input_ids,
+        )
+        mrl_image_features = None
+        mrl_text_features = None
+        if resolution is not None:
+            mrl_image_features = self.matryoshka_layer.forward(
+                outputs["image_features"], resolution
+            )
+            mrl_text_features = self.matryoshka_layer.forward(
+                outputs["text_features"], resolution
+            )
+        return {
+            "image_features": outputs["image_features"],
+            "text_features": outputs["text_features"],
+            "mrl_image_features": mrl_image_features,
+            "mrl_text_features": mrl_text_features,
+            "logit_scale": outputs["logit_scale"],
+            "logit_bias": outputs["logit_bias"],
+        }
+
+    def encode_images(
+        self,
+        images: List[Image.Image],
+        normalize=False,
+        resolution: Union[int, None] = None,
+    ):
+        image_inputs = [self.transform(image) for image in images]
+        image_inputs = torch.stack(image_inputs, dim=0).to(self._device)
+        with torch.inference_mode():
+            features = self.model.visual(image_inputs)
+            if resolution is not None:
+                if resolution not in self.matryoshka_layer.resolutions:
+                    raise ValueError(
+                        f"Resolution {resolution} not in {self.matryoshka_layer.resolutions}"
+                    )
+                features = self.matryoshka_layer.layers[str(resolution)](features)
+            return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_texts(
+        self,
+        texts: List[str],
+        langs: Union[List[str], None] = None,
+        normalize=False,
+        resolution: Union[int, None] = None,
+    ):
+        if langs is None:
+            langs = ["eng_Latn"] * len(texts)
+        texts = [f"{lang}{text}" for lang, text in zip(langs, texts)]
+        input_ids = self.tokenizer.tokenizer.batch_encode_plus(
+            texts, return_tensors="pt", padding="longest", add_special_tokens=False
+        )["input_ids"].to(self._device)
+        with torch.inference_mode():
+            features = self.model.text(input_ids)
+            if resolution is not None:
+                if resolution not in self.matryoshka_layer.resolutions:
+                    raise ValueError(
+                        f"Resolution {resolution} not in {self.matryoshka_layer.resolutions}"
+                    )
+                features = self.matryoshka_layer.layers[str(resolution)](features)
+            return F.normalize(features, dim=-1) if normalize else features
+
+    def get_logits(
+        self,
+        images: List[Image.Image],
+        texts: List[str],
+        langs: Union[List[str], None] = None,
+        resolution: Union[int, None] = None,
+    ):
+        image_features = self.encode_images(
+            images, normalize=True, resolution=resolution
+        )
+        text_features = self.encode_texts(
+            texts, langs, normalize=True, resolution=resolution
+        )
+        with torch.inference_mode():
+            image_logits = (
+                self.model.logit_scale.exp() * image_features @ text_features.T
+            )
+            if self.model.logit_bias is not None:
+                image_logits += self.model.logit_bias
+        text_logits = image_logits.T
+        return image_logits, text_logits