Translation as Language Modeling¶

Goals:

  • Practice getting data into and out of a language model.
    • embeddings (input and output)
    • logits for next words
    • cross-entropy loss
  • Explore different methods of decoding for sequence generation
  • Explain how data flows between the encoder and decoder in a sequence-to-sequence model
  • Interpret attention weights.

Setup¶

Install libraries.

In [ ]:
#%pip install -q datasets transformers[sentencepiece]

Import PyTorch and the HuggingFace Transformers library.

In [ ]:
import torch
import torch.nn as nn
import torch.nn.functional as F
import transformers
import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

Load a Marian Machine Translation model.

Specifically, we're using one that was trained on the OPUS corpus (opus-mt) to translate text in any romance language (roa) to English (en).

In [ ]:
from transformers import MarianMTModel, MarianTokenizer
model_name = 'Helsinki-NLP/opus-mt-roa-en'
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name).to(device)
print(f"The model has {model.num_parameters():,d} parameters.")

/Users/ka37/Courses/cs344/.venv/lib/python3.11/site-packages/transformers/models/marian/tokenization_marian.py:197: UserWarning: Recommended: pip install sacremoses.
  warnings.warn("Recommended: pip install sacremoses.")

The model has 77,943,296 parameters.

Finally, these wrappers will make the code below easier to understand (you should completely ignore them).

In [ ]:
from functools import partial
from transformers.models.marian.modeling_marian import shift_tokens_right
prepend_start_token = partial(
    shift_tokens_right,
    pad_token_id = model.config.pad_token_id, decoder_start_token_id = model.config.decoder_start_token_id)
encoder = model.get_encoder()
decoder = model.get_decoder()
encoder.forward = partial(encoder.forward, output_attentions=True, output_hidden_states=True)
decoder.forward = partial(decoder.forward, output_attentions=True, output_hidden_states=True)

Warm-up¶

Let's practice with the tokenizer. This should be mostly review, but we'll do it in the way that the HuggingFace docs do it.

In [ ]:
spanish_text = "Yo les doy vida eterna."
spanish_batch = tokenizer(spanish_text, return_tensors='pt', padding=True, ).to(device)
spanish_batch
Out[ ]:
{'input_ids': tensor([[ 2554,    29,    73,   131,   860, 21658,     3,     0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1]])}

Since we're only translating one sentence, we can ignore attention_mask (which just helps ignore padding tokens) and the extra initial dimension of the input_ids.

In [ ]:
input_ids = spanish_batch.input_ids
input_ids.shape
Out[ ]:
torch.Size([1, 8])
In [ ]:
tokenizer.convert_ids_to_tokens(input_ids[0])
Out[ ]:
['▁Yo', '▁les', '▁do', 'y', '▁vida', '▁eterna', '.', '</s>']

Now let's ask the model to generate a translation. Lots of magic happens here; we'll peel back the layers shortly.

In [ ]:
translated = model.generate(input_ids = input_ids, num_beams=1, do_sample=False)
translated.shape
Out[ ]:
torch.Size([1, 8])

Decode the result!

In [ ]:
english_text = tokenizer.decode(translated[0])
english_text
Out[ ]:
'<pad> I give them eternal life.</s>'

Scoring a candidate translation¶

The model gives us conditional probabilities of each word: P(word_i | src, word_1, word_2, ..., word_{i-1}). Those underflow, so we actually use the log probabilities, which are conveniently the cross-entropy loss values of each token.

First, let's look at the loss that the model gives. We'll compare the correct translation with an incorrect one:

In [ ]:
def tokenize_target_sentence(sentence):
    return tokenizer(text_target=sentence, return_tensors='pt', padding=True).to(device)['input_ids']
correct_target_ids = tokenize_target_sentence("I give them eternal life.")
wrong_target_ids = tokenize_target_sentence("I give them eternal death.")

Let's run a forward pass through the full model (encoder and decoder) with the complete candidate translation. First, the correct translation:

