<kaggle_start><code># # HW2B: Neural Machine Translation
# In this project, you will build a neural machine translation system using modern techniques for sequence-to-sequence modeling. You will first implement a baseline encoder-decoder architecture, then improve upon the baseline by adding an attention mechanism and implementing beam search. The end result will be a fully functional translation system capable of translating simple German sentences into English.
# ## Setup
# First we install and import the required dependencies. These include:
# * `torch` for modeling and training
# * `torchtext` for data collection
# * `sentencepiece` for subword tokenization
# * `sacrebleu` for BLEU score evaluation
# Standard library imports
import json
import math
import random
import pdb

# Third party imports
import matplotlib.pyplot as plt
import numpy as np
import sacrebleu
import sentencepiece
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchtext
import tqdm.notebook

# Before proceeding, let's verify that we're connected to a GPU runtime and that `torch` can detect the GPU.
# We'll define a variable `device` here to use throughout the code so that we can easily change to run on CPU for debugging.
assert torch.cuda.is_available()
device = torch.device("cuda")
print("Using device:", device)

# ## Data
# The data for this assignment comes from the [Multi30K dataset](https://arxiv.org/abs/1605.00459), which contains English and German captions for images from Flickr. We can download and unpack it using `torchtext`. We use the Multi30K dataset because it is simpler than standard translation benchmark datasets and allows for models to be trained and evaluated in a matter of minutes rather than days.
# We will be translating from German to English in this assignment, but the same techniques apply equally well to any language pair.
#
extensions = [".de", ".en"]
source_field = torchtext.data.Field(tokenize=lambda x: x)
target_field = torchtext.data.Field(tokenize=lambda x: x)
training_data, validation_data, test_data = torchtext.datasets.Multi30k.splits(
    extensions, [source_field, target_field], root="."
)

# Now that we have the data, let's see how large each split is and look at a few examples.
print("Number of training examples:", len(training_data))
print("Number of validation examples:", len(validation_data))
print("Number of test examples:", len(test_data))
print()
for example in training_data[:10]:
    print(example.src)
    print(example.trg)
    print()

# ## Vocabulary
# We can use `sentencepiece` to create a joint German-English subword vocabulary from the training corpus. Because the number of training examples is small, we choose a smaller vocabulary size than would be used for large-scale NMT.
args = {
    "pad_id": 0,
    "bos_id": 1,
    "eos_id": 2,
    "unk_id": 3,
    "input": "multi30k/train.de,multi30k/train.en",
    "vocab_size": 8000,
    "model_prefix": "multi30k",
}
combined_args = " ".join("--{}={}".format(key, value) for key, value in args.items())
sentencepiece.SentencePieceTrainer.Train(combined_args)

# This creates two files: `multi30k.model` and `multi30k.vocab`. The first is a binary file containing the relevant data for the vocabulary. The second is a human-readable listing of each subword and its associated score.
# We can preview the contents of the vocabulary by looking at the first few rows from the human-readable file.

# As we can see, the vocabulary consists of four special tokens (`` for padding, `` for beginning of sentence (BOS), `` for end of sentence (EOS), `` for unknown) and a mixture of German and English words and subwords. In order to ensure reversability, word boundaries are encoded with a special unicode character "▁" (U+2581).
# To use the vocabulary, we first need to load it from the binary file produced above.
vocab = sentencepiece.SentencePieceProcessor()
vocab.Load("multi30k.model")

# The vocabulary object includes a number of methods for working with full sequences or individual pieces. We explore the most relevant ones below. A complete interface can be found on [GitHub](https://github.com/google/sentencepiece/tree/master/python#usage) for reference.
print("Vocabulary size:", vocab.GetPieceSize())
print()
for example in training_data[:3]:
    sentence = example.trg
    pieces = vocab.EncodeAsPieces(sentence)
    indices = vocab.EncodeAsIds(sentence)
    print(sentence)
    print(pieces)
    print(vocab.DecodePieces(pieces))
    print(indices)
    print(vocab.DecodeIds(indices))
    print()
piece = vocab.EncodeAsPieces("the")[0]
index = vocab.PieceToId(piece)
print(piece)
print(index)
print(vocab.IdToPiece(index))

# We define some constants here for the first three special tokens that you may find useful in the following sections.
pad_id = vocab.PieceToId("<pad>")
bos_id = vocab.PieceToId("<s>")
eos_id = vocab.PieceToId("</s>")

# Note that these tokens will be stripped from the output when converting from word pieces to text. This may be helpful when implementing greedy search and beam search.
sentence = training_data[0].trg
indices = vocab.EncodeAsIds(sentence)
indices_augmented = [bos_id] + indices + [eos_id, pad_id, pad_id, pad_id]
print(vocab.DecodeIds(indices))
print(vocab.DecodeIds(indices_augmented))
print(vocab.DecodeIds(indices) == vocab.DecodeIds(indices_augmented))


# ## Baseline sequence-to-sequence model
# With our data and vocabulary loaded, we're now ready to build a baseline sequence-to-sequence model.  Later on we'll add an attention mechanism to the model.
# Let's begin by defining a batch iterator for the training data. Given a dataset and a batch size, it will iterate over the dataset and yield pairs of tensors containing the subword indices for the source and target sentences in the batch, respectively.  Fill in `make_batch` below.
def make_batch(sentences):
    """Convert a list of sentences into a batch of subword indices.
    Args:
      sentences: A list of sentences, each of which is a string.
    Returns:
      A LongTensor of size (max_sequence_length, batch_size) containing the
      subword indices for the sentences, where max_sequence_length is the length
      of the longest sentence as encoded by the subword vocabulary and batch_size
      is the number of sentences in the batch. A beginning-of-sentence token
      should be included before each sequence, and an end-of-sentence token should
      be included after each sequence. Empty slots at the end of shorter sequences
      should be filled with padding tokens. The tensor should be located on the
      device defined at the beginning of the notebook.
    """
    # Implementation tip: You can use the nn.utils.rnn.pad_sequence utility
    # function to combine a list of variable-length sequences with padding.
    # YOUR CODE HERE
    ...
    # BEGIN SOLUTION
    # END SOLUTION


def make_batch_iterator(dataset, batch_size, shuffle=False):
    """Make a batch iterator that yields source-target pairs.
    Args:
      dataset: A torchtext dataset object.
      batch_size: An integer batch size.
      shuffle: A boolean indicating whether to shuffle the examples.
    Yields:
      Pairs of tensors constructed by calling the make_batch function on the
      source and target sentences in the current group of examples. The max
      sequence length can differ between the source and target tensor, but the
      batch size will be the same. The final batch may be smaller than the given
      batch size.
    """
    examples = list(dataset)
    if shuffle:
        random.shuffle(examples)
    for start_index in range(0, len(examples), batch_size):
        example_batch = examples[start_index : start_index + batch_size]
        source_sentences = [example.src for example in example_batch]
        target_sentences = [example.trg for example in example_batch]
        yield make_batch(source_sentences), make_batch(target_sentences)


test_batch = make_batch(["a test input", "a longer input than the first"])
print("Example batch tensor:")
print(test_batch)
assert test_batch.shape[1] == 2
assert test_batch[0, 0] == bos_id
assert test_batch[0, 1] == bos_id
assert test_batch[-1, 0] == pad_id
assert test_batch[-1, 1] == eos_id


