Machine Translation

! pip install datasets
! pip install evaluate
! pip install torchtext

# Download the embedding models
! python -m spacy download en_core_web_sm 
! python -m spacy download de_core_news_sm

import random

import datasets
import evaluate
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import spacy
import torch
import torch.nn as nn

import torchtext; torchtext.disable_torchtext_deprecation_warning()
import torchtext.vocab; torchtext.disable_torchtext_deprecation_warning()

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

######################### DO NOT CHANGE THIS CELL ##########################
seed = 1234

random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.deterministic = True

######################### DO NOT CHANGE THIS CELL ##########################
dataset = datasets.load_dataset("bentrevett/multi30k")

en_nlp = spacy.load("en_core_web_sm")
de_nlp = spacy.load("de_core_news_sm")

sos_token = "<sos>"
eos_token = "<eos>"


def tokenize_example(example, max_length=1000):
    en_tokens = [token.text.lower() for token in en_nlp.tokenizer(example["en"])][:max_length]
    de_tokens = [token.text.lower() for token in de_nlp.tokenizer(example["de"])][:max_length]

    en_tokens = [sos_token] + en_tokens + [eos_token]
    de_tokens = [sos_token] + de_tokens + [eos_token]
    return {"en_tokens": en_tokens, "de_tokens": de_tokens}


train_data = dataset["train"].map(tokenize_example)
valid_data = dataset["validation"].map(tokenize_example)
test_data = dataset["test"].map(tokenize_example)

######################### DO NOT CHANGE THIS CELL ##########################
min_freq = 2
unk_token = "<unk>"
pad_token = "<pad>"

special_tokens = [unk_token, pad_token, sos_token, eos_token]

en_vocab = torchtext.vocab.build_vocab_from_iterator(
    train_data["en_tokens"],
    min_freq=min_freq,
    specials=special_tokens,
)

de_vocab = torchtext.vocab.build_vocab_from_iterator(
    train_data["de_tokens"],
    min_freq=min_freq,
    specials=special_tokens,
)

assert en_vocab[unk_token] == de_vocab[unk_token]
assert en_vocab[pad_token] == de_vocab[pad_token]

unk_index = en_vocab[unk_token]
pad_index = en_vocab[pad_token]

en_vocab.set_default_index(unk_index)
de_vocab.set_default_index(unk_index)

######################### DO NOT CHANGE THIS CELL ##########################
def numericalize_example(example):
    en_ids = en_vocab.lookup_indices(example["en_tokens"])
    de_ids = de_vocab.lookup_indices(example["de_tokens"])
    return {"en_ids": en_ids, "de_ids": de_ids}

train_data = train_data.map(numericalize_example)
valid_data = valid_data.map(numericalize_example)
test_data = test_data.map(numericalize_example)

######################### DO NOT CHANGE THIS CELL ##########################
format_columns = ["en_ids", "de_ids"]

train_data = train_data.with_format(
    type="torch", columns=format_columns, output_all_columns=True
)

valid_data = valid_data.with_format(
    type="torch",
    columns=format_columns,
    output_all_columns=True,
)

test_data = test_data.with_format(
    type="torch",
    columns=format_columns,
    output_all_columns=True,
)

######################### DO NOT CHANGE THIS CELL ##########################
BATCH_SIZE = 128

def get_data_loader(dataset, batch_size, pad_index, shuffle=False):
    def collate_fn(batch):
        batch_en_ids = [example["en_ids"] for example in batch]
        batch_de_ids = [example["de_ids"] for example in batch]
        batch_en_ids = nn.utils.rnn.pad_sequence(batch_en_ids, padding_value=pad_index)
        batch_de_ids = nn.utils.rnn.pad_sequence(batch_de_ids, padding_value=pad_index)
        batch = {
            "en_ids": batch_en_ids,
            "de_ids": batch_de_ids,
        }
        return batch

    data_loader = torch.utils.data.DataLoader(
        dataset=dataset,
        batch_size=batch_size,
        collate_fn=collate_fn,
        shuffle=shuffle,
    )
    return data_loader

train_data_loader = get_data_loader(train_data, BATCH_SIZE, pad_index, shuffle=True)
valid_data_loader = get_data_loader(valid_data, BATCH_SIZE, pad_index)
test_data_loader = get_data_loader(test_data, BATCH_SIZE, pad_index)

class Encoder(nn.Module):
    def __init__(self, input_dim, embedding_dim, encoder_hidden_dim, decoder_hidden_dim, dropout):
        super().__init__()
        # TODO

    def forward(self, src):
        # TODO
        return outputs, hidden

class Attention(nn.Module):
    def __init__(self, encoder_hidden_dim, decoder_hidden_dim):
        super().__init__()
        # TODO

    def forward(self, hidden, encoder_outputs):
        # TODO
        return result

class Decoder(nn.Module):
    def __init__(
        self,
        output_dim,
        embedding_dim,
        encoder_hidden_dim,
        decoder_hidden_dim,
        dropout,
        attention,
    ):
        super().__init__()
        # TODO

    def forward(self, input, hidden, encoder_outputs):
        # TODO
        return prediction, hidden, attention

class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder):
        super().__init__()
        # TODO

    def forward(self, src, trg):
        # TODO
        return outputs

input_dim = len(de_vocab)
output_dim = len(en_vocab)
encoder_embedding_dim = 256
decoder_embedding_dim = 256
encoder_hidden_dim = 512
decoder_hidden_dim = 512
encoder_dropout = 0.5
decoder_dropout = 0.5

attention = ...  # TODO
encoder = ...  # TODO
decoder = ...  # TODO
model = ...  # TODO

# TODO: Write the training loop. Remember to collect training and validation statistics.

# TODO: Train the model on the provided dataset

# TODO: Plot the loss function

# TODO: Evaluate the model on the test set

def translate_sentence(sentence, model):
    # TODO
    return en_tokens, de_tokens, attention


# TODO: Translation examples

def plot_attention(sentence, translation, attention):
    # TODO


# TODO: Attention visualisation

# TODO