title: Word embeddings with BERT and map tensors to words
body: <p>I try to aggregate BERT embeddings on the token level. For each token in the corpus vocabulary, I would like to create a list of all their contextual embeddings and average them to get one representation for each token in the vocabulary.</p>
  <p>The code is pasted below.</p>
  <p><strong>Question:</strong> How to map output tensors (see object <code>token_vecs_sum</code> in the last line of the code below) to particular tokens?</p>
  <p><strong>Preprocess data</strong></p>
  <pre><code>!pip install transformers
  import torch
  from transformers import BertTokenizer
  from nltk import tokenize
  import nltk
  nltk.download('punkt')
  import re
  
  MAX_LEN = 64
  
  sentences = ['Some sentences. Some sentences are.', 'Some sentences are really.', 'Some sentences are really hard.']
  
  tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)
  
  def preprocessing_for_bert(data):
    input_ids = []
    attention_masks = []
    for row in data:
      sents = tokenize.sent_tokenize(row)
      print(sents)
      for sent in sents:
        encoded_sent = tokenizer.encode_plus(text=sent,
                                            add_special_tokens=True,
                                            max_length=MAX_LEN,
                                            pad_to_max_length=True,
                                            return_attention_mask=True,
                                            truncation=True)
        input_ids.append(encoded_sent.get('input_ids'))
        attention_masks.append(encoded_sent.get('attention_mask'))
    # Convert lists to tensors
    input_ids = torch.tensor(input_ids)
    attention_masks = torch.tensor(attention_masks)
  
    return input_ids, attention_masks
  
  
  tokens_tensor, segments_tensor = preprocessing_for_bert(sentences)
  </code></pre>
  <p><strong>Load pretrained model</strong></p>
  <pre><code>import torch
  from transformers import BertTokenizer, BertModel
  
  # Load pre-trained model (weights)
  model = BertModel.from_pretrained('bert-base-uncased', output_hidden_states = True)
  
  # Put the model in &quot;evaluation&quot; mode, meaning feed-forward operation.
  model.eval()
  </code></pre>
  <p><strong>Run through BERT</strong></p>
  <pre><code># Run the text through BERT, and collect all of the hidden states produced
  # from all 12 layers. 
  with torch.no_grad():
  
      outputs = model(tokens_tensor, segments_tensor)
  
      # Evaluating the model will return a different number of objects based on 
      # how it's  configured in the `from_pretrained` call earlier. In this case, 
      # becase we set `output_hidden_states = True`, the third item will be the 
      # hidden states from all layers.
      hidden_states = outputs[2]
  </code></pre>
  <p><strong>The output results</strong></p>
  <pre><code>print (&quot;Number of layers:&quot;, len(hidden_states), &quot;  (initial embeddings + 12 BERT layers)&quot;)
  
  layer_i = 0
  print (&quot;Number of batches:&quot;, len(hidden_states[layer_i]))
  
  batch_i = 0
  print (&quot;Number of tokens:&quot;, len(hidden_states[layer_i][batch_i]))
  
  token_i = 0
  print (&quot;Number of hidden units:&quot;, len(hidden_states[layer_i][batch_i][token_i]))
  
  Number of layers: 13   (initial embeddings + 12 BERT layers)
  Number of batches: 4
  Number of tokens: 64
  Number of hidden units: 768
  </code></pre>
  <p><strong>Aggregation</strong></p>
  <pre><code>token_embeddings = torch.stack(hidden_states, dim=0)
  # Average over batches
  token_embeddings = token_embeddings.mean(1)
  token_embeddings = token_embeddings.permute(1,0,2)
  token_embeddings.size()
  ## -&gt; torch.Size([64, 13, 768])
  </code></pre>
  <p><strong>Prepare token embedding matrix</strong></p>
  <pre><code>token_vecs_sum = []
  # For each token in the sentence...
  for token in token_embeddings:
    # `token` is a [13 x 768] tensor
    # Sum the vectors from the last four layers.
    sum_vec = torch.sum(token[-4:], dim=0)  
    # Use `sum_vec` to represent `token`.
    token_vecs_sum.append(sum_vec)
  
  print ('Shape is: %d x %d' % (len(token_vecs_sum), len(token_vecs_sum[0])))
  </code></pre>