Trjtdtk

I am trying my hand at BERT and I got so far that I can feed a sentence into BertTokenizer, and run that through the BERT model which gives me the output layers back. Modified from PyTorch code over at HuggingFace.

import logging

import torch
from torch.utils.data import TensorDataset, DataLoader, SequentialSampler

from pytorch_pretrained_bert.tokenization import BertTokenizer
from pytorch_pretrained_bert.modeling import BertModel

logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s - %(message)s', 
 datefmt = '%m/%d/%Y %H:%M:%S',
 level = logging.INFO)
logger = logging.getLogger(__name__)

class InputFeatures(object):
 def __init__(self, tokens, input_ids, input_mask, input_type_ids):
 self.tokens = tokens
 self.input_ids = input_ids
 self.input_mask = input_mask
 self.input_type_ids = input_type_ids

def convert_sentences_to_features(sentences, max_seq_length, tokenizer): 
 features = []
 for sentence in sentences:
 # tokenizer will also separate on punctuation
 # see https://github.com/google-research/bert#tokenization
 tokens = tokenizer.tokenize(sentence)

 # limit size of tokens
 if len(tokens) > max_seq_length - 2:
 tokens = tokens[0:(max_seq_length - 2)]

 # add [CLS] and [SEP], as expected in BERT
 tokens = ['[CLS]', *tokens, '[SEP]']


 input_type_ids = [0] * len(tokens)
 input_ids = tokenizer.convert_tokens_to_ids(tokens)

 # The mask has 1 for real tokens and 0 for padding tokens. Only real
 # tokens are attended to.
 input_mask = [1] * len(input_ids)

 # Zero-pad up to the sequence length.
 while len(input_ids) < max_seq_length:
 input_ids.append(0)
 input_mask.append(0)
 input_type_ids.append(0)

 features.append(InputFeatures(tokens=tokens, 
 input_ids=input_ids, 
 input_mask=input_mask, 
 input_type_ids=input_type_ids)
 )
 return features


def main(sentences, layers='-1, -2, -3, -4', max_seq_length=512, bert_model='bert-large-uncased', 
 do_lower_case=True, batch_size=32, no_cuda=False):
 device = torch.device('cuda' if torch.cuda.is_available() and not no_cuda else 'cpu')

 # 'layers' indicates which layers we want to concatenate
 layer_idxs = [int(l) for l in layers.split(',')]

 # init tokenizer
 tokenizer = BertTokenizer.from_pretrained(bert_model, do_lower_case=do_lower_case)

 # returns a list of 'InputFeatures'
 features = convert_sentences_to_features(sentences, max_seq_length, tokenizer)

 # init model and move to device
 model = BertModel.from_pretrained(bert_model)
 model.to(device) 

 # extract IDs and mask from the features
 all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
 all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)

 # prepare dataset and dataloader
 eval_data = TensorDataset(all_input_ids, all_input_mask)
 eval_sampler = SequentialSampler(eval_data)
 eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=batch_size)

 model.eval()

 for input_ids, input_mask in eval_dataloader:
 input_ids = input_ids.to(device)
 input_mask = input_mask.to(device)

 all_encoder_layers, _ = model(input_ids, token_type_ids=None, attention_mask=input_mask)

 # put layers to concatenate in list, and use torch.cat
 layers_to_concat = [all_encoder_layers[idx] for idx in layer_idxs]
 concat = torch.cat(layers_to_concat, dim=-1)

 logger.info(concat.size())
 logger.info(concat)

if __name__ == "__main__":
 proc_args = 
 'sentences': ['I saw Bert today !', 'Do you like bananas ?', 'Some sentences are really horrendous to parse .'],
 'max_seq_length': 32
 
 main(**proc_args)

This works and gives me an output size of 3, 32, 4096 in this case. This comes down to batch_size, seq_length, layers*hidden_size. In practice, this means that I have a representation of each token in its context. But I would like to extract the sentence representation from this. In an (bidirectional) RNN, you would typically take the top leftmost node's output which contained the latest hidden state and latest output, but I am not sure whether this is also true for transformers, as the architecture is different.

What is the best way to extract the sentence representation from a sequence of representations in a transformer? The size would be batch_size, hidden_size.