Spaces:
Sleeping
Sleeping
File size: 1,525 Bytes
08d7644 |
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 |
# coding=utf-8
# Copyleft 2019 project LXRT.
import torch.nn as nn
from ..param import args
from ..lxrt.entry import LXRTEncoder
from ..lxrt.modeling import BertLayerNorm, GeLU
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
# Max length including <bos> and <eos>
MAX_VQA_LENGTH = 20
class VQAModel(nn.Module):
def __init__(self, num_answers):
super().__init__()
# # Build LXRT encoder
# self.lxrt_encoder = LXRTEncoder(
# args,
# max_seq_length=MAX_VQA_LENGTH
# )
# hid_dim = self.lxrt_encoder.dim
#
# # VQA Answer heads
# self.logit_fc = nn.Sequential(
# nn.Linear(hid_dim, hid_dim * 2),
# GeLU(),
# BertLayerNorm(hid_dim * 2, eps=1e-12),
# nn.Linear(hid_dim * 2, num_answers)
# )
# self.logit_fc.apply(self.lxrt_encoder.model.init_bert_weights)
self.tokenizer = AutoTokenizer.from_pretrained("unc-nlp/lxmert-vqa-uncased")
self.model = AutoModelForQuestionAnswering.from_pretrained("unc-nlp/lxmert-vqa-uncased")
def forward(self, feat, pos, sent):
"""
b -- batch_size, o -- object_number, f -- visual_feature_size
:param feat: (b, o, f)
:param pos: (b, o, 4)
:param sent: (b,) Type -- list of string
:param leng: (b,) Type -- int numpy array
:return: (b, num_answer) The logit of each answers.
"""
return self.model(sent, feat, pos)
|