# Now we will define the model itself. It should consist of a bidirectional LSTM encoder that encodes the input sentence into a fixed-size representation, and an LSTM decoder that uses this representation to produce the output sentence.
class Seq2seqBaseline(nn.Module):
    def __init__(self):
        super().__init__()
        # Initialize your model's parameters here. To get started, we suggest
        # setting all embedding and hidden dimensions to 256, using encoder and
        # decoder LSTMs with 2 layers, and using a dropout rate of 0.5.
        # Implementation tip: To create a bidirectional LSTM, you don't need to
        # create two LSTM networks. Instead use nn.LSTM(..., bidirectional=True).
        # YOUR CODE HERE
        ...
        # BEGIN SOLUTION
        # END SOLUTION

    def encode(self, source):
        """Encode the source batch using a bidirectional LSTM encoder.
        Args:
          source: An integer tensor with shape (max_source_sequence_length,
            batch_size) containing subword indices for the source sentences.
        Returns:
          A tuple with three elements:
            encoder_output: The output of the bidirectional LSTM with shape
              (max_source_sequence_length, batch_size, 2 * hidden_size).
            encoder_mask: A boolean tensor with shape (max_source_sequence_length,
              batch_size) indicating which encoder outputs correspond to padding
              tokens. Its elements should be True at positions corresponding to
              padding tokens and False elsewhere.
            encoder_hidden: The final hidden states of the bidirectional LSTM (after
              a suitable projection) that will be used to initialize the decoder.
              This should be a pair of tensors (h_n, c_n), each with shape
              (num_layers, batch_size, hidden_size). Note that the hidden state
              returned by the LSTM cannot be used directly. Its initial dimension is
              twice the required size because it contains state from two directions.
        The first two return values are not required for the baseline model and will
        only be used later in the attention model. If desired, they can be replaced
        with None for the initial implementation.
        """
        # Implementation tip: consider using packed sequences to more easily work
        # with the variable-length sequences represented by the source tensor.
        # See https://pytorch.org/docs/stable/nn.html#torch.nn.utils.rnn.PackedSequence.
        # Implementation tip: there are many simple ways to combine the forward
        # and backward portions of the final hidden state, e.g. addition, averaging,
        # or a linear transformation of the appropriate size. Any of these
        # should let you reach the required performance.
        # Compute a tensor containing the length of each source sequence.
        lengths = torch.sum(source != pad_id, axis=0)
        # YOUR CODE HERE
        ...
        # BEGIN SOLUTION
        # END SOLUTION

    def decode(self, decoder_input, initial_hidden, encoder_output, encoder_mask):
        """Run the decoder LSTM starting from an initial hidden state.
        The third and fourth arguments are not used in the baseline model, but are
        included for compatibility with the attention model in the next section.
        Args:
          decoder_input: An integer tensor with shape (max_decoder_sequence_length,
            batch_size) containing the subword indices for the decoder input. During
            evaluation, where decoding proceeds one step at a time, the initial
            dimension should be 1.
          initial_hidden: A pair of tensors (h_0, c_0) representing the initial
            state of the decoder, each with shape (num_layers, batch_size,
            hidden_size).
          encoder_output: The output of the encoder with shape
            (max_source_sequence_length, batch_size, 2 * hidden_size).
          encoder_mask: The output mask from the encoder with shape
            (max_source_sequence_length, batch_size). Encoder outputs at positions
            with a True value correspond to padding tokens and should be ignored.
        Returns:
          A tuple with three elements:
            logits: A tensor with shape (max_decoder_sequence_length, batch_size,
              vocab_size) containing unnormalized scores for the next-word
              predictions at each position.
            decoder_hidden: A pair of tensors (h_n, c_n) with the same shape as
              initial_hidden representing the updated decoder state after processing
              the decoder input.
            attention_weights: This will be implemented later in the attention
              model, but in order to maintain compatible type signatures, we also
              include it here. This can be None or any other placeholder value.
        """
        # These arguments are not used in the baseline model.
        del encoder_output
        del encoder_mask
        # YOUR CODE HERE
        ...
        # BEGIN SOLUTION
        # END SOLUTION

    def compute_loss(self, source, target):
        """Run the model on the source and compute the loss on the target.
        Args:
          source: An integer tensor with shape (max_source_sequence_length,
            batch_size) containing subword indices for the source sentences.
          target: An integer tensor with shape (max_target_sequence_length,
            batch_size) containing subword indices for the target sentences.
        Returns:
          A scalar float tensor representing cross-entropy loss on the current batch.
        """
        # Implementation tip: don't feed the target tensor directly to the decoder.
        # To see why, note that for a target sequence like <s> A B C </s>, you would
        # want to run the decoder on the prefix <s> A B C and have it predict the
        # suffix A B C </s>.
        # YOUR CODE HERE
        ...
        # BEGIN SOLUTION
        # END SOLUTION


# We define the following functions for training.  This code will run as provided, but you are welcome to modify the training loop to adjust the optimizer settings, add learning rate decay, etc.
def train(model, num_epochs, batch_size, model_file):
    """Train the model and save its best checkpoint.
    Model performance across epochs is evaluated using token-level accuracy on the
    validation set. The best checkpoint obtained during training will be stored on
    disk and loaded back into the model at the end of training.
    """
    optimizer = torch.optim.Adam(model.parameters())
    best_accuracy = 0.0
    for epoch in tqdm.notebook.trange(num_epochs, desc="training", unit="epoch"):
        with tqdm.notebook.tqdm(
            make_batch_iterator(training_data, batch_size, shuffle=True),
            desc="epoch {}".format(epoch + 1),
            unit="batch",
            total=math.ceil(len(training_data) / batch_size),
        ) as batch_iterator:
            model.train()
            total_loss = 0.0
            for i, (source, target) in enumerate(batch_iterator, start=1):
                optimizer.zero_grad()
                loss = model.compute_loss(source, target)
                total_loss += loss.item()
                loss.backward()
                optimizer.step()
                batch_iterator.set_postfix(mean_loss=total_loss / i)
            validation_perplexity, validation_accuracy = evaluate_next_token(
                model, validation_data
            )
            batch_iterator.set_postfix(
                mean_loss=total_loss / i,
                validation_perplexity=validation_perplexity,
                validation_token_accuracy=validation_accuracy,
            )
            if validation_accuracy > best_accuracy:
                print(
                    "Obtained a new best validation accuracy of {:.2f}, saving model "
                    "checkpoint to {}...".format(validation_accuracy, model_file)
                )
                torch.save(model.state_dict(), model_file)
                best_accuracy = validation_accuracy
    print("Reloading best model checkpoint from {}...".format(model_file))
    model.load_state_dict(torch.load(model_file))


def evaluate_next_token(model, dataset, batch_size=64):
    """Compute token-level perplexity and accuracy metrics.
    Note that the perplexity here is over subwords, not words.
    This function is used for validation set evaluation at the end of each epoch
    and should not be modified.
    """
    model.eval()
    total_cross_entropy = 0.0
    total_predictions = 0
    correct_predictions = 0
    with torch.no_grad():
        for source, target in make_batch_iterator(dataset, batch_size):
            encoder_output, encoder_mask, encoder_hidden = model.encode(source)
            decoder_input, decoder_target = target[:-1], target[1:]
            logits, decoder_hidden, attention_weights = model.decode(
                decoder_input, encoder_hidden, encoder_output, encoder_mask
            )
            total_cross_entropy += F.cross_entropy(
                logits.permute(1, 2, 0),
                decoder_target.permute(1, 0),
                ignore_index=pad_id,
                reduction="sum",
            ).item()
            total_predictions += (decoder_target != pad_id).sum().item()
            correct_predictions += (
                ((decoder_target != pad_id) & (decoder_target == logits.argmax(2)))
                .sum()
                .item()
            )
    perplexity = math.exp(total_cross_entropy / total_predictions)
    accuracy = 100 * correct_predictions / total_predictions
    return perplexity, accuracy


