|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
from transformers import ( |
|
BertModel, |
|
BertConfig, |
|
PretrainedConfig, |
|
PreTrainedModel, |
|
) |
|
from transformers.modeling_outputs import SequenceClassifierOutput |
|
|
|
|
|
class BertConfigForWebshop(PretrainedConfig): |
|
model_type = "bert" |
|
|
|
def __init__( |
|
self, |
|
pretrained_bert=True, |
|
image=False, |
|
**kwargs |
|
): |
|
self.pretrained_bert = pretrained_bert |
|
self.image = image |
|
super().__init__(**kwargs) |
|
|
|
|
|
class BiAttention(nn.Module): |
|
def __init__(self, input_size, dropout): |
|
super().__init__() |
|
self.dropout = nn.Dropout(dropout) |
|
self.input_linear = nn.Linear(input_size, 1, bias=False) |
|
self.memory_linear = nn.Linear(input_size, 1, bias=False) |
|
self.dot_scale = nn.Parameter( |
|
torch.zeros(size=(input_size,)).uniform_(1. / (input_size ** 0.5)), |
|
requires_grad=True) |
|
self.init_parameters() |
|
|
|
def init_parameters(self): |
|
return |
|
|
|
def forward(self, context, memory, mask): |
|
bsz, input_len = context.size(0), context.size(1) |
|
memory_len = memory.size(1) |
|
context = self.dropout(context) |
|
memory = self.dropout(memory) |
|
|
|
input_dot = self.input_linear(context) |
|
memory_dot = self.memory_linear(memory).view(bsz, 1, memory_len) |
|
cross_dot = torch.bmm( |
|
context * self.dot_scale, |
|
memory.permute(0, 2, 1).contiguous()) |
|
att = input_dot + memory_dot + cross_dot |
|
att = att - 1e30 * (1 - mask[:, None]) |
|
|
|
weight_one = F.softmax(att, dim=-1) |
|
output_one = torch.bmm(weight_one, memory) |
|
weight_two = (F.softmax(att.max(dim=-1)[0], dim=-1) |
|
.view(bsz, 1, input_len)) |
|
output_two = torch.bmm(weight_two, context) |
|
return torch.cat( |
|
[context, output_one, context * output_one, |
|
output_two * output_one], |
|
dim=-1) |
|
|
|
|
|
class BertModelForWebshop(PreTrainedModel): |
|
|
|
config_class = BertConfigForWebshop |
|
|
|
def __init__(self, config): |
|
super().__init__(config) |
|
bert_config = BertConfig.from_pretrained('bert-base-uncased') |
|
if config.pretrained_bert: |
|
self.bert = BertModel.from_pretrained('bert-base-uncased') |
|
else: |
|
self.bert = BertModel(config) |
|
self.bert.resize_token_embeddings(30526) |
|
self.attn = BiAttention(768, 0.0) |
|
self.linear_1 = nn.Linear(768 * 4, 768) |
|
self.relu = nn.ReLU() |
|
self.linear_2 = nn.Linear(768, 1) |
|
if config.image: |
|
self.image_linear = nn.Linear(512, 768) |
|
else: |
|
self.image_linear = None |
|
|
|
@staticmethod |
|
def get_aggregated(output, lens, method): |
|
""" |
|
Get the aggregated hidden state of the encoder. |
|
B x D |
|
""" |
|
if method == 'mean': |
|
return torch.stack([output[i, :j, :].mean(0) for i, j in enumerate(lens)], dim=0) |
|
elif method == 'last': |
|
return torch.stack([output[i, j-1, :] for i, j in enumerate(lens)], dim=0) |
|
elif method == 'first': |
|
return output[:, 0, :] |
|
|
|
def forward(self, state_input_ids, state_attention_mask, action_input_ids, action_attention_mask, sizes, images=None, labels=None): |
|
sizes = sizes.tolist() |
|
|
|
state_rep = self.bert(state_input_ids, attention_mask=state_attention_mask)[0] |
|
if images is not None and self.image_linear is not None: |
|
images = self.image_linear(images) |
|
state_rep = torch.cat([images.unsqueeze(1), state_rep], dim=1) |
|
state_attention_mask = torch.cat([state_attention_mask[:, :1], state_attention_mask], dim=1) |
|
action_rep = self.bert(action_input_ids, attention_mask=action_attention_mask)[0] |
|
state_rep = torch.cat([state_rep[i:i+1].repeat(j, 1, 1) for i, j in enumerate(sizes)], dim=0) |
|
state_attention_mask = torch.cat([state_attention_mask[i:i+1].repeat(j, 1) for i, j in enumerate(sizes)], dim=0) |
|
act_lens = action_attention_mask.sum(1).tolist() |
|
state_action_rep = self.attn(action_rep, state_rep, state_attention_mask) |
|
state_action_rep = self.relu(self.linear_1(state_action_rep)) |
|
act_values = self.get_aggregated(state_action_rep, act_lens, 'mean') |
|
act_values = self.linear_2(act_values).squeeze(1) |
|
|
|
logits = [F.log_softmax(_, dim=0) for _ in act_values.split(sizes)] |
|
|
|
loss = None |
|
if labels is not None: |
|
loss = - sum([logit[label] for logit, label in zip(logits, labels)]) / len(logits) |
|
|
|
return SequenceClassifierOutput( |
|
loss=loss, |
|
logits=logits, |
|
) |
|
|
