File size: 2,286 Bytes
795f737 | 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 | """Trainable action classifier + candidate pointer; non-autoregressive baseline."""
import torch
from torch import nn
from huggingface_hub import PyTorchModelHubMixin
class PointerPolicy(nn.Module, PyTorchModelHubMixin, library_name='baim', tags=['browser-agent','cpu']):
def __init__(self, vocab_size=4096, width=64, encoder='mean', lexical_features=True):
super().__init__()
self.encoder_kind = encoder
self.lexical_features = lexical_features
self.embedding = nn.Embedding(vocab_size,width,padding_idx=0)
if encoder == 'gru':
self.encoder = nn.GRU(width,width,batch_first=True)
elif encoder == 'transformer':
self.position = nn.Embedding(24,width)
self.encoder = nn.TransformerEncoder(nn.TransformerEncoderLayer(width,4,width*2,
dropout=0.0,batch_first=True),num_layers=1,enable_nested_tensor=False)
elif encoder != 'mean':
raise ValueError('unknown encoder')
self.action = nn.Sequential(nn.Linear(width,width),nn.ReLU(),nn.Linear(width,3))
self.pointer = nn.Sequential(nn.Linear(width*4+5,width),nn.ReLU(),nn.Linear(width,1))
def embed(self, ids):
mask = ids.ne(0)
x = self.embedding(ids)
if self.encoder_kind == 'gru':
x,_ = self.encoder(x)
elif self.encoder_kind == 'transformer':
x = x + self.position(torch.arange(ids.shape[-1],device=ids.device))
# Empty padded candidates need one unmasked token to avoid NaNs.
safe = mask.clone()
safe[:,0] = True
x = self.encoder(x,src_key_padding_mask=~safe)
return (x*mask.unsqueeze(-1)).sum(1)/mask.sum(1,keepdim=True).clamp_min(1)
def forward(self, goal, elements, features, mask):
g = self.embed(goal)
batch,count,length = elements.shape
e = self.embed(elements.reshape(batch*count,length)).reshape(batch,count,-1)
expanded = g.unsqueeze(1).expand_as(e)
lexical = features if self.lexical_features else torch.zeros_like(features)
pairs = torch.cat([expanded,e,expanded*e,torch.abs(expanded-e),lexical],dim=-1)
pointers = self.pointer(pairs).squeeze(-1).masked_fill(~mask,-1e4)
return self.action(g),pointers
|