# We can now train the baseline model.
# Since we haven't yet defined a decoding method to output an entire string, we will measure performance for now by computing perplexity and the accuracy of predicting the next token given a gold prefix of the output. A correct implementation should get a validation token accuracy above 55%. The training code will automatically save the model with the highest validation accuracy and reload that checkpoint's parameters at the end of training.
# You are welcome to adjust these parameters based on your model implementation.
num_epochs = 10
batch_size = 16
baseline_model = Seq2seqBaseline().to(device)
train(baseline_model, num_epochs, batch_size, "baseline_model.pt")


# **Download your baseline model here.** Once you have a model you are happy with, you are encouraged to download it or save it to your Google Drive in case your session disconnects. The best baseline model has been saved to `baseline_model.pt` in the local filesystem. You will need a trained model while implementing inference below and to generate your final predictions. To download session files from Kaggle, please click the data tab on the right side of the screen and expand the `/kaggle/working` folder
# For evaluation, we also need to be able to generate entire strings from the model. We'll first define a greedy inference procedure here. Later on, we'll implement beam search.
# A correct implementation of greedy decoding should get above 20 BLEU on the validation set.
def predict_greedy(model, sentences, max_length=100):
    """Make predictions for the given inputs using greedy inference.
    Args:
      model: A sequence-to-sequence model.
      sentences: A list of input sentences, represented as strings.
      max_length: The maximum length at which to truncate outputs in order to
        avoid non-terminating inference.
    Returns:
      A list of predicted translations, represented as strings.
    """
    # Requirement: your implementation must be batched. This means that you should
    # make only one call to model.encode() at the start of the function, and make
    # only one call to model.decode() per inference step.
    # Implementation tip: once an EOS token has been generated, force the output
    # for that example to be padding tokens in all subsequent time steps by
    # adding a large positive number like 1e9 to the appropriate logits.
    # YOUR CODE HERE
    ...
    # BEGIN SOLUTION
    # END SOLUTION


def evaluate(model, dataset, batch_size=64, method="greedy"):
    assert method in {"greedy", "beam"}
    source_sentences = [example.src for example in dataset]
    target_sentences = [example.trg for example in dataset]
    model.eval()
    predictions = []
    with torch.no_grad():
        for start_index in range(0, len(source_sentences), batch_size):
            if method == "greedy":
                prediction_batch = predict_greedy(
                    model, source_sentences[start_index : start_index + batch_size]
                )
            else:
                prediction_batch = predict_beam(
                    model, source_sentences[start_index : start_index + batch_size]
                )
                prediction_batch = [candidates[0] for candidates in prediction_batch]
            predictions.extend(prediction_batch)
    return sacrebleu.corpus_bleu(predictions, [target_sentences]).score


print(
    "Baseline model validation BLEU using greedy search:",
    evaluate(baseline_model, validation_data),
)


def show_predictions(model, num_examples=4, include_beam=False):
    for example in validation_data[:num_examples]:
        print("Input:")
        print(" ", example.src)
        print("Target:")
        print(" ", example.trg)
        print("Greedy prediction:")
        print(" ", predict_greedy(model, [example.src])[0])
        if include_beam:
            print("Beam predictions:")
            for candidate in predict_beam(model, [example.src])[0]:
                print(" ", candidate)
        print()


print("Baseline model sample predictions:")
print()
show_predictions(baseline_model)


# ## Sequence-to-sequence model with attention
# Next, we extend the baseline model to include an attention mechanism in the decoder. This circumvents the need to store all information about the source sentence in a fixed-size representation, and should substantially improve performance and convergence time.
# Your implementation should use bilinear attention, where the attention distribution over the encoder outputs $e_1, \dots, e_n$ given a decoder LSTM output $d$ is obtained via a softmax of the dot products after a suitable projection to get them to the same size: $w_i \propto \exp ( d^\top W e_i )$. The unnormalized attention logits for encoder outputs corresponding to padding tokens should be offset with a large negative value to ensure that the corresponding attention weights are $0$.
# After computing the attention distribution, take a weighted sum of the projected encoder outputs to obtain the attention context $c = \sum_i w_i We_i$, and add this to the decoder output $d$ to obtain the final representation to be passed to the vocabulary projection layer.
class Seq2seqAttention(Seq2seqBaseline):
    def __init__(self):
        super().__init__()
        # Initialize any additional parameters needed for this model that are not
        # already included in the baseline model.
        # YOUR CODE HERE
        ...
        # BEGIN SOLUTION
        # END SOLUTION

    def decode(self, decoder_input, initial_hidden, encoder_output, encoder_mask):
        """Run the decoder LSTM starting from an initial hidden state.
        The third and fourth arguments are not used in the baseline model, but are
        included for compatibility with the attention model in the next section.
        Args:
          decoder_input: An integer tensor with shape (max_decoder_sequence_length,
            batch_size) containing the subword indices for the decoder input. During
            evaluation, where decoding proceeds one step at a time, the initial
            dimension should be 1.
          initial_hidden: A pair of tensors (h_0, c_0) representing the initial
            state of the decoder, each with shape (num_layers, batch_size,
            hidden_size).
          encoder_output: The output of the encoder with shape
            (max_source_sequence_length, batch_size, 2 * hidden_size).
          encoder_mask: The output mask from the encoder with shape
            (max_source_sequence_length, batch_size). Encoder outputs at positions
            with a True value correspond to padding tokens and should be ignored.
        Returns:
          A tuple with three elements:
            logits: A tensor with shape (max_decoder_sequence_length, batch_size,
              vocab_size) containing unnormalized scores for the next-word
              predictions at each position.
            decoder_hidden: A pair of tensors (h_n, c_n) with the same shape as
              initial_hidden representing the updated decoder state after processing
              the decoder input.
            attention_weights: A tensor with shape (max_decoder_sequence_length,
              batch_size, max_source_sequence_length) representing the normalized
              attention weights. This should sum to 1 along the last dimension.
        """
        # Implementation tip: use a large negative number like -1e9 instead of
        # float("-inf") when masking logits to avoid numerical issues.
        # Implementation tip: the function torch.einsum may be useful here.
        # See https://rockt.github.io/2018/04/30/einsum for a tutorial.
        # YOUR CODE HERE
        ...
        # BEGIN SOLUTION
        # END SOLUTION


# As before, we can train an attention model using the provided training code.
# A correct implementation should get a validation token accuracy above 64 and a validation BLEU above 36 with greedy search.
# You are welcome to adjust these parameters based on your model implementation.
num_epochs = 10
batch_size = 16
attention_model = Seq2seqAttention().to(device)
train(attention_model, num_epochs, batch_size, "attention_model.pt")
print(
    "Attention model validation BLEU using greedy search:",
    evaluate(attention_model, validation_data),
)

# **Download your attention model here.** Once you have a model you are happy with, you are encouraged to download it or save it to your Google Drive in case your session disconnects. The best attention model has been saved to `attention_model.pt` in the local filesystem. You will need a trained model while implementing beam search below and to generate your final predictions.
print(
    "Attention model validation BLEU using greedy search:",
    evaluate(attention_model, validation_data),
)
print()
print("Attention model sample predictions:")
print()
show_predictions(attention_model)


