| | from torch import nn |
| | import transformers |
| | import torch |
| | from PIL import Image |
| |
|
| |
|
| | class CLIPModel(nn.Module): |
| | def __init__(self, model_name: str = "openai/clip-vit-base-patch32", processor_name = None): |
| | super(CLIPModel, self).__init__() |
| |
|
| | if processor_name is None: |
| | processor_name = model_name |
| |
|
| | self.model = transformers.CLIPModel.from_pretrained(model_name) |
| | self.processor = transformers.CLIPProcessor.from_pretrained(processor_name) |
| |
|
| | def __repr__(self): |
| | return "CLIPModel()" |
| |
|
| | def forward(self, features): |
| | image_embeds = [] |
| | text_embeds = [] |
| |
|
| | if 'pixel_values' in features: |
| | vision_outputs = self.model.vision_model(pixel_values=features['pixel_values']) |
| | image_embeds = self.model.visual_projection(vision_outputs[1]) |
| |
|
| | if 'input_ids' in features: |
| | text_outputs = self.model.text_model( |
| | input_ids=features.get('input_ids'), |
| | attention_mask=features.get('attention_mask', None), |
| | position_ids=features.get('position_ids', None), |
| | output_attentions=features.get('output_attentions', None), |
| | output_hidden_states=features.get('output_hidden_states', None), |
| | ) |
| | text_embeds = self.model.text_projection(text_outputs[1]) |
| |
|
| | sentence_embedding = [] |
| | image_features = iter(image_embeds) |
| | text_features = iter(text_embeds) |
| |
|
| | for idx, input_type in enumerate(features['image_text_info']): |
| | if input_type == 0: |
| | sentence_embedding.append(next(image_features)) |
| | else: |
| | sentence_embedding.append(next(text_features)) |
| |
|
| | features['sentence_embedding'] = torch.stack(sentence_embedding).float() |
| |
|
| | return features |
| |
|
| |
|
| | def tokenize(self, texts): |
| | images = [] |
| | texts_values = [] |
| | image_text_info = [] |
| |
|
| | for idx, data in enumerate(texts): |
| | if isinstance(data, Image.Image): |
| | images.append(data) |
| | image_text_info.append(0) |
| | else: |
| | texts_values.append(data) |
| | image_text_info.append(1) |
| |
|
| | if len(texts_values) == 0: |
| | texts_values = None |
| | if len(images) == 0: |
| | images = None |
| |
|
| | inputs = self.processor(text=texts_values, images=images, return_tensors="pt", padding=True) |
| | inputs['image_text_info'] = image_text_info |
| | return inputs |
| |
|
| |
|
| | def save(self, output_path: str): |
| | self.model.save_pretrained(output_path) |
| | self.processor.save_pretrained(output_path) |
| |
|
| | @staticmethod |
| | def load(input_path: str): |
| | return CLIPModel(model_name=input_path) |
| |
|
| |
|
| |
|
