| 293 | |
| 294 | |
| 295 | class CLIPModel(nn.Module): |
| 296 | def __init__(self, config: CLIPConfig): |
| 297 | self.text_model = ClipTextModel(config.text_config) |
| 298 | self.vision_model = ClipVisionModel(config.vision_config) |
| 299 | |
| 300 | text_embed_dim = config.text_config.hidden_size |
| 301 | vision_embed_dim = config.vision_config.hidden_size |
| 302 | projection_dim = config.projection_dim |
| 303 | |
| 304 | self.visual_projection = nn.Linear(vision_embed_dim, projection_dim, bias=False) |
| 305 | self.text_projection = nn.Linear(text_embed_dim, projection_dim, bias=False) |
| 306 | self.logit_scale = mx.array(0.0) |
| 307 | |
| 308 | def get_text_features(self, x: mx.array) -> mx.array: |
| 309 | return self.text_projection(self.text_model(x).pooler_output) |
| 310 | |
| 311 | def get_image_features(self, x: mx.array) -> mx.array: |
| 312 | return self.visual_projection(self.vision_model(x).pooler_output) |
| 313 | |
| 314 | def __call__( |
| 315 | self, |
| 316 | input_ids: Optional[mx.array] = None, |
| 317 | pixel_values: Optional[mx.array] = None, |
| 318 | return_loss=False, |
| 319 | ) -> CLIPModelOutput: |
| 320 | if input_ids is not None: |
| 321 | text_model_output = self.text_model(input_ids) |
| 322 | text_embeds = self.text_projection(text_model_output.pooler_output) |
| 323 | text_embeds = text_embeds / LA.norm(text_embeds, axis=-1, keepdims=True) |
| 324 | else: |
| 325 | text_embeds = None |
| 326 | text_model_output = None |
| 327 | |
| 328 | if pixel_values is not None: |
| 329 | vision_model_output = self.vision_model(pixel_values) |
| 330 | image_embeds = self.visual_projection(vision_model_output.pooler_output) |
| 331 | image_embeds = image_embeds / LA.norm(image_embeds, axis=-1, keepdims=True) |
| 332 | else: |
| 333 | image_embeds = None |
| 334 | vision_model_output = None |
| 335 | |
| 336 | if return_loss and (input_ids is None or pixel_values is None): |
| 337 | raise ValueError("Must provide text and image inputs to compute loss.") |
| 338 | |
| 339 | if return_loss: |
| 340 | logit_scale = mx.exp(self.logit_scale) |
| 341 | logits = (text_embeds @ image_embeds.T) * logit_scale |
| 342 | loss = clip_loss(logits) |
| 343 | else: |
| 344 | loss = None |
| 345 | |
| 346 | return CLIPModelOutput( |
| 347 | loss=loss, |
| 348 | text_embeds=text_embeds, |
| 349 | image_embeds=image_embeds, |
| 350 | vision_model_output=vision_model_output, |
| 351 | text_model_output=text_model_output, |
| 352 | ) |