# ## Beam Search
# Now it's time to implement beam search.
# Similar to greedy search, beam search generates one token at a time. However, rather than keeping only the single best hypothesis, we instead keep the top $k$ candidates at each time step. This is accomplished by computing the set of next-token extensions for each item on the beam and finding the top $k$ across all candidates according to total log-probability.
# Candidates that are finished should stay on the beam through the end of inference. The search process concludes once all $k$ items on the beam are complete.
# With beam search, you should get an improvement of at least 0.5 BLEU over greedy search, and should reach above 21 BLEU without attention and above 37 BLEU with attention.
# **Tips:**
# 1) A good general strategy when doing complex code like this is to carefully annotate each line with a comment saying what each dimension represents.
# 2) You should only need one call to topk per step. You do not need to have a topk just over vocabulary first, you can directly go from vocab_size*beam_size to beam_size items.
# 3) Be sure you are correctly keeping track of which beam item a candidate is selected from and updating the beam states, such as LSTM hidden state, accordingly. A single state from the previous time step may need to be used for multiple new beam items or not at all. This includes all state associated with a beam, including all past tokens output by the beam and any extra tensors such as ones remembering when a beam is finished.
# 4) Pay attention to how you interleave things when using a single dimension to represent multiple things.  It will make a difference when you start reshaping to separate them out.  It may be easier to start with everything separate, then temporarily combine as needed.
# 5) For efficiency, we suggest that you implement all beam manipulations using batched PyTorch computations rather than Python for-loops.
# 6) Once an EOS token has been generated, force the output for that candidate to be padding tokens in all subsequent time steps by adding a large positive number like 1e9 to the appropriate logits. This will ensure that the candidate stays on the beam, as its probability will be very close to 1 and its score will effectively remain the same as when it was first completed.  All other (invalid) token continuations will have extremely low log probability and will not make it onto the beam.
# 7) While you are encouraged to keep your tensor dimensions constant for simplicity (aside from the sequence length), some special care will need to be taken on the first iteration to ensure that your beam doesn't fill up with k identical copies of the same candidate.
#
def predict_beam(model, sentences, k=5, max_length=100):
    """Make predictions for the given inputs using beam search.
    Args:
      model: A sequence-to-sequence model.
      sentences: A list of input sentences, represented as strings.
      k: The size of the beam.
      max_length: The maximum length at which to truncate outputs in order to
        avoid non-terminating inference.
    Returns:
      A list of beam predictions. Each element in the list should be a list of k
      strings corresponding to the top k predictions for the corresponding input,
      sorted in descending order by score.
    """
    # Requirement: your implementation must be batched. This means that you should
    # make only one call to model.encode() at the start of the function, and make
    # only one call to model.decode() per inference step.
    # Does top-k return relative ordering, if not how to return at end of method?
    # EOS candidate getting knocked.
    # YOUR CODE HERE
    ...
    # BEGIN SOLUTION
    # END SOLUTION


print(
    "Baseline model validation BLEU using beam search:",
    evaluate(baseline_model, validation_data, method="beam"),
)
print()
print("Baseline model sample predictions:")
print()
show_predictions(baseline_model, include_beam=True)
print(
    "Attention model validation BLEU using beam search:",
    evaluate(attention_model, validation_data, method="beam"),
)
print()
print("Attention model sample predictions:")
print()
show_predictions(attention_model, include_beam=True)


# ## Attention visualization: 1-Page Analysis
# Once you have everything working in the sections above, add some code here to visualize the decoder attention learned by the attention model using `matplotlib`.
# You may visualize decoder attention on gold source-target pairs from the validation data. You do not need to run any inference.
# For this section, you will submit a write-up interpreting attention maps generated by your model. Your write-up should be 1-page maximum in length and should be submitted in PDF format. You may use any editor you like, but we recommend using LaTeX and working in an environment like Overleaf. For full credit, your write-up should include:
# * A figure with attention map plots for 4 sentence pairs from the validation set (the method `imshow`, or equivalent, will likely be useful here). We encourage you to look through more maps to aid your analysis, but please only include 4 representative plots in the figure.
# * A brief discussion over trends you discover in the plots. Do the maps line up with your intuition, are there any surprising alignments? Are there any many-to-one or many-to-many alignments, or mainly one-to-one? Using a tool like Google Translate on substrings may help give some insight into this.
# When you submit the file, please name it report.pdf.
# You may find the following annotated heatmap tutorial helpful:
# https://matplotlib.org/3.1.3/gallery/images_contours_and_fields/image_annotated_heatmap.html.
# YOUR CODE HERE
...
# BEGIN SOLUTION
# END SOLUTION


# \* Describe your findings in this text cell. \*
# ## Submission
# Turn in the following files on Gradescope:
# * proj_2.ipynb (this file; please rename to match)
# * predictions.json (the predictions file generated by running the cell below)
# * report.pdf
# Be sure to check the output of the autograder after it runs.  It should confirm that no files are missing and that the output files have the correct format.
# The code below will generate the required predictions file.  **Note** that it is designed to create the file even if some required elements are missing so that you can submit for partial credit.  If you want full credit, you should check the output to make sure there are no warnings indicating missing portions.
# Run this cell to generate the predictions.json file required for submission.
def get_raw_predictions(model, dataset, method, batch_size=64):
    assert method in {"greedy", "beam"}
    source_sentences = [example.src for example in dataset]
    target_sentences = [example.trg for example in dataset]
    model.eval()
    predictions = []
    with torch.no_grad():
        for start_index in range(0, len(source_sentences), batch_size):
            if method == "greedy":
                prediction_batch = predict_greedy(
                    model, source_sentences[start_index : start_index + batch_size]
                )
            else:
                prediction_batch = predict_beam(
                    model, source_sentences[start_index : start_index + batch_size]
                )
            predictions.extend(prediction_batch)
    return predictions


def generate_predictions_file_for_submission(filepath):
    models = {"baseline": baseline_model, "attention": attention_model}
    datasets = {"validation": validation_data, "test": test_data}
    methods = ["greedy", "beam"]
    predictions = {}
    for model_name, model in models.items():
        for dataset_name, dataset in datasets.items():
            for method in methods:
                print(
                    "Getting predictions for {} model on {} set using {} "
                    "search...".format(model_name, dataset_name, method)
                )
                if model_name not in predictions:
                    predictions[model_name] = {}
                if dataset_name not in predictions[model_name]:
                    predictions[model_name][dataset_name] = {}
                try:
                    predictions[model_name][dataset_name][method] = get_raw_predictions(
                        model, dataset, method
                    )
                except:
                    print(
                        "!!! WARNING: An exception was raised, setting predictions to None !!!"
                    )
                    predictions[model_name][dataset_name][method] = None
    print("Writing predictions to {}...".format(filepath))
    with open(filepath, "w") as outfile:
        json.dump(predictions, outfile, indent=2)
    print("Finished writing predictions to {}.".format(filepath))