In [ ]:
@torch.no_grad() # We don't need to compute gradients 
def get_logprob_of_translation(src_ids, tgt_ids):
    model_outputs = model(
        input_ids = src_ids,
        labels = tgt_ids
    )
    return -model_outputs.loss # TODO: multiply by num tokens? Replace by manually doing cross_entropy_loss?
get_logprob_of_translation(spanish_batch.input_ids, correct_target_ids)
Out[ ]:
tensor(-0.2088)

Now (your turn) the incorrect translation:

In [ ]:
# your code here
Out[ ]:
tensor(-1.3203)

Dig In!¶

Ok now how did it do that?

You may find it helpful to have the documentation for the MarianMT model in HuggingFace Transformers open. But you can do all of this without referring to it.

The guts of the model¶

I've ripped out all the plumbing code and things you only need in special situations to just show the guts of the model below. Study this code carefully with the help of the questions below it. Add comments to describe what each line does. Include, where applicable, the shape of the tensors involved.

In [ ]:
encoder_input_ids = spanish_batch.input_ids
target_ids = correct_target_ids
decoder_input_ids = prepend_start_token(target_ids)

with torch.no_grad():
    encoder_outputs = encoder(input_ids = encoder_input_ids)

    decoder_outputs = decoder(
        input_ids = decoder_input_ids,
        encoder_hidden_states = encoder_outputs.last_hidden_state
    )

    output_embedding = decoder_outputs.last_hidden_state
    token_embeddings = model.lm_head.weight
    logits = output_embedding @ token_embeddings.t()
    logits += model.final_logits_bias

    # ignore the batch dimension.
    logits = logits[0]

nlls_of_correct_tokens = F.cross_entropy(logits, target_ids[0], reduction='none')
nlls_of_correct_tokens.mean()
Out[ ]:
tensor(0.2088)

Explain logits.shape.

In [ ]:
logits.shape
Out[ ]:
torch.Size([7, 65001])

your narrative answer here

In [ ]:
tokenizer.convert_ids_to_tokens(logits.argmax(dim=1))
Out[ ]:
['▁I', '▁give', '▁them', '▁eternal', '▁life', '.', '</s>']
In [ ]:
tokenizer.convert_ids_to_tokens(target_ids[0])
Out[ ]:
['▁I', '▁give', '▁them', '▁eternal', '▁life', '.', '</s>']

What tensor contains all of the information from the Spanish sentence that is used to generate the English sentence? Explain each element of the shape of that tensor.

