models/NeMo/slm/slm_pretraining_sft.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building Your Own Small Language Model from Scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Author: [Yoshi Suhara](https://github.com/suhara)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8CYLKLLkhgLV"
   },
   "source": [
    "*_NOTE_**: This notebook has been tested in the following environment:\n",
    "- [NeMo Framework](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo) (`nvcr.io/nvidia/nemo:24.01.framework`) \n",
    "\n",
    "## Overview\n",
    "\n",
    "Small Language Models (SLMs) are a type of LLM family with fewer number of parameters (~3B) so that they can be easily used with consumer GPUs (e.g., NVIDIA RTX GPUs) and embedded systems (e.g., NVIDIA Jetson Orin). Recently, many open-source SLMs such as [phi2](https://huggingface.co/microsoft/phi-2), [Gemma](https://huggingface.co/google/gemma-2b), and [TinyLlama](https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0) have been released. Training SLMs does not require as much computational resources as larger models. This helps us consider training our own SLM for specific purposes. \n",
    "\n",
    "In this tutorial, you will learn how to train an SLM from scratch using pre-training and supervised fine-tuning techniques. You'll also learn how to generate text using your SLM. The exact number of GPUs needed will depend on the amount of training data and the model architecture you use. This notebook was tested on 1x A100/H100-80GB GPU.\n",
    "\n",
    "\n",
    "### Objective\n",
    "\n",
    "In this notebook, you will be learning how to:\n",
    "- Pre-train an SLM from scratch\n",
    "- Generate responses from the trained SLM\n",
    "- Fine-tune the pre-trained SLM\n",
    "\n",
    "This tutorial uses the following Nvidia services and resources:\n",
    "- [NeMo Framework](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo) \n",
    "- [NeMo Toolkit](https://github.com/NVIDIA/NeMo)\n",
    "- [NeMo Aligner](https://github.com/NVIDIA/NeMo-Aligner)\n",
    "\n",
    "This tutorial uses the following opensource components and resources:\n",
    "- [Hugging Face Tokenizers](https://huggingface.co/docs/tokenizers/en/indexhttps://huggingface.co/docs/tokenizers/en/index)\n",
    "- [Hugging Face Models](https://huggingface.co/modelss)\n",
    "- [Hugging Face Datasets](https://huggingface.co/datasets)\n",
    "- [cosmopedia-100k](https://huggingface.co/datasets/HuggingFaceTB/cosmopedia-100k) - Hugging Face\n",
    "- [databricks-dolly-15k](https://huggingface.co/datasets/databricks/databricks-dolly-15k)\n",
    "\n",
    "\n",
    "The steps performed include:\n",
    "- Step 1: Downloading Pre-training Data and Tokenizer\n",
    "- Step 2: Preprocessing Pre-training Data\n",
    "- Step 3: Pre-training SLM from Scratch\n",
    "- Step 4: Generating Text using Your SLM\n",
    "- Step 5: Improving Trained SLM to Talk Better (Optional)\n",
    "\n",
    "\n",
    "## Background: Training Techniques for SLMs\n",
    "\n",
    "\n",
    "### Pre-training\n",
    "\n",
    "**Pre-training** is an important step for LLMs to learn knowledge from a large amount of textual data. The exact same techniques as LLMs can be used to train SLMs. A recent trend is to use high-quality data including synthetically generated by larger LMs. For example, Microsoft's phi series claim that they used a small amount of \"text-book quality\" pre-training data generated by GPT-3.5/4 to train the models [paper](https://arxiv.org/abs/2306.11644).\n",
    "\n",
    "In this tutorial, we will use a subset of [Cosmopedia](https://huggingface.co/datasets/HuggingFaceTB/cosmopedia), which is generated by Mixtral-8x7B model---one of the best open-source LLMs.\n",
    "\n",
    "\n",
    "### Supervised Fine-tuning / Alignment (Optional) \n",
    "\n",
    "**Supervised Fine-Tuning (SFT)** is the process of training all of a model’s parameters on supervised data of inputs and outputs to help acquire **instruction-following** capability. This notebook describes the steps involved in fine-tuning your custom SLM. Another type of training is to use human preference data to teach the model to output generations that would be preferred by human. This type of alignment is often called **Reinforcement Learning from Human Feedback (RLHF)**. [NeMo-Aligner](https://github.com/NVIDIA/NeMo-Aligner) implements major techniques including PPO, DPO and NVIDIA's preference-based aligment techninque SteerLM.\n",
    "\n",
    "\n",
    "## Before you begin\n",
    "\n",
    "### Getting NeMo Framework\n",
    "\n",
    "[NeMo Framework](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo)  is a generative AI framework built for researchers and PyTorch developers working on large language models (LLMs), multimodal models (MM), automatic speech recognition (ASR), and text-to-speech synthesis (TTS). The primary objective of NeMo is to provide a scalable framework for researchers and developers from industry and academia to more easily implement and design new generative AI models by being able to leverage existing code and pretrained models.\n",
    "\n",
    "If you haven't already, you can pull a container that includes the version of NeMo Framework and all dependencies needed for this notebook with the following:\n",
    "\n",
    "```bash\n",
    "docker pull nvcr.io/nvidia/nemo:24.01.framework\n",
    "```\n",
    "\n",
    "The best way to run this notebook is from within the container. You can do that by launching the container with the following command\n",
    "\n",
    "```bash\n",
    "docker run -it --rm --gpus all --ipc=host --network host -v $(pwd):/workspace nvcr.io/nvidia/nemo:24.01.framework\n",
    "```\n",
    "\n",
    "Then, from within the container, start the jupyter server with\n",
    "\n",
    "```bash\n",
    "jupyter lab --no-browser --port=8080 --allow-root --ip 0.0.0.0\n",
    "```\n",
    "\n",
    "Alternatively, you can combine the two steps and can directly launch a Jupyter Lab server with the following Docker command.\n",
    "\n",
    "```bash\n",
    "docker run -it --rm --gpus all --ipc=host --network host -v $(pwd):/workspace nvcr.io/nvidia/nemo:24.01.framework \"jupyter lab --no-browser --port=8080 --ip 0.0.0.0\"\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Installation\n",
    "\n",
    "Install the following packages required to execute this notebook. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --quiet ipywidgets\n",
    "!pip install --upgrade --quiet huggingface_hub"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Restart the kernel after installing packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "\n",
    "app = IPython.Application.instance()\n",
    "app.kernel.do_shutdown(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import random\n",
    "import requests\n",
    "\n",
    "from datasets import load_dataset\n",
    "import numpy as np\n",
    "from transformers import AutoTokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1W7nSFP-hgLX"
   },
   "source": [
    "## Step 1: Download Pre-training Data and Tokenizer\n",
    "\n",
    "Collecting high-quality data and blending them to create pre-training data is crucial for the quality of SLMs. It is getting more common to use powerful LLMs such as GPT-3.5/4 and/or Mixtral 8x7B to generate synthetic data to train SLMs. Synthetic data generation is active research area and different types of recipes have been created.\n",
    "\n",
    "In this example, we will use a small portion of [Cosmopedia](https://huggingface.co/datasets/HuggingFaceTB/cosmopedia) to quickly walk through the training process with the NeMo Framework.\n",
    "\n",
    "\n",
    "### Cosmopedia 100k data\n",
    "\n",
    "Run the following command to download [Cosmopedia-100k](https://huggingface.co/datasets/HuggingFaceTB/cosmopedia-100k)---a subset of Cosmopedia."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 318,
     "referenced_widgets": [
      "2968599b0ac047ee8abb591068b0fc51",
      "dbe565396e3e45a0ad5cb59f4cbc43e2",
      "fe6a62e2e0e44cfd8ce8fdea0406a8aa",
      "e3925a733c924a439a2454ffa7ca2363",
      "59dec006ca154f0e833bf02fb1886626",
      "7ea553939bf741ee84174c935e3d6e6e",
      "f713d501a82d4fa0960318ad96ffcfa2",
      "70fd0842944d44a08357823c6bb60c3e",
      "98f4bf45612a41f2b50ff7b903ca00a6",
      "e38684570b8441d0a460ece94645cb96",
      "6c9bd7937a334be1bc71be7bcfeb831d",
      "48ff4d9ee0284a278715b449e112f1fb",
      "3b7d28f07f054038a85d77810bbc3ff5",
      "c6fa9476c7824f75ba78560173cf2c9a",
      "7da6f7f14ce24986baa2ef47adbabbe3",
      "e437ed2842e9450ca5b4cf5561e5996c",
      "6e39f41f44564a3682fc864fbc57233e",
      "25ce70abc10c4f48a916cdfe0cf4568c",
      "57ce5c49a1bd404393ca3f42ebad8e6b",
      "2e0c5716c7fd40cd854896ee086150e7",
      "353667831390430ca9e1026132101bbd",
      "4f66ede3ff6c451385a25042631f26c1",
      "71a1d619be59486da7efa887d586b383",
      "7d943d6fff8b42a4b405ad02df20d6cf",
      "f48d53599c5e4a718607969904835a49",
      "ab1c423de1e64e7d8d02eea888161f3f",
      "662d1617c46c4640a9eb286307b9f855",
      "e92a7766829949559d9030153a0a7256",
      "59df1fbe27e04924900c585a635e238e",
      "a0fd28b6ff574dd38d12339333336ce6",
      "7bb077eb9ccd42f28bcc1d245f1904cb",
      "ad8636fc0c8e4e128aa8efa19a0e8623",
      "9341dca351074f4db14a6b76cf8c722b",
      "d0d3e4b8f28c43aaa9d505aab38c5b1b",
      "bda815dcbb2a4e6b9ff535cdb9debaf2",
      "1b1b4f0f0e0e4757a86bc30919f244d4",
      "8320287538c040d091caa681ea745992",
      "0104e084968f4690a400894611e937fa",
      "6ff7f6b154584ad09a31ce9c779f6a73",
      "3eaa1290e8034ba2bf3aee98555f2218",
      "5730cbe985484ca19f666d475b3561be",
      "989c2e90ec5a46c498fc897bc540eed8",
      "fba8de14bde342b0ad29476924fb601f",
      "cb418b13f15c42fa8fa63c9587aaf7bf",
      "3b032d2dbf244641ae0a49fecff36abe",
      "de5a5da9c9514e89bbb9ed2ce0cdebb3",
      "85ffda2602e745faace8e22c889cb0e9",
      "3a29e607fb7e4591bd35267dde360667",
      "18fb08af07c241028002709d7a52a125",
      "bbd037941f6845bab9b03be741695ad7",
      "0efe12220814483aa2266bab884fd838",
      "154ea9b75c47405fb2b5cbca23444939",
      "0010996f02ed4662a07291713924b12b",
      "b7e935d02f204f4a8eff7534188bc6c5",
      "1529be25ce264c28b794cb75b2975a6a"
     ]
    },
    "id": "xyRm42q7h1bZ",
    "outputId": "fe3fdba2-b97d-4251-b73c-bf69dcd74fc0"
   },
   "outputs": [],
   "source": [
    "# Download Cosmopedia 100k data\n",
    "TEXT_FILE = \"cosmopedia-100k.jsonl\"\n",
    "\n",
    "if not os.path.exists(TEXT_FILE):\n",
    "    dataset = load_dataset(\"HuggingFaceTB/cosmopedia-100k\", split=\"train\")\n",
    "    dataset.to_json(TEXT_FILE, orient=\"records\", lines=True)\n",
    "else:\n",
    "    print(\"Dataset already downloaded. Skip.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tokenizer\n",
    "\n",
    "You will need a pre-trained tokenizer to preprocess the pre-training dataset. In this tutorial, we will use a pre-trained Byte Pair Encoding (BPE) tokenizer trained for GPT-2. More details about BPE and other types of tokenizers can be found in [Hugging Face's documentation](https://huggingface.co/docs/transformers/en/tokenizer_summary). \n",
    "\n",
    "The pre-trained tokenizer needs the following two files \n",
    "\n",
    "- `vocab.json` contains mapping between tokens and token IDs.\n",
    "- `merges.txt` stores the rules on what tokens will be merged at each iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "SsFl49-fhgLY",
    "outputId": "ab0c768b-17dd-47bf-cc58-b74abd8ea37b"
   },
   "outputs": [],
   "source": [
    "VOCAB_FILE = \"vocab.json\"\n",
    "MERGE_FILE = \"merges.txt\"\n",
    "\n",
    "# Download the tokenizer files\n",
    "!wget -nc https://huggingface.co/openai-community/gpt2/resolve/main/vocab.json -O {VOCAB_FILE}\n",
    "!wget -nc https://huggingface.co/openai-community/gpt2/resolve/main/merges.txt -O {MERGE_FILE}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jpBdsOIIhgLY",
    "outputId": "5a3332c2-99de-48ed-8c44-11b3c660ce35"
   },
   "outputs": [],
   "source": [
    "# Confirm if NeMo directory, merges.txt, and vocab.json exist\n",
    "!ls"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Preprocessing Pre-training Data\n",
    "\n",
    "Next, you will need to tokenize the pre-training data using the tokenizer that you just downloaded to convert the pre-training data into the format that is ready to use for training. More specifically, we want each data to be converted into a sequence of token IDs, which is the input format for the LM.\n",
    "\n",
    "[The NeMo Toolkit](https://github.com/NVIDIA/NeMo/tree/main) has [a preprocessing script](https://github.com/NVIDIA/NeMo/blob/main/scripts/nlp_language_modeling/preprocess_data_for_megatron.py) for this purpose. The script will create two binary files (.idx and .bin), which can be efficiently loaded during the pre-training step.\n",
    "\n",
    "Run the command below to preprocess the cosmopedia data. It will take a few minutes or more depending on the environment. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "tq0-DyOFhgLY",
    "outputId": "11dce117-d62a-4d57-823b-261b5cae9d53"
   },
   "outputs": [],
   "source": [
    "DATA_PREFIX = \"cosmopedia-100k\"\n",
    "\n",
    "idx_file = \"{}_text_document.idx\".format(DATA_PREFIX)\n",
    "bin_file = \"{}_text_document.bin\".format(DATA_PREFIX)\n",
    "if not (os.path.exists(idx_file) and os.path.exists(bin_file)):\n",
    "    !python /opt/NeMo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py \\\n",
    "        --input={TEXT_FILE} \\\n",
    "        --json-keys=text \\\n",
    "        --tokenizer-library=megatron \\\n",
    "        --tokenizer-type=GPT2BPETokenizer \\\n",
    "        --dataset-impl=mmap \\\n",
    "        --merge-file={MERGE_FILE} \\\n",
    "        --vocab-file={VOCAB_FILE} \\\n",
    "        --output-prefix={DATA_PREFIX} \\\n",
    "        --append-eod \\\n",
    "        --workers=4\n",
    "else:\n",
    "    print(\"Files already exist. Skip.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After the preprocess, confirm if the following files were created\n",
    "- `cosmopedia-100k_text_document.idx`\n",
    "- `cosmopedia-100k_text_document.bin`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fxPFztekKnvg",
    "outputId": "507dab8c-995a-4029-d85b-93ba5fd450a7"
   },
   "outputs": [],
   "source": [
    "!ls"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Pre-training SLM from Scratch\n",
    "\n",
    "Now we are ready to train a custom SLM using the pre-training data we just preprocessed. To start training, we need to configure training settings including the model architecture. In this tutorial, we train a 200M GPT model. The training process might take more than 30 minutes for 1000 steps. You can use a larger value for `MAX_STEPS` if you'd like the model to learn more. At the end of training, the script will create a NeMo checkpoint `results/megatron_gpt/checkpoints/megatron_gpt.nemo`, which we weill use for text generation later.\n",
    "\n",
    "### Model Architecture / Optimization Settings (Advanced)\n",
    "\n",
    "Here are tips for deciding SLM architecture design. The following descriptions are for advanced users who want to have a better understanding in how those hyperparameters determine the final model architecture.\n",
    "\n",
    "#### Architecture configuration\n",
    "\n",
    "- Increasing the depth of the model (= a higher value for `model.num_layers`) is a good way to have a larger model with fewer risks.\n",
    "- Increasing the value for `model.hidden_size` may make the model take longer time to converge.\n",
    "- Make sure `model.ffn_hidden_size` is always larger than `model.hidden_size` (`4 * model.hidden_size` is a commonly used setting)\n",
    "\n",
    "Below is the default setting for this tutorial.\n",
    "\n",
    "- `model.num_layers: 12`\n",
    "- `model.hidden_size: 768`\n",
    "- `model.ffn_hidden_size: 3072`\n",
    "- `model.num_attention_heads: 12`\n",
    "\n",
    "#### Optimization settings (Advanced)\n",
    "\n",
    "- A higher learning rate (LR) helps the model converge faster with a risk of loss divergence.\n",
    "- If you observe unstable training/validation loss values over steps, consider decreasing `model.optim.lr`.\n",
    "- `model.optim.sched.min_lr` should be set to 1/10 of `model.optim.lr`.\n",
    "\n",
    "Below is the default setting for this tutorial.\n",
    "\n",
    "- `model.optim.lr: 2e-3`\n",
    "- `model.optim.sched.min_lr: 2e-4`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Fav1wQGthgLZ",
    "outputId": "16b41f68-6324-4ec4-8aed-9f7bfbe67caa"
   },
   "outputs": [],
   "source": [
    "# ======================================\n",
    "MAX_STEPS = 1000  # Change to higher values such as 2124 (= 1 epoch) if you'd like to get a better model :)\n",
    "NUM_LAYERS = 12\n",
    "NUM_GPUS = 1\n",
    "HIDDEN_SIZE = 768\n",
    "FFN_HIDDEN_SIZE = 3072\n",
    "NUM_ATTENTION_HEADS = 12\n",
    "SEQ_LENGTH = 1024\n",
    "\n",
    "MAX_LR = 2e-3\n",
    "MIN_LR = 2e-4\n",
    "WARMUP_STEPS = 100\n",
    "\n",
    "TENSOR_MP_SIZE = 1\n",
    "PIPELINE_MP_SIZE = 1\n",
    "INPUT_DATA_PREFIX = \"cosmopedia-100k_text_document\"\n",
    "INDEX_MAPPING_DIR = \"index_cache\"\n",
    "EXP_DIR = \"results\"\n",
    "# =======================================    \n",
    "    \n",
    "!python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py \\\n",
    "        trainer.devices=1 \\\n",
    "        trainer.accelerator=gpu \\\n",
    "        trainer.log_every_n_steps=100 \\\n",
    "        trainer.val_check_interval=100 \\\n",
    "        trainer.max_steps={MAX_STEPS} \\\n",
    "        trainer.precision=16 \\\n",
    "        trainer.gradient_clip_val=1.0 \\\n",
    "        exp_manager.exp_dir={EXP_DIR} \\\n",
    "        model.global_batch_size=64 \\\n",
    "        model.tensor_model_parallel_size={TENSOR_MP_SIZE} \\\n",
    "        model.pipeline_model_parallel_size={PIPELINE_MP_SIZE} \\\n",
    "        model.optim.name=fused_adam \\\n",
    "        model.optim.lr={MAX_LR} \\\n",
    "        model.optim.sched.warmup_steps={WARMUP_STEPS} \\\n",
    "        model.optim.sched.min_lr={MIN_LR} \\\n",
    "        model.optim.sched.constant_steps=0 \\\n",
    "        model.max_position_embeddings={SEQ_LENGTH} \\\n",
    "        model.encoder_seq_length={SEQ_LENGTH} \\\n",
    "        model.data.seq_length={SEQ_LENGTH} \\\n",
    "        model.tokenizer.type=GPT2BPETokenizer \\\n",
    "        model.tokenizer.library=megatron \\\n",
    "        model.tokenizer.vocab_file={VOCAB_FILE} \\\n",
    "        model.tokenizer.merge_file={MERGE_FILE} \\\n",
    "        model.data.eod_mask_loss=True \\\n",
    "        model.data.splits_string=\\'98,1,1\\' \\\n",
    "        model.num_layers={NUM_LAYERS} \\\n",
    "        model.hidden_size={HIDDEN_SIZE} \\\n",
    "        model.num_attention_heads={NUM_ATTENTION_HEADS} \\\n",
    "        model.ffn_hidden_size={FFN_HIDDEN_SIZE} \\\n",
    "        model.data.data_prefix=[{INPUT_DATA_PREFIX}] \\\n",
    "        model.data.index_mapping_dir={INDEX_MAPPING_DIR} \\\n",
    "        exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### After training\n",
    "\n",
    "After the training, please confirm if the following .nemo checkpoint has been created.\n",
    "- `results/megatron_gpt/checkpoints/megatron_gpt.nemo`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3lDEtaIyhgLZ",
    "outputId": "57196d10-c4da-4416-db6b-9127e2f6ef0c"
   },
   "outputs": [],
   "source": [
    "# Check if `results/megatron_gpt/checkpoints/megatron_gpt.nemo` has been created.\n",
    "!ls results/megatron_gpt/checkpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Generating Text using Your SLM\n",
    "\n",
    "Let's play with this freshly  made SLM!\n",
    "\n",
    "Open [megatron_gpt_eval_server.ipynb](./megatron_gpt_eval_server.ipynb) in another tab on the Jupyter server to launch a text generating server. This step has to be done after Step 3. In case the port is in use, change the port number accordingly. \n",
    "\n",
    "- [megatron_gpt_eval_server.ipynb](./megatron_gpt_eval_server.ipynb)\n",
    "\n",
    "```\n",
    "!python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py \\\n",
    "  gpt_model_file=results/megatron_gpt/checkpoints/megatron_gpt.nemo \\\n",
    "  server=True \\\n",
    "  port=55555\n",
    "```\n",
    "\n",
    "You can communicate with the server with the following Python script to get a response to prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 686
    },
    "id": "knaeLXJMhgLZ",
    "outputId": "34fa632a-383f-45ec-93db-94135c22f0a5"
   },
   "outputs": [],
   "source": [
    "class MegatronGPTEvalClient:\n",
    "    def __init__(self,\n",
    "                 batch_size: int = 1,\n",
    "                 port_num: int = 55555,\n",
    "                 headers = {\"Content-Type\": \"application/json\"}):\n",
    "        self.batch_size = batch_size\n",
    "        self.port_num = port_num\n",
    "        self.headers = headers\n",
    "\n",
    "    def generate(self, text: str):\n",
    "        data = {\"sentences\": [text] * self.batch_size,\n",
    "                \"tokens_to_generate\": 32,\n",
    "                \"temperature\": 1.0,\n",
    "                \"add_BOS\": True,\n",
    "                \"top_k\": 0,\n",
    "                \"top_p\": 0.9,\n",
    "                \"greedy\": False,\n",
    "                \"all_probs\": False,\n",
    "                \"repetition_penalty\": 1.2,\n",
    "                \"min_tokens_to_generate\": 2}\n",
    "        resp = requests.put('http://localhost:{}/generate'.format(self.port_num),\n",
    "                            data=json.dumps(data),\n",
    "                            headers=self.headers)\n",
    "        sentences = resp.json()['sentences']\n",
    "        generation = sentences[0]\n",
    "        return generation[len(text):].lstrip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = MegatronGPTEvalClient()\n",
    "client.generate(\"How are you?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-rzS9ikzhgLZ",
    "outputId": "53640ff7-72ee-42ad-f77e-362042cfc43c"
   },
   "source": [
    "## Step 5: Improving Trained SLM to Talk Better (Optional)\n",
    "\n",
    "Are you happy with the responses that the SLM generated? Probably not. You must have got something that looks like English but the SLM did not answer your question and/or follow your instruction. A trial run got a response like below\n",
    "\n",
    "\n",
    "```\n",
    ">>> client.generate(\"How are you?\")\n",
    "\n",
    "There's what it's all like complex AI is how data theft can help organizations save inventory and continuous products that make they need to their chances.\n",
    "```\n",
    "\n",
    "Why did it happen? The pre-training step teaches the SLM how to generate natural English sentences but may not always help the SLM learn to answer questions and follow instructions, especially when only a limited amount of data is used. It seems that the SLM still does not have a sufficient instruction-following capability.\n",
    "\n",
    "As mentioned above, **Supervised Fine-Tuning (SFT)** (or often called as **Instruction Tuning**) is a solution for this. With SFT, the model can learn to follow instructions and have a better conversational ability. We will use [NeMo-Aligner](https://github.com/NVIDIA/NeMo-Aligner), which implements major alignment techniques including SFT and Reinforcement Learning from Human Feedback (RLHF). Please see [the NeMo-Aligner official GitHub repo](https://github.com/NVIDIA/NeMo-Aligner) for more details.\n",
    "\n",
    "\n",
    "In this optional step, we will further train the SLM by Instruction Tuning to talk better. Please make sure if you've completed the pre-training step and you have a .nemo checkpoint file (`megatron_gpt.nemo`).\n",
    "\n",
    "### Step 5-1: Download Dolly 15k dataset\n",
    "\n",
    "We need to download training data for Insturction Tuning. Typical Insutruction Tuning data contain pairs of **prompts** and **human generated responses**, which help the model to learn how to follow instructions. \n",
    "\n",
    "In this example, we will use DataBricks' [Dolly 15k dataset](https://huggingface.co/datasets/databricks/databricks-dolly-15k)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "m_fgYJQohgLa"
   },
   "outputs": [],
   "source": [
    "!wget -nc https://huggingface.co/datasets/databricks/databricks-dolly-15k/resolve/main/databricks-dolly-15k.jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Preparation: Format Conversion\n",
    "\n",
    "Let's take a look at one example in the JSONL file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!head -1 databricks-dolly-15k.jsonl | jq ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since NeMo-Aligner's SFT script expects the input data to follow the format below, we need to concatenate `instruction` and `context` into a single `input` and use `response` as `output` in the converted file.\n",
    "\n",
    "```\n",
    "{\"input\": str,\n",
    " \"output\": str,\n",
    " \"category\": str}\n",
    "```\n",
    "\n",
    "Note that randomly shuffling the order of `instruction` and `context` is a commonly used technique that helps stabilize the model. Run the following Python script to convert into the NeMo-Aligner format. This script also removes examples that are longer than the context length (the default value for this notebook example is `512`.) Update the context length value accordingly if you'd like to use longer context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"openai-community/gpt2\")\n",
    "\n",
    "max_token = 512\n",
    "input_file = \"databricks-dolly-15k.jsonl\"\n",
    "output_file = input_file.replace(\".jsonl\", \"-output.jsonl\")\n",
    "with open(input_file, \"r\") as fin, open(output_file, \"w\") as fout:\n",
    "    for line in fin:\n",
    "        # Read JSONL line in original format\n",
    "        line = json.loads(line)\n",
    "        context = line[\"context\"].strip()\n",
    "        instruction = line[\"instruction\"].strip()\n",
    "        if context == \"\" or instruction == \"\":\n",
    "            continue\n",
    "        output = line[\"response\"]\n",
    "        # Randomly shuffle to make context/instruction comes first \n",
    "        fmt = random.choice([r'f\"{context}\\n\\n{instruction}\"',\n",
    "                             r'f\"{instruction}\\n\\n{context}\"'])\n",
    "        input = eval(fmt)\n",
    "        if len(tokenizer.encode(\"{} {}\".format(input, output))) < max_token:\n",
    "            fout.write(\"{}\\n\".format(\n",
    "                json.dumps({\"input\": input, \"output\": output, \"category\": line[\"category\"]})))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a quick look at how many examples were filtered out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wc -l databricks-dolly-15k.jsonl \n",
    "!wc -l databricks-dolly-15k-output.jsonl "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Only 3762 examples were remained after filtering. We should definitely consider increasing the context length for the next trial. Let's keep moving forward with this option for now.\n",
    "\n",
    "Confirm if the converted file follows the format described above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!head -1 databricks-dolly-15k-output.jsonl | jq ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5-2: Fine-Tuning SLM with NeMo-Aligner\n",
    "\n",
    "We are ready to use a subset of the Dolly 15k dataset to fine-tune the SLM to make it a better model. :)\n",
    "This training step will consume 3762 examples and it should take about 10+ minutes. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAIN_FILEPATH = \"databricks-dolly-15k-output.jsonl\"\n",
    "VALID_FILEPATH = \"databricks-dolly-15k-output.jsonl\"\n",
    "NEMO_CKPT_PATH = \"results/megatron_gpt/checkpoints/megatron_gpt.nemo\"\n",
    "EXP_DIR = \"results\"\n",
    "\n",
    "!python /opt/NeMo-Aligner/examples/nlp/gpt/train_gpt_sft.py \\\n",
    "   trainer.precision=bf16 \\\n",
    "   trainer.num_nodes=1 \\\n",
    "   trainer.devices=1 \\\n",
    "   trainer.sft.max_epochs=1 \\\n",
    "   trainer.sft.max_steps=-1 \\\n",
    "   trainer.sft.val_check_interval=20 \\\n",
    "   trainer.sft.limit_val_batches=10 \\\n",
    "   model.tensor_model_parallel_size=1 \\\n",
    "   model.pipeline_model_parallel_size=1 \\\n",
    "   model.megatron_amp_O2=False \\\n",
    "   model.restore_from_path={NEMO_CKPT_PATH} \\\n",
    "   model.optim.lr=5e-6 \\\n",
    "   model.answer_only_loss=True \\\n",
    "   ++model.bias_activation_fusion=true \\\n",
    "   model.data.num_workers=8 \\\n",
    "   model.data.train_ds.micro_batch_size=1 \\\n",
    "   model.data.train_ds.global_batch_size=64 \\\n",
    "   model.data.train_ds.file_path={TRAIN_FILEPATH} \\\n",
    "   model.data.train_ds.drop_last=True \\\n",
    "   model.data.validation_ds.micro_batch_size=1 \\\n",
    "   model.data.validation_ds.global_batch_size=64 \\\n",
    "   model.data.validation_ds.file_path={VALID_FILEPATH} \\\n",
    "   model.data.validation_ds.drop_last=True \\\n",
    "   exp_manager.exp_dir={EXP_DIR} \\\n",
    "   exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if megatron_gpt_sft.nemo has been created\n",
    "!ls results/megatron_gpt_sft/checkpoints/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5-3: Generating Responses from Fine-tuned SLM\n",
    "\n",
    "Open the tab for [megatron_gpt_eval_server.ipynb](./megatron_gpt_eval_server.ipynb) again (open a new tab if you haven't) to launch a text generating server. You will need to do an additional fix on the .nemo checkpoint for this time.\n",
    "\n",
    "- [megatron_gpt_eval_server.ipynb](./megatron_gpt_eval_server.ipynb)\n",
    "\n",
    "```\n",
    "!python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py \\\n",
    "  gpt_model_file=results/megatron_gpt_sft/checkpoints/megatron_gpt_sft.nemo \\\n",
    "  server=True \\\n",
    "  port=55555\n",
    "```\n",
    "\n",
    "You can communicate with the server with the following Python script to get a response to prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = MegatronGPTEvalClient()\n",
    "client.generate(\"How are you?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "I hope you got better responses with the fine-tuned SLM! However, just fine-tuning with 3762 examples may not be sufficient.\n",
    "\n",
    "Now that you know how to pre-traini and fine-tune your own SLM, you should be able to explore different model architectures as well as different/more training data to make smarter SLMs using the NeMo Framework!"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "machine_shape": "hm",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "0010996f02ed4662a07291713924b12b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "0104e084968f4690a400894611e937fa": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "0efe12220814483aa2266bab884fd838": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "1529be25ce264c28b794cb75b2975a6a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "154ea9b75c47405fb2b5cbca23444939": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "18fb08af07c241028002709d7a52a125": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "1b1b4f0f0e0e4757a86bc30919f244d4": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5730cbe985484ca19f666d475b3561be",
      "max": 100000,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_989c2e90ec5a46c498fc897bc540eed8",
      "value": 100000
     }
    },
    "25ce70abc10c4f48a916cdfe0cf4568c": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "2968599b0ac047ee8abb591068b0fc51": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_dbe565396e3e45a0ad5cb59f4cbc43e2",
       "IPY_MODEL_fe6a62e2e0e44cfd8ce8fdea0406a8aa",
       "IPY_MODEL_e3925a733c924a439a2454ffa7ca2363"
      ],
      "layout": "IPY_MODEL_59dec006ca154f0e833bf02fb1886626"
     }
    },
    "2e0c5716c7fd40cd854896ee086150e7": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "353667831390430ca9e1026132101bbd": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "3a29e607fb7e4591bd35267dde360667": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_b7e935d02f204f4a8eff7534188bc6c5",
      "placeholder": "​",
      "style": "IPY_MODEL_1529be25ce264c28b794cb75b2975a6a",
      "value": " 100/100 [00:06&lt;00:00, 14.78ba/s]"
     }
    },
    "3b032d2dbf244641ae0a49fecff36abe": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_de5a5da9c9514e89bbb9ed2ce0cdebb3",
       "IPY_MODEL_85ffda2602e745faace8e22c889cb0e9",
       "IPY_MODEL_3a29e607fb7e4591bd35267dde360667"
      ],
      "layout": "IPY_MODEL_18fb08af07c241028002709d7a52a125"
     }
    },
    "3b7d28f07f054038a85d77810bbc3ff5": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6e39f41f44564a3682fc864fbc57233e",
      "placeholder": "​",
      "style": "IPY_MODEL_25ce70abc10c4f48a916cdfe0cf4568c",
      "value": "Downloading data: 100%"
     }
    },
    "3eaa1290e8034ba2bf3aee98555f2218": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "48ff4d9ee0284a278715b449e112f1fb": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_3b7d28f07f054038a85d77810bbc3ff5",
       "IPY_MODEL_c6fa9476c7824f75ba78560173cf2c9a",
       "IPY_MODEL_7da6f7f14ce24986baa2ef47adbabbe3"
      ],
      "layout": "IPY_MODEL_e437ed2842e9450ca5b4cf5561e5996c"
     }
    },
    "4f66ede3ff6c451385a25042631f26c1": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "5730cbe985484ca19f666d475b3561be": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "57ce5c49a1bd404393ca3f42ebad8e6b": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "59dec006ca154f0e833bf02fb1886626": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "59df1fbe27e04924900c585a635e238e": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "662d1617c46c4640a9eb286307b9f855": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "6c9bd7937a334be1bc71be7bcfeb831d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "6e39f41f44564a3682fc864fbc57233e": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "6ff7f6b154584ad09a31ce9c779f6a73": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "70fd0842944d44a08357823c6bb60c3e": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "71a1d619be59486da7efa887d586b383": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_7d943d6fff8b42a4b405ad02df20d6cf",
       "IPY_MODEL_f48d53599c5e4a718607969904835a49",
       "IPY_MODEL_ab1c423de1e64e7d8d02eea888161f3f"
      ],
      "layout": "IPY_MODEL_662d1617c46c4640a9eb286307b9f855"
     }
    },
    "7bb077eb9ccd42f28bcc1d245f1904cb": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "7d943d6fff8b42a4b405ad02df20d6cf": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_e92a7766829949559d9030153a0a7256",
      "placeholder": "​",
      "style": "IPY_MODEL_59df1fbe27e04924900c585a635e238e",
      "value": "Downloading data: 100%"
     }
    },
    "7da6f7f14ce24986baa2ef47adbabbe3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_353667831390430ca9e1026132101bbd",
      "placeholder": "​",
      "style": "IPY_MODEL_4f66ede3ff6c451385a25042631f26c1",
      "value": " 153M/153M [00:08&lt;00:00, 20.0MB/s]"
     }
    },
    "7ea553939bf741ee84174c935e3d6e6e": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "8320287538c040d091caa681ea745992": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_fba8de14bde342b0ad29476924fb601f",
      "placeholder": "​",
      "style": "IPY_MODEL_cb418b13f15c42fa8fa63c9587aaf7bf",
      "value": " 100000/100000 [00:01&lt;00:00, 62919.64 examples/s]"
     }
    },
    "85ffda2602e745faace8e22c889cb0e9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_154ea9b75c47405fb2b5cbca23444939",
      "max": 100,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_0010996f02ed4662a07291713924b12b",
      "value": 100
     }
    },
    "9341dca351074f4db14a6b76cf8c722b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "989c2e90ec5a46c498fc897bc540eed8": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "98f4bf45612a41f2b50ff7b903ca00a6": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "a0fd28b6ff574dd38d12339333336ce6": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "ab1c423de1e64e7d8d02eea888161f3f": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_ad8636fc0c8e4e128aa8efa19a0e8623",
      "placeholder": "​",
      "style": "IPY_MODEL_9341dca351074f4db14a6b76cf8c722b",
      "value": " 153M/153M [00:10&lt;00:00, 16.4MB/s]"
     }
    },
    "ad8636fc0c8e4e128aa8efa19a0e8623": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "b7e935d02f204f4a8eff7534188bc6c5": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "bbd037941f6845bab9b03be741695ad7": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "bda815dcbb2a4e6b9ff535cdb9debaf2": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6ff7f6b154584ad09a31ce9c779f6a73",
      "placeholder": "​",
      "style": "IPY_MODEL_3eaa1290e8034ba2bf3aee98555f2218",
      "value": "Generating train split: 100%"
     }
    },
    "c6fa9476c7824f75ba78560173cf2c9a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_57ce5c49a1bd404393ca3f42ebad8e6b",
      "max": 153421593,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_2e0c5716c7fd40cd854896ee086150e7",
      "value": 153421593
     }
    },
    "cb418b13f15c42fa8fa63c9587aaf7bf": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "d0d3e4b8f28c43aaa9d505aab38c5b1b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_bda815dcbb2a4e6b9ff535cdb9debaf2",
       "IPY_MODEL_1b1b4f0f0e0e4757a86bc30919f244d4",
       "IPY_MODEL_8320287538c040d091caa681ea745992"
      ],
      "layout": "IPY_MODEL_0104e084968f4690a400894611e937fa"
     }
    },
    "dbe565396e3e45a0ad5cb59f4cbc43e2": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_7ea553939bf741ee84174c935e3d6e6e",
      "placeholder": "​",
      "style": "IPY_MODEL_f713d501a82d4fa0960318ad96ffcfa2",
      "value": "Downloading readme: 100%"
     }
    },
    "de5a5da9c9514e89bbb9ed2ce0cdebb3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_bbd037941f6845bab9b03be741695ad7",
      "placeholder": "​",
      "style": "IPY_MODEL_0efe12220814483aa2266bab884fd838",
      "value": "Creating json from Arrow format: 100%"
     }
    },
    "e38684570b8441d0a460ece94645cb96": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "e3925a733c924a439a2454ffa7ca2363": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_e38684570b8441d0a460ece94645cb96",
      "placeholder": "​",
      "style": "IPY_MODEL_6c9bd7937a334be1bc71be7bcfeb831d",
      "value": " 944/944 [00:00&lt;00:00, 70.9kB/s]"
     }
    },
    "e437ed2842e9450ca5b4cf5561e5996c": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "e92a7766829949559d9030153a0a7256": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "f48d53599c5e4a718607969904835a49": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_a0fd28b6ff574dd38d12339333336ce6",
      "max": 153206051,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_7bb077eb9ccd42f28bcc1d245f1904cb",
      "value": 153206051
     }
    },
    "f713d501a82d4fa0960318ad96ffcfa2": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "fba8de14bde342b0ad29476924fb601f": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "fe6a62e2e0e44cfd8ce8fdea0406a8aa": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_70fd0842944d44a08357823c6bb60c3e",
      "max": 944,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_98f4bf45612a41f2b50ff7b903ca00a6",
      "value": 944
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}