File size: 7,196 Bytes
955949b
 
 
 
 
89d694e
 
955949b
89d694e
955949b
 
 
89d694e
 
 
 
 
 
 
 
 
 
 
955949b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
89d694e
 
 
955949b
89d694e
955949b
 
 
89d694e
 
 
 
 
 
 
d91b3a3
89d694e
955949b
 
 
 
 
 
 
 
89d694e
 
955949b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2f07d8b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
955949b
 
 
 
 
 
89d694e
955949b
 
 
 
 
 
 
 
 
 
2f07d8b
955949b
 
 
 
 
 
 
 
89d694e
955949b
 
89d694e
955949b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2f07d8b
955949b
 
2f07d8b
955949b
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
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 PIL import Image
from transformers import PretrainedConfig, PreTrainedModel


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


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):
    config_class = MatryoshkaNllbClipConfig

    def __init__(self, config: MatryoshkaNllbClipConfig, device):
        super().__init__(config)
        if isinstance(device, str):
            device = torch.device(device)
        self.config = config
        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.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_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)
        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 text_features(
        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.image_features(
            images, normalize=True, resolution=resolution
        )
        text_features = self.text_features(
            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