generate_predictions_file_for_submission("predictions.json")
<|endoftext|>
[{"Id": 87597153, "ScriptId": 24604634, "ParentScriptVersionId": NaN, "ScriptLanguageId": 9, "AuthorUserId": 9564005, "CreationDate": "02/11/2022 21:51:40", "VersionNumber": 1.0, "Title": "cs288-hw2b-public", "EvaluationDate": "02/11/2022", "IsChange": false, "TotalLines": 733.0, "LinesInsertedFromPrevious": 0.0, "LinesChangedFromPrevious": 0.0, "LinesUnchangedFromPrevious": 733.0, "LinesInsertedFromFork": 0.0, "LinesDeletedFromFork": 0.0, "LinesChangedFromFork": 0.0, "LinesUnchangedFromFork": 733.0, "TotalVotes": 0}]
# # HW2B: Neural Machine Translation
# In this project, you will build a neural machine translation system using modern techniques for sequence-to-sequence modeling. You will first implement a baseline encoder-decoder architecture, then improve upon the baseline by adding an attention mechanism and implementing beam search. The end result will be a fully functional translation system capable of translating simple German sentences into English.
# ## Setup
# First we install and import the required dependencies. These include:
# * `torch` for modeling and training
# * `torchtext` for data collection
# * `sentencepiece` for subword tokenization
# * `sacrebleu` for BLEU score evaluation
# Standard library imports
import json
import math
import random
import pdb

# Third party imports
import matplotlib.pyplot as plt
import numpy as np
import sacrebleu
import sentencepiece
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchtext
import tqdm.notebook

# Before proceeding, let's verify that we're connected to a GPU runtime and that `torch` can detect the GPU.
# We'll define a variable `device` here to use throughout the code so that we can easily change to run on CPU for debugging.
assert torch.cuda.is_available()
device = torch.device("cuda")
print("Using device:", device)

# ## Data
# The data for this assignment comes from the [Multi30K dataset](https://arxiv.org/abs/1605.00459), which contains English and German captions for images from Flickr. We can download and unpack it using `torchtext`. We use the Multi30K dataset because it is simpler than standard translation benchmark datasets and allows for models to be trained and evaluated in a matter of minutes rather than days.
# We will be translating from German to English in this assignment, but the same techniques apply equally well to any language pair.
#
extensions = [".de", ".en"]
source_field = torchtext.data.Field(tokenize=lambda x: x)
target_field = torchtext.data.Field(tokenize=lambda x: x)
training_data, validation_data, test_data = torchtext.datasets.Multi30k.splits(
    extensions, [source_field, target_field], root="."
)

# Now that we have the data, let's see how large each split is and look at a few examples.
print("Number of training examples:", len(training_data))
print("Number of validation examples:", len(validation_data))
print("Number of test examples:", len(test_data))
print()
for example in training_data[:10]:
    print(example.src)
    print(example.trg)
    print()

# ## Vocabulary
# We can use `sentencepiece` to create a joint German-English subword vocabulary from the training corpus. Because the number of training examples is small, we choose a smaller vocabulary size than would be used for large-scale NMT.
args = {
    "pad_id": 0,
    "bos_id": 1,
    "eos_id": 2,
    "unk_id": 3,
    "input": "multi30k/train.de,multi30k/train.en",
    "vocab_size": 8000,
    "model_prefix": "multi30k",
}
combined_args = " ".join("--{}={}".format(key, value) for key, value in args.items())
sentencepiece.SentencePieceTrainer.Train(combined_args)

# This creates two files: `multi30k.model` and `multi30k.vocab`. The first is a binary file containing the relevant data for the vocabulary. The second is a human-readable listing of each subword and its associated score.
# We can preview the contents of the vocabulary by looking at the first few rows from the human-readable file.

# As we can see, the vocabulary consists of four special tokens (`` for padding, `` for beginning of sentence (BOS), `` for end of sentence (EOS), `` for unknown) and a mixture of German and English words and subwords. In order to ensure reversability, word boundaries are encoded with a special unicode character "▁" (U+2581).
# To use the vocabulary, we first need to load it from the binary file produced above.
vocab = sentencepiece.SentencePieceProcessor()
vocab.Load("multi30k.model")

# The vocabulary object includes a number of methods for working with full sequences or individual pieces. We explore the most relevant ones below. A complete interface can be found on [GitHub](https://github.com/google/sentencepiece/tree/master/python#usage) for reference.
print("Vocabulary size:", vocab.GetPieceSize())
print()
for example in training_data[:3]:
    sentence = example.trg
    pieces = vocab.EncodeAsPieces(sentence)
    indices = vocab.EncodeAsIds(sentence)
    print(sentence)
    print(pieces)
    print(vocab.DecodePieces(pieces))
    print(indices)
    print(vocab.DecodeIds(indices))
    print()
piece = vocab.EncodeAsPieces("the")[0]
index = vocab.PieceToId(piece)
print(piece)
print(index)
print(vocab.IdToPiece(index))

# We define some constants here for the first three special tokens that you may find useful in the following sections.
pad_id = vocab.PieceToId("<pad>")
bos_id = vocab.PieceToId("<s>")
eos_id = vocab.PieceToId("</s>")

# Note that these tokens will be stripped from the output when converting from word pieces to text. This may be helpful when implementing greedy search and beam search.
sentence = training_data[0].trg
indices = vocab.EncodeAsIds(sentence)
indices_augmented = [bos_id] + indices + [eos_id, pad_id, pad_id, pad_id]
print(vocab.DecodeIds(indices))
print(vocab.DecodeIds(indices_augmented))
print(vocab.DecodeIds(indices) == vocab.DecodeIds(indices_augmented))


# ## Baseline sequence-to-sequence model
# With our data and vocabulary loaded, we're now ready to build a baseline sequence-to-sequence model.  Later on we'll add an attention mechanism to the model.
# Let's begin by defining a batch iterator for the training data. Given a dataset and a batch size, it will iterate over the dataset and yield pairs of tensors containing the subword indices for the source and target sentences in the batch, respectively.  Fill in `make_batch` below.
def make_batch(sentences):
    """Convert a list of sentences into a batch of subword indices.
    Args:
      sentences: A list of sentences, each of which is a string.
    Returns:
      A LongTensor of size (max_sequence_length, batch_size) containing the
      subword indices for the sentences, where max_sequence_length is the length
      of the longest sentence as encoded by the subword vocabulary and batch_size
      is the number of sentences in the batch. A beginning-of-sentence token
      should be included before each sequence, and an end-of-sentence token should
      be included after each sequence. Empty slots at the end of shorter sequences
      should be filled with padding tokens. The tensor should be located on the
      device defined at the beginning of the notebook.
    """
    # Implementation tip: You can use the nn.utils.rnn.pad_sequence utility
    # function to combine a list of variable-length sequences with padding.
    # YOUR CODE HERE
    ...
    # BEGIN SOLUTION
    # END SOLUTION


def make_batch_iterator(dataset, batch_size, shuffle=False):
    """Make a batch iterator that yields source-target pairs.
    Args:
      dataset: A torchtext dataset object.
      batch_size: An integer batch size.
      shuffle: A boolean indicating whether to shuffle the examples.
    Yields:
      Pairs of tensors constructed by calling the make_batch function on the
      source and target sentences in the current group of examples. The max
      sequence length can differ between the source and target tensor, but the
      batch size will be the same. The final batch may be smaller than the given
      batch size.
    """
    examples = list(dataset)
    if shuffle:
        random.shuffle(examples)
    for start_index in range(0, len(examples), batch_size):
        example_batch = examples[start_index : start_index + batch_size]
        source_sentences = [example.src for example in example_batch]
        target_sentences = [example.trg for example in example_batch]
        yield make_batch(source_sentences), make_batch(target_sentences)


