remote_bert.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "remote_bert.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "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/ratmcu/wikipiifed/blob/master/remote_bert.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pc6U5qZZIWMO"
      },
      "source": [
        "##Loading the final dataset\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y0ylg_3uhplH"
      },
      "source": [
        "import tarfile\n",
        "!pip install wget\n",
        "import os\n",
        "import wget"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JKjWXKs5v73O"
      },
      "source": [
        "#experiment_code\n",
        "wget.download('https://archive.org/download/wikiner_dataset_csv.tar/wikiner_dataset_txt.tar.gz')\n",
        "tar = tarfile.open('wikiner_dataset_txt.tar.gz', mode='r')\n",
        "tar.extractall('./dataset_txt')\n",
        "tar.close()"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kcOVFvRV2pUu"
      },
      "source": [
        "## Installing PySyft"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cU7pqJpXkjtj"
      },
      "source": [
        "import sys\n",
        "# !pip install syft\n",
        "!git clone \"https://github.com/ratmcu/PySyft.git\"\n",
        "!cd PySyft && pip install -e .\n",
        "!cd PySyft &&  python setup.py install\n",
        "sys.path.append(\"PySyft\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_T9VNpab-_oc"
      },
      "source": [
        "## run everything below after restarting the runtime"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kQXFlb55vbi4"
      },
      "source": [
        "#run the above cell and restart the runtime then run below\n",
        "\n",
        "import syft as sy\n",
        "import torch\n",
        "import torch.optim as optim\n",
        "import os\n",
        "import numpy as np\n",
        "import argparse\n",
        "from torchvision import datasets, transforms\n",
        "import pandas as pd\n",
        "from torch.utils import data\n",
        "\n",
        "!pip install pytorch-pretrained-bert\n",
        "from pytorch_pretrained_bert import BertTokenizer\n",
        "tokenizer = BertTokenizer.from_pretrained('bert-base-cased', do_lower_case=False)\n",
        "import pytorch_pretrained_bert\n",
        "\n",
        "nn = torch.nn"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BZOIb-HFcxmV"
      },
      "source": [
        "#experiment_code\n",
        "hook = sy.TorchHook(torch)  # <-- NEW: hook PyTorch ie add extra functionalities to support Federated Learning\n",
        "bob = sy.VirtualWorker(hook, id=\"bob\")  # <-- NEW: define remote worker bob\n",
        "alice = sy.VirtualWorker(hook, id=\"alice\")  # <-- NEW: and alice"
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ismGxQ-x3Kv"
      },
      "source": [
        "## Dataset Class"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4uJ1neHLx-wG"
      },
      "source": [
        "utility class for managing dataset and padding functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b0Ev0aJhF5Py"
      },
      "source": [
        "tags = ['BD', 'BP', 'PR', 'SP', 'CH', 'ED']\n",
        "VOCAB_list = ['<PAD>', 'O',]\n",
        "for tag in tags:\n",
        "    VOCAB_list.append('I-'+tag)\n",
        "    VOCAB_list.append('B-'+tag)\n",
        "VOCAB = tuple(VOCAB_list)\n",
        "tag2idx = {tag: idx for idx, tag in enumerate(VOCAB)}\n",
        "idx2tag = {idx: tag for idx, tag in enumerate(VOCAB)}\n",
        "\n",
        "class NerDataset(data.Dataset):\n",
        "    def __init__(self, fpath):\n",
        "        \"\"\"\n",
        "        fpath: [train|valid|test].txt\n",
        "        \"\"\"\n",
        "        entries = open(fpath, 'r').read().strip().split(\"\\n\\n\")\n",
        "        sents, tags_li = [], [] # list of lists\n",
        "        for entry in entries:\n",
        "            lines = entry.splitlines()\n",
        "            words = [line.split()[0] for line in entry.splitlines() if len(line.split()) > 1]\n",
        "            tags = ([line.split()[-1] for line in entry.splitlines() if len(line.split()) > 1])\n",
        "            if not (len(tags) != 0 and tags.count(tags[0]) == len(tags)):\n",
        "                sents.append([\"[CLS]\"] + words + [\"[SEP]\"])\n",
        "                tags_li.append([\"<PAD>\"] + tags + [\"<PAD>\"])\n",
        "        self.sents, self.tags_li = sents, tags_li\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.sents)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        words, tags = self.sents[idx], self.tags_li[idx] # words, tags: string list\n",
        "        # We give credits only to the first piece.\n",
        "        x, y = [], [] # list of ids\n",
        "        is_heads = [] # list. 1: the token is the first piece of a word\n",
        "        for w, t in zip(words, tags):\n",
        "            if ord(w[0]) in [65533, 8206, 150, 61656, 128, 157] : #bad tokens that caused miss matches in the token and is_head legths\n",
        "                continue \n",
        "            tokens = tokenizer.tokenize(w) if w not in (\"[CLS]\", \"[SEP]\") else [w]\n",
        "            xx = tokenizer.convert_tokens_to_ids(tokens)\n",
        "\n",
        "            is_head = [1] + [0]*(len(tokens) - 1)\n",
        "\n",
        "            t = [t] + [\"<PAD>\"] * (len(tokens) - 1)  # <PAD>: no decision\n",
        "            yy = [tag2idx[each] for each in t]  # (T,)\n",
        "\n",
        "            x.extend(xx)\n",
        "            is_heads.extend(is_head)\n",
        "            y.extend(yy)\n",
        "        try:\n",
        "          assert len(x)==len(y)==len(is_heads), f\"len(x)={len(x)}, len(y)={len(y)}, len(is_heads)={len(is_heads)}\"\n",
        "        except AssertionError:\n",
        "          print(tags)\n",
        "          for tag in words:\n",
        "              for c in tag:\n",
        "                  print(c, ord(c))\n",
        "          print(words)\n",
        "          print(x)\n",
        "          print(y)\n",
        "          print(is_heads)\n",
        "          raise BaseException(f\"len(x)={len(x)}, len(y)={len(y)}, len(is_heads)={len(is_heads)}\")\n",
        "        # seqlen\n",
        "        seqlen = len(y)\n",
        "\n",
        "        # to string\n",
        "        words = \" \".join(words)\n",
        "        tags = \" \".join(tags)\n",
        "        return words, x, is_heads, tags, y, seqlen\n",
        "    \n",
        "    def append(self, other):\n",
        "        self.sents.extend(other.sents)\n",
        "        self.tags_li.extend(other.tags_li)\n",
        "\n",
        "def pad(batch):\n",
        "    '''Pads to 50'''\n",
        "    f = lambda x: [sample[x] for sample in batch]\n",
        "    g = lambda x, seqlen: [sample[x] + [\" #!#!\",[0],[0],\" <PAD>\"][x] * (seqlen - len(sample[x])) if len(sample[x]) < seqlen else sample[x][:seqlen] for sample in batch]  \n",
        "    seqlens = f(-1)\n",
        "    maxlen = min(50, np.array(seqlens).max())\n",
        "    words = f(0)\n",
        "    is_heads = g(2, maxlen)\n",
        "    tags = [sample[3] for sample in batch] #g(3, maxlen)\n",
        "    g = lambda x, seqlen: [ sample[x] + [0] * (seqlen - len(sample[x])) if len(sample[x]) < seqlen else sample[x][:seqlen] for sample in batch]  # 0: <pad>\n",
        "    x = g(1, maxlen)\n",
        "    y = g(-2, maxlen)\n",
        "    f = torch.LongTensor\n",
        "    return words, f(x), is_heads, tags, f(y), [maxlen for sample in batch]\n",
        "\n",
        "def pad_max(batch):\n",
        "    '''Pads to the longest sample'''\n",
        "    f = lambda x: [sample[x] for sample in batch]\n",
        "    words = f(0)\n",
        "    is_heads = f(2)\n",
        "    tags = f(3)\n",
        "    seqlens = f(-1)\n",
        "    maxlen = np.array(seqlens).max()\n",
        "\n",
        "    f = lambda x, seqlen: [sample[x] + [0] * (seqlen - len(sample[x])) for sample in batch] # 0: <pad>\n",
        "    x = f(1, maxlen)\n",
        "    y = f(-2, maxlen)\n",
        "    f = torch.LongTensor\n",
        "    return words, f(x), is_heads, tags, f(y), seqlens\n",
        "\n",
        "def pad_fed(batch):\n",
        "    '''Pads always to 50 since there will be a fixed size in the embedding layer'''\n",
        "    f = lambda x: [sample[x] for sample in batch]\n",
        "    g = lambda x, seqlen: [sample[x] + [\" #!#!\",[0],[0],\" <PAD>\"][x] * (seqlen - len(sample[x])) if len(sample[x]) < seqlen else sample[x][:seqlen] for sample in batch]  \n",
        "    seqlens = f(-1)\n",
        "    maxlen = 50\n",
        "    words = f(0)\n",
        "    is_heads = g(2, maxlen)\n",
        "    tags = [sample[3] for sample in batch] \n",
        "    g = lambda x, seqlen: [ sample[x] + [0] * (seqlen - len(sample[x])) if len(sample[x]) < seqlen else sample[x][:seqlen] for sample in batch]  # 0: <pad>\n",
        "\n",
        "    x = g(1, maxlen)\n",
        "    y = g(-2, maxlen)\n",
        "\n",
        "    f = torch.LongTensor\n",
        "    return words, f(x), is_heads, tags, f(y), [maxlen for sample in batch]"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9lG1cWaZSM2o"
      },
      "source": [
        "## BERT Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SAFb5ZCHn6GN"
      },
      "source": [
        "# coding=utf-8\n",
        "# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.\n",
        "# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\"\"\"PyTorch BERT model.\"\"\"\n",
        "\n",
        "from __future__ import absolute_import, division, print_function, unicode_literals\n",
        "\n",
        "import copy\n",
        "import json\n",
        "import logging\n",
        "import math\n",
        "import os\n",
        "import shutil\n",
        "import tarfile\n",
        "import tempfile\n",
        "import sys\n",
        "from io import open\n",
        "\n",
        "# import torch\n",
        "# from torch import nn\n",
        "# from torch.nn import CrossEntropyLoss\n",
        "# import pytorch_pretrained_bert\n",
        "from pytorch_pretrained_bert.file_utils import cached_path, WEIGHTS_NAME, CONFIG_NAME\n",
        "\n",
        "logger = logging.getLogger(__name__)\n",
        "\n",
        "PRETRAINED_MODEL_ARCHIVE_MAP = {\n",
        "    'bert-base-uncased': \"https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased.tar.gz\",\n",
        "    'bert-large-uncased': \"https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased.tar.gz\",\n",
        "    'bert-base-cased': \"https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased.tar.gz\",\n",
        "    'bert-large-cased': \"https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased.tar.gz\",\n",
        "    'bert-base-multilingual-uncased': \"https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased.tar.gz\",\n",
        "    'bert-base-multilingual-cased': \"https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased.tar.gz\",\n",
        "    'bert-base-chinese': \"https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese.tar.gz\",\n",
        "}\n",
        "BERT_CONFIG_NAME = 'bert_config.json'\n",
        "TF_WEIGHTS_NAME = 'model.ckpt'\n",
        "\n",
        "def load_tf_weights_in_bert(model, tf_checkpoint_path):\n",
        "    \"\"\" Load tf checkpoints in a pytorch model\n",
        "    \"\"\"\n",
        "    try:\n",
        "        import re\n",
        "        import numpy as np\n",
        "        import tensorflow as tf\n",
        "    except ImportError:\n",
        "        print(\"Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see \"\n",
        "            \"https://www.tensorflow.org/install/ for installation instructions.\")\n",
        "        raise\n",
        "    tf_path = os.path.abspath(tf_checkpoint_path)\n",
        "    print(\"Converting TensorFlow checkpoint from {}\".format(tf_path))\n",
        "    # Load weights from TF model\n",
        "    init_vars = tf.train.list_variables(tf_path)\n",
        "    names = []\n",
        "    arrays = []\n",
        "    for name, shape in init_vars:\n",
        "        print(\"Loading TF weight {} with shape {}\".format(name, shape))\n",
        "        array = tf.train.load_variable(tf_path, name)\n",
        "        names.append(name)\n",
        "        arrays.append(array)\n",
        "\n",
        "    for name, array in zip(names, arrays):\n",
        "        name = name.split('/')\n",
        "        # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v\n",
        "        # which are not required for using pretrained model\n",
        "        if any(n in [\"adam_v\", \"adam_m\", \"global_step\"] for n in name):\n",
        "            print(\"Skipping {}\".format(\"/\".join(name)))\n",
        "            continue\n",
        "        pointer = model\n",
        "        for m_name in name:\n",
        "            if re.fullmatch(r'[A-Za-z]+_\\d+', m_name):\n",
        "                l = re.split(r'_(\\d+)', m_name)\n",
        "            else:\n",
        "                l = [m_name]\n",
        "            if l[0] == 'kernel' or l[0] == 'gamma':\n",
        "                pointer = getattr(pointer, 'weight')\n",
        "            elif l[0] == 'output_bias' or l[0] == 'beta':\n",
        "                pointer = getattr(pointer, 'bias')\n",
        "            elif l[0] == 'output_weights':\n",
        "                pointer = getattr(pointer, 'weight')\n",
        "            elif l[0] == 'squad':\n",
        "                pointer = getattr(pointer, 'classifier')\n",
        "            else:\n",
        "                try:\n",
        "                    pointer = getattr(pointer, l[0])\n",
        "                except AttributeError:\n",
        "                    print(\"Skipping {}\".format(\"/\".join(name)))\n",
        "                    continue\n",
        "            if len(l) >= 2:\n",
        "                num = int(l[1])\n",
        "                pointer = pointer[num]\n",
        "        if m_name[-11:] == '_embeddings':\n",
        "            pointer = getattr(pointer, 'weight')\n",
        "        elif m_name == 'kernel':\n",
        "            array = np.transpose(array)\n",
        "        try:\n",
        "            assert pointer.shape == array.shape\n",
        "        except AssertionError as e:\n",
        "            e.args += (pointer.shape, array.shape)\n",
        "            raise\n",
        "        print(\"Initialize PyTorch weight {}\".format(name))\n",
        "        pointer.data = torch.from_numpy(array)\n",
        "    return model\n",
        "\n",
        "\n",
        "def gelu(x):\n",
        "    \"\"\"Implementation of the gelu activation function.\n",
        "        For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):\n",
        "        0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))\n",
        "        Also see https://arxiv.org/abs/1606.08415\n",
        "    \"\"\"\n",
        "    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))\n",
        "\n",
        "\n",
        "def swish(x):\n",
        "    return x * torch.sigmoid(x)\n",
        "\n",
        "\n",
        "ACT2FN = {\"gelu\": gelu, \"relu\": torch.nn.functional.relu, \"swish\": swish}\n",
        "\n",
        "\n",
        "class BertConfig(object):\n",
        "    \"\"\"Configuration class to store the configuration of a `BertModel`.\n",
        "    \"\"\"\n",
        "    def __init__(self,\n",
        "                 vocab_size_or_config_json_file,\n",
        "                 hidden_size=768,\n",
        "                 num_hidden_layers=12,\n",
        "                 num_attention_heads=12,\n",
        "                 intermediate_size=3072,\n",
        "                 hidden_act=\"gelu\",\n",
        "                 hidden_dropout_prob=0.1,\n",
        "                 attention_probs_dropout_prob=0.1,\n",
        "                 max_position_embeddings=512,\n",
        "                 type_vocab_size=2,\n",
        "                 initializer_range=0.02):\n",
        "        \"\"\"Constructs BertConfig.\n",
        "\n",
        "        Args:\n",
        "            vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `BertModel`.\n",
        "            hidden_size: Size of the encoder layers and the pooler layer.\n",
        "            num_hidden_layers: Number of hidden layers in the Transformer encoder.\n",
        "            num_attention_heads: Number of attention heads for each attention layer in\n",
        "                the Transformer encoder.\n",
        "            intermediate_size: The size of the \"intermediate\" (i.e., feed-forward)\n",
        "                layer in the Transformer encoder.\n",
        "            hidden_act: The non-linear activation function (function or string) in the\n",
        "                encoder and pooler. If string, \"gelu\", \"relu\" and \"swish\" are supported.\n",
        "            hidden_dropout_prob: The dropout probabilitiy for all fully connected\n",
        "                layers in the embeddings, encoder, and pooler.\n",
        "            attention_probs_dropout_prob: The dropout ratio for the attention\n",
        "                probabilities.\n",
        "            max_position_embeddings: The maximum sequence length that this model might\n",
        "                ever be used with. Typically set this to something large just in case\n",
        "                (e.g., 512 or 1024 or 2048).\n",
        "            type_vocab_size: The vocabulary size of the `token_type_ids` passed into\n",
        "                `BertModel`.\n",
        "            initializer_range: The sttdev of the truncated_normal_initializer for\n",
        "                initializing all weight matrices.\n",
        "        \"\"\"\n",
        "        if isinstance(vocab_size_or_config_json_file, str) or (sys.version_info[0] == 2\n",
        "                        and isinstance(vocab_size_or_config_json_file, unicode)):\n",
        "            with open(vocab_size_or_config_json_file, \"r\", encoding='utf-8') as reader:\n",
        "                json_config = json.loads(reader.read())\n",
        "            for key, value in json_config.items():\n",
        "                self.__dict__[key] = value\n",
        "        elif isinstance(vocab_size_or_config_json_file, int):\n",
        "            self.vocab_size = vocab_size_or_config_json_file\n",
        "            self.hidden_size = hidden_size\n",
        "            self.num_hidden_layers = num_hidden_layers\n",
        "            self.num_attention_heads = num_attention_heads\n",
        "            self.hidden_act = hidden_act\n",
        "            self.intermediate_size = intermediate_size\n",
        "            self.hidden_dropout_prob = hidden_dropout_prob\n",
        "            self.attention_probs_dropout_prob = attention_probs_dropout_prob\n",
        "            self.max_position_embeddings = max_position_embeddings\n",
        "            self.type_vocab_size = type_vocab_size\n",
        "            self.initializer_range = initializer_range\n",
        "        else:\n",
        "            raise ValueError(\"First argument must be either a vocabulary size (int)\"\n",
        "                             \"or the path to a pretrained model config file (str)\")\n",
        "\n",
        "    @classmethod\n",
        "    def from_dict(cls, json_object):\n",
        "        \"\"\"Constructs a `BertConfig` from a Python dictionary of parameters.\"\"\"\n",
        "        config = BertConfig(vocab_size_or_config_json_file=-1)\n",
        "        for key, value in json_object.items():\n",
        "            config.__dict__[key] = value\n",
        "        return config\n",
        "\n",
        "    @classmethod\n",
        "    def from_json_file(cls, json_file):\n",
        "        \"\"\"Constructs a `BertConfig` from a json file of parameters.\"\"\"\n",
        "        with open(json_file, \"r\", encoding='utf-8') as reader:\n",
        "            text = reader.read()\n",
        "        return cls.from_dict(json.loads(text))\n",
        "\n",
        "    def __repr__(self):\n",
        "        return str(self.to_json_string())\n",
        "\n",
        "    def to_dict(self):\n",
        "        \"\"\"Serializes this instance to a Python dictionary.\"\"\"\n",
        "        output = copy.deepcopy(self.__dict__)\n",
        "        return output\n",
        "\n",
        "    def to_json_string(self):\n",
        "        \"\"\"Serializes this instance to a JSON string.\"\"\"\n",
        "        return json.dumps(self.to_dict(), indent=2, sort_keys=True) + \"\\n\"\n",
        "\n",
        "    def to_json_file(self, json_file_path):\n",
        "        \"\"\" Save this instance to a json file.\"\"\"\n",
        "        with open(json_file_path, \"w\", encoding='utf-8') as writer:\n",
        "            writer.write(self.to_json_string())\n",
        "\n",
        "try:\n",
        "    from apex.normalization.fused_layer_norm import FusedLayerNorm as BertLayerNorm\n",
        "except ImportError:\n",
        "    logger.info(\"Better speed can be achieved with apex installed from https://www.github.com/nvidia/apex .\")\n",
        "    class BertLayerNorm(nn.Module):\n",
        "        def __init__(self, hidden_size, eps=1e-12):\n",
        "            \"\"\"Construct a layernorm module in the TF style (epsilon inside the square root).\n",
        "            \"\"\"\n",
        "            super(BertLayerNorm, self).__init__()\n",
        "            self.weight = nn.Parameter(torch.ones(hidden_size))\n",
        "            self.bias = nn.Parameter(torch.zeros(hidden_size))\n",
        "            self.variance_epsilon = eps\n",
        "\n",
        "        def forward(self, x):\n",
        "            u = x.mean(-1, keepdim=True)\n",
        "            s = (x - u).pow(2).mean(-1, keepdim=True)\n",
        "            x = (x - u) / torch.sqrt(s + self.variance_epsilon)\n",
        "            return self.weight * x + self.bias\n",
        "\n",
        "class BertEmbeddings(nn.Module):\n",
        "    \"\"\"Construct the embeddings from word, position and token_type embeddings.\n",
        "    \"\"\"\n",
        "    def __init__(self, config):\n",
        "        super(BertEmbeddings, self).__init__()\n",
        "        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0)\n",
        "        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)\n",
        "        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)\n",
        "\n",
        "        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load\n",
        "        # any TensorFlow checkpoint file\n",
        "        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)\n",
        "        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n",
        "\n",
        "        self.register_buffer('position_ids', torch.arange(50, dtype=torch.long, device='cuda').unsqueeze(0).expand(torch.Size([128, 50])))\n",
        "        self.register_buffer('token_type_ids', torch.zeros(torch.Size([128, 50]), dtype=torch.long, device='cuda'))\n",
        "        \n",
        "    def forward(self, input_ids, token_type_ids=None):\n",
        "\n",
        "        words_embeddings = self.word_embeddings(input_ids)\n",
        "        position_embeddings = self.position_embeddings(self.position_ids)\n",
        "        token_type_embeddings = self.token_type_embeddings(self.token_type_ids) \n",
        "        embeddings = words_embeddings + position_embeddings + token_type_embeddings\n",
        "        embeddings = self.LayerNorm(embeddings)\n",
        "        embeddings = self.dropout(embeddings)\n",
        "        return embeddings\n",
        "\n",
        "\n",
        "class BertSelfAttention(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super(BertSelfAttention, self).__init__()\n",
        "        if config.hidden_size % config.num_attention_heads != 0:\n",
        "            raise ValueError(\n",
        "                \"The hidden size (%d) is not a multiple of the number of attention \"\n",
        "                \"heads (%d)\" % (config.hidden_size, config.num_attention_heads))\n",
        "        self.num_attention_heads = config.num_attention_heads\n",
        "        self.attention_head_size = int(config.hidden_size / config.num_attention_heads) # 768/12 = 64\n",
        "        self.all_head_size = self.num_attention_heads * self.attention_head_size  # 12 * 64\n",
        "\n",
        "        self.query = nn.Linear(config.hidden_size, self.all_head_size)\n",
        "        self.key = nn.Linear(config.hidden_size, self.all_head_size)\n",
        "        self.value = nn.Linear(config.hidden_size, self.all_head_size)\n",
        "\n",
        "        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)\n",
        "\n",
        "    def transpose_for_scores(self, x):\n",
        "        x = x.view(128, 50, 12, 64)\n",
        "        x = x.permute(0, 2, 1, 3)\n",
        "        return x\n",
        "\n",
        "    def forward(self, hidden_states, attention_mask):\n",
        "        mixed_query_layer = self.query(hidden_states)\n",
        "        mixed_key_layer = self.key(hidden_states)\n",
        "        mixed_value_layer = self.value(hidden_states)\n",
        "\n",
        "        query_layer = self.transpose_for_scores(mixed_query_layer)\n",
        "        key_layer = self.transpose_for_scores(mixed_key_layer)\n",
        "        value_layer = self.transpose_for_scores(mixed_value_layer)\n",
        "\n",
        "        # Take the dot product between \"query\" and \"key\" to get the raw attention scores.\n",
        "        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))\n",
        "        attention_scores = attention_scores / math.sqrt(self.attention_head_size)\n",
        "        # Apply the attention mask is (precomputed for all layers in BertModel forward() function)\n",
        "        attention_scores = attention_scores + attention_mask\n",
        "\n",
        "        # Normalize the attention scores to probabilities.\n",
        "        attention_probs = nn.Softmax(dim=-1)(attention_scores)\n",
        "\n",
        "        # This is actually dropping out entire tokens to attend to, which might\n",
        "        # seem a bit unusual, but is taken from the original Transformer paper.\n",
        "        attention_probs = self.dropout(attention_probs)\n",
        "\n",
        "        context_layer = torch.matmul(attention_probs, value_layer)        \n",
        "        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()\n",
        "        loc = context_layer.location\n",
        "        context_layer = context_layer.get()\n",
        "        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)\n",
        "        context_layer = context_layer.view(*new_context_layer_shape)\n",
        "        context_layer = context_layer.send(loc)\n",
        "        return context_layer\n",
        "\n",
        "\n",
        "class BertSelfOutput(nn.Module):\n",
        "    def __init__(self, config):  \n",
        "        super(BertSelfOutput, self).__init__()\n",
        "        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n",
        "        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)\n",
        "        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n",
        "\n",
        "    def forward(self, hidden_states, input_tensor):\n",
        "        hidden_states = self.dense(hidden_states)\n",
        "        hidden_states = self.dropout(hidden_states)\n",
        "        hidden_states = self.LayerNorm(hidden_states + input_tensor)\n",
        "        return hidden_states\n",
        "\n",
        "\n",
        "class BertAttention(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super(BertAttention, self).__init__()\n",
        "        self.self = BertSelfAttention(config)\n",
        "        self.output = BertSelfOutput(config)\n",
        "\n",
        "    def forward(self, input_tensor, attention_mask):\n",
        "        self_output = self.self(input_tensor, attention_mask)\n",
        "        attention_output = self.output(self_output, input_tensor)\n",
        "        return attention_output\n",
        "\n",
        "\n",
        "class BertIntermediate(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super(BertIntermediate, self).__init__()\n",
        "        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)\n",
        "        if isinstance(config.hidden_act, str) or (sys.version_info[0] == 2 and isinstance(config.hidden_act, unicode)):\n",
        "            self.intermediate_act_fn = ACT2FN[config.hidden_act]\n",
        "        else:\n",
        "            self.intermediate_act_fn = config.hidden_act\n",
        "\n",
        "    def forward(self, hidden_states):\n",
        "        hidden_states = self.dense(hidden_states)\n",
        "        hidden_states = self.intermediate_act_fn(hidden_states)\n",
        "        return hidden_states\n",
        "\n",
        "\n",
        "class BertOutput(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super(BertOutput, self).__init__()\n",
        "        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)\n",
        "        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)\n",
        "        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n",
        "\n",
        "    def forward(self, hidden_states, input_tensor):\n",
        "        hidden_states = self.dense(hidden_states)\n",
        "        hidden_states = self.dropout(hidden_states)\n",
        "        hidden_states = self.LayerNorm(hidden_states + input_tensor)\n",
        "        return hidden_states\n",
        "\n",
        "\n",
        "class BertLayer(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super(BertLayer, self).__init__()\n",
        "        self.attention = BertAttention(config)\n",
        "        self.intermediate = BertIntermediate(config)\n",
        "        self.output = BertOutput(config)\n",
        "\n",
        "    def forward(self, hidden_states, attention_mask):\n",
        "        attention_output = self.attention(hidden_states, attention_mask)\n",
        "        intermediate_output = self.intermediate(attention_output)\n",
        "        layer_output = self.output(intermediate_output, attention_output)\n",
        "        return layer_output\n",
        "\n",
        "\n",
        "class BertEncoder(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super(BertEncoder, self).__init__()\n",
        "        layer = BertLayer(config)\n",
        "        self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])\n",
        "\n",
        "    def forward(self, hidden_states, attention_mask, output_all_encoded_layers=True):\n",
        "        all_encoder_layers = []\n",
        "        for layer_module in self.layer:\n",
        "            hidden_states = layer_module(hidden_states, attention_mask)\n",
        "            if output_all_encoded_layers:\n",
        "                all_encoder_layers.append(hidden_states)\n",
        "        if not output_all_encoded_layers:\n",
        "            all_encoder_layers.append(hidden_states)\n",
        "        return all_encoder_layers\n",
        "\n",
        "\n",
        "class BertPooler(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super(BertPooler, self).__init__()\n",
        "        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n",
        "        self.activation = nn.Tanh()\n",
        "\n",
        "    def forward(self, hidden_states):\n",
        "        # We \"pool\" the model by simply taking the hidden state corresponding\n",
        "        # to the first token.\n",
        "        first_token_tensor = hidden_states[:, 0]\n",
        "        pooled_output = self.dense(first_token_tensor)\n",
        "        pooled_output = self.activation(pooled_output)\n",
        "        return pooled_output\n",
        "\n",
        "\n",
        "class BertPredictionHeadTransform(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super(BertPredictionHeadTransform, self).__init__()\n",
        "        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n",
        "        if isinstance(config.hidden_act, str) or (sys.version_info[0] == 2 and isinstance(config.hidden_act, unicode)):\n",
        "            self.transform_act_fn = ACT2FN[config.hidden_act]\n",
        "        else:\n",
        "            self.transform_act_fn = config.hidden_act\n",
        "        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)\n",
        "\n",
        "    def forward(self, hidden_states):\n",
        "        hidden_states = self.dense(hidden_states)\n",
        "        hidden_states = self.transform_act_fn(hidden_states)\n",
        "        hidden_states = self.LayerNorm(hidden_states)\n",
        "        return hidden_states\n",
        "\n",
        "\n",
        "class BertLMPredictionHead(nn.Module):\n",
        "    def __init__(self, config, bert_model_embedding_weights):\n",
        "        super(BertLMPredictionHead, self).__init__()\n",
        "        self.transform = BertPredictionHeadTransform(config)\n",
        "\n",
        "        # The output weights are the same as the input embeddings, but there is\n",
        "        # an output-only bias for each token.\n",
        "        self.decoder = nn.Linear(bert_model_embedding_weights.size(1),\n",
        "                                 bert_model_embedding_weights.size(0),\n",
        "                                 bias=False)\n",
        "        self.decoder.weight = bert_model_embedding_weights\n",
        "        self.bias = nn.Parameter(torch.zeros(bert_model_embedding_weights.size(0)))\n",
        "\n",
        "    def forward(self, hidden_states):\n",
        "        hidden_states = self.transform(hidden_states)\n",
        "        hidden_states = self.decoder(hidden_states) + self.bias\n",
        "        return hidden_states\n",
        "\n",
        "\n",
        "class BertOnlyMLMHead(nn.Module):\n",
        "    def __init__(self, config, bert_model_embedding_weights):\n",
        "        super(BertOnlyMLMHead, self).__init__()\n",
        "        self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights)\n",
        "\n",
        "    def forward(self, sequence_output):\n",
        "        prediction_scores = self.predictions(sequence_output)\n",
        "        return prediction_scores\n",
        "\n",
        "\n",
        "class BertOnlyNSPHead(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super(BertOnlyNSPHead, self).__init__()\n",
        "        self.seq_relationship = nn.Linear(config.hidden_size, 2)\n",
        "\n",
        "    def forward(self, pooled_output):\n",
        "        seq_relationship_score = self.seq_relationship(pooled_output)\n",
        "        return seq_relationship_score\n",
        "\n",
        "\n",
        "class BertPreTrainingHeads(nn.Module):\n",
        "    def __init__(self, config, bert_model_embedding_weights):\n",
        "        super(BertPreTrainingHeads, self).__init__()\n",
        "        self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights)\n",
        "        self.seq_relationship = nn.Linear(config.hidden_size, 2)\n",
        "\n",
        "    def forward(self, sequence_output, pooled_output):\n",
        "        prediction_scores = self.predictions(sequence_output)\n",
        "        seq_relationship_score = self.seq_relationship(pooled_output)\n",
        "        return prediction_scores, seq_relationship_score\n",
        "\n",
        "class BertPreTrainedModel(nn.Module):\n",
        "    \"\"\" An abstract class to handle weights initialization and\n",
        "        a simple interface for dowloading and loading pretrained models.\n",
        "    \"\"\"\n",
        "    def __init__(self, config, *inputs, **kwargs):\n",
        "        super(BertPreTrainedModel, self).__init__()\n",
        "        if not isinstance(config, BertConfig):\n",
        "            raise ValueError(\n",
        "                \"Parameter config in `{}(config)` should be an instance of class `BertConfig`. \"\n",
        "                \"To create a model from a Google pretrained model use \"\n",
        "                \"`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`\".format(\n",
        "                    self.__class__.__name__, self.__class__.__name__\n",
        "                ))\n",
        "        self.config = config\n",
        "\n",
        "    def init_bert_weights(self, module):\n",
        "        \"\"\" Initialize the weights.\n",
        "        \"\"\"\n",
        "        if isinstance(module, (nn.Linear, nn.Embedding)):\n",
        "            # Slightly different from the TF version which uses truncated_normal for initialization\n",
        "            # cf https://github.com/pytorch/pytorch/pull/5617\n",
        "            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)\n",
        "        elif isinstance(module, BertLayerNorm):\n",
        "            module.bias.data.zero_()\n",
        "            module.weight.data.fill_(1.0)\n",
        "        if isinstance(module, nn.Linear) and module.bias is not None:\n",
        "            module.bias.data.zero_()\n",
        "\n",
        "    @classmethod\n",
        "    def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):\n",
        "        \"\"\"\n",
        "        Instantiate a BertPreTrainedModel from a pre-trained model file or a pytorch state dict.\n",
        "        Download and cache the pre-trained model file if needed.\n",
        "\n",
        "        Params:\n",
        "            pretrained_model_name_or_path: either:\n",
        "                - a str with the name of a pre-trained model to load selected in the list of:\n",
        "                    . `bert-base-uncased`\n",
        "                    . `bert-large-uncased`\n",
        "                    . `bert-base-cased`\n",
        "                    . `bert-large-cased`\n",
        "                    . `bert-base-multilingual-uncased`\n",
        "                    . `bert-base-multilingual-cased`\n",
        "                    . `bert-base-chinese`\n",
        "                - a path or url to a pretrained model archive containing:\n",
        "                    . `bert_config.json` a configuration file for the model\n",
        "                    . `pytorch_model.bin` a PyTorch dump of a BertForPreTraining instance\n",
        "                - a path or url to a pretrained model archive containing:\n",
        "                    . `bert_config.json` a configuration file for the model\n",
        "                    . `model.chkpt` a TensorFlow checkpoint\n",
        "            from_tf: should we load the weights from a locally saved TensorFlow checkpoint\n",
        "            cache_dir: an optional path to a folder in which the pre-trained models will be cached.\n",
        "            state_dict: an optional state dictionnary (collections.OrderedDict object) to use instead of Google pre-trained models\n",
        "            *inputs, **kwargs: additional input for the specific Bert class\n",
        "                (ex: num_labels for BertForSequenceClassification)\n",
        "        \"\"\"\n",
        "        state_dict = kwargs.get('state_dict', None)\n",
        "        kwargs.pop('state_dict', None)\n",
        "        cache_dir = kwargs.get('cache_dir', None)\n",
        "        kwargs.pop('cache_dir', None)\n",
        "        from_tf = kwargs.get('from_tf', False)\n",
        "        kwargs.pop('from_tf', None)\n",
        "\n",
        "        if pretrained_model_name_or_path in PRETRAINED_MODEL_ARCHIVE_MAP:\n",
        "            archive_file = PRETRAINED_MODEL_ARCHIVE_MAP[pretrained_model_name_or_path]\n",
        "        else:\n",
        "            archive_file = pretrained_model_name_or_path\n",
        "        # redirect to the cache, if necessary\n",
        "        try:\n",
        "            resolved_archive_file = cached_path(archive_file, cache_dir=cache_dir)\n",
        "        except EnvironmentError:\n",
        "            logger.error(\n",
        "                \"Model name '{}' was not found in model name list ({}). \"\n",
        "                \"We assumed '{}' was a path or url but couldn't find any file \"\n",
        "                \"associated to this path or url.\".format(\n",
        "                    pretrained_model_name_or_path,\n",
        "                    ', '.join(PRETRAINED_MODEL_ARCHIVE_MAP.keys()),\n",
        "                    archive_file))\n",
        "            return None\n",
        "        if resolved_archive_file == archive_file:\n",
        "            logger.info(\"loading archive file {}\".format(archive_file))\n",
        "        else:\n",
        "            logger.info(\"loading archive file {} from cache at {}\".format(\n",
        "                archive_file, resolved_archive_file))\n",
        "        tempdir = None\n",
        "        if os.path.isdir(resolved_archive_file) or from_tf:\n",
        "            serialization_dir = resolved_archive_file\n",
        "        else:\n",
        "            # Extract archive to temp dir\n",
        "            tempdir = tempfile.mkdtemp()\n",
        "            logger.info(\"extracting archive file {} to temp dir {}\".format(\n",
        "                resolved_archive_file, tempdir))\n",
        "            with tarfile.open(resolved_archive_file, 'r:gz') as archive:\n",
        "                archive.extractall(tempdir)\n",
        "            serialization_dir = tempdir\n",
        "        # Load config\n",
        "        config_file = os.path.join(serialization_dir, CONFIG_NAME)\n",
        "        if not os.path.exists(config_file):\n",
        "            # Backward compatibility with old naming format\n",
        "            config_file = os.path.join(serialization_dir, BERT_CONFIG_NAME)\n",
        "        config = BertConfig.from_json_file(config_file)\n",
        "        logger.info(\"Model config {}\".format(config))\n",
        "        # Instantiate model.\n",
        "        model = cls(config, *inputs, **kwargs)\n",
        "        if state_dict is None and not from_tf:\n",
        "            weights_path = os.path.join(serialization_dir, WEIGHTS_NAME)\n",
        "            state_dict = torch.load(weights_path, map_location='cpu')\n",
        "        if tempdir:\n",
        "            # Clean up temp dir\n",
        "            shutil.rmtree(tempdir)\n",
        "        if from_tf:\n",
        "            # Directly load from a TensorFlow checkpoint\n",
        "            weights_path = os.path.join(serialization_dir, TF_WEIGHTS_NAME)\n",
        "            return load_tf_weights_in_bert(model, weights_path)\n",
        "        # Load from a PyTorch state_dict\n",
        "        old_keys = []\n",
        "        new_keys = []\n",
        "        for key in state_dict.keys():\n",
        "            new_key = None\n",
        "            if 'gamma' in key:\n",
        "                new_key = key.replace('gamma', 'weight')\n",
        "            if 'beta' in key:\n",
        "                new_key = key.replace('beta', 'bias')\n",
        "            if new_key:\n",
        "                old_keys.append(key)\n",
        "                new_keys.append(new_key)\n",
        "        for old_key, new_key in zip(old_keys, new_keys):\n",
        "            state_dict[new_key] = state_dict.pop(old_key)\n",
        "\n",
        "        missing_keys = []\n",
        "        unexpected_keys = []\n",
        "        error_msgs = []\n",
        "        # copy state_dict so _load_from_state_dict can modify it\n",
        "        metadata = getattr(state_dict, '_metadata', None)\n",
        "        state_dict = state_dict.copy()\n",
        "        if metadata is not None:\n",
        "            state_dict._metadata = metadata\n",
        "\n",
        "        def load(module, prefix=''):\n",
        "            local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})\n",
        "            module._load_from_state_dict(\n",
        "                state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)\n",
        "            for name, child in module._modules.items():\n",
        "                if child is not None:\n",
        "                    load(child, prefix + name + '.')\n",
        "        start_prefix = ''\n",
        "        if not hasattr(model, 'bert') and any(s.startswith('bert.') for s in state_dict.keys()):\n",
        "            start_prefix = 'bert.'\n",
        "        load(model, prefix=start_prefix)\n",
        "        if len(missing_keys) > 0:\n",
        "            logger.info(\"Weights of {} not initialized from pretrained model: {}\".format(\n",
        "                model.__class__.__name__, missing_keys))\n",
        "        if len(unexpected_keys) > 0:\n",
        "            logger.info(\"Weights from pretrained model not used in {}: {}\".format(\n",
        "                model.__class__.__name__, unexpected_keys))\n",
        "        if len(error_msgs) > 0:\n",
        "            raise RuntimeError('Error(s) in loading state_dict for {}:\\n\\t{}'.format(\n",
        "                               model.__class__.__name__, \"\\n\\t\".join(error_msgs)))\n",
        "        return model\n",
        "\n",
        "\n",
        "class BertModel(BertPreTrainedModel):\n",
        "    \"\"\"BERT model (\"Bidirectional Embedding Representations from a Transformer\").\n",
        "\n",
        "    Params:\n",
        "        config: a BertConfig class instance with the configuration to build a new model\n",
        "\n",
        "    Inputs:\n",
        "        `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]\n",
        "            with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts\n",
        "            `extract_features.py`, `run_classifier.py` and `run_squad.py`)\n",
        "        `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token\n",
        "            types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to\n",
        "            a `sentence B` token (see BERT paper for more details).\n",
        "        `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices\n",
        "            selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max\n",
        "            input sequence length in the current batch. It's the mask that we typically use for attention when\n",
        "            a batch has varying length sentences.\n",
        "        `output_all_encoded_layers`: boolean which controls the content of the `encoded_layers` output as described below. Default: `True`.\n",
        "\n",
        "    Outputs: Tuple of (encoded_layers, pooled_output)\n",
        "        `encoded_layers`: controled by `output_all_encoded_layers` argument:\n",
        "            - `output_all_encoded_layers=True`: outputs a list of the full sequences of encoded-hidden-states at the end\n",
        "                of each attention block (i.e. 12 full sequences for BERT-base, 24 for BERT-large), each\n",
        "                encoded-hidden-state is a torch.FloatTensor of size [batch_size, sequence_length, hidden_size],\n",
        "            - `output_all_encoded_layers=False`: outputs only the full sequence of hidden-states corresponding\n",
        "                to the last attention block of shape [batch_size, sequence_length, hidden_size],\n",
        "        `pooled_output`: a torch.FloatTensor of size [batch_size, hidden_size] which is the output of a\n",
        "            classifier pretrained on top of the hidden state associated to the first character of the\n",
        "            input (`CLS`) to train on the Next-Sentence task (see BERT's paper).\n",
        "\n",
        "    Example usage:\n",
        "    ```python\n",
        "    # Already been converted into WordPiece token ids\n",
        "    input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])\n",
        "    input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])\n",
        "    token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])\n",
        "\n",
        "    config = modeling.BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,\n",
        "        num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)\n",
        "\n",
        "    model = modeling.BertModel(config=config)\n",
        "    all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask)\n",
        "    ```\n",
        "    \"\"\"\n",
        "    def __init__(self, config):\n",
        "        super(BertModel, self).__init__(config)\n",
        "        self.embeddings = BertEmbeddings(config)\n",
        "        self.encoder = BertEncoder(config)\n",
        "        self.pooler = BertPooler(config)\n",
        "        self.apply(self.init_bert_weights)\n",
        "        extended_attention_mask = torch.ones(torch.Size([128, 1, 1, 50])).to('cuda')\n",
        "        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0\n",
        "        self.register_buffer('extended_attention_mask', extended_attention_mask)\n",
        "\n",
        "        print('this is the model implemeted on the nb')\n",
        "          \n",
        "    def forward(self, input_ids, token_type_ids=None, attention_mask=None, output_all_encoded_layers=True):\n",
        "        if attention_mask is None:\n",
        "            attention_mask = torch.ones_like(input_ids)\n",
        "        if token_type_ids is None:\n",
        "            token_type_ids = torch.zeros_like(input_ids)\n",
        "\n",
        "        embedding_output = self.embeddings(input_ids)\n",
        "        encoded_layers = self.encoder(embedding_output,\n",
        "                                      self.extended_attention_mask,\n",
        "                                      output_all_encoded_layers=output_all_encoded_layers)\n",
        "        sequence_output = encoded_layers[-1]\n",
        "        pooled_output = self.pooler(sequence_output)\n",
        "        if not output_all_encoded_layers:\n",
        "            encoded_layers = encoded_layers[-1]\n",
        "        return encoded_layers, pooled_output\n"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qnKZh68zF5SR"
      },
      "source": [
        "class Net(nn.Module):\n",
        "    def __init__(self, top_rnns=False, vocab_size=None, device='cpu', finetuning=False):\n",
        "        super().__init__()\n",
        "        self.bert = BertModel.from_pretrained('bert-base-cased')\n",
        "\n",
        "        self.top_rnns=top_rnns\n",
        "        if top_rnns:\n",
        "            self.rnn = nn.LSTM(bidirectional=True, num_layers=2, input_size=768, hidden_size=768//2, batch_first=True)\n",
        "        self.fc = nn.Linear(768, vocab_size)\n",
        "\n",
        "        self.device = device\n",
        "        self.finetuning = finetuning\n",
        "\n",
        "    def forward(self, x, y, ):\n",
        "        '''\n",
        "        x: (N, T). int64\n",
        "        y: (N, T). int64\n",
        "        Returns\n",
        "        enc: (N, T, VOCAB)\n",
        "        '''\n",
        "        x = x.to(self.device)\n",
        "        y = y.to(self.device)\n",
        "\n",
        "        if self.training and self.finetuning:\n",
        "            # print(\"->bert.train()\")\n",
        "            self.bert.train()\n",
        "            # self.bert.send(x.location)\n",
        "            encoded_layers, _ = self.bert(x)\n",
        "            enc = encoded_layers[-1]\n",
        "            # self.bert.get()\n",
        "        else:\n",
        "            self.bert.eval()\n",
        "            with torch.no_grad():\n",
        "                encoded_layers, _ = self.bert(x)\n",
        "                enc = encoded_layers[-1]\n",
        "\n",
        "        if self.top_rnns:\n",
        "            enc, _ = self.rnn(enc)\n",
        "        logits = self.fc(enc)\n",
        "        y_hat = logits.argmax(-1)\n",
        "        return logits, y, y_hat"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YMh_ApScra5T"
      },
      "source": [
        "## Data Iterators"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "acL2OzyNpArT"
      },
      "source": [
        "#experiment_code\n",
        "paths_annot = sorted([os.path.join(f[0], name) for f in os.walk('./dataset_txt') \n",
        "                if len(f[2])!=0 for name in f[2] if os.path.splitext(name)[-1] == '.txt' and name.split('_')[0]=='annot'],\n",
        "               key=lambda path: int(path.split('_')[-1].split('.')[0]))\n",
        "training_files = [path for path in paths_annot if 'training' in path.split('/')]\n",
        "eval_files = [path for path in paths_annot if 'eval' in path.split('/')]\n",
        "test_files = [path for path in paths_annot if 'test' in path.split('/')]"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "INJK415_emDK"
      },
      "source": [
        "#experiment_code\n",
        "eval_dataset = NerDataset(eval_files[0])\n",
        "_ = [eval_dataset.append(NerDataset(eval_file)) for eval_file in eval_files[1:]]\n",
        "\n",
        "test_dataset = NerDataset(test_files[0])\n",
        "_ = [test_dataset.append(NerDataset(test_file)) for test_file in test_files[1:]]\n",
        "\n",
        "eval_iter = data.DataLoader(dataset=eval_dataset,\n",
        "                              batch_size=1,\n",
        "                              shuffle=True,\n",
        "                              num_workers=16,\n",
        "                              collate_fn=pad)\n",
        "test_iter = data.DataLoader(dataset=test_dataset,\n",
        "                              batch_size=32,\n",
        "                              shuffle=True,\n",
        "                              num_workers=16,\n",
        "                              collate_fn=pad)\n",
        "print(len(eval_iter))\n",
        "print(len(test_iter))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qjF9fet1G7N9"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jbQB93qwltG5"
      },
      "source": [
        "## Training function"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yUgSDE_AJqRb"
      },
      "source": [
        "def federated_train(model, device, iterator, criterion):\n",
        "    model.train()\n",
        "    loc = ''\n",
        "    for i, batch in enumerate(iterator):\n",
        "        # words, x, is_heads, tags, y, seqlens = batch\n",
        "        x, y = batch\n",
        "        if loc != x.location:\n",
        "            print(f'data location {x.location}')\n",
        "            if model.location != None:\n",
        "                model.get()\n",
        "            model.send(x.location) # <-- NEW: send the model to the right location\n",
        "            optimizer = optim.Adam(model.parameters(), lr = lr)\n",
        "            loc = x.location\n",
        "        x, y = x.to(device), y.to(device)\n",
        "        _y = y # for monitoring\n",
        "        optimizer.zero_grad()\n",
        "        logits, y, _ = model(x, y) # logits: (N, T, VOCAB), y: (N, T)\n",
        "        logits = logits.view(-1, len(VOCAB_list)) # (N*T, VOCAB), vocab len = #tags*2<I,N> + 0 + <PAD> \n",
        "        y = y.view(-1)  # (N*T,)\n",
        "        loss = criterion(logits, y)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        if i%10==0: # monitoring\n",
        "            print(f\"step: {i}, loss: {loss.get().item()}\")\n",
        "    model.get() # <-- NEW: get the model back    "
      ],
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S8hpX44WLC3L"
      },
      "source": [
        "## Loading model for training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VxrHwSyPlqIr"
      },
      "source": [
        "#experiment_code\n",
        "print(torch.cuda.max_memory_allocated(device=None))\n",
        "lr = 0.0001\n",
        "finetuning = True\n",
        "top_rnns = False\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "model = Net(top_rnns=top_rnns, vocab_size=len(VOCAB), device=device, finetuning=finetuning).to(device)\n",
        "print(torch.cuda.max_memory_allocated(device=None))\n",
        "criterion = nn.CrossEntropyLoss(ignore_index=0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jggSzl1Nf6oM"
      },
      "source": [
        " ## Distributing the dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fBv4x3ZjIXPx"
      },
      "source": [
        "import math\n",
        "def dataset_federate(dataset, workers):\n",
        "    \"\"\"\n",
        "    Add a method to easily transform a torch.Dataset or a sy.BaseDataset\n",
        "    into a sy.FederatedDataset. The dataset given is split in len(workers)\n",
        "    part and sent to each workers\n",
        "    \"\"\"\n",
        "    # logger.info(\"Scanning and sending data to {}...\".format(\", \".join([w.id for w in workers])))\n",
        "    # take ceil to have exactly len(workers) sets after splitting\n",
        "    data_size = math.ceil(len(dataset) / len(workers))\n",
        "    datasets = []\n",
        "    data_loader = torch.utils.data.DataLoader(dataset, batch_size=data_size, collate_fn=pad_fed)\n",
        "    for dataset_idx, (words, x, is_heads, tags, y, seqlen) in enumerate(data_loader):\n",
        "        worker = workers[dataset_idx % len(workers)]\n",
        "        data_batches = x.size(0)//128\n",
        "        x =  x.narrow(0, 0, int(data_batches*128))        \n",
        "        y =  y.narrow(0, 0, int(data_batches*128))        \n",
        "        data = x.send(worker)\n",
        "        targets = y.send(worker)\n",
        "        datasets.append(sy.BaseDataset(data, targets))  # .send(worker)\n",
        "    return sy.FederatedDataset(datasets)\n",
        "NerDataset.federate = dataset_federate"
      ],
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tzE3TyLUgCOZ"
      },
      "source": [
        "#experiment_code\n",
        "train_dataset = NerDataset(training_files[0])\n",
        "_ = [train_dataset.append(NerDataset(train_file)) for train_file in training_files[1:200]]"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xz7P3OVdLz35"
      },
      "source": [
        "#experiment_code\n",
        "federated_train_loader = sy.FederatedDataLoader(train_dataset.federate((bob, alice)),\n",
        "                              batch_size=128,\n",
        "                              shuffle=True)\n",
        "federated_train_iter = federated_train_loader"
      ],
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QurWsl_yFZQe"
      },
      "source": [
        "## Federated training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rgJ4XQGaFYpH"
      },
      "source": [
        "#experiment_code\n",
        "for epoch in range(1, 1+1):\n",
        "    federated_train(model, device, federated_train_iter, criterion)\n",
        "    print(f\"=========eval at epoch={epoch}=========\")\n",
        "print(torch.cuda.max_memory_allocated(device=None))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yw_CtHT0t7JL",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "41ed3dcc-e098-4a44-d215-9b345e76ca9b"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Thu Apr 22 23:45:34 2021       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 465.19.01    Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   57C    P8    10W /  70W |      0MiB / 15109MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}