(The leading "1" is the batch dimension; you can ignore this unless you're translating multiple sentence simultaneously.)

What is the "shape" of this model? Specifically:

  1. What is the dimensionality of the hidden vectors it uses to represent everything? (How does this relate to the dimensionality of the token embeddings?)
  2. How many internal layers does the model have?
In [ ]:
encoder_outputs.last_hidden_state.shape
Out[ ]:
torch.Size([1, 8, 512])
In [ ]:
model.config.num_hidden_layers
Out[ ]:
6

Visualize attentions¶

Read these as: the row token looks at the column token.

There are actually 8 attention heads for each of the 6 layers, so to visualize simply, we take the mean over the attention weights (which are all positive).

In [ ]:
decoder_outputs.cross_attentions[0].shape
Out[ ]:
torch.Size([1, 8, 7, 8])
In [ ]:
layer = -1
if layer < 0: layer += len(encoder_outputs.attentions)

fig, axs = plt.subplots(2, 2, figsize=(12, 12))

encoder_tokens = tokenizer.convert_ids_to_tokens(encoder_input_ids[0])
decoder_tokens = tokenizer.convert_ids_to_tokens(decoder_input_ids[0])
encoder_ticks = torch.arange(len(encoder_tokens)) + 0.5
decoder_ticks = torch.arange(len(decoder_tokens)) + 0.5

ax = axs[0, 0]
ax.pcolormesh(encoder_outputs.attentions[layer][0].mean(dim=0).cpu().numpy())
ax.set_title(f"Encoder Self-Attention Weights for layer {layer} (avg over all {model.config.num_attention_heads} heads)")
ax.set_xticks(encoder_ticks, encoder_tokens)
ax.set_yticks(encoder_ticks, encoder_tokens)

ax = axs[0, 1]
ax.pcolormesh(decoder_outputs.cross_attentions[layer][0].mean(dim=0).cpu().numpy(), vmin=0, vmax=1)
ax.set_title(f"Cross-Attention Weights for layer {layer} (avg over all {model.config.num_attention_heads} heads)")
ax.set_xticks(encoder_ticks, encoder_tokens)
ax.set_yticks(decoder_ticks, decoder_tokens);

ax = axs[1, 0]
ax.pcolormesh(decoder_outputs.attentions[layer][0].mean(dim=0).cpu().numpy())
ax.set_title(f"Decoder Self-Attention Weights for layer {layer} (avg over all {model.config.num_attention_heads} heads)")
ax.set_xticks(torch.arange(7)+.5, tokenizer.convert_ids_to_tokens(decoder_input_ids[0]))
ax.set_yticks(torch.arange(7)+.5, tokenizer.convert_ids_to_tokens(decoder_input_ids[0]));
No description has been provided for this image

Similarity¶

Notice that the last step of the model is a dot product with all the token embeddings. Recall that a dot product is a measure of similarity. Let's look at similarity in embedding space.

In [ ]:
normalized_token_embeddings = token_embeddings / token_embeddings.norm(p=2, dim=1, keepdim=True)
In [ ]:
query_word = "London"
query_ids = tokenizer(text_target=query_word, add_special_tokens=False)['input_ids']
print(query_ids)
query = token_embeddings[query_ids].mean(dim=0)
similarities = query @ normalized_token_embeddings.t()
most_similar_indices = similarities.topk(50).indices
tokenizer.convert_ids_to_tokens(most_similar_indices)

[5226]
Out[ ]:
['<pad>',
 '▁London',
 '▁Moscow',
 '▁Cambridge',
 '▁Kingston',
 '▁Bremen',
 '▁Windsor',
 '▁Philadelphia',
 '▁Melbourne',
 '▁Baltimore',
 '▁Bristol',
 '▁Cleveland',
 '▁Houston',
 '▁Belfast',
 '▁Denver',
 '▁Baghdad',
 '▁Liverpool',
 '▁Oregon',
 '▁England',
 '▁Edinburgh',
 '▁Tripoli',
 '▁Missouri',
 '▁Flanders',
 '▁Mumbai',
 '▁Churchill',
 '▁Istanbul',
 '▁Bermuda',
 '▁Barcelona',
 '▁Kentucky',
 '▁Detroit',
 '▁Honda',
 '▁Lorraine',
 '▁Tibet',
 '▁Brussels',
 '▁Lusaka',
 '▁Honduran',
 '▁Madison',
 '▁Bordeaux',
 '▁Mormon',
 '▁Maryland',
 '▁Alabama',
 '▁Damascus',
 '▁Tibetan',
 '▁Versailles',
 '▁Iowa',
 '▁Orleans',
 '▁Burgundy',
 '▁Naples',
 '▁Murcia',
 '▁Glasgow']

Your turn: now, take query vectors from the output_embeddings that were calculated above and find the most similar token embeddings.

Compare the results with the translation output you saw from the model earlier.

In [ ]:
# your code here
Out[ ]:
['▁them',
 '▁you',
 '▁eternal',
 "▁'",
 '▁to',
 '▁everlasting',
 '▁it',
 '▁the',
 "'",
 ',',
 '▁these',
 '▁unto',
 '▁[',
 '▁him',
 '▁forever',
 '▁all',
 '▁that',
 '▁their',
 '▁those',
 '▁up',
 '▁life',
 '▁they',
 '▁for',
 '▁y',
 '▁You',
 '▁ye',
 '▁out',
 '▁Oh',
 '▁I',
 '▁your',
 '▁an',
 '▁people',
 '▁-',
 '▁eternity',
 '▁"',
 '▁(',
 '▁YOU',
 '▁such',
 '▁her',
 '▁birth',
 '▁us',
 '▁perpetual',
 '▁forth',
 '▁of',
 '▁this',
 '▁a',
 '▁lasting',
 '▁Eternal',
 '▁lifelong',
 '▁Him']