Spaces:
Sleeping
Sleeping
| # Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved. | |
| # | |
| # This work is made available under the Nvidia Source Code License-NC. | |
| # To view a copy of this license, visit | |
| # https://github.com/NVlabs/prismer/blob/main/LICENSE | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from einops.einops import rearrange | |
| from model.prismer import Prismer | |
| class PrismerVQA(Prismer): | |
| def forward(self, experts, question, answer=None, weights=None, train=True, inference='rank', k_test=128): | |
| device = experts['rgb'].device | |
| question = ['<s>' + ques.capitalize() for ques in question] | |
| question = self.tokenizer(question, padding='longest', truncation=True, max_length=35, | |
| add_special_tokens=False, return_tensors="pt").to(device) | |
| if train: | |
| experts_train = self.expert_encoder(experts) | |
| experts_train = rearrange(experts_train, 'l b d -> b l d') # batch_size, num_latents, output_dim | |
| answer = [' ' + ans.capitalize() + '</s>' for ans in answer] | |
| answer = self.tokenizer(answer, padding='longest', return_tensors="pt", add_special_tokens=False).to(device) | |
| input_ids = torch.cat([question.input_ids, answer.input_ids], dim=1).long() | |
| attention_mask = torch.cat([question.attention_mask, answer.attention_mask], dim=1) | |
| answer_targets = input_ids.masked_fill(input_ids == self.tokenizer.pad_token_id, -100) | |
| answer_targets[:, :-answer.input_ids.shape[1]] = -100 | |
| answer_output = self.text_decoder(input_ids, | |
| attention_mask=attention_mask, | |
| encoder_hidden_states=experts_train, | |
| labels=answer_targets, | |
| return_dict=True) | |
| loss = weights * answer_output.loss | |
| loss = loss.mean() | |
| return loss | |
| else: | |
| if inference == 'generate': | |
| num_beams = 3 | |
| input_ids = question.input_ids | |
| attention_masks = question.attention_mask | |
| experts_train = self.expert_encoder(experts) | |
| experts_train = rearrange(experts_train, 'l b d -> b l d') # batch_size, num_latents, output_dim | |
| experts_train = experts_train.repeat_interleave(num_beams, dim=0) | |
| outputs = self.text_decoder.generate(input_ids=input_ids, | |
| encoder_hidden_states=experts_train, | |
| attention_mask=attention_masks, | |
| max_length=input_ids.shape[1] + 10, | |
| min_length=input_ids.shape[1] + 2, | |
| num_beams=num_beams, | |
| length_penalty=-1) | |
| answers = [] | |
| for i in range(len(outputs)): | |
| answer = self.tokenizer.decode(outputs[i, len(input_ids[i]):], skip_special_tokens=True) | |
| answers.append(answer.lower().strip()) | |
| return answers | |
| elif inference == 'rank': | |
| experts_train = self.expert_encoder(experts) | |
| experts_train = rearrange(experts_train, 'l b d -> b l d') | |
| answer = [' ' + ans.capitalize() + '</s>' for ans in answer] | |
| answer = self.tokenizer(answer, padding='longest', return_tensors='pt', add_special_tokens=False).to(device) | |
| start_ids = question.input_ids | |
| attention_masks = question.attention_mask | |
| start_output = self.text_decoder(start_ids, | |
| attention_mask=attention_masks, | |
| encoder_hidden_states=experts_train, | |
| return_dict=True) | |
| logits = start_output.logits[:, -1, :] | |
| answer_first_token = answer.input_ids[:, 0] | |
| prob_first_token = F.softmax(logits, dim=1).index_select(dim=1, index=answer_first_token) | |
| _, topk_ids = prob_first_token.topk(k_test, dim=1) | |
| # answer input: [num_question * k, answer_len] | |
| answer_input_ids = [] | |
| answer_input_atts = [] | |
| for b, topk_id in enumerate(topk_ids): | |
| answer_input_ids.append(answer.input_ids.index_select(dim=0, index=topk_id)) | |
| answer_input_atts.append(answer.attention_mask.index_select(dim=0, index=topk_id)) | |
| answer_input_ids = torch.cat(answer_input_ids, dim=0) | |
| answer_input_atts = torch.cat(answer_input_atts, dim=0) | |
| # repeat encoder's output for top-k answers | |
| input_ids = torch.cat([tile(start_ids, 0, k_test), answer_input_ids], dim=1).long() | |
| attention_masks = torch.cat([tile(attention_masks, 0, k_test), answer_input_atts], dim=1) | |
| experts_train = tile(experts_train, 0, k_test) | |
| answer_targets = input_ids.masked_fill(input_ids == self.tokenizer.pad_token_id, -100) | |
| answer_targets[:, :-answer.input_ids.shape[1]] = -100 | |
| output = self.text_decoder(input_ids, | |
| attention_mask=attention_masks, | |
| encoder_hidden_states=experts_train, | |
| labels=answer_targets, | |
| return_dict=True) | |
| log_probs_sum = -output.loss / torch.sum(answer_targets != -100, dim=-1) | |
| log_probs_sum = log_probs_sum.view(-1, k_test) | |
| max_topk_ids = log_probs_sum.argmax(dim=1) | |
| max_ids = topk_ids[max_topk_ids >= 0, max_topk_ids] | |
| return max_ids | |
| def tile(x, dim, n_tile): | |
| init_dim = x.size(dim) | |
| repeat_idx = [1] * x.dim() | |
| repeat_idx[dim] = n_tile | |
| x = x.repeat(*repeat_idx) | |
| order_index = torch.LongTensor(np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)])) | |
| return torch.index_select(x, dim, order_index.to(x.device)) | |