Update nllb_mrl.py

nllb_mrl.py

@@ -1,26 +1,21 @@
+from dataclasses import dataclass
 from typing import List, Union
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from open_clip import create_model, get_tokenizer
-from open_clip.transform import PreprocessCfg, image_transform_v2
+from huggingface_hub import PyTorchModelHubMixin
+from open_clip import create_model_and_transforms, get_tokenizer
 from PIL import Image
-from transformers import PretrainedConfig, PreTrainedModel
+from transformers import PretrainedConfig
 
 
+@dataclass
 class MatryoshkaNllbClipConfig(PretrainedConfig):
-    def __init__(
-        self,
-        clip_model_name: str = "",
-        target_resolution: int = -1,
-        mrl_resolutions: List[int] = [],
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.clip_model_name = clip_model_name
-        self.target_resolution = target_resolution
-        self.mrl_resolutions = mrl_resolutions
+    clip_model_name: str
+    clip_model_version: str
+    target_resolution: int
+    mrl_resolutions: List[int]
 
 
 class MatryoshkaLayer(nn.Module):
@@ -42,23 +37,16 @@ class MatryoshkaLayer(nn.Module):
         return outputs
 
 
-class MatryoshkaNllbClip(PreTrainedModel):
-    config_class = MatryoshkaNllbClipConfig
-
+class MatryoshkaNllbClip(nn.Module, PyTorchModelHubMixin):
     def __init__(self, config: MatryoshkaNllbClipConfig, device):
-        super().__init__(config)
+        super().__init__()
         if isinstance(device, str):
             device = torch.device(device)
         self.config = config
-        self.model = create_model(
-            config.clip_model_name, output_dict=True
+        self.model, _, self.transform = create_model_and_transforms(
+            config.clip_model_name, config.clip_model_version, output_dict=True
         )
-        pp_cfg = PreprocessCfg(**self.model.visual.preprocess_cfg)
-        self.transform = image_transform_v2(
-            pp_cfg,
-            is_train=False,
-        )
-        self._device = device
+        self.device = device
         self.model.to(device)
         self.matryoshka_layer = MatryoshkaLayer(
             config.mrl_resolutions, config.target_resolution
@@ -67,8 +55,8 @@ class MatryoshkaNllbClip(PreTrainedModel):
         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)
+        image_inputs = image_inputs.to(self.device)
+        input_ids = input_ids.to(self.device)
         outputs = self.model(
             image=image_inputs,
             text=input_ids,
@@ -91,14 +79,46 @@ class MatryoshkaNllbClip(PreTrainedModel):
             "logit_bias": outputs["logit_bias"],
         }
 
-    def encode_images(
+    def encode_image(
+        self,
+        image,
+        normalize=False,
+        resolution: Union[int, None] = None,
+    ):
+        with torch.inference_mode():
+            features = self.model.visual(image)
+            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_text(
+        self,
+        text,
+        normalize=False,
+        resolution: Union[int, None] = None,
+    ):
+        with torch.inference_mode():
+            features = self.model.text(text)
+            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 image_features(
         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)
+        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:
@@ -109,7 +129,7 @@ class MatryoshkaNllbClip(PreTrainedModel):
                 features = self.matryoshka_layer.layers[str(resolution)](features)
             return F.normalize(features, dim=-1) if normalize else features
 
-    def encode_texts(
+    def text_features(
         self,
         texts: List[str],
         langs: Union[List[str], None] = None,
@@ -118,10 +138,10 @@ class MatryoshkaNllbClip(PreTrainedModel):
     ):
         if langs is None:
             langs = ["eng_Latn"] * len(texts)
-        texts = [f"{lang}{text}" for lang, text in zip(langs, 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)
+        )["input_ids"].to(self.device)
         with torch.inference_mode():
             features = self.model.text(input_ids)
             if resolution is not None:
@@ -139,10 +159,10 @@ class MatryoshkaNllbClip(PreTrainedModel):
         langs: Union[List[str], None] = None,
         resolution: Union[int, None] = None,
     ):
-        image_features = self.encode_images(
+        image_features = self.image_features(
             images, normalize=True, resolution=resolution
         )
-        text_features = self.encode_texts(
+        text_features = self.text_features(
             texts, langs, normalize=True, resolution=resolution
         )
         with torch.inference_mode():