Update nllb_mrl.py

nllb_mrl.py

@@ -1,21 +1,26 @@
-from dataclasses import dataclass
 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 open_clip import create_model, get_tokenizer
+from open_clip.transform import PreprocessCfg, image_transform_v2
 from PIL import Image
-from transformers import PretrainedConfig
+from transformers import PretrainedConfig, PreTrainedModel
 
 
-@dataclass
 class MatryoshkaNllbClipConfig(PretrainedConfig):
-    clip_model_name: str
-    clip_model_version: str
-    target_resolution: int
-    mrl_resolutions: List[int]
+    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
 
 
 class MatryoshkaLayer(nn.Module):
@@ -37,16 +42,23 @@ class MatryoshkaLayer(nn.Module):
         return outputs
 
 
-class MatryoshkaNllbClip(nn.Module, PyTorchModelHubMixin):
+class MatryoshkaNllbClip(PreTrainedModel):
+    config_class = MatryoshkaNllbClipConfig
+
     def __init__(self, config: MatryoshkaNllbClipConfig, device):
-        super().__init__()
+        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.model = create_model(
+            config.clip_model_name, 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
@@ -55,8 +67,8 @@ class MatryoshkaNllbClip(nn.Module, PyTorchModelHubMixin):
         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,
@@ -118,7 +130,7 @@ class MatryoshkaNllbClip(nn.Module, PyTorchModelHubMixin):
         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:
@@ -138,10 +150,10 @@ class MatryoshkaNllbClip(nn.Module, PyTorchModelHubMixin):
     ):
         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: