| """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)) |
| |
| 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 |
|
|
