Spaces:
Sleeping
Sleeping
File size: 3,007 Bytes
3662ec7 |
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 |
import torch
import torch.nn as nn
from transformers import CLIPModel, AutoTokenizer, AutoModelForSeq2SeqLM
from huggingface_hub import hf_hub_download
class CLIP2MT5_CrossAttention(nn.Module):
def __init__(self, clip_name='openai/clip-vit-base-patch32',
t5_name='mukayese/mt5-base-turkish-summarization'):
super().__init__()
self.clip = CLIPModel.from_pretrained(clip_name)
self.tokenizer = AutoTokenizer.from_pretrained(t5_name)
self.t5 = AutoModelForSeq2SeqLM.from_pretrained(t5_name)
self.vis_proj = nn.Linear(
self.clip.config.vision_config.hidden_size,
self.t5.config.d_model
)
def forward(self, images, input_ids, attention_mask, labels=None):
vision_outputs = self.clip.vision_model(pixel_values=images).last_hidden_state
vision_embeds = self.vis_proj(vision_outputs)
text_embeds = self.t5.encoder.embed_tokens(input_ids)
extended_input_embeds = torch.cat([vision_embeds, text_embeds], dim=1)
extended_attention_mask = torch.cat([
torch.ones(vision_embeds.size(0), vision_embeds.size(1),
dtype=attention_mask.dtype, device=attention_mask.device),
attention_mask
], dim=1)
if labels is not None:
labels = labels.clone()
labels[labels == self.tokenizer.pad_token_id] = -100
return self.t5(
inputs_embeds=extended_input_embeds,
attention_mask=extended_attention_mask,
labels=labels,
return_dict=True
)
@torch.no_grad()
def generate(self, images, input_ids, attention_mask, **gen_kwargs):
vision_outputs = self.clip.vision_model(pixel_values=images).last_hidden_state
vision_embeds = self.vis_proj(vision_outputs)
text_embeds = self.t5.encoder.embed_tokens(input_ids)
extended_input_embeds = torch.cat([vision_embeds, text_embeds], dim=1)
extended_attention_mask = torch.cat([
torch.ones(vision_embeds.size(0), vision_embeds.size(1),
dtype=attention_mask.dtype, device=attention_mask.device),
attention_mask
], dim=1)
return self.t5.generate(
inputs_embeds=extended_input_embeds,
attention_mask=extended_attention_mask,
**gen_kwargs
)
# HF Loader for STATE_DICT
def load_model(
repo_id: str,
filename: str = "model.pt",
clip_name="openai/clip-vit-base-patch32",
t5_name="mukayese/mt5-base-turkish-summarization",
device=None
):
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_path = hf_hub_download(repo_id=repo_id, filename=filename)
model = CLIP2MT5_CrossAttention(clip_name=clip_name, t5_name=t5_name)
state = torch.load(model_path, map_location=device)
model.load_state_dict(state)
model.to(device)
model.eval()
return model
|