Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
from transformers import AutoTokenizer, AutoModel | |
import clip | |
import skimage.io as io | |
import PIL.Image | |
from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup | |
import os | |
import pickle | |
import torch | |
import torch.nn as nn | |
from torch.utils.data import Dataset, DataLoader | |
from torch.nn import functional as F | |
import pandas as pd | |
from tqdm import tqdm | |
from PIL import Image | |
from typing import Tuple | |
import numpy as np | |
import time | |
import json | |
import nltk | |
nltk.download('punkt') | |
class Adapter(nn.Module): | |
def forward(self, x): | |
return self.model(x) | |
def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh): | |
super(Adapter, self).__init__() | |
layers = [] | |
for i in range(len(sizes) -1): | |
layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias)) | |
if i < len(sizes) - 2: | |
layers.append(act()) | |
self.model = nn.Sequential(*layers) | |
class ClipGPT2Model(nn.Module): | |
def __init__(self, img_feature_length, img_feature_size = 512): | |
super(ClipGPT2Model, self).__init__() | |
self.img_feature_length = img_feature_length | |
self.gpt = GPT2LMHeadModel.from_pretrained('gpt2') | |
self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1] | |
self.clip_project = Adapter((img_feature_size, | |
(self.gpt_embedding_size * img_feature_length) // 2, | |
self.gpt_embedding_size * img_feature_length)) | |
def get_dummy_token(self, | |
batch_size: int, | |
device: torch.device) -> torch.Tensor: | |
return torch.zeros(batch_size, self.img_feature_length, dtype=torch.int64, device=device) | |
def forward(self, | |
tokens: torch.Tensor, | |
feature: torch.Tensor, | |
mask = None, | |
labels = None): | |
embedding_text = self.gpt.transformer.wte(tokens) | |
feature_projections = self.clip_project(feature).view(-1, self.img_feature_length, self.gpt_embedding_size) | |
embedding_cat = torch.cat((feature_projections, embedding_text), dim=1) | |
if labels is not None: | |
dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device) | |
labels = torch.cat((dummy_token, tokens), dim=1) | |
out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask) | |
return out | |
def generate_beam( | |
model, | |
tokenizer, | |
entry_length, | |
temperature, | |
embed=None, | |
beam_size: int = 10, | |
prompt=None, | |
stop_token: str = ".", | |
): | |
model.eval() | |
stop_token_index = tokenizer.encode(stop_token)[0] | |
tokens = None | |
scores = None | |
device = next(model.parameters()).device | |
seq_lengths = torch.ones(beam_size, device=device) | |
is_stopped = torch.zeros(beam_size, device=device, dtype=torch.bool) | |
with torch.no_grad(): | |
if embed is not None: | |
generated = embed | |
else: | |
if tokens is None: | |
tokens = torch.tensor(tokenizer.encode(prompt)) | |
tokens = tokens.unsqueeze(0).to(device) | |
generated = model.gpt.transformer.wte(tokens) | |
for i in range(entry_length): | |
outputs = model.gpt(inputs_embeds=generated) | |
logits = outputs.logits | |
logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0) | |
logits = logits.softmax(-1).log() | |
if scores is None: | |
scores, next_tokens = logits.topk(beam_size, -1) | |
generated = generated.expand(beam_size, *generated.shape[1:]) | |
next_tokens, scores = next_tokens.permute(1, 0), scores.squeeze(0) | |
if tokens is None: | |
tokens = next_tokens | |
else: | |
tokens = tokens.expand(beam_size, *tokens.shape[1:]) | |
tokens = torch.cat((tokens, next_tokens), dim=1) | |
else: | |
logits[is_stopped] = -float(np.inf) | |
logits[is_stopped, 0] = 0 | |
scores_sum = scores[:, None] + logits | |
seq_lengths[~is_stopped] += 1 | |
scores_sum_average = scores_sum / seq_lengths[:, None] | |
scores_sum_average, next_tokens = scores_sum_average.view(-1).topk( | |
beam_size, -1 | |
) | |
next_tokens_source = next_tokens // scores_sum.shape[1] | |
seq_lengths = seq_lengths[next_tokens_source] | |
next_tokens = next_tokens % scores_sum.shape[1] | |
next_tokens = next_tokens.unsqueeze(1) | |
tokens = tokens[next_tokens_source] | |
tokens = torch.cat((tokens, next_tokens), dim=1) | |
generated = generated[next_tokens_source] | |
scores = scores_sum_average * seq_lengths | |
is_stopped = is_stopped[next_tokens_source] | |
next_token_embed = model.gpt.transformer.wte(next_tokens.squeeze()).view( | |
generated.shape[0], 1, -1 | |
) | |
generated = torch.cat((generated, next_token_embed), dim=1) | |
is_stopped = is_stopped + next_tokens.eq(stop_token_index).squeeze() | |
if is_stopped.all(): | |
break | |
scores = scores / seq_lengths | |
output_list = tokens.cpu().numpy() | |
output_texts = [ | |
tokenizer.decode(output[: int(length)]) | |
for output, length in zip(output_list, seq_lengths) | |
] | |
order = scores.argsort(descending=True) | |
output_texts = [output_texts[i] for i in order] | |
return output_texts | |
def generate_caption_clipgpt(img, entry_length, temperature): | |
prefix_length = 10 | |
model = ClipGPT2Model(prefix_length) | |
model.load_state_dict(torch.load('model/model_train_best_run_clipGPT.pt', map_location=torch.device('cpu'))) | |
model = model.eval() | |
device=torch.device('cpu') | |
model = model.to(device) | |
clip_model, preprocess = clip.load('ViT-B/32', device, jit=False) | |
tokenizer = GPT2Tokenizer.from_pretrained("gpt2") | |
start_time = time.time() | |
# pil_image = PIL.Image.fromarray(img) | |
image = preprocess(img).unsqueeze(0).to(device) | |
with torch.no_grad(): | |
prefix = clip_model.encode_image(image).to(device, dtype=torch.float32) | |
prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1) | |
beam_caption = generate_beam(model, tokenizer, entry_length, temperature, embed=prefix_embed)[0] | |
end_time = time.time() | |
print("--- Time taken to generate: %s seconds ---" % (end_time - start_time)) | |
return beam_caption | |