test_batch = make_batch(["a test input", "a longer input than the first"])
print("Example batch tensor:")
print(test_batch)
assert test_batch.shape[1] == 2
assert test_batch[0, 0] == bos_id
assert test_batch[0, 1] == bos_id
assert test_batch[-1, 0] == pad_id
assert test_batch[-1, 1] == eos_id


# Now we will define the model itself. It should consist of a bidirectional LSTM encoder that encodes the input sentence into a fixed-size representation, and an LSTM decoder that uses this representation to produce the output sentence.
class Seq2seqBaseline(nn.Module):
    def __init__(self):
        super().__init__()
        # Initialize your model's parameters here. To get started, we suggest
        # setting all embedding and hidden dimensions to 256, using encoder and
        # decoder LSTMs with 2 layers, and using a dropout rate of 0.5.
        # Implementation tip: To create a bidirectional LSTM, you don't need to
        # create two LSTM networks. Instead use nn.LSTM(..., bidirectional=True).
        # YOUR CODE HERE
        ...
        # BEGIN SOLUTION
        # END SOLUTION

    def encode(self, source):
        """Encode the source batch using a bidirectional LSTM encoder.
        Args:
          source: An integer tensor with shape (max_source_sequence_length,
            batch_size) containing subword indices for the source sentences.
        Returns:
          A tuple with three elements:
            encoder_output: The output of the bidirectional LSTM with shape
              (max_source_sequence_length, batch_size, 2 * hidden_size).
            encoder_mask: A boolean tensor with shape (max_source_sequence_length,
              batch_size) indicating which encoder outputs correspond to padding
              tokens. Its elements should be True at positions corresponding to
              padding tokens and False elsewhere.
            encoder_hidden: The final hidden states of the bidirectional LSTM (after
              a suitable projection) that will be used to initialize the decoder.
              This should be a pair of tensors (h_n, c_n), each with shape
              (num_layers, batch_size, hidden_size). Note that the hidden state
              returned by the LSTM cannot be used directly. Its initial dimension is
              twice the required size because it contains state from two directions.
        The first two return values are not required for the baseline model and will
        only be used later in the attention model. If desired, they can be replaced
        with None for the initial implementation.
        """
        # Implementation tip: consider using packed sequences to more easily work
        # with the variable-length sequences represented by the source tensor.
        # See https://pytorch.org/docs/stable/nn.html#torch.nn.utils.rnn.PackedSequence.
        # Implementation tip: there are many simple ways to combine the forward
        # and backward portions of the final hidden state, e.g. addition, averaging,
        # or a linear transformation of the appropriate size. Any of these
        # should let you reach the required performance.
        # Compute a tensor containing the length of each source sequence.
        lengths = torch.sum(source != pad_id, axis=0)
        # YOUR CODE HERE
        ...
        # BEGIN SOLUTION
        # END SOLUTION

    def decode(self, decoder_input, initial_hidden, encoder_output, encoder_mask):
        """Run the decoder LSTM starting from an initial hidden state.
        The third and fourth arguments are not used in the baseline model, but are
        included for compatibility with the attention model in the next section.
        Args:
          decoder_input: An integer tensor with shape (max_decoder_sequence_length,
            batch_size) containing the subword indices for the decoder input. During
            evaluation, where decoding proceeds one step at a time, the initial
            dimension should be 1.
          initial_hidden: A pair of tensors (h_0, c_0) representing the initial
            state of the decoder, each with shape (num_layers, batch_size,
            hidden_size).
          encoder_output: The output of the encoder with shape
            (max_source_sequence_length, batch_size, 2 * hidden_size).
          encoder_mask: The output mask from the encoder with shape
            (max_source_sequence_length, batch_size). Encoder outputs at positions
            with a True value correspond to padding tokens and should be ignored.
        Returns:
          A tuple with three elements:
            logits: A tensor with shape (max_decoder_sequence_length, batch_size,
              vocab_size) containing unnormalized scores for the next-word
              predictions at each position.
            decoder_hidden: A pair of tensors (h_n, c_n) with the same shape as
              initial_hidden representing the updated decoder state after processing
              the decoder input.
            attention_weights: This will be implemented later in the attention
              model, but in order to maintain compatible type signatures, we also
              include it here. This can be None or any other placeholder value.
        """
        # These arguments are not used in the baseline model.
        del encoder_output
        del encoder_mask
        # YOUR CODE HERE
        ...
        # BEGIN SOLUTION
        # END SOLUTION

    def compute_loss(self, source, target):
        """Run the model on the source and compute the loss on the target.
        Args:
          source: An integer tensor with shape (max_source_sequence_length,
            batch_size) containing subword indices for the source sentences.
          target: An integer tensor with shape (max_target_sequence_length,
            batch_size) containing subword indices for the target sentences.
        Returns:
          A scalar float tensor representing cross-entropy loss on the current batch.
        """
        # Implementation tip: don't feed the target tensor directly to the decoder.
        # To see why, note that for a target sequence like <s> A B C </s>, you would
        # want to run the decoder on the prefix <s> A B C and have it predict the
        # suffix A B C </s>.
        # YOUR CODE HERE
        ...
        # BEGIN SOLUTION
        # END SOLUTION


# We define the following functions for training.  This code will run as provided, but you are welcome to modify the training loop to adjust the optimizer settings, add learning rate decay, etc.
def train(model, num_epochs, batch_size, model_file):
    """Train the model and save its best checkpoint.
    Model performance across epochs is evaluated using token-level accuracy on the
    validation set. The best checkpoint obtained during training will be stored on
    disk and loaded back into the model at the end of training.
    """
    optimizer = torch.optim.Adam(model.parameters())
    best_accuracy = 0.0
    for epoch in tqdm.notebook.trange(num_epochs, desc="training", unit="epoch"):
        with tqdm.notebook.tqdm(
            make_batch_iterator(training_data, batch_size, shuffle=True),
            desc="epoch {}".format(epoch + 1),
            unit="batch",
            total=math.ceil(len(training_data) / batch_size),
        ) as batch_iterator:
            model.train()
            total_loss = 0.0
            for i, (source, target) in enumerate(batch_iterator, start=1):
                optimizer.zero_grad()
                loss = model.compute_loss(source, target)
                total_loss += loss.item()
                loss.backward()
                optimizer.step()
                batch_iterator.set_postfix(mean_loss=total_loss / i)
            validation_perplexity, validation_accuracy = evaluate_next_token(
                model, validation_data
            )
            batch_iterator.set_postfix(
                mean_loss=total_loss / i,
                validation_perplexity=validation_perplexity,
                validation_token_accuracy=validation_accuracy,
            )
            if validation_accuracy > best_accuracy:
                print(
                    "Obtained a new best validation accuracy of {:.2f}, saving model "
                    "checkpoint to {}...".format(validation_accuracy, model_file)
                )
                torch.save(model.state_dict(), model_file)
                best_accuracy = validation_accuracy
    print("Reloading best model checkpoint from {}...".format(model_file))
    model.load_state_dict(torch.load(model_file))


