NLLB-200-Fine-tuning-One-Shot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ol4zOqWnGQn3",
    "outputId": "08a48784-d7fc-4590-ab30-baf899397a92"
   },
   "outputs": [],
   "source": [
    "!pip3 install --upgrade transformers accelerate datasets ctranslate2 -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p /workspace/models/cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "73a3ff0ee0f64df299770dbf1e7df420",
      "2c83e70f249348ed82634066b2f49ec8",
      "299c95c8b346460d9cff8c226174ee37",
      "c2b2907d370f4cd08f09b736e4db656d",
      "8e52b7ad068d4c03a42a63c5559ef3ee",
      "955026f52f704484b1901cf96f255c15",
      "d6c224becab2466f931c6f70732914ba",
      "7f40268b51074b13b764f3b9857b8bf6",
      "c8e1b54c220c43699145cf110fc279cb",
      "2afb39cc79b341bc87e6146fa4ee0b58",
      "91461605886d4526be4b7729c1b42dc7"
     ]
    },
    "id": "PvLu8UPoGQ53",
    "outputId": "69a2fa00-9fcc-4613-af1e-d054cbfb355a"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM\n",
    "import torch\n",
    "\n",
    "cache_dir = \"/workspace/models/cache\"\n",
    "\n",
    "model_name = \"facebook/nllb-200-3.3B\"\n",
    "\n",
    "\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(\n",
    "                                            model_name,\n",
    "                                            device_map='auto',\n",
    "                                            use_cache=False,\n",
    "                                            cache_dir=cache_dir\n",
    "                                            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oQLz0iJ4KD5S"
   },
   "outputs": [],
   "source": [
    "# Language codes supported by NLLB-200\n",
    "\n",
    "SRC_LANG = \"eng_Latn\"\n",
    "TGT_LANG = \"fra_Latn\"\n",
    "\n",
    "# SRC_LANG = \"eng_Latn\"\n",
    "# TGT_LANG = \"por_Latn\"\n",
    "\n",
    "# SRC_LANG = \"eng_Latn\"\n",
    "# TGT_LANG = \"swh_Latn\"\n",
    "\n",
    "# SRC_LANG = \"swh_Latn\"\n",
    "# TGT_LANG = \"eng_Latn\"\n",
    "\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name,\n",
    "                                          cache_dir=cache_dir,\n",
    "                                          src_lang=SRC_LANG,\n",
    "                                          tgt_lang=TGT_LANG\n",
    "                                          )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hnh5nkfxQsyb"
   },
   "source": [
    "# Loading the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4EznpQIHQhGo",
    "outputId": "7267b6dc-6486-4a59-ee5c-3428aa0cbcd9"
   },
   "outputs": [],
   "source": [
    "# Load the training dataset\n",
    "\n",
    "import os\n",
    "\n",
    "trainsmall = True\n",
    "trainmedium = False if trainsmall else True\n",
    "\n",
    "src = \"en\"\n",
    "tgt = \"fr\"\n",
    "\n",
    "# src = \"en\"\n",
    "# tgt = \"pt\"\n",
    "\n",
    "# src = \"en\"\n",
    "# tgt = \"sw\"\n",
    "\n",
    "# src = \"sw\"\n",
    "# tgt = \"en\"\n",
    "\n",
    "# Change the path to your datasets\n",
    "directory = f\"/workspace/data/{src}-{tgt}/train\"\n",
    "\n",
    "if trainsmall and (tgt == \"fr\" or tgt == \"pt\"):\n",
    "    source_train_file = os.path.join(directory, f\"all-filtered.en.real.trainsmall\")\n",
    "    target_train_file = os.path.join(directory, f\"all-filtered.{tgt}.real.trainsmall\")\n",
    "elif trainmedium and (tgt == \"fr\" or tgt == \"pt\"):\n",
    "    source_train_file = os.path.join(directory, f\"all-filtered.en.real.trainmedium\")\n",
    "    target_train_file = os.path.join(directory, f\"all-filtered.{tgt}.real.trainmedium\")\n",
    "else:\n",
    "    source_train_file = os.path.join(directory, f\"mixed.filtered.{src}.real.train\")\n",
    "    target_train_file = os.path.join(directory, f\"mixed.filtered.{tgt}.real.train\")\n",
    "    \n",
    "\n",
    "with open(source_train_file, encoding=\"utf-8\") as source, open(target_train_file, encoding=\"utf-8\") as target:\n",
    "    source_train_sentences = [sent.strip() for sent in source]\n",
    "    target_train_sentences = [sent.strip() for sent in target]\n",
    "\n",
    "print(source_train_file, target_train_file, sep=\"\\n\")\n",
    "print(len(source_train_sentences))\n",
    "print(source_train_sentences[10])\n",
    "print(target_train_sentences[10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the test dataset\n",
    "\n",
    "directory = f\"/workspace/data/{src}-{tgt}/test\"\n",
    "\n",
    "if tgt == \"fr\" or tgt == \"pt\":\n",
    "    source_test_file = os.path.join(directory, f\"all-filtered.en.real.test\")\n",
    "    target_test_file = os.path.join(directory, f\"all-filtered.{tgt}.real.test\")\n",
    "else:\n",
    "    source_test_file = os.path.join(directory, f\"medical.filtered.{src}.real.test\")\n",
    "    target_test_file = os.path.join(directory, f\"medical.filtered.{tgt}.real.test\")\n",
    "\n",
    "with open(source_test_file, encoding=\"utf-8\") as source, open(target_test_file, encoding=\"utf-8\") as target:\n",
    "    source_test_sentences = [sent.strip() for sent in source][:1000]\n",
    "    target_test_sentences = [sent.strip() for sent in target][:1000]\n",
    "\n",
    "print(source_test_file, target_test_file, sep=\"\\n\")\n",
    "print(len(source_test_sentences))\n",
    "print(source_test_sentences[0])\n",
    "print(target_test_sentences[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 91
    },
    "id": "hnhRUrh1LnoU",
    "outputId": "ccb1e69b-534a-404a-d3e0-1565ce965f40"
   },
   "outputs": [],
   "source": [
    "# # Test inference\n",
    "\n",
    "# from transformers import pipeline\n",
    "\n",
    "# translator = pipeline(\"translation\",\n",
    "#                       model=model,\n",
    "#                       tokenizer=tokenizer,\n",
    "#                       src_lang=SRC_LANG,\n",
    "#                       tgt_lang=TGT_LANG)\n",
    "\n",
    "# translator(source_sentences[0])[0][\"translation_text\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Concatenate Fuzzy Matches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the fuzzy matches from unique Context Dataset\n",
    "\n",
    "directory = f\"/workspace/data/{src}-{tgt}/train\"\n",
    "\n",
    "if trainsmall and (tgt == \"fr\" or tgt == \"pt\"):\n",
    "    context_train_file = os.path.join(directory, f\"all-filtered.{src}-{tgt}.ms-multi-12.online.trainsmall\")\n",
    "elif trainmedium and (tgt == \"fr\" or tgt == \"pt\"):\n",
    "    context_train_file = os.path.join(directory, f\"all-filtered.{src}-{tgt}.ms-multi-12.online.trainmedium\")\n",
    "else:\n",
    "    context_train_file = os.path.join(directory, f\"mixed.filtered.{src}{tgt}.ms-multi-12.online.train\")\n",
    "\n",
    "with open(context_train_file, encoding=\"utf-8\") as context:\n",
    "    lines = [line.strip().split(\" ||| \") for line in context.readlines()]\n",
    "    scores = [float(line[0].strip()) for line in lines]\n",
    "    fuzzy_source_sentences = [line[1].strip() for line in lines]\n",
    "    online_source_sentences = [line[2].strip() for line in lines]\n",
    "    fuzzy_target_prefixes = [line[3].strip() for line in lines]\n",
    "\n",
    "n = 0\n",
    "print(\"Context dataset:\",context_train_file)\n",
    "print(\"Number of lines:\", len(lines), end=\"\\n\\n\")\n",
    "print(fuzzy_source_sentences[n])\n",
    "print(online_source_sentences[n])\n",
    "print(fuzzy_target_prefixes[n])\n",
    "print(target_train_sentences[n])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to create zero-shot and one-shot prompts\n",
    "\n",
    "def create_prompt(source_lang,\n",
    "                  target_lang,\n",
    "                  fuzzy_sources,\n",
    "                  fuzzy_targets,\n",
    "                  new_sources,\n",
    "                  new_targets,\n",
    "                  one_shot=True\n",
    "                  ):\n",
    "\n",
    "    prompts = []\n",
    "\n",
    "    if one_shot:\n",
    "        for fuzzy_src, fuzzy_tgt, new_src, new_tgt in zip(fuzzy_sources, fuzzy_targets, new_sources, new_targets):\n",
    "            src = fuzzy_src + SRC_LANG + \"•\" + new_src\n",
    "            tgt = fuzzy_tgt + TGT_LANG + \"•\" + new_tgt\n",
    "            prompts.append((src, tgt))\n",
    "    else:\n",
    "        for new_src, new_tgt in zip(new_sources, new_targets):\n",
    "            prompts.append((new_src, new_tgt))\n",
    "    \n",
    "    return prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zero_shot_prompts = create_prompt(SRC_LANG,\n",
    "                                  TGT_LANG,\n",
    "                                  fuzzy_source_sentences,\n",
    "                                  fuzzy_target_prefixes,\n",
    "                                  online_source_sentences,\n",
    "                                  target_train_sentences,\n",
    "                                  one_shot=False\n",
    "                                  )\n",
    "\n",
    "one_shot_prompts = create_prompt(SRC_LANG,\n",
    "                                 TGT_LANG,\n",
    "                                 fuzzy_source_sentences,\n",
    "                                 fuzzy_target_prefixes,\n",
    "                                 online_source_sentences,\n",
    "                                 target_train_sentences,\n",
    "                                 one_shot=True\n",
    "                                 )\n",
    "\n",
    "print(zero_shot_prompts[0], end=\"\\n\\n\")\n",
    "print(one_shot_prompts[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "prompts = zero_shot_prompts + one_shot_prompts\n",
    "# prompts = random.sample(prompts, len(prompts))  # shuffle\n",
    "prompts[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_train_sentences, target_train_sentences = zip(*prompts)\n",
    "\n",
    "print(len(source_train_sentences))\n",
    "print(source_train_sentences[0])\n",
    "print(target_train_sentences[0])\n",
    "print(source_train_sentences[-1])\n",
    "print(target_train_sentences[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tokenizer.tokenize(target_train_sentences[-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hkUC0AC6MSoS"
   },
   "source": [
    "# Fine-tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "rzK-hBjhMR2Z",
    "outputId": "6f9c0f25-b469-4de8-cdae-cabfaaf7af1b"
   },
   "outputs": [],
   "source": [
    "from datasets import Dataset, DatasetDict, load_dataset\n",
    "\n",
    "src_key = \"sentence_\" + SRC_LANG\n",
    "tgt_key = \"sentence_\" + TGT_LANG\n",
    "\n",
    "data_train = []\n",
    "\n",
    "for src_sent, tgt_sent in zip(source_train_sentences, target_train_sentences):\n",
    "    data_train.append({src_key: src_sent, tgt_key: tgt_sent})\n",
    "\n",
    "data_test = []\n",
    "\n",
    "for src_sent, tgt_sent in zip(source_test_sentences, target_test_sentences):\n",
    "    data_test.append({src_key: src_sent, tgt_key: tgt_sent})\n",
    "\n",
    "data_finetune = Dataset.from_list(data_train)\n",
    "data_validate = Dataset.from_list(data_test)\n",
    "\n",
    "print(data_finetune)\n",
    "print(data_validate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "fe35f61eaa9d4df3935b8391d7354376",
      "43a3eb02d1aa4c299d71709454eb29ab",
      "d156c5171fa347218661b0910579d037",
      "213a88c75ec4453da9deeada952a3658",
      "05c396da268546108feace21824b7bb0",
      "6f4fa3a400d64db682206a5fcfcea5b0",
      "3eb504acdfda4111a85606874f2c7eaa",
      "d72fdee6dac64f108a2cfb3927e4b731",
      "5e546e437e1847f0a93a28df5588b4a9",
      "496db2e272cb4d498a9632904c84ea43",
      "444a4bd0c71b4f6f99ae53a56ea9936e"
     ]
    },
    "id": "cY9keNuEODmW",
    "outputId": "30b8dbd1-d580-403a-e2c0-20a5646b3883"
   },
   "outputs": [],
   "source": [
    "def tokenize_fn(examples):\n",
    "    return tokenizer(examples[src_key], text_target=examples[tgt_key], padding=\"max_length\", truncation=True, max_length=512)\n",
    "\n",
    "tokenized_finetune = data_finetune.map(tokenize_fn, batched=True)\n",
    "tokenized_validate = data_validate.map(tokenize_fn, batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "sWrggllBrm3s",
    "outputId": "21d2c531-c02b-49ab-ac11-4bcffdc6b736"
   },
   "outputs": [],
   "source": [
    "print(tokenized_finetune)\n",
    "print(tokenized_validate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tokenizer.decode(tokenized_finetune[0][\"input_ids\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XTmfHYCMlckB"
   },
   "outputs": [],
   "source": [
    "from transformers import Seq2SeqTrainingArguments, Seq2SeqTrainer, logging\n",
    "\n",
    "epochs = 1\n",
    "\n",
    "learning_rate = 5e-5\n",
    "\n",
    "# Define training arguments\n",
    "training_args = Seq2SeqTrainingArguments(\n",
    "    output_dir=f\"/workspace/models/{src}-{tgt}\",\n",
    "    num_train_epochs=epochs,\n",
    "    per_device_train_batch_size=8,\n",
    "    per_device_eval_batch_size=8,\n",
    "    gradient_accumulation_steps=4,\n",
    "    eval_accumulation_steps=4,\n",
    "    #gradient_checkpointing=True,\n",
    "    fp16=True,\n",
    "    fp16_full_eval=True,\n",
    "\n",
    "    learning_rate=learning_rate,\n",
    "    lr_scheduler_type='constant',  # \"constant\", \"linear\", \"cosine\"\n",
    "    \n",
    "    eval_strategy=\"steps\",  # or \"epoch\"\n",
    "    eval_steps=100,\n",
    "    save_strategy=\"epoch\",\n",
    "    logging_steps=50,\n",
    "    report_to='none',    \n",
    ")\n",
    "\n",
    "# Initialize the trainer\n",
    "trainer = Seq2SeqTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_finetune,\n",
    "    eval_dataset=tokenized_validate,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# training_args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Start training\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert to CTranslate2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch = 1\n",
    "\n",
    "output_dir = f\"/workspace/models/{src}-{tgt}/saved_model_oneshot_medium_{learning_rate}_constant_epoch-{epoch}\"\n",
    "\n",
    "trainer.save_model(output_dir)\n",
    "tokenizer.save_pretrained(output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "quantization = \"float16\" # or \"int8\"\n",
    "\n",
    "ct2_output_dir = f\"/workspace/models/{src}-{tgt}/ct2_model_oneshot_medium_{learning_rate}_constant_epoch-{epoch}_{quantization}\" \n",
    "\n",
    "!ct2-transformers-converter --model {output_dir} \\\n",
    "--output_dir {ct2_output_dir} \\\n",
    "--quantization {quantization} --force \\\n",
    "&& echo \"CTranslate2 model saved at: {ct2_output_dir}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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": {
    "05c396da268546108feace21824b7bb0": {
     "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
     }
    },
    "213a88c75ec4453da9deeada952a3658": {
     "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_496db2e272cb4d498a9632904c84ea43",
      "placeholder": "​",
      "style": "IPY_MODEL_444a4bd0c71b4f6f99ae53a56ea9936e",
      "value": " 10000/10000 [00:06&lt;00:00, 1626.29 examples/s]"
     }
    },
    "299c95c8b346460d9cff8c226174ee37": {
     "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_7f40268b51074b13b764f3b9857b8bf6",
      "max": 3,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_c8e1b54c220c43699145cf110fc279cb",
      "value": 3
     }
    },
    "2afb39cc79b341bc87e6146fa4ee0b58": {
     "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
     }
    },
    "2c83e70f249348ed82634066b2f49ec8": {
     "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_955026f52f704484b1901cf96f255c15",
      "placeholder": "​",
      "style": "IPY_MODEL_d6c224becab2466f931c6f70732914ba",
      "value": "Loading checkpoint shards: 100%"
     }
    },
    "3eb504acdfda4111a85606874f2c7eaa": {
     "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": ""
     }
    },
    "43a3eb02d1aa4c299d71709454eb29ab": {
     "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_6f4fa3a400d64db682206a5fcfcea5b0",
      "placeholder": "​",
      "style": "IPY_MODEL_3eb504acdfda4111a85606874f2c7eaa",
      "value": "Map: 100%"
     }
    },
    "444a4bd0c71b4f6f99ae53a56ea9936e": {
     "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": ""
     }
    },
    "496db2e272cb4d498a9632904c84ea43": {
     "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
     }
    },
    "5e546e437e1847f0a93a28df5588b4a9": {
     "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": ""
     }
    },
    "6f4fa3a400d64db682206a5fcfcea5b0": {
     "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
     }
    },
    "73a3ff0ee0f64df299770dbf1e7df420": {
     "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_2c83e70f249348ed82634066b2f49ec8",
       "IPY_MODEL_299c95c8b346460d9cff8c226174ee37",
       "IPY_MODEL_c2b2907d370f4cd08f09b736e4db656d"
      ],
      "layout": "IPY_MODEL_8e52b7ad068d4c03a42a63c5559ef3ee"
     }
    },
    "7f40268b51074b13b764f3b9857b8bf6": {
     "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
     }
    },
    "8e52b7ad068d4c03a42a63c5559ef3ee": {
     "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
     }
    },
    "91461605886d4526be4b7729c1b42dc7": {
     "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": ""
     }
    },
    "955026f52f704484b1901cf96f255c15": {
     "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
     }
    },
    "c2b2907d370f4cd08f09b736e4db656d": {
     "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_2afb39cc79b341bc87e6146fa4ee0b58",
      "placeholder": "​",
      "style": "IPY_MODEL_91461605886d4526be4b7729c1b42dc7",
      "value": " 3/3 [00:13&lt;00:00,  4.12s/it]"
     }
    },
    "c8e1b54c220c43699145cf110fc279cb": {
     "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": ""
     }
    },
    "d156c5171fa347218661b0910579d037": {
     "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_d72fdee6dac64f108a2cfb3927e4b731",
      "max": 10000,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_5e546e437e1847f0a93a28df5588b4a9",
      "value": 10000
     }
    },
    "d6c224becab2466f931c6f70732914ba": {
     "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": ""
     }
    },
    "d72fdee6dac64f108a2cfb3927e4b731": {
     "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
     }
    },
    "fe35f61eaa9d4df3935b8391d7354376": {
     "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_43a3eb02d1aa4c299d71709454eb29ab",
       "IPY_MODEL_d156c5171fa347218661b0910579d037",
       "IPY_MODEL_213a88c75ec4453da9deeada952a3658"
      ],
      "layout": "IPY_MODEL_05c396da268546108feace21824b7bb0"
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}