add model

bert.py

@@ -0,0 +1,118 @@
+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
+from .modules import EncoderRNN, BiAttention, get_aggregated
+
+
+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 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
+
+        # for state value prediction, used in RL
+        self.linear_3 = nn.Sequential(
+                nn.Linear(768, 128),
+                nn.LeakyReLU(),
+                nn.Linear(128, 1),
+            )
+
+    def forward(self, state_input_ids, state_attention_mask, action_input_ids, action_attention_mask, sizes, images=None, labels=None):
+        sizes = sizes.tolist()
+        # print(state_input_ids.shape, action_input_ids.shape)
+        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 = 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,
+        )
+
+    def rl_forward(self, state_batch, act_batch, value=False, q=False, act=False):
+        act_values = []
+        act_sizes = []
+        values = []
+        for state, valid_acts in zip(state_batch, act_batch):
+            with torch.set_grad_enabled(not act):
+                state_ids = torch.tensor([state.obs]).cuda()
+                state_mask = (state_ids > 0).int()
+                act_lens = [len(_) for _ in valid_acts]
+                act_ids = [torch.tensor(_) for _ in valid_acts]
+                act_ids = nn.utils.rnn.pad_sequence(act_ids, batch_first=True).cuda()
+                act_mask = (act_ids > 0).int()
+                act_size = torch.tensor([len(valid_acts)]).cuda()
+                if self.image_linear is not None:
+                    images = [state.image_feat]
+                    images = [torch.zeros(512) if _ is None else _ for _ in images] 
+                    images = torch.stack(images).cuda()  # BS x 512
+                else:
+                    images = None
+                logits = self.forward(state_ids, state_mask, act_ids, act_mask, act_size, images=images).logits[0]
+                act_values.append(logits)
+                act_sizes.append(len(valid_acts))
+            if value:
+                v = self.bert(state_ids, state_mask)[0]
+                values.append(self.linear_3(v[0][0]))
+        act_values = torch.cat(act_values, dim=0)
+        act_values = torch.cat([F.log_softmax(_, dim=0) for _ in act_values.split(act_sizes)], dim=0)
+        # Optionally, output state value prediction
+        if value:
+            values = torch.cat(values, dim=0)
+            return act_values, act_sizes, values
+        else:
+            return act_values, act_sizes

config.json

@@ -0,0 +1,14 @@
+{
+  "architectures": [
+    "BertModelForWebshop"
+  ],
+  "auto_map": {
+    "AutoConfig": "bert.BertConfigForWebshop",
+    "AutoModel": "bert.BertModelForWebshop"
+  },
+  "image": true,
+  "model_type": "bert",
+  "pretrained_bert": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.17.0"
+}

modules.py

@@ -0,0 +1,156 @@
+import itertools
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import rnn
+
+
+def duplicate(output, mask, lens, act_sizes):
+    """
+    Duplicate the output based on the action sizes.
+    """
+    output = torch.cat([output[i:i+1].repeat(j, 1, 1) for i, j in enumerate(act_sizes)], dim=0)
+    mask = torch.cat([mask[i:i+1].repeat(j, 1) for i, j in enumerate(act_sizes)], dim=0)
+    lens = list(itertools.chain.from_iterable([lens[i:i+1] * j for i, j in enumerate(act_sizes)]))
+    return output, mask, lens
+
+
+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, :]
+
+
+class EncoderRNN(nn.Module):
+    def __init__(self, input_size, num_units, nlayers, concat,
+                 bidir, layernorm, return_last):
+        super().__init__()
+        self.layernorm = (layernorm == 'layer')
+        if layernorm:
+            self.norm = nn.LayerNorm(input_size)
+
+        self.rnns = []
+        for i in range(nlayers):
+            if i == 0:
+                input_size_ = input_size
+                output_size_ = num_units
+            else:
+                input_size_ = num_units if not bidir else num_units * 2
+                output_size_ = num_units
+            self.rnns.append(
+                nn.GRU(input_size_, output_size_, 1,
+                       bidirectional=bidir, batch_first=True))
+
+        self.rnns = nn.ModuleList(self.rnns)
+        self.init_hidden = nn.ParameterList(
+            [nn.Parameter(
+                torch.zeros(size=(2 if bidir else 1, 1, num_units)),
+                requires_grad=True) for _ in range(nlayers)])
+        self.concat = concat
+        self.nlayers = nlayers
+        self.return_last = return_last
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        with torch.no_grad():
+            for rnn_layer in self.rnns:
+                for name, p in rnn_layer.named_parameters():
+                    if 'weight_ih' in name:
+                        torch.nn.init.xavier_uniform_(p.data)
+                    elif 'weight_hh' in name:
+                        torch.nn.init.orthogonal_(p.data)
+                    elif 'bias' in name:
+                        p.data.fill_(0.0)
+                    else:
+                        p.data.normal_(std=0.1)
+
+    def get_init(self, bsz, i):
+        return self.init_hidden[i].expand(-1, bsz, -1).contiguous()
+
+    def forward(self, inputs, input_lengths=None):
+        bsz, slen = inputs.size(0), inputs.size(1)
+        if self.layernorm:
+            inputs = self.norm(inputs)
+        output = inputs
+        outputs = []
+        lens = 0
+        if input_lengths is not None:
+            lens = input_lengths  # .data.cpu().numpy()
+        for i in range(self.nlayers):
+            hidden = self.get_init(bsz, i)
+            # output = self.dropout(output)
+            if input_lengths is not None:
+                output = rnn.pack_padded_sequence(output, lens,
+                                                  batch_first=True,
+                                                  enforce_sorted=False)
+            output, hidden = self.rnns[i](output, hidden)
+            if input_lengths is not None:
+                output, _ = rnn.pad_packed_sequence(output, batch_first=True)
+                if output.size(1) < slen:
+                    # used for parallel
+                    # padding = Variable(output.data.new(1, 1, 1).zero_())
+                    padding = torch.zeros(
+                        size=(1, 1, 1), dtype=output.type(),
+                        device=output.device())
+                    output = torch.cat(
+                        [output,
+                         padding.expand(
+                             output.size(0),
+                             slen - output.size(1),
+                             output.size(2))
+                         ], dim=1)
+            if self.return_last:
+                outputs.append(
+                    hidden.permute(1, 0, 2).contiguous().view(bsz, -1))
+            else:
+                outputs.append(output)
+        if self.concat:
+            return torch.cat(outputs, dim=2)
+        return outputs[-1]
+
+
+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)

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32c00421372d607ffbd15f317f0040569dbe3cc7843f8885d2b54ffd2db9d0a8
+size 449449751