def evaluate_next_token(model, dataset, batch_size=64):
    """Compute token-level perplexity and accuracy metrics.
    Note that the perplexity here is over subwords, not words.
    This function is used for validation set evaluation at the end of each epoch
    and should not be modified.
    """
    model.eval()
    total_cross_entropy = 0.0
    total_predictions = 0
    correct_predictions = 0
    with torch.no_grad():
        for source, target in make_batch_iterator(dataset, batch_size):
            encoder_output, encoder_mask, encoder_hidden = model.encode(source)
            decoder_input, decoder_target = target[:-1], target[1:]
            logits, decoder_hidden, attention_weights = model.decode(
                decoder_input, encoder_hidden, encoder_output, encoder_mask
            )
            total_cross_entropy += F.cross_entropy(
                logits.permute(1, 2, 0),
                decoder_target.permute(1, 0),
                ignore_index=pad_id,
                reduction="sum",
            ).item()
            total_predictions += (decoder_target != pad_id).sum().item()
            correct_predictions += (
                ((decoder_target != pad_id) & (decoder_target == logits.argmax(2)))
                .sum()
                .item()
            )
    perplexity = math.exp(total_cross_entropy / total_predictions)
    accuracy = 100 * correct_predictions / total_predictions
    return perplexity, accuracy


# We can now train the baseline model.
# Since we haven't yet defined a decoding method to output an entire string, we will measure performance for now by computing perplexity and the accuracy of predicting the next token given a gold prefix of the output. A correct implementation should get a validation token accuracy above 55%. The training code will automatically save the model with the highest validation accuracy and reload that checkpoint's parameters at the end of training.
# You are welcome to adjust these parameters based on your model implementation.
num_epochs = 10
batch_size = 16
baseline_model = Seq2seqBaseline().to(device)
train(baseline_model, num_epochs, batch_size, "baseline_model.pt")


# **Download your baseline model here.** Once you have a model you are happy with, you are encouraged to download it or save it to your Google Drive in case your session disconnects. The best baseline model has been saved to `baseline_model.pt` in the local filesystem. You will need a trained model while implementing inference below and to generate your final predictions. To download session files from Kaggle, please click the data tab on the right side of the screen and expand the `/kaggle/working` folder
# For evaluation, we also need to be able to generate entire strings from the model. We'll first define a greedy inference procedure here. Later on, we'll implement beam search.
# A correct implementation of greedy decoding should get above 20 BLEU on the validation set.
def predict_greedy(model, sentences, max_length=100):
    """Make predictions for the given inputs using greedy inference.
    Args:
      model: A sequence-to-sequence model.
      sentences: A list of input sentences, represented as strings.
      max_length: The maximum length at which to truncate outputs in order to
        avoid non-terminating inference.
    Returns:
      A list of predicted translations, represented as strings.
    """
    # Requirement: your implementation must be batched. This means that you should
    # make only one call to model.encode() at the start of the function, and make
    # only one call to model.decode() per inference step.
    # Implementation tip: once an EOS token has been generated, force the output
    # for that example to be padding tokens in all subsequent time steps by
    # adding a large positive number like 1e9 to the appropriate logits.
    # YOUR CODE HERE
    ...
    # BEGIN SOLUTION
    # END SOLUTION


def evaluate(model, dataset, batch_size=64, method="greedy"):
    assert method in {"greedy", "beam"}
    source_sentences = [example.src for example in dataset]
    target_sentences = [example.trg for example in dataset]
    model.eval()
    predictions = []
    with torch.no_grad():
        for start_index in range(0, len(source_sentences), batch_size):
            if method == "greedy":
                prediction_batch = predict_greedy(
                    model, source_sentences[start_index : start_index + batch_size]
                )
            else:
                prediction_batch = predict_beam(
                    model, source_sentences[start_index : start_index + batch_size]
                )
                prediction_batch = [candidates[0] for candidates in prediction_batch]
            predictions.extend(prediction_batch)
    return sacrebleu.corpus_bleu(predictions, [target_sentences]).score


print(
    "Baseline model validation BLEU using greedy search:",
    evaluate(baseline_model, validation_data),
)


def show_predictions(model, num_examples=4, include_beam=False):
    for example in validation_data[:num_examples]:
        print("Input:")
        print(" ", example.src)
        print("Target:")
        print(" ", example.trg)
        print("Greedy prediction:")
        print(" ", predict_greedy(model, [example.src])[0])
        if include_beam:
            print("Beam predictions:")
            for candidate in predict_beam(model, [example.src])[0]:
                print(" ", candidate)
        print()


print("Baseline model sample predictions:")
print()
show_predictions(baseline_model)


# ## Sequence-to-sequence model with attention
# Next, we extend the baseline model to include an attention mechanism in the decoder. This circumvents the need to store all information about the source sentence in a fixed-size representation, and should substantially improve performance and convergence time.
# Your implementation should use bilinear attention, where the attention distribution over the encoder outputs $e_1, \dots, e_n$ given a decoder LSTM output $d$ is obtained via a softmax of the dot products after a suitable projection to get them to the same size: $w_i \propto \exp ( d^\top W e_i )$. The unnormalized attention logits for encoder outputs corresponding to padding tokens should be offset with a large negative value to ensure that the corresponding attention weights are $0$.
# After computing the attention distribution, take a weighted sum of the projected encoder outputs to obtain the attention context $c = \sum_i w_i We_i$, and add this to the decoder output $d$ to obtain the final representation to be passed to the vocabulary projection layer.
class Seq2seqAttention(Seq2seqBaseline):
    def __init__(self):
        super().__init__()
        # Initialize any additional parameters needed for this model that are not
        # already included in the baseline model.
        # YOUR CODE HERE
        ...
        # BEGIN SOLUTION
        # END SOLUTION

    def decode(self, decoder_input, initial_hidden, encoder_output, encoder_mask):
        """Run the decoder LSTM starting from an initial hidden state.
        The third and fourth arguments are not used in the baseline model, but are
        included for compatibility with the attention model in the next section.
        Args:
          decoder_input: An integer tensor with shape (max_decoder_sequence_length,
            batch_size) containing the subword indices for the decoder input. During
            evaluation, where decoding proceeds one step at a time, the initial
            dimension should be 1.
          initial_hidden: A pair of tensors (h_0, c_0) representing the initial
            state of the decoder, each with shape (num_layers, batch_size,
            hidden_size).
          encoder_output: The output of the encoder with shape
            (max_source_sequence_length, batch_size, 2 * hidden_size).
          encoder_mask: The output mask from the encoder with shape
            (max_source_sequence_length, batch_size). Encoder outputs at positions
            with a True value correspond to padding tokens and should be ignored.
        Returns:
          A tuple with three elements:
            logits: A tensor with shape (max_decoder_sequence_length, batch_size,
              vocab_size) containing unnormalized scores for the next-word
              predictions at each position.
            decoder_hidden: A pair of tensors (h_n, c_n) with the same shape as
              initial_hidden representing the updated decoder state after processing
              the decoder input.
            attention_weights: A tensor with shape (max_decoder_sequence_length,
              batch_size, max_source_sequence_length) representing the normalized
              attention weights. This should sum to 1 along the last dimension.
        """
        # Implementation tip: use a large negative number like -1e9 instead of
        # float("-inf") when masking logits to avoid numerical issues.
        # Implementation tip: the function torch.einsum may be useful here.
        # See https://rockt.github.io/2018/04/30/einsum for a tutorial.
        # YOUR CODE HERE
        ...
        # BEGIN SOLUTION
        # END SOLUTION


# As before, we can train an attention model using the provided training code.
# A correct implementation should get a validation token accuracy above 64 and a validation BLEU above 36 with greedy search.
# You are welcome to adjust these parameters based on your model implementation.
num_epochs = 10
batch_size = 16
attention_model = Seq2seqAttention().to(device)
train(attention_model, num_epochs, batch_size, "attention_model.pt")
print(
    "Attention model validation BLEU using greedy search:",
    evaluate(attention_model, validation_data),
)

