Italian_Pre_training_AlBERTo_from_scratch_with_cloud_TPU.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Italian Pre-training AlBERTo from scratch with cloud TPU",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E2XB_l-Hgzq_",
        "colab_type": "text"
      },
      "source": [
        "# Pre-training AlBERTo from scratch with cloud TPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vZPgpRl5g2e2",
        "colab_type": "text"
      },
      "source": [
        "In this notebook It is described how a pre-training a state-of-the-art Natural Language Understanding model <b>AlBERTo</b> based on [BERT](https://arxiv.org/abs/1810.04805.) for Italian Language using Google Cloud infrastructure.\n",
        "\n",
        "For persistent storage of training data and model, you will require a Google Cloud Storage bucket [Google Cloud Storage](https://cloud.google.com/products/storage/). \n",
        "The notebook can not be run without a GCS bucket. In that case, however, you will not be able to train the model.\n",
        "\n",
        "\n",
        "<br>\n",
        "<b>**Note** </b>\n",
        "Pre-training a BERT-Base model on a TPUv2 will take about 50 hours. Google Colab is not designed for executing such long-running jobs and will interrupt the training process every 12 hours or so. For uninterrupted training, consider using a preemptible TPUv2 instance. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8wVeP7SmOmTe",
        "colab_type": "text"
      },
      "source": [
        " <span> </span>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ODimOhBR05yR",
        "colab_type": "text"
      },
      "source": [
        "MIT License for the **original notebook sourcecode**\n",
        "\n",
        "Copyright (c) [2019] [Antyukhov Denis Olegovich]\n",
        "\n",
        "*Permission is hereby granted, free of charge, to any person obtaining a copy\n",
        "of this software and associated documentation files (the \"Software\"), to deal\n",
        "in the Software without restriction, including without limitation the rights\n",
        "to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
        "copies of the Software, and to permit persons to whom the Software is\n",
        "furnished to do so, subject to the following conditions:*\n",
        "\n",
        "The above copyright notice and this permission notice shall be included in all\n",
        "copies or substantial portions of the Software.\n",
        "\n",
        "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
        "IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
        "FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
        "AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
        "LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
        "OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fjad5jsr9YaM",
        "colab_type": "text"
      },
      "source": [
        "## Step 1: setting up training environment\n",
        "First and foremost, we get the packages required to train the model. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9S4CiOh3RzFW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#INSTALL SENTENCES ANALYZERS\n",
        "!pip install sentencepiece\n",
        "!pip install ekphrasis\n",
        "\n",
        "#CLONE THE BERT GIT REPOSITORY\n",
        "!git clone https://github.com/google-research/bert\n",
        "\n",
        "import os\n",
        "import sys\n",
        "import json\n",
        "import nltk\n",
        "import random\n",
        "import logging\n",
        "import tensorflow as tf\n",
        "import sentencepiece as spm\n",
        "\n",
        "from glob import glob\n",
        "from google.colab import auth, drive\n",
        "from tensorflow.keras.utils import Progbar\n",
        "\n",
        "sys.path.append(\"bert\")\n",
        "from bert import modeling, optimization, tokenization\n",
        "from bert.run_pretraining import input_fn_builder, model_fn_builder\n",
        "\n",
        "#LOGIN TO GOOGLE CLOUD \n",
        "auth.authenticate_user()\n",
        "  \n",
        "# configure logging\n",
        "log = logging.getLogger('tensorflow')\n",
        "log.setLevel(logging.INFO)\n",
        "\n",
        "# create formatter and add it to the handlers\n",
        "formatter = logging.Formatter('%(asctime)s :  %(message)s')\n",
        "sh = logging.StreamHandler()\n",
        "sh.setLevel(logging.INFO)\n",
        "sh.setFormatter(formatter)\n",
        "log.handlers = [sh]\n",
        "\n",
        "#INIZIALIZE TPUV2\n",
        "if 'COLAB_TPU_ADDR' in os.environ:\n",
        "  log.info(\"Using TPU runtime\")\n",
        "  USE_TPU = True\n",
        "  TPU_ADDRESS = 'grpc://' + os.environ['COLAB_TPU_ADDR']\n",
        "\n",
        "  with tf.Session(TPU_ADDRESS) as session:\n",
        "    log.info('TPU address is ' + TPU_ADDRESS)\n",
        "    # Upload credentials to TPU.\n",
        "    with open('/content/adc.json', 'r') as f:\n",
        "      auth_info = json.load(f)\n",
        "    tf.contrib.cloud.configure_gcs(session, credentials=auth_info)\n",
        "    \n",
        "else:\n",
        "  log.warning('Not connected to TPU runtime')\n",
        "  USE_TPU = False"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oRQd4-v0nQqH",
        "colab_type": "text"
      },
      "source": [
        "## Step 2: preprocessing text\n",
        "\n",
        "The raw text data we will use for training purpose contains punсtuation, uppercase letters and non-UTF symbols which we will remove before proceeding. During inference, we will apply the same normalization procedure to new data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iCi2oSdInRkC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from ekphrasis.classes.preprocessor import TextPreProcessor\n",
        "from ekphrasis.classes.tokenizer import SocialTokenizer\n",
        "from ekphrasis.dicts.emoticons import emoticons\n",
        "\n",
        "text_processor = TextPreProcessor (\n",
        "    # terms that will be normalized\n",
        "    normalize=[ 'url' , 'email', 'user', 'percent', 'money', 'phone', 'time', 'date', 'number'] ,\n",
        "    # terms that will be annotated\n",
        "    annotate={\"hashtag\"} ,\n",
        "    fix_html=True ,  # fix HTML tokens\n",
        "\n",
        "    unpack_hashtags=True ,  # perform word segmentation on hashtags\n",
        "\n",
        "    # select a tokenizer. You can use SocialTokenizer, or pass your own\n",
        "    # the tokenizer, should take as input a string and return a list of tokens\n",
        "    tokenizer=SocialTokenizer ( lowercase=True ).tokenize,\n",
        "    dicts = [ emoticons ]\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0gngtEZWqVhY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#LOAD TWITA FROM STORAGE\n",
        "data = open(\"twita/download/twita_200M.txt\", \"r\")\n",
        "f_data = open(\"twita/download/processed_dataset_final.txt\", \"w\")\n",
        "\n",
        "#PROCESS DATA\n",
        "import re\n",
        "for s in data:\n",
        "    s = s.lower()\n",
        "    s = str(\" \".join(text_processor.pre_process_doc(s)))\n",
        "    s = re.sub(r\"[^a-zA-ZÀ-ú</>!?♥♡\\s\\U00010000-\\U0010ffff]\", ' ', s)\n",
        "    s = re.sub(r\"\\s+\", ' ', s)\n",
        "    s = re.sub(r'(\\w)\\1{2,}',r'\\1\\1', s)\n",
        "    s = re.sub ( r'^\\s' , '' , s )\n",
        "    s = re.sub ( r'\\s$' , '' , s )\n",
        "\n",
        "    #print(s)\n",
        "    f_data.write(s+\"\\n\")\n",
        "    f_data.flush()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zVO9EnUwrluQ",
        "colab_type": "text"
      },
      "source": [
        "## Step 3: building the vocabulary\n",
        "\n",
        "For the next step, we will learn a new vocabulary that we will use to represent our dataset. \n",
        "\n",
        "We will be using SentencePiece tokenizer in unigram mode. While it is not directly compatible with BERT, with a small hack we can make it work.\n",
        "\n",
        "SentencePiece requires quite a lot of RAM, so running it on the full dataset in Colab will crash the kernel. To avoid this, we will randomly subsample a fraction of the dataset for building the vocabulary. \n",
        "\n",
        "The typical values for VOC_SIZE are somewhere in between 32000 and 128000. We reserve NUM_PLACEHOLDERS tokens in case one wants to update the vocabulary and fine-tune the model after the pre-training phase is finished."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "18nn6eW_s-fV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "MODEL_PREFIX = \"alberto_lower_case_128\" #@param {type: \"string\"}\n",
        "VOC_SIZE = 128000 #@param {type:\"integer\"}\n",
        "SUBSAMPLE_SIZE = 10000000 #@param {type:\"integer\"}\n",
        "NUM_PLACEHOLDERS = 256 #@param {type:\"integer\"}\n",
        "\n",
        "SPM_COMMAND = ('--input={} --model_prefix={} '\n",
        "               '--vocab_size={} --input_sentence_size={} '\n",
        "               '--shuffle_input_sentence=true ' \n",
        "               '--bos_id=-1 --eos_id=-1').format(\n",
        "               PRC_DATA_FPATH, MODEL_PREFIX, \n",
        "               VOC_SIZE - NUM_PLACEHOLDERS, SUBSAMPLE_SIZE)\n",
        "\n",
        "spm.SentencePieceTrainer.Train(SPM_COMMAND)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BNiLdWXTh9cj",
        "colab_type": "text"
      },
      "source": [
        "As we can see, the WordPiece tokenizer prepends the subwords which occur in the middle of words with '##'. The subwords occurring at the beginning of words are unchanged. If the subword occurs both in the beginning and in the middle of words, both versions (with and without '##') are added to the vocabulary."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nYlBqiv5UD-j",
        "colab_type": "text"
      },
      "source": [
        "SentencePiece has created two files: tokenizer.model and tokenizer.vocab"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yBsURk_h5jw4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def read_sentencepiece_vocab(filepath):\n",
        "  voc = []\n",
        "  with open(filepath, encoding='utf-8') as fi:\n",
        "    for line in fi:\n",
        "      voc.append(line.split(\"\\t\")[0])\n",
        "  # skip the first <unk> token\n",
        "  voc = voc[1:]\n",
        "  return voc"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3YjqHVRpmlKq",
        "colab_type": "text"
      },
      "source": [
        "As we may observe, SentencePiece does quite the opposite to WordPiece. From the [documentation](https://github.com/google/sentencepiece/blob/master/README.md):\n",
        "\n",
        "\n",
        "SentencePiece first escapes the whitespace with a meta-symbol \"▁\" (U+2581) as follows:\n",
        "\n",
        "`Hello▁World`.\n",
        "\n",
        "Then, this text is segmented into small pieces, for example:\n",
        "\n",
        "`[Hello] [▁Wor] [ld] [.]`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OD-28l_0p0PQ",
        "colab_type": "text"
      },
      "source": [
        "Subwords which occur after whitespace (which are also those that most words begin with) are prepended with '▁', while others are unchanged. This excludes subwords which only occur at the beginning of sentences and nowhere else. These cases should be quite rare, however. \n",
        "\n",
        "So, in order to obtain a vocabulary analogous to WordPiece, we need to perform a simple conversion, removing \"▁\" from the tokens that contain it and adding \"##\"  to the ones that don't."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5QJGFjzOMbfN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def parse_sentencepiece_token(token):\n",
        "    if token.startswith(\"▁\"):\n",
        "        return token[1:]\n",
        "    else:\n",
        "        return \"##\" + token"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "64dcVgD98S28",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "bert_vocab = list(map(parse_sentencepiece_token, snt_vocab))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mL9ZR3RN9IMA",
        "colab_type": "text"
      },
      "source": [
        "We also add some special control symbols which are required by the BERT architecture. By convention, we put those at the beginning of the vocabulary."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VdTlXDPL8cHE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "ctrl_symbols = [\"[PAD]\",\"[UNK]\",\"[CLS]\",\"[SEP]\",\"[MASK]\"]\n",
        "bert_vocab = ctrl_symbols + bert_vocab"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ry1jvEr9EIdd",
        "colab_type": "text"
      },
      "source": [
        "We also append some placeholder tokens to the vocabulary. Those are useful if one wishes to update the pre-trained model with new, task-specific tokens. \n",
        "\n",
        "In that case, the placeholder tokens are replaced with new real ones, the pre-training data is re-generated, and the model is fine-tuned on new data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XLYSTil4E0Dm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "bert_vocab += [\"[UNUSED_{}]\".format(i) for i in range(VOC_SIZE - len(bert_vocab))]\n",
        "print(len(bert_vocab))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qJKk_7JI-MtW",
        "colab_type": "text"
      },
      "source": [
        "Finally, we write the obtained vocabulary to file."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8G1jg0cj9Duf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "VOC_FNAME = \"vocabulary_lower_case_128.txt\" #@param {type:\"string\"}\n",
        "\n",
        "with open(VOC_FNAME, \"w\") as fo:\n",
        "  for token in bert_vocab:\n",
        "    fo.write(token+\"\\n\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EwFtStCo__QX",
        "colab_type": "text"
      },
      "source": [
        "## Step 4: generating pre-training data\n",
        "\n",
        "With the vocabulary at hand, we are ready to generate pre-training data for the BERT model. Since our dataset might be quite large, we will split it into shards:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xyN1nI04-uKV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "PRC_DATA_FPATH = \"twita/download/twita_200M.txt\"\n",
        "!mkdir ./shards\n",
        "!split -a 4 -l 256000 -d $PRC_DATA_FPATH ./shards/shard_\n",
        "!ls ./shards/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w-FSq3zNFvLs",
        "colab_type": "text"
      },
      "source": [
        "Before we start generating, we need to set some model-specific parameters.  "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UnZcD0yIBGPd",
        "colab_type": "code",
        "cellView": "both",
        "colab": {}
      },
      "source": [
        "MAX_SEQ_LENGTH = 128 #@param {type:\"integer\"}\n",
        "MASKED_LM_PROB = 0.15 #@param\n",
        "MAX_PREDICTIONS = 20 #@param {type:\"integer\"}\n",
        "DO_LOWER_CASE = True #@param {type:\"boolean\"}\n",
        "PROCESSES = 2 #@param {type:\"integer\"}\n",
        "PRETRAINING_DIR = \"pretraining_data\" #@param {type:\"string\"}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k-MibOBkFam2",
        "colab_type": "text"
      },
      "source": [
        "Now, for each shard we need to call *create_pretraining_data.py* script. To that end, we will employ the  *xargs* command. \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YbZjIeVP0T36",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "XARGS_CMD = (\"ls ./shards/ | \"\n",
        "             \"xargs -n 1 -P {} -I{} \"\n",
        "             \"python3 bert/create_pretraining_data.py \"\n",
        "             \"--input_file=./shards/{} \"\n",
        "             \"--output_file={}/{}.tfrecord \"\n",
        "             \"--vocab_file={} \"\n",
        "             \"--do_lower_case={} \"\n",
        "             \"--max_predictions_per_seq={} \"\n",
        "             \"--max_seq_length={} \"\n",
        "             \"--masked_lm_prob={} \"\n",
        "             \"--random_seed=34 \"\n",
        "             \"--dupe_factor=5\")\n",
        "\n",
        "XARGS_CMD = XARGS_CMD.format(PROCESSES, '{}', '{}', PRETRAINING_DIR, '{}', \n",
        "                             VOC_FNAME, DO_LOWER_CASE, \n",
        "                             MAX_PREDICTIONS, MAX_SEQ_LENGTH, MASKED_LM_PROB)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u5XzaY8xCdiV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tf.gfile.MkDir(PRETRAINING_DIR)\n",
        "!$XARGS_CMD"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZKnW_hVxPoP6",
        "colab_type": "text"
      },
      "source": [
        "## Step 5: setting up persistent storage\n",
        "\n",
        "To preserve our hard-earned assets, we will persist them to Google Cloud Storage. Provided that you have created the GCS bucket, this should be simple.\n",
        "\n",
        "We will create two directories in GCS, one for the data and one for the model.\n",
        "In the model directory, we will put the model vocabulary and configuration file.\n",
        "\n",
        "**Configure your BUCKET_NAME variable here before proceeding, otherwise the model and data will not be saved.**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gDtrt68QQIHs",
        "colab_type": "code",
        "cellView": "both",
        "colab": {}
      },
      "source": [
        "BUCKET_NAME = \"alberto_files\" #@param {type:\"string\"}\n",
        "MODEL_DIR = \"alberto_model\" #@param {type:\"string\"}\n",
        "tf.gfile.MkDir(MODEL_DIR)\n",
        "\n",
        "if not BUCKET_NAME:\n",
        "  log.warning(\"WARNING: BUCKET_NAME is not set. \"\n",
        "              \"You will not be able to train the model.\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YigCcV-hSHVH",
        "colab_type": "text"
      },
      "source": [
        "Below is the sample hyperparameter configuration for BERT-base. Change at your own risk."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SEpSGpUKReKW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# use this for BERT-base\n",
        "\n",
        "bert_base_config = {\n",
        "  \"attention_probs_dropout_prob\": 0.1, \n",
        "  \"directionality\": \"bidi\", \n",
        "  \"hidden_act\": \"gelu\", \n",
        "  \"hidden_dropout_prob\": 0.1, \n",
        "  \"hidden_size\": 768, \n",
        "  \"initializer_range\": 0.02, \n",
        "  \"intermediate_size\": 3072, \n",
        "  \"max_position_embeddings\": 512, \n",
        "  \"num_attention_heads\": 12, \n",
        "  \"num_hidden_layers\": 12, \n",
        "  \"pooler_fc_size\": 768, \n",
        "  \"pooler_num_attention_heads\": 12, \n",
        "  \"pooler_num_fc_layers\": 3, \n",
        "  \"pooler_size_per_head\": 128, \n",
        "  \"pooler_type\": \"first_token_transform\", \n",
        "  \"type_vocab_size\": 2, \n",
        "  \"vocab_size\": VOC_SIZE\n",
        "}\n",
        "\n",
        "with open(\"{}/bert_config.json\".format(MODEL_DIR), \"w\") as fo:\n",
        "  json.dump(bert_base_config, fo, indent=2)\n",
        "  \n",
        "with open(\"{}/{}\".format(MODEL_DIR, VOC_FNAME), \"w\") as fo:\n",
        "  for token in bert_vocab:\n",
        "    fo.write(token+\"\\n\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "txgrEDugRG48",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "if BUCKET_NAME:\n",
        "  !gsutil -m cp -r $MODEL_DIR $PRETRAINING_DIR gs://$BUCKET_NAME"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3gdQEOzhYmSh",
        "colab_type": "text"
      },
      "source": [
        "## Step 6: training the model\n",
        "\n",
        "We are almost ready to begin training our model. <br>\n",
        "IN CASE OF KERNEL STOP, RESTART FROM HERE!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yXAuzsJfYrio",
        "colab_type": "code",
        "cellView": "both",
        "colab": {}
      },
      "source": [
        "BUCKET_NAME = \"alberto_files\" #@param {type:\"string\"}\n",
        "MODEL_DIR = \"alberto_model\" #@param {type:\"string\"}\n",
        "PRETRAINING_DIR = \"pretraining_data\" #@param {type:\"string\"}\n",
        "VOC_FNAME = \"vocabulary_lower_case_128.txt\" #@param {type:\"string\"}\n",
        "\n",
        "# Input data pipeline config\n",
        "TRAIN_BATCH_SIZE = 128 #@param {type:\"integer\"}\n",
        "MAX_PREDICTIONS = 20 #@param {type:\"integer\"}\n",
        "MAX_SEQ_LENGTH = 128 #@param {type:\"integer\"}\n",
        "MASKED_LM_PROB = 0.15 #@param\n",
        "\n",
        "# Training procedure config\n",
        "EVAL_BATCH_SIZE = 64\n",
        "LEARNING_RATE = 2e-5\n",
        "TRAIN_STEPS = 1000000 #@param {type:\"integer\"}\n",
        "SAVE_CHECKPOINTS_STEPS = 2500 #@param {type:\"integer\"}\n",
        "NUM_TPU_CORES = 8\n",
        "\n",
        "if BUCKET_NAME:\n",
        "  BUCKET_PATH = \"gs://{}\".format(BUCKET_NAME)\n",
        "else:\n",
        "  BUCKET_PATH = \".\"\n",
        "\n",
        "BERT_GCS_DIR = \"{}/{}\".format(BUCKET_PATH, MODEL_DIR)\n",
        "DATA_GCS_DIR = \"{}/{}\".format(BUCKET_PATH, PRETRAINING_DIR)\n",
        "\n",
        "VOCAB_FILE = os.path.join(BERT_GCS_DIR, VOC_FNAME)\n",
        "CONFIG_FILE = os.path.join(BERT_GCS_DIR, \"bert_config.json\")\n",
        "\n",
        "INIT_CHECKPOINT = tf.train.latest_checkpoint(BERT_GCS_DIR)\n",
        "\n",
        "bert_config = modeling.BertConfig.from_json_file(CONFIG_FILE)\n",
        "input_files = tf.gfile.Glob(os.path.join(DATA_GCS_DIR,'*tfrecord'))\n",
        "\n",
        "log.info(\"Using checkpoint: {}\".format(INIT_CHECKPOINT))\n",
        "log.info(\"Using {} data shards\".format(len(input_files)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MQTFjITdd53F",
        "colab_type": "text"
      },
      "source": [
        "Prepare the training run configuration, build the estimator and input function."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CMahsqUnZ55z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model_fn = model_fn_builder(\n",
        "      bert_config=bert_config,\n",
        "      init_checkpoint=INIT_CHECKPOINT,\n",
        "      learning_rate=LEARNING_RATE,\n",
        "      num_train_steps=TRAIN_STEPS,\n",
        "      num_warmup_steps=10,\n",
        "      use_tpu=USE_TPU,\n",
        "      use_one_hot_embeddings=True)\n",
        "\n",
        "tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)\n",
        "\n",
        "run_config = tf.contrib.tpu.RunConfig(\n",
        "    cluster=tpu_cluster_resolver,\n",
        "    model_dir=BERT_GCS_DIR,\n",
        "    save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS,\n",
        "    tpu_config=tf.contrib.tpu.TPUConfig(\n",
        "        iterations_per_loop=SAVE_CHECKPOINTS_STEPS,\n",
        "        num_shards=NUM_TPU_CORES,\n",
        "        per_host_input_for_training=tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2))\n",
        "\n",
        "estimator = tf.contrib.tpu.TPUEstimator(\n",
        "    use_tpu=USE_TPU,\n",
        "    model_fn=model_fn,\n",
        "    config=run_config,\n",
        "    train_batch_size=TRAIN_BATCH_SIZE,\n",
        "    eval_batch_size=EVAL_BATCH_SIZE)\n",
        "  \n",
        "train_input_fn = input_fn_builder(\n",
        "        input_files=input_files,\n",
        "        max_seq_length=MAX_SEQ_LENGTH,\n",
        "        max_predictions_per_seq=MAX_PREDICTIONS,\n",
        "        is_training=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JrCuEbr6dv8U",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#RUN THE TRAINING\n",
        "estimator.train(input_fn=train_input_fn, max_steps=TRAIN_STEPS)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N_OeXod-fHMT",
        "colab_type": "text"
      },
      "source": [
        "Training the model with 1 million steps will take ~50 hours. \n",
        "\n",
        "In case the kernel is restarted, you may always continue training from the latest checkpoint. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "77ZIAAATfzdF",
        "colab_type": "text"
      },
      "source": [
        "This concludes the guide to pre-training AlBERTo from scratch on a cloud TPU."
      ]
    }
  ]
}