Transformer_playground.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Transformer playground",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/littell/NamedTransformer/blob/master/Transformer_playground.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bD1L4FEJVUeh",
        "colab_type": "text"
      },
      "source": [
        "# Transformer in NamedTensor\n",
        "\n",
        "Based on this PyTorch Transformer code: https://blog.floydhub.com/the-transformer-in-pytorch/\n",
        "\n",
        "Reimplemented in NamedTensor: https://github.com/harvardnlp/namedtensor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cWBNCxwQowDM",
        "colab_type": "text"
      },
      "source": [
        "## Install the necessary software\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uDe6LN_xorLj",
        "colab_type": "code",
        "outputId": "af0f81a9-e276-45c9-f005-3c42ac361500",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "# -*- coding: utf-8 -*-\n",
        "\n",
        "from __future__ import unicode_literals, print_function\n",
        "\n",
        "from io import open\n",
        "import spacy\n",
        "import torch\n",
        "import numpy as np\n",
        "\n",
        "import math\n",
        "import random\n",
        "import logging\n",
        "import hashlib\n",
        "from collections import defaultdict\n",
        "import yaml\n",
        "import time\n",
        "\n",
        "from copy import deepcopy\n",
        "\n",
        "# bidict\n",
        "!pip install bidict\n",
        "from bidict import bidict\n",
        "\n",
        "# namedtensor\n",
        "!pip install -qq git+https://github.com/harvardnlp/namedtensor\n",
        "from namedtensor import ntorch\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "print(f\"Device = {device}\")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: bidict in /usr/local/lib/python3.6/dist-packages (0.18.0)\n",
            "  Building wheel for namedtensor (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Device = cuda\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j6lhC9dNo4zz",
        "colab_type": "text"
      },
      "source": [
        "## Load some data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o1EPdVRWo5P2",
        "colab_type": "code",
        "outputId": "d2af60e2-8a4b-42c8-bf91-b0429f5e23b0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 306
        }
      },
      "source": [
        "!wget 'http://opus.nlpl.eu/download.php?f=GlobalVoices/v2017q3/moses/en-sw.txt.zip' -O swahili_data.zip\n",
        "!unzip -qo swahili_data.zip\n",
        "!wc -l GlobalVoices.en-sw.en"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2019-07-08 20:47:19--  http://opus.nlpl.eu/download.php?f=GlobalVoices/v2017q3/moses/en-sw.txt.zip\n",
            "Resolving opus.nlpl.eu (opus.nlpl.eu)... 193.166.25.9\n",
            "Connecting to opus.nlpl.eu (opus.nlpl.eu)|193.166.25.9|:80... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: https://object.pouta.csc.fi/OPUS-GlobalVoices/v2017q3/moses/en-sw.txt.zip [following]\n",
            "--2019-07-08 20:47:20--  https://object.pouta.csc.fi/OPUS-GlobalVoices/v2017q3/moses/en-sw.txt.zip\n",
            "Resolving object.pouta.csc.fi (object.pouta.csc.fi)... 86.50.254.0, 86.50.254.1\n",
            "Connecting to object.pouta.csc.fi (object.pouta.csc.fi)|86.50.254.0|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 2890397 (2.8M) [application/zip]\n",
            "Saving to: ‘swahili_data.zip’\n",
            "\n",
            "swahili_data.zip    100%[===================>]   2.76M  1.39MB/s    in 2.0s    \n",
            "\n",
            "2019-07-08 20:47:23 (1.39 MB/s) - ‘swahili_data.zip’ saved [2890397/2890397]\n",
            "\n",
            "29698 GlobalVoices.en-sw.en\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MyPfmPcYR14o",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#!wget 'http://www.statmt.org/europarl/v7/fr-en.tgz' -O fr_en.tgz\n",
        "#!tar xvzf fr_en.tgz\n",
        "#!wc -l europarl-v7.fr-en.fr"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GRhkgv9MHp-L",
        "colab_type": "text"
      },
      "source": [
        "## Some utility functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zEaewlpOHnea",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def load_yaml(path):\n",
        "    with open(path, \"r\", encoding=\"utf-8\") as fin:\n",
        "        return yaml.safe_load(fin)\n",
        "      \n",
        "def save_yaml(path, obj):\n",
        "    with open(path, \"w\", encoding=\"utf-8\") as fout:\n",
        "        fout.write(yaml.dump(obj, allow_unicode=True, default_flow_style=False))\n",
        "\n",
        "def hash_to_int(obj):\n",
        "    s = str(obj).encode('utf-8')\n",
        "    h = hashlib.md5(s).hexdigest()\n",
        "    return int(h, 16)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6JgB_OUTIk5v",
        "colab_type": "text"
      },
      "source": [
        "## Describe the dataset\n",
        "\n",
        "The data pipeline below requires metadata about the files (where they are, what they represent, what format they're in) rather than the filenames themselves.  This is to help simplify things when our data sources are heterogeneous (e.g. when different sources cover different pairs of languages).\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MX2tAeoA8Cau",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "unsharded_fr_en_dataset = '''\n",
        "metadata: \n",
        "    name: Europarl English-French\n",
        "    url: http://www.statmt.org/europarl/v7/fr-en.tgz\n",
        "    downloaded_on: 2019-06-22\n",
        "format: multiple_file_txt\n",
        "data:\n",
        "    en:\n",
        "        lang: eng\n",
        "        path: europarl-v7.fr-en.en\n",
        "    fr:\n",
        "        lang: fra\n",
        "        path: europarl-v7.fr-en.fr\n",
        "'''\n",
        "\n",
        "unsharded_sw_en_dataset = '''\n",
        "metadata: \n",
        "    name: GlobalVoices Swahili-English\n",
        "    url: http://opus.nlpl.eu/download.php?f=GlobalVoices/v2017q3/moses/en-sw.txt.zip\n",
        "    downloaded_on: 2019-06-22\n",
        "    desc: The GlobalVoices Swahili-English news corpus, as processed by the OPUS project.\n",
        "format: multiple_file_txt\n",
        "data:\n",
        "    en:\n",
        "        lang: eng\n",
        "        path: GlobalVoices.en-sw.en\n",
        "    sw:\n",
        "        lang: swh\n",
        "        path: GlobalVoices.en-sw.sw\n",
        "'''"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BonGOLY9LykI",
        "colab_type": "text"
      },
      "source": [
        "## Split into shards\n",
        "\n",
        "We split the corpus into ```num_shards``` shards, by hashing each (L1, L2, ...) sentence tuple into an integer ```i```, and assigning it to the ```i % num_shards``` shard.  \n",
        "\n",
        "Each of these shards gets a dataset metadata file, too, automatically generated from the parent file.\n",
        "\n",
        "Then, when we design the experiment, we'll describe the test/dev/train sets in terms of these shards."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WXhahNaVK94F",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "def make_shards(dataset, num_shards=10):\n",
        "  \n",
        "    #logging.debug(f\"Splitting {paths.values()} into {num_shards} shards\")\n",
        "    \n",
        "    ids = list(dataset[\"data\"])   # just get the keys\n",
        "    ids.sort()       # Since Python doesn't hash dictionary keys things the same\n",
        "                     # way each time, this would lead to different sentence\n",
        "                     # tuples, and then when we hash *those*, we would get \n",
        "                     # different results and then different shards.  So for\n",
        "                     # reproducability, we sort the keys alphabetically first.\n",
        "                                  \n",
        "    sentences_per_id = {}\n",
        "    for id in dataset[\"data\"]:\n",
        "        path = dataset[\"data\"][id][\"path\"]\n",
        "        with open(path, \"r\", encoding=\"utf-8\") as fin:\n",
        "            sentences_per_id[id] = fin.readlines()\n",
        "    \n",
        "    sentences = [ sentences_per_id[id] for id in ids ]     # keep these in the same order\n",
        "    \n",
        "    # make sure input files have same number of sentences\n",
        "    num_sentences = -1\n",
        "    for ss in sentences:\n",
        "      if num_sentences != -1 and len(ss) != num_sentences:\n",
        "        logging.warning(\"Parallel source files have different number of \"\n",
        "                f\"sentences in dataset {dataset['metadata']['name']}\")\n",
        "    \n",
        "    shards = { i:defaultdict(list) for i in range(num_shards) }\n",
        "    \n",
        "    for sentence_tuple in zip(*sentences):\n",
        "        h = hash_to_int(sentence_tuple) % num_shards\n",
        "        for id, s in zip(ids, sentence_tuple):\n",
        "            shards[h][id].append(s)\n",
        "\n",
        "    for h in range(num_shards):\n",
        "        new_dataset = deepcopy(dataset)\n",
        "        new_dataset[\"metadata\"][\"sharded\"] = True\n",
        "        new_dataset[\"metadata\"][\"shard\"] = h\n",
        "        new_dataset[\"metadata\"][\"num_shards\"] = num_shards\n",
        "        for id in ids:\n",
        "            filename = f\"shard{h}.{id}.txt\"\n",
        "            new_dataset[\"data\"][id][\"path\"] = filename\n",
        "            lines = shards[h][id]\n",
        "            with open(filename, \"w\", encoding=\"utf-8\") as fout:\n",
        "                fout.write(\"\".join(lines))\n",
        "        filename = f\"shard{h}.yaml\"\n",
        "        save_yaml(filename, new_dataset)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y2rtdCu75Ir8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "dataset = yaml.load(unsharded_sw_en_dataset)\n",
        "make_shards(dataset)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T744O5Rn6QzD",
        "colab_type": "code",
        "outputId": "ec2bbe0d-2b45-44a8-f083-78f7ec6d2a05",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 513
        }
      },
      "source": [
        "!ls\n",
        "!head -n 1 shard8.sw.txt\n",
        "!head -n 1 shard8.en.txt\n",
        "!cat shard8.yaml"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "GlobalVoices.en-sw.en\tshard1.sw.txt  shard4.yaml    shard8.en.txt\n",
            "GlobalVoices.en-sw.ids\tshard1.yaml    shard5.en.txt  shard8.sw.txt\n",
            "GlobalVoices.en-sw.sw\tshard2.en.txt  shard5.sw.txt  shard8.yaml\n",
            "LICENSE\t\t\tshard2.sw.txt  shard5.yaml    shard9.en.txt\n",
            "README\t\t\tshard2.yaml    shard6.en.txt  shard9.sw.txt\n",
            "sample_data\t\tshard3.en.txt  shard6.sw.txt  shard9.yaml\n",
            "shard0.en.txt\t\tshard3.sw.txt  shard6.yaml    swahili_data.zip\n",
            "shard0.sw.txt\t\tshard3.yaml    shard7.en.txt\n",
            "shard0.yaml\t\tshard4.en.txt  shard7.sw.txt\n",
            "shard1.en.txt\t\tshard4.sw.txt  shard7.yaml\n",
            "Angola na Brazil zina uhusiano maalum baina yao, kutokana na matumizi ya lugha moja na pamoja na historia yao ya kutawaliwa huko nyuma - nchi hizi mbili zilikuwa sehemu ya himaya ya Wareno – na utamaduni wa pamoja unaotoka na historia ya asili moja.\n",
            "Angola and Brazil have a special relationship towards each other, partially because of their common language and their shared colonial past - both countries were part of the Portuguese Empire - and the cultural ties that stem from this shared history.\n",
            "data:\n",
            "  en:\n",
            "    lang: eng\n",
            "    path: shard8.en.txt\n",
            "  sw:\n",
            "    lang: swh\n",
            "    path: shard8.sw.txt\n",
            "format: multiple_file_txt\n",
            "metadata:\n",
            "  desc: The GlobalVoices Swahili-English news corpus, as processed by the OPUS project.\n",
            "  downloaded_on: 2019-06-22\n",
            "  name: GlobalVoices Swahili-English\n",
            "  num_shards: 10\n",
            "  shard: 8\n",
            "  sharded: true\n",
            "  url: http://opus.nlpl.eu/download.php?f=GlobalVoices/v2017q3/moses/en-sw.txt.zip\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BGx8KHxyLxZk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "\n",
        "class DataIterator:\n",
        "  \n",
        "    def __init__(self, inner_iter=None, make_cache=True, deterministic=True):\n",
        "        self.inner_iter = inner_iter\n",
        "        if not deterministic:\n",
        "            self.deterministic = False\n",
        "        elif inner_iter is None:\n",
        "            self.deterministic = True\n",
        "        else:\n",
        "            self.deterministic = inner_iter.deterministic\n",
        "        self.itr = None\n",
        "        self.make_cache = make_cache\n",
        "        self.cache = []\n",
        "    \n",
        "      \n",
        "    def __iter__(self):\n",
        "        if self.deterministic and self.cache:\n",
        "            return iter(self.cache)\n",
        "        self.itr = self.go()\n",
        "        return self\n",
        "      \n",
        "    def __next__(self):\n",
        "        result = next(self.itr)\n",
        "        if self.deterministic:\n",
        "            self.cache.append(result)\n",
        "        return result\n",
        "        \n",
        "        \n",
        "\n",
        "class SplitIterator(DataIterator):\n",
        "    ''' iterates through a split, according to an\n",
        "        experiment spec '''\n",
        "  \n",
        "    def __init__(self, experiment_spec, split_name):\n",
        "        super(SplitIterator, self).__init__(None, True)\n",
        "        self.experiment_spec = experiment_spec\n",
        "        self.split_name = split_name\n",
        "    \n",
        "    def go(self):\n",
        "        for filename in self.experiment_spec[self.split_name]:\n",
        "            dataset = load_yaml(filename)\n",
        "            ids = list(dataset[\"data\"])\n",
        "\n",
        "\n",
        "            sentences = [] \n",
        "            for id in ids:\n",
        "                file_path = dataset[\"data\"][id][\"path\"]\n",
        "                with open(file_path, \"r\", encoding=\"utf-8\") as fin:\n",
        "                    sentences.append([s.rstrip(\"\\n\") for s in fin.readlines()])\n",
        "\n",
        "            num_sentences = -1\n",
        "            for ss in sentences:\n",
        "                if num_sentences != -1 and len(ss) != num_sentences:\n",
        "                    logging.warning(\"Parallel source files have different \"\n",
        "                                f\"number of sentences: {paths_dict.values()}\")\n",
        "\n",
        "            for ss in zip(*sentences):\n",
        "                yield dict(zip(ids, ss))\n",
        "\n",
        "class TokenizingIterator(DataIterator):\n",
        "  \n",
        "    def __init__(self, inner_iter):\n",
        "        super(TokenizingIterator, self).__init__(inner_iter, True)\n",
        "        self.basic_tokenizer = spacy.load('en').tokenizer\n",
        "\n",
        "    def tokenize(self, sentence):\n",
        "        return [token.text for token in self.basic_tokenizer(sentence)]\n",
        "      \n",
        "    def go(self):\n",
        "        for entry in iter(self.inner_iter):\n",
        "            yield { id : self.tokenize(s) for id, s in entry.items() } \n",
        "    \n",
        "class LowercasingIterator(DataIterator):\n",
        "  \n",
        "    def __init__(self, inner_iter):\n",
        "        super(LowercasingIterator, self).__init__(inner_iter, True)\n",
        "        \n",
        "    def go(self):\n",
        "        for entry in iter(self.inner_iter):\n",
        "            yield { id : s.lower() for id, s in entry.items() } \n",
        "    \n",
        "class StartEndIterator(DataIterator):\n",
        "  \n",
        "    def __init__(self, inner_iter, vocab):\n",
        "        super(StartEndIterator, self).__init__(inner_iter, True)\n",
        "        self.sos = vocab.SOS_TOKEN\n",
        "        self.eos = vocab.EOS_TOKEN\n",
        "  \n",
        "    def go(self):\n",
        "        for entry in iter(self.inner_iter):\n",
        "            yield { id : [self.sos] + s + [self.eos]\n",
        "                                    for id, s in entry.items() }\n",
        "  \n",
        "  \n",
        "class IndexingIterator(DataIterator):\n",
        "  \n",
        "    def __init__(self, inner_iter, vocab, max_vocab=60000):\n",
        "        super(IndexingIterator, self).__init__(inner_iter, True)\n",
        "        self.vocab = vocab\n",
        "        self.max_vocab = max_vocab\n",
        "  \n",
        "      \n",
        "    def go(self):\n",
        "        for entry in iter(self.inner_iter):\n",
        "            result = defaultdict(list)\n",
        "            for id, tokens in entry.items():\n",
        "                for token in tokens:\n",
        "                    index = self.vocab.stoi(token)\n",
        "                    result[id].append(index)\n",
        "            yield result\n",
        "\n",
        "class NonBucketingIterator(DataIterator):\n",
        "  \n",
        "    def __init__(self, inner_iter, tokens_per_bucket=1000, bucket_size_iter=5):\n",
        "        super(NonBucketingIterator, self).__init__(inner_iter, True)\n",
        "        \n",
        "    def go(self):\n",
        "        for entry in iter(self.inner_iter):\n",
        "            result = {}\n",
        "            for id, indices in entry.items():\n",
        "                result[id] = ntorch.tensor([indices], names=('batch','tokens'))\n",
        "            yield result\n",
        "\n",
        "class BucketingIterator(DataIterator):\n",
        "  \n",
        "    def __init__(self, inner_iter, tokens_per_bucket=3000, bucket_size_iter=5):\n",
        "        super(BucketingIterator, self).__init__(inner_iter, True)\n",
        "        self.tokens_per_bucket = tokens_per_bucket\n",
        "        self.bucket_size_iter = 5\n",
        "        \n",
        "\n",
        "    def go(self):\n",
        "        #if shuffle:\n",
        "        #  inner_iter = list(inner_iter)  # collect all sentences\n",
        "        #  random.shuffle(inner_iter)\n",
        "\n",
        "        buckets = {}  # buckets contain torch.Tensors of various sizes.  When one is ready we\n",
        "                      # send it down the line and empty it out.\n",
        "        vacancies = {}  # how many remaining sentences can fit in each bucket\n",
        "        for entry in iter(self.inner_iter):\n",
        "            # find whichever sentence is longest; that determines our bucket size\n",
        "            length = self.bucket_size_iter  # no point starting smaller than the smallest bucket\n",
        "            for id, indices in entry.items():\n",
        "                length = max(length, len(indices))\n",
        "            length = math.ceil(length / self.bucket_size_iter) \\\n",
        "                                            * self.bucket_size_iter\n",
        "\n",
        "            # make the bucket if it doesn't exist\n",
        "            if length not in buckets:\n",
        "                buckets[length] = defaultdict(list)\n",
        "                vacancies[length] = math.floor(self.tokens_per_bucket / length)\n",
        "\n",
        "            for id, indices in entry.items():\n",
        "                len_indices = len(indices)\n",
        "                padded_indices = indices + [0] * (length - len(indices))\n",
        "                assert(len(padded_indices) == length)\n",
        "                buckets[length][id].append(padded_indices)\n",
        "\n",
        "            vacancies[length] -= 1\n",
        "\n",
        "            if vacancies[length] <= 0:\n",
        "                result = {}\n",
        "                for id, sentences in buckets[length].items():\n",
        "                    result[id] = ntorch.tensor(sentences, names=('batch','tokens'))\n",
        "                del buckets[length]\n",
        "                yield result\n",
        "\n",
        "        # once we've run out of sentences, yield any remaining buckets\n",
        "        for length, bucket in list(buckets.items()):\n",
        "            result = {}\n",
        "            for id, sentences in bucket.items():\n",
        "              result[id] = ntorch.tensor(sentences, names=('batch', 'tokens'))\n",
        "            del buckets[length]\n",
        "            yield result\n",
        "\n",
        "            \n",
        "class Vocab:\n",
        "   \n",
        "    def __init__(self, max_vocab=60000):\n",
        "      \n",
        "        self.PAD_TOKEN = \"<pad>\"\n",
        "        self.PAD_INDEX = 0\n",
        "        self.UNK_TOKEN = \"<unk>\"\n",
        "        self.UNK_INDEX = 1\n",
        "        self.SOS_TOKEN = \"<sos>\"\n",
        "        self.SOS_INDEX = 2\n",
        "        self.EOS_TOKEN = \"<eos>\"\n",
        "        self.EOS_INDEX = 3\n",
        "        self.vocab = bidict({ self.PAD_TOKEN: self.PAD_INDEX, \n",
        "                              self.UNK_TOKEN: self.UNK_INDEX, \n",
        "                              self.SOS_TOKEN: self.SOS_INDEX, \n",
        "                              self.EOS_TOKEN: self.EOS_INDEX })\n",
        "        self.max_vocab = max_vocab\n",
        "        self.is_frozen = False\n",
        "        \n",
        "    def __len__(self):\n",
        "        return len(self.vocab)\n",
        "      \n",
        "    def freeze(self):\n",
        "        self.is_frozen = True\n",
        "        \n",
        "    def stoi(self, s):\n",
        "      \n",
        "        if s not in self.vocab:\n",
        "            if len(self.vocab) < self.max_vocab and not self.is_frozen:\n",
        "                index = len(self.vocab)\n",
        "                self.vocab[s] = index\n",
        "                return index\n",
        "            return self.vocab[self.UNK_TOKEN]\n",
        "          \n",
        "        return self.vocab[s]\n",
        "      \n",
        "    def itos(self, i):\n",
        "        if i not in self.vocab.inverse:\n",
        "            return self.UNK_TOKEN\n",
        "        return self.vocab.inverse[i]\n",
        "    \n",
        "    \n",
        "def setup_data_iterators(experiment_spec, vocab, max_vocab=60000):\n",
        "  \n",
        "    \n",
        "    splits = { \n",
        "        \"train\": SplitIterator(experiment_spec, \"train\"),\n",
        "        \"dev\": SplitIterator(experiment_spec, \"dev\"),\n",
        "        \"test\": SplitIterator(experiment_spec, \"test\")\n",
        "    }\n",
        "    \n",
        "    for split_name in splits:\n",
        "        itr = splits[split_name]\n",
        "        \n",
        "        itr = LowercasingIterator(itr)\n",
        "        itr = TokenizingIterator(itr)\n",
        "        itr = StartEndIterator(itr, vocab)\n",
        "        itr = IndexingIterator(itr, vocab, max_vocab)\n",
        "        \n",
        "        if split_name == \"train\":      \n",
        "            itr = BucketingIterator(itr)\n",
        "        else:\n",
        "            itr = NonBucketingIterator(itr)\n",
        "            \n",
        "        splits[split_name] = itr\n",
        "        \n",
        "    for i, e in enumerate(iter(splits[\"train\"])):\n",
        "        pass # run it once through train to learn the vocabulary\n",
        "    \n",
        "    vocab.freeze()\n",
        "    \n",
        "    return splits[\"train\"], splits[\"dev\"], splits[\"test\"]\n",
        "   \n",
        "        \n",
        "        \n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-OSqKctoulIA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Embedder(ntorch.nn.Module):\n",
        "  \n",
        "    def __init__(self, vocab_size, model_size=512):\n",
        "        super().__init__()\n",
        "        self.embed = ntorch.nn.Embedding(vocab_size, model_size)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        return self.embed(x)\n",
        "      "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UWumB3-4uzKG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "class PositionEmbedder(ntorch.nn.Module):\n",
        "    \n",
        "    def __init__(self, model_size=512, max_seq_len = 200):\n",
        "        super().__init__()\n",
        "        self.model_size = model_size\n",
        "        \n",
        "        model_size_half = math.floor(float(model_size) / 2)\n",
        "        model_size_half2 = math.ceil(float(model_size) / 2)\n",
        "        \n",
        "        # gotta do this all in vanilla pytorch because\n",
        "        # one of the versions of torch.pow isn't yet\n",
        "        # implemented in namedtensor\n",
        "        \n",
        "        i = torch.arange(model_size_half, requires_grad=False)\n",
        "        i = i.unsqueeze(0).expand(max_seq_len, model_size_half).float()\n",
        "        \n",
        "        i2 = torch.arange(model_size_half2, requires_grad=False)\n",
        "        i2 = i2.unsqueeze(0).expand(max_seq_len, model_size_half2).float()\n",
        "        \n",
        "        pos = torch.arange(max_seq_len, requires_grad=False)\n",
        "        pos = pos.unsqueeze(1).float()    \n",
        "        pe1 = torch.sin(pos / torch.pow(10000.0, i / model_size))\n",
        "        pe2 = torch.cos(pos / torch.pow(10000.0, i2 / model_size))\n",
        "        pe = torch.cat((pe1, pe2), dim=1)\n",
        "        self.pe = ntorch.tensor(pe, names=(\"tokens\", \"embedding\")).to(device)\n",
        "               \n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x * math.sqrt(self.model_size)\n",
        "        narrowed_pe = self.pe.narrow(\"tokens\", 0, x.size(\"tokens\"))\n",
        "        x = x + narrowed_pe\n",
        "        return x\n",
        "      \n",
        "\n",
        "#plt.imshow(pe.pe.cpu().numpy())\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bdZ9-BukNe3G",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "\n",
        "def make_source_mask(source_batch, vocab):\n",
        "  \n",
        "    source_mask = (source_batch != vocab.PAD_INDEX).float().to(device)\n",
        "    source_mask = source_mask.rename(\"tokens\", \"tokens2\")\n",
        "    return source_mask\n",
        "  \n",
        "def make_target_mask(target_batch, vocab):\n",
        "\n",
        "    target_mask = (target_batch != vocab.PAD_INDEX).float()\n",
        "    size = target_batch.size(\"tokens\")\n",
        "    nopeak_mask = np.triu(np.ones((1, size, size)), k=1)\n",
        "    nopeak_mask = ntorch.tensor(nopeak_mask == 0, \n",
        "                       names=(\"batch\", \"tokens\", \"tokens2\")).float().to(device)\n",
        "    \n",
        "    target_mask = target_mask * nopeak_mask\n",
        "    return target_mask\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "48Y-YMQacNv6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "\n",
        "\n",
        "class MultiHeadAttention(ntorch.nn.Module):\n",
        "  \n",
        "    def __init__(self, num_heads=8, model_size=512, dropout_rate = 0.1):\n",
        "        super().__init__()\n",
        "        \n",
        "        self.model_size = model_size\n",
        "        assert(model_size % num_heads == 0)\n",
        "        self.d_key = model_size // num_heads\n",
        "        self.num_heads = num_heads\n",
        "        \n",
        "        self.query_linear = ntorch.nn.Linear(model_size, model_size).spec(\"embedding\")\n",
        "        self.value_linear = ntorch.nn.Linear(model_size, model_size).spec(\"embedding\")\n",
        "        self.key_linear = ntorch.nn.Linear(model_size, model_size).spec(\"embedding\")\n",
        "        self.dropout = ntorch.nn.Dropout(dropout_rate)\n",
        "        self.out_linear = ntorch.nn.Linear(model_size, model_size).spec(\"embedding\")\n",
        "        \n",
        "      \n",
        "    def forward(self, query, key, value, mask=None):\n",
        "        \n",
        "        # make keys/queries/values (for all \"heads\" at once)\n",
        "        key = self.key_linear(key)\n",
        "        query = self.query_linear(query)\n",
        "        value = self.value_linear(value)\n",
        "        \n",
        "        # then actually split them into separate heads\n",
        "        key = key.split(\"embedding\", (\"heads\", \"embedding\"), heads=self.num_heads).rename(\"tokens\", \"tokens2\")\n",
        "        query = query.split(\"embedding\", (\"heads\", \"embedding\"), heads=self.num_heads)\n",
        "        value = value.split(\"embedding\", (\"heads\", \"embedding\"), heads=self.num_heads).rename(\"tokens\", \"tokens2\")\n",
        "        \n",
        "        # attention\n",
        "        scores = ntorch.dot(\"embedding\", query, key) /  math.sqrt(self.d_key)   # compare the embeddings for each pair of tokens\n",
        "        scores = scores * mask                                                  # mask out scores you shouldn't know\n",
        "        scores = scores.softmax(\"tokens2\")                               \n",
        "        scores = self.dropout(scores)                                  \n",
        "        output = ntorch.dot(\"tokens2\", scores, value)                            # calculate the new token embeddings\n",
        "      \n",
        "        # put the heads back together\n",
        "        output = output.stack((\"heads\", \"embedding\"), \"embedding\")\n",
        "        output = self.out_linear(output)\n",
        "        return output\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eKaxN3y2GKg2",
        "colab_type": "code",
        "outputId": "6811bd01-e3d3-417b-8f3e-461747af5f10",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "e = torch.randn(10, 80, 8*64)\n",
        "\n",
        "query = torch.randn(10, 8, 80, 64)\n",
        "key = torch.randn(10, 8, 80, 64)\n",
        "value = torch.randn(10, 8, 80, 64)\n",
        "\n",
        "\n",
        "\n",
        "def attention(q, k, v, d_k, mask=None, dropout=None):\n",
        "    \n",
        "  scores = torch.matmul(q, k.transpose(-2, -1)) /  math.sqrt(d_k)\n",
        "\n",
        "  if mask is not None:\n",
        "    mask = mask.unsqueeze(1)\n",
        "    scores = scores.masked_fill(mask == 0, -1e9)\n",
        "  scores = torch.nn.functional.softmax(scores, dim=-1)\n",
        "  print(scores.shape)\n",
        "  \n",
        "  if dropout is not None:\n",
        "    scores = dropout(scores)\n",
        "        \n",
        "  output = torch.matmul(scores, v)\n",
        "  return output\n",
        "\n",
        "x = attention(query, key, value, 64)\n",
        "print(x.shape)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "torch.Size([10, 8, 80, 80])\n",
            "torch.Size([10, 8, 80, 64])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pAdT823EqWyG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uIPLbg_vDUw5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class LayerNorm(ntorch.nn.Module):\n",
        "    def __init__(self, model_size=512, eps = 1e-6):\n",
        "        super().__init__()\n",
        "    \n",
        "        self.model_size = model_size\n",
        "        # create two learnable parameters to calibrate normalisation\n",
        "        alpha = ntorch.ones(model_size, names=(\"embedding\"))\n",
        "        bias = ntorch.zeros(model_size, names=(\"embedding\"))\n",
        "        \n",
        "        self.register_parameter(\"alpha\", alpha)\n",
        "        self.register_parameter(\"bias\", bias)\n",
        "        \n",
        "        self.eps = eps\n",
        "        \n",
        "        #if torch.cuda.is_available():\n",
        "        #    self.alpha = self.alpha.cuda()\n",
        "        #    self.bias = self.bias.cuda()\n",
        "        \n",
        "        #self.register_parameter(\"alpha\", self.alpha)\n",
        "        #self.register_parameter(\"bias\", self.bias)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        norm = self.alpha * (x - x.mean(\"embedding\")) \n",
        "        norm = norm / (x.std(\"embedding\") + self.eps) \n",
        "        norm = norm + self.bias\n",
        "        return norm\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mW9xAA2xM86M",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class SelfAttentionLayer(ntorch.nn.Module):\n",
        "  \n",
        "    def __init__(self, model_size=512, num_heads=8, dropout_rate=0.1):\n",
        "        super().__init__()\n",
        "        self.norm = LayerNorm(model_size)\n",
        "        self.attn = MultiHeadAttention(num_heads, model_size)\n",
        "        self.dropout = ntorch.nn.Dropout(dropout_rate)\n",
        "        \n",
        "    def forward(self, x, mask):\n",
        "        x_norm = self.norm(x)\n",
        "        x_attn = self.attn(x_norm, x_norm, x_norm, mask)\n",
        "        x = x + self.dropout(x_attn)\n",
        "        return x\n",
        "      \n",
        "class OtherAttentionLayer(ntorch.nn.Module):\n",
        "  \n",
        "    def __init__(self, model_size=512, num_heads=8, dropout_rate=0.1):\n",
        "        super().__init__()\n",
        "        self.norm = LayerNorm(model_size)\n",
        "        self.attn = MultiHeadAttention(num_heads, model_size)\n",
        "        self.dropout = ntorch.nn.Dropout(dropout_rate)\n",
        "        \n",
        "    def forward(self, x, other_x, mask):\n",
        "        x_norm = self.norm(x)\n",
        "        x_attn = self.attn(x_norm, other_x, other_x, mask)\n",
        "        x = x + self.dropout(x_attn)\n",
        "        return x\n",
        "      \n",
        "class FeedForward(ntorch.nn.Module):\n",
        "  \n",
        "    def __init__(self, model_size=512, hidden_size=512, dropout_rate=0.1):\n",
        "        super().__init__() \n",
        "        self.norm = LayerNorm(model_size)\n",
        "        self.linear1 = ntorch.nn.Linear(model_size, hidden_size)\n",
        "        self.dropout1 = ntorch.nn.Dropout(dropout_rate)\n",
        "        self.linear2 = ntorch.nn.Linear(hidden_size, model_size)\n",
        "        self.dropout2 = ntorch.nn.Dropout(dropout_rate)\n",
        "        \n",
        "    def forward(self, x):\n",
        "      \n",
        "        x_norm = self.norm(x)\n",
        "        x_norm = self.linear1(x_norm)\n",
        "        x_norm = ntorch.relu(x_norm)\n",
        "        x_norm = self.dropout1(x_norm)\n",
        "        x_norm = self.linear2(x_norm)\n",
        "        x = x + self.dropout2(x_norm)\n",
        "        return x      \n",
        "      \n",
        "class DecoderLayer(ntorch.nn.Module):\n",
        "  \n",
        "    def __init__(self, model_size=512, num_heads=8, dropout_rate=0.1):\n",
        "        super().__init__()\n",
        "        \n",
        "        self.self_attn = SelfAttentionLayer(model_size, num_heads, dropout_rate)\n",
        "        self.other_attn = OtherAttentionLayer(model_size, num_heads, dropout_rate)\n",
        "        #self.norm1 = LayerNorm(model_size)\n",
        "        #self.norm2 = LayerNorm(model_size)\n",
        "        #self.norm3 = LayerNorm(model_size)\n",
        "        \n",
        "        #self.dropout1 = ntorch.nn.Dropout(dropout_rate)\n",
        "        #self.dropout2 = ntorch.nn.Dropout(dropout_rate)\n",
        "        #self.dropout3 = ntorch.nn.Dropout(dropout_rate)\n",
        "        \n",
        "        #self.attn1 = MultiHeadAttention(num_heads, model_size)\n",
        "        #self.attn2 = MultiHeadAttention(num_heads, model_size)\n",
        "        \n",
        "        self.ff = FeedForward(model_size, model_size, dropout_rate)\n",
        "        \n",
        "    def forward(self, x, e_outputs, src_mask, trg_mask):\n",
        "        \n",
        "        x = self.self_attn(x, trg_mask)\n",
        "        x = self.other_attn(x, e_outputs, src_mask)\n",
        "        x = self.ff(x)\n",
        "        \n",
        "        #x_norm = self.norm1(x)\n",
        "        #x_attn = self.attn1(x_norm, x_norm, x_norm, trg_mask)\n",
        "        #x = x + self.dropout1(x_attn)\n",
        "        \n",
        "        #x_norm = self.norm2(x)\n",
        "        #x_attn = self.attn2(x_norm, e_outputs, e_outputs, src_mask)\n",
        "        #x = x + self.dropout2(x_attn)\n",
        "        \n",
        "        #x_norm = self.norm3(x)\n",
        "        #x_ff = self.ff(x_norm)\n",
        "        #x = x + self.dropout3(x_ff)\n",
        "        \n",
        "        return x\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QguYJFVlL4rR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class EncoderLayer(ntorch.nn.Module):\n",
        "    \n",
        "    def __init__(self, model_size=512, num_heads=8, dropout_rate=0.1):\n",
        "      \n",
        "        super().__init__()\n",
        "        \n",
        "        self.self_attn = SelfAttentionLayer(model_size, num_heads, dropout_rate)\n",
        "        \n",
        "        #self.norm_1 = LayerNorm(model_size)\n",
        "        #self.norm_2 = LayerNorm(model_size)\n",
        "        #self.attn = MultiHeadAttention(num_heads, model_size)\n",
        "        self.ff = FeedForward(model_size, model_size, dropout_rate)\n",
        "        #self.dropout1 = ntorch.nn.Dropout(dropout_rate)\n",
        "        #self.dropout2 = ntorch.nn.Dropout(dropout_rate)\n",
        "        \n",
        "    def forward(self, x, mask):\n",
        "        \n",
        "        # Block 1: attention\n",
        "        #x_norm = self.norm_1(x)\n",
        "        #x_attn = self.attn(x_norm,x_norm,x_norm,mask)\n",
        "        #x = x + self.dropout1(x_attn)\n",
        "        \n",
        "        x = self.self_attn(x, mask)\n",
        "        x = self.ff(x)\n",
        "        \n",
        "        # Block 2: feed-forward\n",
        "        #x_norm = self.norm_2(x)\n",
        "        #x_ff = self.ff(x_norm)\n",
        "        #x = x + self.dropout2(x_ff)\n",
        "        return x\n",
        "      \n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A3_V_ptgndkE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "class Encoder(ntorch.nn.Module):\n",
        "  \n",
        "    def __init__(self, vocab_size, model_size=512, num_layers=6, num_heads=8):\n",
        "      \n",
        "        super().__init__()\n",
        "        self.num_layers = num_layers\n",
        "        self.embed = Embedder(vocab_size, model_size)\n",
        "        self.pe = PositionEmbedder(model_size)\n",
        "        self.layers = ntorch.nn.ModuleList([\n",
        "              EncoderLayer(model_size, num_heads) \n",
        "                  for l in range(num_layers)])\n",
        "        self.norm = LayerNorm(model_size)\n",
        "        \n",
        "    def forward(self, src, src_mask):\n",
        "        x = self.embed(src)\n",
        "        x = self.pe(x)\n",
        "        for i in range(self.num_layers):\n",
        "            x = self.layers[i](x, src_mask)\n",
        "        return self.norm(x)\n",
        "      \n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ImQ0PLE9tMMH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Decoder(ntorch.nn.Module):\n",
        "  \n",
        "    def __init__(self, vocab_size, model_size=512, num_layers=6, num_heads=8):\n",
        "      \n",
        "        super().__init__()\n",
        "        self.num_layers = num_layers\n",
        "        self.embed = Embedder(vocab_size, model_size)\n",
        "        self.pe = PositionEmbedder(model_size)\n",
        "        self.layers = ntorch.nn.ModuleList([\n",
        "              DecoderLayer(model_size, num_heads) \n",
        "                  for l in range(num_layers)])\n",
        "        self.norm = LayerNorm(model_size)\n",
        "        self.out_linear = ntorch.nn.Linear(model_size, vocab_size) \\\n",
        "                                   .spec(\"embedding\", name_out=\"logits\")\n",
        "        \n",
        "    def forward(self, trg, e_outputs, src_mask, trg_mask):\n",
        "        x = self.embed(trg)\n",
        "        x = self.pe(x)\n",
        "        for i in range(self.num_layers):\n",
        "            x = self.layers[i](x, e_outputs, src_mask, trg_mask)\n",
        "        x = self.norm(x)\n",
        "        x = self.out_linear(x)\n",
        "        return x\n",
        "      \n",
        "      "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fgVowzITuTcm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Transformer(ntorch.nn.Module):\n",
        "  \n",
        "    def __init__(self, vocab, model_size=512, num_layers=6, num_heads=8):\n",
        "        super().__init__()\n",
        "        self.vocab = vocab\n",
        "        vocab_size = len(vocab)\n",
        "        self.encoder = Encoder(vocab_size, model_size, num_layers, num_heads)\n",
        "        self.decoder = Decoder(vocab_size, model_size, num_layers, num_heads)\n",
        "        \n",
        "    def forward(self, src, trg):\n",
        "        src_mask = make_source_mask(src, self.vocab)\n",
        "        trg_mask = make_target_mask(trg, self.vocab)\n",
        "        src_outputs = self.encoder(src, src_mask)\n",
        "        #src_mask = src_mask.rename(\"tokens\", \"tokens2\")\n",
        "        trg_output = self.decoder(trg, src_outputs, src_mask, trg_mask)\n",
        "        return trg_output\n",
        "      "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZVbZftVw2SNq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "NUM_LAYERS_DEFAULT = 6\n",
        "NUM_HEADS_DEFAULT = 8\n",
        "MODEL_SIZE_DEFAULT = 512\n",
        "LEARNING_RATE_DEFAULT = 0.0001\n",
        "ADAM_BETA1_DEFAULT = 0.9\n",
        "ADAM_BETA2_DEFAULT = 0.98\n",
        "\n",
        "class Experiment:\n",
        "  \n",
        "    def __init__(self, exp_spec):\n",
        "        self.exp_spec = exp_spec\n",
        "        self.setup_data()\n",
        "        self.setup_model()\n",
        "            \n",
        "    def setup_data(self):\n",
        "        self.vocab = Vocab(100000)\n",
        "        logging.info(\"Loading data and making vocab\")\n",
        "        self.train_iter, self.dev_iter, self.test_iter = \\\n",
        "            setup_data_iterators(self.exp_spec, self.vocab)\n",
        "\n",
        "    def setup_model(self):\n",
        "        \n",
        "        logging.info(\"Setting up model\")\n",
        "        \n",
        "        if \"hyperparams\" not in self.exp_spec:\n",
        "            logging.error(\"Experiment specification contains no hyperparameters.\")\n",
        "            sys.exit()\n",
        "\n",
        "        hyperparams = self.exp_spec[\"hyperparams\"]        \n",
        "        self.num_layers = hyperparams.get(\"num_layers\", NUM_LAYERS_DEFAULT)\n",
        "        self.num_heads = hyperparams.get(\"num_heads\", NUM_HEADS_DEFAULT)\n",
        "        self.model_size = hyperparams.get(\"model_size\", MODEL_SIZE_DEFAULT)\n",
        "        self.learning_rate = hyperparams.get(\"learning_rate\", LEARNING_RATE_DEFAULT)\n",
        "        self.adam_beta1 = hyperparams.get(\"adam_beta1\", ADAM_BETA1_DEFAULT)\n",
        "        self.adam_beta2 = hyperparams.get(\"adam_beta2\", ADAM_BETA2_DEFAULT)\n",
        "            \n",
        "                \n",
        "        self.model = Transformer(self.vocab, \n",
        "                                 num_layers=self.num_layers,\n",
        "                                 num_heads=self.num_heads,\n",
        "                                 model_size=self.model_size).to(device)\n",
        "\n",
        "        for p in self.model.parameters():\n",
        "            if p.dim() > 1:\n",
        "                torch.nn.init.xavier_uniform_(p)        \n",
        "      \n",
        "        self.optim = torch.optim.Adam(self.model.parameters(), \n",
        "                                      lr=self.learning_rate, \n",
        "                                      betas=(self.adam_beta1, self.adam_beta2), \n",
        "                                      eps=1e-9)\n",
        "        self.loss_function = ntorch.nn.CrossEntropyLoss(\n",
        "                    ignore_index=self.vocab.PAD_INDEX)\\\n",
        "                      .spec(\"logits\").to(device)\n",
        "\n",
        "\n",
        "\n",
        "    def train(self, epochs):\n",
        "\n",
        "        self.model.train()\n",
        "\n",
        "        start = time.time()\n",
        "        temp = start\n",
        "\n",
        "        total_loss = 0\n",
        "\n",
        "        for i, epoch in enumerate(range(epochs)):\n",
        "\n",
        "            for j, batch in enumerate(iter(self.train_iter)):\n",
        "                src = batch[\"sw\"].to(device)\n",
        "                trg = batch[\"en\"].to(device)\n",
        "\n",
        "\n",
        "                trg_num_tokens = trg.size(\"tokens\")\n",
        "                trg_input = trg[{\"tokens\":slice(0, trg_num_tokens-1)}]\n",
        "                trg_targets = trg[{\"tokens\":slice(1, trg_num_tokens)}]\n",
        "\n",
        "                preds = self.model(src, trg_input)\n",
        "\n",
        "                self.optim.zero_grad()\n",
        "\n",
        "                loss = self.loss_function(preds, trg_targets)\n",
        "                loss.backward()\n",
        "                self.optim.step()\n",
        "                \n",
        "                total_loss += float(loss.detach().cpu().numpy())\n",
        "\n",
        "                if j % 100 == 0:\n",
        "                    print(\".\", end=\"\")\n",
        "\n",
        "            print()\n",
        "            print(\"time = %dm, epoch %d, loss = %.4f, %ds per epoch\" % ((time.time() - start) // 60,\n",
        "                        epoch + 1, total_loss, time.time() - temp))\n",
        "            total_loss = 0\n",
        "            temp = time.time()\n",
        "                \n",
        "    def translate(self, itr, max_len=120):\n",
        "\n",
        "        self.model.eval()\n",
        "        \n",
        "        result_sentences = []\n",
        "        \n",
        "        for batch_idx, batch in enumerate(iter(itr)):\n",
        "          \n",
        "            if batch_idx >= 10:\n",
        "                break\n",
        "                \n",
        "            src_batch = batch[\"sw\"].to(device)\n",
        "            \n",
        "            for src_sent_idx in range(src_batch.size(\"batch\")):\n",
        "                src = src_batch[{\"batch\":src_sent_idx}]\n",
        "                src_mask = make_source_mask(src, self.vocab)\n",
        "                src_outputs = self.model.encoder(src, src_mask)\n",
        "\n",
        "                trg_outputs = ntorch.zeros(1, max_len, \n",
        "                                           names=(\"batch\", \"tokens\"), \n",
        "                                           dtype=torch.long, device=device)\n",
        "                trg_outputs[{\"tokens\":0}] = self.vocab.SOS_INDEX\n",
        "\n",
        "                sentence = []\n",
        "                \n",
        "                for i in range(1, max_len):    \n",
        "\n",
        "                    trg_mask = np.triu(np.ones((1, i, i)), k=1).astype('uint8')\n",
        "                    trg_mask = ntorch.tensor(trg_mask == 0, \n",
        "                           names=(\"batch\", \"tokens\", \"tokens2\")).float().to(device)\n",
        "\n",
        "                    outputs_so_far = trg_outputs[{\"tokens\":slice(0,i)}]\n",
        "                    dec_outputs = self.model.decoder(outputs_so_far, \n",
        "                                                     src_outputs,\n",
        "                                                     src_mask,\n",
        "                                                     trg_mask)\n",
        "\n",
        "                    logits = dec_outputs.softmax(\"logits\")\n",
        "\n",
        "                    val, ix = logits[{\"batch\":0,\"tokens\":i-1}].topk(\"logits\", 1)\n",
        "                    output_token = int(ix[{\"logits\":0}].cpu().numpy())\n",
        "                    #print(output_token)\n",
        "                    sentence.append(self.vocab.itos(output_token))\n",
        "                    trg_outputs[{\"tokens\":i}] = output_token\n",
        "                    \n",
        "                    if output_token in [self.vocab.EOS_INDEX, self.vocab.PAD_INDEX]:\n",
        "                        break\n",
        "                                \n",
        "                result_sentences.append(\" \".join(sentence))\n",
        "                \n",
        "        return result_sentences\n",
        "      \n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jrkAiyMxsNUq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "experiment_sw_en = '''\n",
        "metadata:\n",
        "    id: sw_en_dev8_test9_3_8_256\n",
        "    name: GlobalVoices Swahili-English\n",
        "    created_on: 2019-06-23\n",
        "hyperparams:\n",
        "    num_layers: 2\n",
        "    num_heads: 8\n",
        "    model_size: 128\n",
        "    learning_rate: 0.0001\n",
        "    beta1: 0.9\n",
        "    beta2: 0.98\n",
        "train:\n",
        "    - shard0.yaml\n",
        "    - shard1.yaml\n",
        "    - shard2.yaml\n",
        "    - shard3.yaml\n",
        "    - shard4.yaml\n",
        "    - shard5.yaml\n",
        "    - shard6.yaml\n",
        "    - shard7.yaml\n",
        "dev:\n",
        "    - shard8.yaml\n",
        "test:\n",
        "    - shard9.yaml\n",
        "'''\n",
        "\n",
        "experiment_fr_en = '''\n",
        "metadata:\n",
        "    id: fr_en_dev8_test9_4_8_256\n",
        "    name: Europarl French-English\n",
        "    created_on: 2019-06-23\n",
        "hyperparams:\n",
        "    num_layers: 2\n",
        "    num_heads: 8\n",
        "    model_size: 128\n",
        "    learning_rate: 0.0001\n",
        "    beta1: 0.9\n",
        "    beta2: 0.98\n",
        "train:\n",
        "    - shard0.yaml\n",
        "    - shard1.yaml\n",
        "    - shard2.yaml\n",
        "    - shard3.yaml\n",
        "    - shard4.yaml\n",
        "    - shard5.yaml\n",
        "    - shard6.yaml\n",
        "    - shard7.yaml\n",
        "dev:\n",
        "    - shard8.yaml\n",
        "test:\n",
        "    - shard9.yaml\n",
        "'''\n",
        "\n",
        "exp_spec = yaml.load(experiment_sw_en)\n",
        "exp = Experiment(exp_spec)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mM_yFC0DsT-n",
        "colab_type": "code",
        "outputId": "5a84ef11-8f8a-475a-86ff-06c7865f12c8",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 544
        }
      },
      "source": [
        "exp.train(60)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "...\n",
            "time = 0m, epoch 1, loss = 2177.2138, 58s per epoch\n",
            "...\n",
            "time = 1m, epoch 2, loss = 1708.2491, 57s per epoch\n",
            "...\n",
            "time = 2m, epoch 3, loss = 1606.9725, 57s per epoch\n",
            "...\n",
            "time = 3m, epoch 4, loss = 1554.2680, 57s per epoch\n",
            "...\n",
            "time = 4m, epoch 5, loss = 1501.5090, 57s per epoch\n",
            "...\n",
            "time = 5m, epoch 6, loss = 1457.1978, 57s per epoch\n",
            "...\n",
            "time = 6m, epoch 7, loss = 1417.4508, 57s per epoch\n",
            "...\n",
            "time = 7m, epoch 8, loss = 1382.2333, 57s per epoch\n",
            "...\n",
            "time = 8m, epoch 9, loss = 1349.7589, 57s per epoch\n",
            "...\n",
            "time = 9m, epoch 10, loss = 1320.1320, 57s per epoch\n",
            "...\n",
            "time = 10m, epoch 11, loss = 1292.7378, 57s per epoch\n",
            "...\n",
            "time = 11m, epoch 12, loss = 1267.1807, 57s per epoch\n",
            "...\n",
            "time = 12m, epoch 13, loss = 1244.0011, 57s per epoch\n",
            "...\n",
            "time = 13m, epoch 14, loss = 1221.3804, 57s per epoch\n",
            "...\n",
            "time = 14m, epoch 15, loss = 1198.1323, 57s per epoch\n",
            "."
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nteO5jMqq8gA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "result_sentences = exp.translate(exp.dev_iter)\n",
        "print(result_sentences[:10])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KrwChJNQPpNw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(\"Model's state_dict:\")\n",
        "num_params = 0\n",
        "for p, t in exp.model.state_dict().items():\n",
        "    print(p, \"\\t\", t.shape)\n",
        "    num_params += np.prod(t.shape)\n",
        "print(num_params)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}