# **Download your attention model here.** Once you have a model you are happy with, you are encouraged to download it or save it to your Google Drive in case your session disconnects. The best attention model has been saved to `attention_model.pt` in the local filesystem. You will need a trained model while implementing beam search below and to generate your final predictions.
print(
    "Attention model validation BLEU using greedy search:",
    evaluate(attention_model, validation_data),
)
print()
print("Attention model sample predictions:")
print()
show_predictions(attention_model)


# ## Beam Search
# Now it's time to implement beam search.
# Similar to greedy search, beam search generates one token at a time. However, rather than keeping only the single best hypothesis, we instead keep the top $k$ candidates at each time step. This is accomplished by computing the set of next-token extensions for each item on the beam and finding the top $k$ across all candidates according to total log-probability.
# Candidates that are finished should stay on the beam through the end of inference. The search process concludes once all $k$ items on the beam are complete.
# With beam search, you should get an improvement of at least 0.5 BLEU over greedy search, and should reach above 21 BLEU without attention and above 37 BLEU with attention.
# **Tips:**
# 1) A good general strategy when doing complex code like this is to carefully annotate each line with a comment saying what each dimension represents.
# 2) You should only need one call to topk per step. You do not need to have a topk just over vocabulary first, you can directly go from vocab_size*beam_size to beam_size items.
# 3) Be sure you are correctly keeping track of which beam item a candidate is selected from and updating the beam states, such as LSTM hidden state, accordingly. A single state from the previous time step may need to be used for multiple new beam items or not at all. This includes all state associated with a beam, including all past tokens output by the beam and any extra tensors such as ones remembering when a beam is finished.
# 4) Pay attention to how you interleave things when using a single dimension to represent multiple things.  It will make a difference when you start reshaping to separate them out.  It may be easier to start with everything separate, then temporarily combine as needed.
# 5) For efficiency, we suggest that you implement all beam manipulations using batched PyTorch computations rather than Python for-loops.
# 6) Once an EOS token has been generated, force the output for that candidate to be padding tokens in all subsequent time steps by adding a large positive number like 1e9 to the appropriate logits. This will ensure that the candidate stays on the beam, as its probability will be very close to 1 and its score will effectively remain the same as when it was first completed.  All other (invalid) token continuations will have extremely low log probability and will not make it onto the beam.
# 7) While you are encouraged to keep your tensor dimensions constant for simplicity (aside from the sequence length), some special care will need to be taken on the first iteration to ensure that your beam doesn't fill up with k identical copies of the same candidate.
#
def predict_beam(model, sentences, k=5, max_length=100):
    """Make predictions for the given inputs using beam search.
    Args:
      model: A sequence-to-sequence model.
      sentences: A list of input sentences, represented as strings.
      k: The size of the beam.
      max_length: The maximum length at which to truncate outputs in order to
        avoid non-terminating inference.
    Returns:
      A list of beam predictions. Each element in the list should be a list of k
      strings corresponding to the top k predictions for the corresponding input,
      sorted in descending order by score.
    """
    # Requirement: your implementation must be batched. This means that you should
    # make only one call to model.encode() at the start of the function, and make
    # only one call to model.decode() per inference step.
    # Does top-k return relative ordering, if not how to return at end of method?
    # EOS candidate getting knocked.
    # YOUR CODE HERE
    ...
    # BEGIN SOLUTION
    # END SOLUTION


print(
    "Baseline model validation BLEU using beam search:",
    evaluate(baseline_model, validation_data, method="beam"),
)
print()
print("Baseline model sample predictions:")
print()
show_predictions(baseline_model, include_beam=True)
print(
    "Attention model validation BLEU using beam search:",
    evaluate(attention_model, validation_data, method="beam"),
)
print()
print("Attention model sample predictions:")
print()
show_predictions(attention_model, include_beam=True)


# ## Attention visualization: 1-Page Analysis
# Once you have everything working in the sections above, add some code here to visualize the decoder attention learned by the attention model using `matplotlib`.
# You may visualize decoder attention on gold source-target pairs from the validation data. You do not need to run any inference.
# For this section, you will submit a write-up interpreting attention maps generated by your model. Your write-up should be 1-page maximum in length and should be submitted in PDF format. You may use any editor you like, but we recommend using LaTeX and working in an environment like Overleaf. For full credit, your write-up should include:
# * A figure with attention map plots for 4 sentence pairs from the validation set (the method `imshow`, or equivalent, will likely be useful here). We encourage you to look through more maps to aid your analysis, but please only include 4 representative plots in the figure.
# * A brief discussion over trends you discover in the plots. Do the maps line up with your intuition, are there any surprising alignments? Are there any many-to-one or many-to-many alignments, or mainly one-to-one? Using a tool like Google Translate on substrings may help give some insight into this.
# When you submit the file, please name it report.pdf.
# You may find the following annotated heatmap tutorial helpful:
# https://matplotlib.org/3.1.3/gallery/images_contours_and_fields/image_annotated_heatmap.html.
# YOUR CODE HERE
...
# BEGIN SOLUTION
# END SOLUTION


# \* Describe your findings in this text cell. \*
# ## Submission
# Turn in the following files on Gradescope:
# * proj_2.ipynb (this file; please rename to match)
# * predictions.json (the predictions file generated by running the cell below)
# * report.pdf
# Be sure to check the output of the autograder after it runs.  It should confirm that no files are missing and that the output files have the correct format.
# The code below will generate the required predictions file.  **Note** that it is designed to create the file even if some required elements are missing so that you can submit for partial credit.  If you want full credit, you should check the output to make sure there are no warnings indicating missing portions.
# Run this cell to generate the predictions.json file required for submission.
def get_raw_predictions(model, dataset, method, batch_size=64):
    assert method in {"greedy", "beam"}
    source_sentences = [example.src for example in dataset]
    target_sentences = [example.trg for example in dataset]
    model.eval()
    predictions = []
    with torch.no_grad():
        for start_index in range(0, len(source_sentences), batch_size):
            if method == "greedy":
                prediction_batch = predict_greedy(
                    model, source_sentences[start_index : start_index + batch_size]
                )
            else:
                prediction_batch = predict_beam(
                    model, source_sentences[start_index : start_index + batch_size]
                )
            predictions.extend(prediction_batch)
    return predictions


def generate_predictions_file_for_submission(filepath):
    models = {"baseline": baseline_model, "attention": attention_model}
    datasets = {"validation": validation_data, "test": test_data}
    methods = ["greedy", "beam"]
    predictions = {}
    for model_name, model in models.items():
        for dataset_name, dataset in datasets.items():
            for method in methods:
                print(
                    "Getting predictions for {} model on {} set using {} "
                    "search...".format(model_name, dataset_name, method)
                )
                if model_name not in predictions:
                    predictions[model_name] = {}
                if dataset_name not in predictions[model_name]:
                    predictions[model_name][dataset_name] = {}
                try:
                    predictions[model_name][dataset_name][method] = get_raw_predictions(
                        model, dataset, method
                    )
                except:
                    print(
                        "!!! WARNING: An exception was raised, setting predictions to None !!!"
                    )
                    predictions[model_name][dataset_name][method] = None
    print("Writing predictions to {}...".format(filepath))
    with open(filepath, "w") as outfile:
        json.dump(predictions, outfile, indent=2)
    print("Finished writing predictions to {}.".format(filepath))


generate_predictions_file_for_submission("predictions.json")
