BARK_AI_Text_to_Speech_Training_Sementic_jonathan (2).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 107536,
     "status": "ok",
     "timestamp": 1721844171310,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "d0g-XYlWeVIx",
    "outputId": "84195687-7025-4a7e-ce23-f54d054d092a"
   },
   "outputs": [],
   "source": [
    "!pip install ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 5,
     "status": "ok",
     "timestamp": 1721844171311,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "h8nS6BdaY726",
    "outputId": "e4aa612c-e41e-4595-da78-c2fbf293a914"
   },
   "outputs": [],
   "source": [
    "%ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 81978,
     "status": "ok",
     "timestamp": 1721844253286,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "CK5ibQS5UrGw",
    "outputId": "34302a25-c3d2-4e07-cec6-3b3da981a747"
   },
   "outputs": [],
   "source": [
    "!pip install pip==24.0\n",
    "!pip install audiolm_pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 58717,
     "status": "ok",
     "timestamp": 1721844311994,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "Uo2ZyPEpB9tU",
    "outputId": "018ccec0-77bd-41d8-d3c3-a50216854bdb"
   },
   "outputs": [],
   "source": [
    "!pip install encodec\n",
    "!pip install accelerate\n",
    "!pip install transformers\n",
    "!pip install huggingface_hub\n",
    "!pip install diffusers\n",
    "!pip install funcy\n",
    "!pip install fairseq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 9330,
     "status": "ok",
     "timestamp": 1721844321316,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "G29ydE3TBW-H"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import os\n",
    "import re\n",
    "import gc\n",
    "import json\n",
    "import math\n",
    "import hashlib\n",
    "import numpy as np\n",
    "import logging\n",
    "import torchaudio\n",
    "from tqdm.auto import tqdm\n",
    "import torch.nn.functional as F\n",
    "from encodec.utils import convert_audio\n",
    "from accelerate import Accelerator\n",
    "from accelerate.utils import set_seed\n",
    "from transformers import BertTokenizer\n",
    "from huggingface_hub import hf_hub_download\n",
    "from packaging import version\n",
    "from diffusers.optimization import get_scheduler\n",
    "\n",
    "from utils.bitsandbytes import BitsAndBytesConfig, importlib_metadata, get_keys_to_not_convert, replace_with_bnb_linear, set_module_quantized_tensor_to_device\n",
    "from utils.lora import convert_linear_layer_to_lora, only_optimize_lora_parameters, convert_lora_to_linear_layer\n",
    "from bark.model import GPTConfig, GPT\n",
    "from bark.model_fine import FineGPT, FineGPTConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4,
     "status": "ok",
     "timestamp": 1721844321316,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "kYqLomzKdIfc",
    "outputId": "737eda63-0d93-444c-8581-d8937f25dcf7"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "print(torch.__version__)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 1944,
     "status": "ok",
     "timestamp": 1721844323257,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "kVj15Ty_rMfb",
    "outputId": "735d29b3-ec5a-49b5-fbea-69a7e1b1ab2e"
   },
   "outputs": [],
   "source": [
    "!wget https://huggingface.co/junwchina/bark-voice-cloning-japanese-HuBERT-quantizer/resolve/main/japanese-HuBERT-quantizer_24_epoch.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xCtu1szcre-l"
   },
   "outputs": [],
   "source": [
    "# !wget https://huggingface.co/rinna/japanese-hubert-base/resolve/main/fairseq/model.pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 3,
     "status": "ok",
     "timestamp": 1721844323257,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "2Ly_KxhWCknu"
   },
   "outputs": [],
   "source": [
    "train_batch_size = 3\n",
    "eval_batch_size = 3\n",
    "grad_accum = 3\n",
    "ckpt_path = 'models/text_2.pt'\n",
    "model_type = \"text\"\n",
    "# dataset_path = 'datasets/joe_biden_state_of_union/'\n",
    "dataset_path = 'datasets/'\n",
    "\n",
    "project_dir = 'logs/'\n",
    "log_with = None\n",
    "hubert_path = 'data/models/hubert/hubert.pt'\n",
    "# hubert_path = '/content/bark-with-voice-clone/model.pt'\n",
    "\n",
    "hubert_tokenizer_path = 'japanese-HuBERT-quantizer_24_epoch.pth'\n",
    "\n",
    "output_dir = 'semantic_output/'\n",
    "resume_from_checkpoint = None\n",
    "\n",
    "checkpointing_steps = 500\n",
    "\n",
    "mixed_precision = 'fp16'\n",
    "bits = 16 #4 4 and 8 bit are a work in progress\n",
    "compute_dtype = torch.bfloat16\n",
    "double_quant = True\n",
    "quant_type = 'nf4'\n",
    "\n",
    "lora_dim = 0\n",
    "lora_scaling = 1\n",
    "lora_dropout = 0.1\n",
    "lora_module_name = 'transformer.h'\n",
    "optimize_lora_params_only = False\n",
    "\n",
    "learning_rate = 1e-4\n",
    "scale_lr = False\n",
    "use_8bit_adam = False\n",
    "adam_beta1 = 0.9\n",
    "adam_beta2 = 0.999\n",
    "adam_epsilon = 1e-8\n",
    "weight_decay = 0.01\n",
    "\n",
    "llm_int8_skip_modules = None\n",
    "keep_in_fp32_modules = ['lm_head']\n",
    "\n",
    "lr_scheduler_type = 'linear'\n",
    "lr_warmup_steps = 60\n",
    "num_train_epochs = 250\n",
    "max_train_steps = None\n",
    "max_grad_norm = 1.0\n",
    "\n",
    "seed = 741"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 365,
     "status": "ok",
     "timestamp": 1721844323620,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "stbBYtVDCtxd"
   },
   "outputs": [],
   "source": [
    "train_batch_size = 1\n",
    "CONTEXT_WINDOW_SIZE = 2048\n",
    "\n",
    "MAX_TEXT_LEN = 256\n",
    "\n",
    "SEMANTIC_RATE_HZ = 49.9\n",
    "SEMANTIC_VOCAB_SIZE = 10_000\n",
    "\n",
    "TEXT_ENCODING_OFFSET = 10_048\n",
    "SEMANTIC_PAD_TOKEN = 10_000\n",
    "TEXT_PAD_TOKEN = 129_595\n",
    "SEMANTIC_INFER_TOKEN = 129_599\n",
    "\n",
    "MAX_SEMANTIC_LEN = 511\n",
    "\n",
    "SAMPLE_RATE = 24_000\n",
    "CHANNELS = 1\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "USE_SMALL_MODELS = os.environ.get(\"SERP_USE_SMALL_MODELS\", False)\n",
    "\n",
    "default_cache_dir = os.path.join(os.path.expanduser(\"~\"), \".cache\")\n",
    "CACHE_DIR = os.path.join(os.getenv(\"XDG_CACHE_HOME\", default_cache_dir), \"serp\", \"bark_v0\")\n",
    "\n",
    "\n",
    "def _clear_cuda_cache():\n",
    "    if torch.cuda.is_available():\n",
    "        torch.cuda.empty_cache()\n",
    "        torch.cuda.synchronize()\n",
    "\n",
    "\n",
    "def _md5(fname):\n",
    "    hash_md5 = hashlib.md5()\n",
    "    with open(fname, \"rb\") as f:\n",
    "        for chunk in iter(lambda: f.read(4096), b\"\"):\n",
    "            hash_md5.update(chunk)\n",
    "    return hash_md5.hexdigest()\n",
    "\n",
    "\n",
    "def _download(from_hf_path, file_name, to_local_path):\n",
    "    to_local_path = to_local_path.replace(\"\\\\\", \"/\")\n",
    "    path = '/'.join(to_local_path.split(\"/\")[:-1])\n",
    "    os.makedirs(path, exist_ok=True)\n",
    "    hf_hub_download(repo_id=from_hf_path, filename=file_name, local_dir=path)\n",
    "    os.replace(os.path.join(path, file_name), to_local_path)\n",
    "\n",
    "\n",
    "def _tokenize(tokenizer, text):\n",
    "    return tokenizer.encode(text, add_special_tokens=False)\n",
    "\n",
    "\n",
    "def _detokenize(tokenizer, enc_text):\n",
    "    return tokenizer.decode(enc_text)\n",
    "\n",
    "\n",
    "def _normalize_whitespace(text):\n",
    "    return re.sub(r\"\\s+\", \" \", text).strip()\n",
    "\n",
    "\n",
    "REMOTE_MODEL_PATHS = {\n",
    "    \"text_small\": {\n",
    "        \"repo_id\": \"suno/bark\",\n",
    "        \"file_name\": \"text.pt\",\n",
    "        \"checksum\": \"b3e42bcbab23b688355cd44128c4cdd3\",\n",
    "    },\n",
    "    \"coarse_small\": {\n",
    "        \"repo_id\": \"suno/bark\",\n",
    "        \"file_name\": \"coarse.pt\",\n",
    "        \"checksum\": \"5fe964825e3b0321f9d5f3857b89194d\",\n",
    "    },\n",
    "    \"fine_small\": {\n",
    "        \"repo_id\": \"suno/bark\",\n",
    "        \"file_name\": \"fine.pt\",\n",
    "        \"checksum\": \"5428d1befe05be2ba32195496e58dc90\",\n",
    "    },\n",
    "    \"text\": {\n",
    "        \"repo_id\": \"suno/bark\",\n",
    "        \"file_name\": \"text_2.pt\",\n",
    "        \"checksum\": \"54afa89d65e318d4f5f80e8e8799026a\",\n",
    "    },\n",
    "    \"coarse\": {\n",
    "        \"repo_id\": \"suno/bark\",\n",
    "        \"file_name\": \"coarse_2.pt\",\n",
    "        \"checksum\": \"8a98094e5e3a255a5c9c0ab7efe8fd28\",\n",
    "    },\n",
    "    \"fine\": {\n",
    "        \"repo_id\": \"suno/bark\",\n",
    "        \"file_name\": \"fine_2.pt\",\n",
    "        \"checksum\": \"59d184ed44e3650774a2f0503a48a97b\",\n",
    "    },\n",
    "}\n",
    "\n",
    "\n",
    "def _load_model(ckpt_path, device, use_small=False, model_type=\"text\", load_pretrained=True):\n",
    "    if model_type == \"text\":\n",
    "        ConfigClass = GPTConfig\n",
    "        ModelClass = GPT\n",
    "    elif model_type == \"coarse\":\n",
    "        ConfigClass = GPTConfig\n",
    "        ModelClass = GPT\n",
    "    elif model_type == \"fine\":\n",
    "        ConfigClass = FineGPTConfig\n",
    "        ModelClass = FineGPT\n",
    "    else:\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    model_key = f\"{model_type}_small\" if use_small or USE_SMALL_MODELS else model_type\n",
    "    model_info = REMOTE_MODEL_PATHS[model_key]\n",
    "\n",
    "    if load_pretrained:\n",
    "        # Load checkpoint if specified or default path\n",
    "        if ckpt_path in [None, '']:\n",
    "            ckpt_path = os.path.join(CACHE_DIR, model_info[\"file_name\"])\n",
    "        if not os.path.exists(ckpt_path):\n",
    "            logger.info(f\"{model_type} model not found, downloading into `{CACHE_DIR}`.\")\n",
    "            _download(model_info[\"repo_id\"], model_info[\"file_name\"], ckpt_path)\n",
    "\n",
    "        checkpoint = torch.load(ckpt_path, map_location=device)\n",
    "\n",
    "        # Configure model with checkpoint arguments\n",
    "        model_args = checkpoint[\"model_args\"]\n",
    "        if \"input_vocab_size\" not in model_args:\n",
    "            model_args[\"input_vocab_size\"] = model_args[\"vocab_size\"]\n",
    "            model_args[\"output_vocab_size\"] = model_args[\"vocab_size\"]\n",
    "            del model_args[\"vocab_size\"]\n",
    "\n",
    "        gptconf = ConfigClass(**model_args)\n",
    "        model = ModelClass(gptconf)\n",
    "        state_dict = checkpoint[\"model\"]\n",
    "\n",
    "        # Fix checkpoint state dict keys\n",
    "        unwanted_prefix = \"_orig_mod.\"\n",
    "        for k, v in list(state_dict.items()):\n",
    "            if k.startswith(unwanted_prefix):\n",
    "                state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)\n",
    "\n",
    "        extra_keys = set(state_dict.keys()) - set(model.state_dict().keys())\n",
    "        extra_keys = set([k for k in extra_keys if not k.endswith(\".attn.bias\")])\n",
    "        missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())\n",
    "        missing_keys = set([k for k in missing_keys if not k.endswith(\".attn.bias\")])\n",
    "\n",
    "        if len(extra_keys) != 0:\n",
    "            raise ValueError(f\"extra keys found: {extra_keys}\")\n",
    "        if len(missing_keys) != 0:\n",
    "            raise ValueError(f\"missing keys: {missing_keys}\")\n",
    "\n",
    "        model.load_state_dict(state_dict, strict=False)\n",
    "        n_params = model.get_num_params()\n",
    "        val_loss = checkpoint[\"best_val_loss\"].item()\n",
    "        print(f\"Loaded {model_type} model with {n_params} params, val_loss={val_loss:.4f}.\")\n",
    "        del checkpoint, state_dict\n",
    "        _clear_cuda_cache()\n",
    "    else:\n",
    "        # Initialize a new model with default or custom configurations\n",
    "        model_args = {\n",
    "            \"input_vocab_size\": 10048,\n",
    "            \"output_vocab_size\": 10048,\n",
    "            \"n_positions\": 1024,\n",
    "            \"n_layer\": 12,\n",
    "            \"n_head\": 12,\n",
    "            \"d_model\": 768,\n",
    "        }\n",
    "        gptconf = ConfigClass(**model_args)\n",
    "        model = ModelClass(gptconf)\n",
    "        print(f\"Initialized new {model_type} model with custom configuration.\")\n",
    "\n",
    "    if model_type == \"text\":\n",
    "        tokenizer = BertTokenizer.from_pretrained(\"bert-base-multilingual-cased\")\n",
    "        return model, tokenizer\n",
    "    return model\n",
    "\n",
    "\n",
    "def load_filepaths_and_text(filename, split=\"|\"):\n",
    "    with open(filename, encoding='utf-8', errors='ignore') as f:\n",
    "        filepaths_and_text = [line.strip().split(split) for line in f]\n",
    "        base = os.path.dirname(filename)\n",
    "        for j in range(len(filepaths_and_text)):\n",
    "            filepaths_and_text[j][0] = os.path.join(base, filepaths_and_text[j][0])\n",
    "            # read .lyrics file\n",
    "            with open(filepaths_and_text[j][0].replace('.wav', '.wav.txt'), encoding='utf-8', errors='ignore') as f:\n",
    "                filepaths_and_text[j][1] = f.read().strip()\n",
    "    return filepaths_and_text\n",
    "\n",
    "class TtsDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, opt):\n",
    "        self.path = os.path.dirname(opt['path'])\n",
    "        self.mode = opt['mode']\n",
    "        self.audiopaths_and_text = load_filepaths_and_text(os.path.join(opt['path'] , opt['mode'] + '.txt'))\n",
    "        self.tokenizer = opt['tokenizer']\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        audiopath_and_text = self.audiopaths_and_text[index]\n",
    "        audiopath, text = audiopath_and_text[0], audiopath_and_text[1]\n",
    "\n",
    "        input_ids = np.array(_tokenize(self.tokenizer, text)) + TEXT_ENCODING_OFFSET\n",
    "        input_ids = torch.from_numpy(input_ids).long()\n",
    "        tokens = np.load(audiopath.replace('.wav', '.npz').replace('wavs', 'tokens'))\n",
    "        semantic_tokens = tokens['semantic']\n",
    "        semantic_tokens = torch.from_numpy(semantic_tokens).long()\n",
    "\n",
    "        return input_ids, semantic_tokens\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.audiopaths_and_text)\n",
    "\n",
    "\n",
    "class TtsCollater():\n",
    "    def __init__(self):\n",
    "        pass\n",
    "    def __call__(self, batch):\n",
    "        max_text_len = MAX_TEXT_LEN\n",
    "        max_semantic_tokens_len = MAX_SEMANTIC_LEN\n",
    "        texts = []\n",
    "        semantic_tokens = []\n",
    "\n",
    "        for b in batch:\n",
    "            text, semantic_tokens_ = b\n",
    "            text = F.pad(text, (0, max_text_len-len(text)), value=TEXT_PAD_TOKEN)\n",
    "            semantic_history = torch.from_numpy(np.array([SEMANTIC_PAD_TOKEN] * 256))\n",
    "            text = torch.cat([text, semantic_history, torch.tensor([SEMANTIC_INFER_TOKEN])])\n",
    "            texts.append(text)\n",
    "            semantic_tokens_ = semantic_tokens_[:max_semantic_tokens_len]\n",
    "            semantic_tokens.append(F.pad(semantic_tokens_, (0, max_semantic_tokens_len-len(semantic_tokens_)), value=SEMANTIC_PAD_TOKEN))\n",
    "\n",
    "        return {\n",
    "            'input_ids': torch.stack(texts).contiguous(),\n",
    "            'semantic_tokens': torch.stack(semantic_tokens).contiguous()\n",
    "        }\n",
    "\n",
    "\n",
    "accelerator = Accelerator(\n",
    "    gradient_accumulation_steps=grad_accum,\n",
    "    mixed_precision=mixed_precision,\n",
    "    log_with=log_with,\n",
    "    project_dir=project_dir,\n",
    ")\n",
    "device = accelerator.device\n",
    "\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "set_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 14824,
     "status": "ok",
     "timestamp": 1721844338441,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "9h03qv-qgljC",
    "outputId": "ccf2bf07-4e92-47e8-ae25-15551be41f02"
   },
   "outputs": [],
   "source": [
    "!pip install audiolm_pytorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PAoL6pAFSfHT"
   },
   "source": [
    "Setup Dataset (only need to do this once)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-aNgtjvcSzfI"
   },
   "outputs": [],
   "source": [
    "# prin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 356
    },
    "executionInfo": {
     "elapsed": 5046,
     "status": "error",
     "timestamp": 1712494249239,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "U5LTia7GDHOm",
    "outputId": "9a6dabb2-8314-4030-91ec-1b13ef110e77"
   },
   "outputs": [],
   "source": [
    "max_duration_sec = 50.00 # the maximum allowed duration in seconds\n",
    "\n",
    "path = dataset_path + \"/\"\n",
    "\n",
    "# From https://github.com/gitmylo/bark-voice-cloning-HuBERT-quantizer\n",
    "from hubert.hubert_manager import HuBERTManager\n",
    "hubert_manager = HuBERTManager()\n",
    "from hubert.pre_kmeans_hubert import CustomHubert\n",
    "from hubert.customtokenizer import CustomTokenizer\n",
    "hubert_manager.make_sure_hubert_installed()\n",
    "hubert_manager.make_sure_tokenizer_installed()\n",
    "\n",
    "# Load the HuBERT model\n",
    "hubert_model = CustomHubert(checkpoint_path=hubert_path).to(device)\n",
    "hubert_model.eval()\n",
    "for param in hubert_model.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "# Load the CustomTokenizer model\n",
    "hubert_tokenizer = CustomTokenizer.load_from_checkpoint(hubert_tokenizer_path).to(device)  # Automatically uses the right layers\n",
    "\n",
    "from bark.generation import load_codec_model\n",
    "codec_model = load_codec_model(use_gpu=True)\n",
    "codec_model.eval()\n",
    "for param in codec_model.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "\n",
    "def get_duration(wav, sr):\n",
    "    return wav.shape[1] / sr\n",
    "\n",
    "valid_lines_train = []\n",
    "# convert wavs to semantic tokens\n",
    "for wav_path, txt in load_filepaths_and_text(path + 'train.txt'):\n",
    "    wav, sr = torchaudio.load(wav_path)\n",
    "    if not get_duration(wav, sr) > max_duration_sec:\n",
    "        valid_lines_train.append((wav_path, txt))\n",
    "    wav = convert_audio(wav, sr, SAMPLE_RATE, CHANNELS).to(device)\n",
    "\n",
    "    semantic_vectors = hubert_model.forward(wav, input_sample_hz=SAMPLE_RATE)\n",
    "    semantic_tokens = hubert_tokenizer.get_token(semantic_vectors)\n",
    "\n",
    "    # save semantic tokens\n",
    "    os.makedirs(os.path.join(path, 'tokens'), exist_ok=True)\n",
    "    semantic_tokens = semantic_tokens.cpu().numpy()\n",
    "\n",
    "    # Extract discrete codes from EnCodec\n",
    "    with torch.no_grad():\n",
    "        encoded_frames = codec_model.encode(wav.unsqueeze(0))\n",
    "    codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1).squeeze()  # [n_q, T]\n",
    "\n",
    "    # move codes to cpu\n",
    "    codes = codes.cpu().numpy()\n",
    "\n",
    "    # save tokens\n",
    "    np.savez_compressed(os.path.join(path, 'tokens', os.path.basename(wav_path).replace('.wav', '.npz')), fine=codes, coarse=codes[:2, :], semantic=semantic_tokens)\n",
    "\n",
    "# rewrite train.txt with valid lines\n",
    "with open(path + 'train.txt', 'w', encoding='utf-8') as f:\n",
    "    for wav_path, txt in valid_lines_train:\n",
    "        wav_path = os.path.relpath(wav_path, dataset_path).replace('\\\\', '/')\n",
    "        f.write(f'{wav_path}|{txt}\\n')\n",
    "\n",
    "valid_lines_valid = []\n",
    "for wav_path, txt in load_filepaths_and_text(path + 'valid.txt'):\n",
    "    wav, sr = torchaudio.load(wav_path)\n",
    "    if not get_duration(wav, sr) > max_duration_sec:\n",
    "        valid_lines_valid.append((wav_path, txt))\n",
    "    wav = convert_audio(wav, sr, SAMPLE_RATE, CHANNELS).to(device)\n",
    "\n",
    "    semantic_vectors = hubert_model.forward(wav, input_sample_hz=SAMPLE_RATE)\n",
    "    semantic_tokens = hubert_tokenizer.get_token(semantic_vectors)\n",
    "\n",
    "    # save semantic tokens\n",
    "    os.makedirs(os.path.join(path, 'tokens'), exist_ok=True)\n",
    "    semantic_tokens = semantic_tokens.cpu().numpy()\n",
    "\n",
    "    # Extract discrete codes from EnCodec\n",
    "    with torch.no_grad():\n",
    "        encoded_frames = codec_model.encode(wav.unsqueeze(0))\n",
    "    codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1).squeeze()  # [n_q, T]\n",
    "\n",
    "    # move codes to cpu\n",
    "    codes = codes.cpu().numpy()\n",
    "\n",
    "    # save tokens\n",
    "    np.savez_compressed(os.path.join(path, 'tokens', os.path.basename(wav_path).replace('.wav', '.npz')), fine=codes, coarse=codes[:2, :], semantic=semantic_tokens)\n",
    "\n",
    "# rewrite valid.txt with valid lines\n",
    "with open(path + 'valid.txt', 'w', encoding='utf-8') as f:\n",
    "    for wav_path, txt in valid_lines_valid:\n",
    "        wav_path = os.path.relpath(wav_path, dataset_path).replace('\\\\', '/')\n",
    "        f.write(f'{wav_path}|{txt}\\n')\n",
    "\n",
    "del hubert_model\n",
    "del hubert_tokenizer\n",
    "del codec_model\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 194,
     "referenced_widgets": [
      "8167c8c4fbfb4269bfa7af77ec077408",
      "1f1bd137a3cf430c8ca78fc0b3a3021b",
      "32919d67aed742809f66ae64c5079acd",
      "c275d3e607e64079a303039928bcf035",
      "1d0e586a34f74cac967844ef2bce71f1",
      "eabcea8b31b04a88a63fd996cf90c14b",
      "2e3a3d758f6f4ebd9c8c9907b1a1f8d7",
      "12692df6cea04d0db9be202c2a296790",
      "fc1c06851cc34694a223b1320f98af5b",
      "60e2f5ec5c934f1ba33d362760a28282",
      "22abab758f7e4e0a826b20a83c408916",
      "d4b0fdf3a7a945619c6905a48725a8ca",
      "ad350891065d4e29b73b5bf29e494b09",
      "df8d89932919459996b19a875b453dda",
      "b84a9d3290784b57b0a2af9215992c9d",
      "778a064b9d7c4a8ea92186810264dc08",
      "a203d336509d4f74999af5b973959fef",
      "4c56b7a8dbac40319087ad029ae72a1c",
      "3f22d50b47e8447a89fd4333cf870582",
      "cce2a28fb6e0455889c04575fbb68564",
      "03cc91e0c3da4e978951506b78633661",
      "26ffdcdba6544f0cb815c83e7466932c",
      "91e23b7000104e77b7d3a88ba7662818",
      "8032cb188c9f41f388588cbfe2e755b1",
      "a2060bd8e3e04cfd8a6b9b7ee3ab8822",
      "4106fe16851c4d9d8b84a005ba7bde46",
      "eb5b52d870ac484eb8ea1e194b4ef33b",
      "42541f3e57ac4d8cafe6218c7c476ee5",
      "a4a60902c35f4c9680e1821a02d3d289",
      "2040f19e32cd48289fb3912754f521f2",
      "02e03aa708cc4f6fae60796c40b8c237",
      "0283835fe96f4b5ba0400cd64ea29b75",
      "acbc0736f23147e19a05e63a1f1ed31d",
      "99c57c8ce08146e4b2787abb65015480",
      "d997312ea272454a98774230b731b748",
      "c3fbf88ac32f4674a4c857a74014aaa4",
      "c416ccbfd6674c1b9e72a62ecd7532d7",
      "4067cbb94b4a4726aeb6ed1c3be9e851",
      "afaf1cb0ab0343e9a933c1a6359c05b9",
      "02a335585b324d2499ce59979d773703",
      "ca4a2016a7ad4ba9a0fb43eadb30c90d",
      "c22396edfa6a45bd9ddc25da8293a5f9",
      "734f52c9430a424683d91dffbac1347e",
      "9118bb97895a4cd79613e692c8c74b3a",
      "352b805f204e4c6d83b21219e6038f29",
      "6c2af273159547a392089c5834c9cd8e",
      "5ca2e215df8e4158b2c8d9efa10765c2",
      "b430967240c24db39e4d92da89bec23a",
      "7f616d3fcd2c42539d9d7787213e2c57",
      "93f49a957d674ccc8ef402b39e9a7d91",
      "951846e55fba4a6290c769c9d35645f9",
      "df24a18ed26d4bf3a5e4f249e8cb9c78",
      "226b40ea6ce446bf88f9292e8216803d",
      "08639571c1224db687cfaefb7a80fa6f",
      "5f71cb7313124cf8bafd6ce70a887df4"
     ]
    },
    "executionInfo": {
     "elapsed": 83526,
     "status": "ok",
     "timestamp": 1721844421958,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "GQK2RiR9CuEk",
    "outputId": "89008958-028f-407a-97f9-3e87b9186407"
   },
   "outputs": [],
   "source": [
    "model, tokenizer = _load_model(ckpt_path, device, use_small=False, model_type=model_type, load_pretrained=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 9,
     "status": "ok",
     "timestamp": 1721844421958,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "x_xXOrBmDS9a"
   },
   "outputs": [],
   "source": [
    "if scale_lr:\n",
    "    learning_rate = (\n",
    "        learning_rate * grad_accum * train_batch_size * accelerator.num_processes\n",
    "    )\n",
    "\n",
    "if use_8bit_adam:\n",
    "    try:\n",
    "        import bitsandbytes as bnb\n",
    "    except ImportError:\n",
    "        raise ImportError(\n",
    "            \"To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`.\"\n",
    "        )\n",
    "\n",
    "    optimizer_class = bnb.optim.AdamW8bit\n",
    "else:\n",
    "    optimizer_class = torch.optim.AdamW"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 8,
     "status": "ok",
     "timestamp": 1721844421958,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "wFsANM2vP886"
   },
   "outputs": [],
   "source": [
    "quantization_config=BitsAndBytesConfig(\n",
    "    load_in_4bit=bits == 4,\n",
    "    load_in_8bit=bits == 8,\n",
    "    llm_int8_threshold=6.0,\n",
    "    llm_int8_has_fp16_weight=False,\n",
    "    bnb_4bit_compute_dtype=compute_dtype,\n",
    "    bnb_4bit_use_double_quant=double_quant,\n",
    "    bnb_4bit_quant_type=quant_type # {'fp4', 'nf4'}\n",
    ")\n",
    "\n",
    "if quantization_config.load_in_8bit or quantization_config.load_in_4bit:\n",
    "    if quantization_config.load_in_8bit:\n",
    "        logger.info(\"Detected 8-bit loading: activating 8-bit loading for this model\")\n",
    "    elif quantization_config.load_in_4bit:\n",
    "        logger.info(\"Detected 4-bit loading: activating 4-bit loading for this model\")\n",
    "\n",
    "    # We keep some modules such as the lm_head in their original dtype for numerical stability reasons\n",
    "    if llm_int8_skip_modules is None or len(llm_int8_skip_modules) == 0:\n",
    "        modules_to_not_convert = [] # get_keys_to_not_convert(model)\n",
    "    else:\n",
    "        modules_to_not_convert = llm_int8_skip_modules\n",
    "\n",
    "    if not isinstance(modules_to_not_convert, list):\n",
    "        modules_to_not_convert = [modules_to_not_convert]\n",
    "\n",
    "    modules_to_not_convert.extend(keep_in_fp32_modules)\n",
    "\n",
    "    supports_4bit = version.parse(importlib_metadata.version(\"bitsandbytes\")) >= version.parse(\"0.39.0\")\n",
    "\n",
    "    if quantization_config.load_in_4bit and not supports_4bit:\n",
    "        raise ValueError(\n",
    "            \"You have a version of `bitsandbytes` that is not compatible with 4bit inference and training\"\n",
    "            \" make sure you have the latest version of `bitsandbytes` installed\"\n",
    "        )\n",
    "\n",
    "    if len(modules_to_not_convert) == 0:\n",
    "        modules_to_not_convert = None\n",
    "\n",
    "    model = replace_with_bnb_linear(\n",
    "        model, modules_to_not_convert=modules_to_not_convert, quantization_config=quantization_config\n",
    "    )\n",
    "\n",
    "    # training in 8-bit is only available in 0.37.0+\n",
    "    model._is_kbit_training_enabled = version.parse(\n",
    "        importlib_metadata.version(\"bitsandbytes\")\n",
    "    ) >= version.parse(\"0.37.0\")\n",
    "\n",
    "    model.config.quantization_config = quantization_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 375,
     "status": "ok",
     "timestamp": 1721844422325,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "wyU7y1dIP8x_"
   },
   "outputs": [],
   "source": [
    "if bits == 4:\n",
    "    from accelerate.utils import CustomDtype\n",
    "    target_dtype = CustomDtype.INT4\n",
    "elif bits == 8:\n",
    "    target_dtype = torch.int8\n",
    "\n",
    "if lora_dim > 0:\n",
    "    for param in model.parameters():\n",
    "        if param.ndim == 1:\n",
    "            # cast the small parameters (e.g. layernorm) to fp32 for stability\n",
    "            param.data = param.data.to(torch.float32)\n",
    "\n",
    "    class CastOutputToFloat(nn.Sequential):\n",
    "        def forward(self, x):\n",
    "            return super().forward(x).to(torch.float32)\n",
    "\n",
    "    model.lm_head = CastOutputToFloat(model.lm_head)\n",
    "\n",
    "    model = convert_linear_layer_to_lora(model, lora_module_name,\n",
    "                                            lora_dim=lora_dim, lora_scaling=lora_scaling,\n",
    "                                            lora_dropout=lora_dropout)\n",
    "    if optimize_lora_params_only:\n",
    "        model = only_optimize_lora_parameters(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 3186,
     "status": "ok",
     "timestamp": 1721844425509,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "NiupgNENP8vX"
   },
   "outputs": [],
   "source": [
    "params_to_optimize = (\n",
    "        param for param in model.parameters() if param.requires_grad\n",
    "    )\n",
    "\n",
    "optimizer = optimizer_class(\n",
    "    params_to_optimize,\n",
    "    lr=learning_rate,\n",
    "    betas=(adam_beta1, adam_beta2),\n",
    "    weight_decay=weight_decay,\n",
    "    eps=adam_epsilon,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 35337,
     "status": "ok",
     "timestamp": 1721844460844,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "PACQx5UnP8sh"
   },
   "outputs": [],
   "source": [
    "opt_train = {\n",
    "    'path': dataset_path,\n",
    "    'tokenizer': tokenizer,\n",
    "    'mode': 'train',\n",
    "}\n",
    "\n",
    "opt_val = {\n",
    "    'path': dataset_path,\n",
    "    'tokenizer': tokenizer,\n",
    "    'mode': 'valid',\n",
    "}\n",
    "\n",
    "train_dataset = TtsDataset(opt_train)\n",
    "validation_dataset = TtsDataset(opt_val)\n",
    "\n",
    "train_dataloader = torch.utils.data.DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=train_batch_size,\n",
    "    collate_fn=TtsCollater(),\n",
    ")\n",
    "\n",
    "validation_dataloader = torch.utils.data.DataLoader(\n",
    "    validation_dataset,\n",
    "    batch_size=eval_batch_size,\n",
    "    collate_fn=TtsCollater(),\n",
    ")\n",
    "\n",
    "criterion = torch.nn.CrossEntropyLoss() #ignore_index=SEMANTIC_PAD_TOKEN)\n",
    "\n",
    "# Scheduler and math around the number of training steps.\n",
    "overrode_max_train_steps = False\n",
    "num_update_steps_per_epoch = math.ceil(len(train_dataloader) / grad_accum)\n",
    "if max_train_steps is None:\n",
    "    max_train_steps = num_train_epochs * num_update_steps_per_epoch\n",
    "    overrode_max_train_steps = True\n",
    "\n",
    "lr_scheduler = get_scheduler(\n",
    "    lr_scheduler_type,\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=lr_warmup_steps * grad_accum,\n",
    "    num_training_steps=max_train_steps * grad_accum,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 1446,
     "status": "ok",
     "timestamp": 1721844462282,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "1FiX9bcTP8pp"
   },
   "outputs": [],
   "source": [
    "model, optimizer, train_dataloader, validation_dataloader, lr_scheduler = accelerator.prepare(\n",
    "    model, optimizer, train_dataloader, validation_dataloader, lr_scheduler\n",
    ")\n",
    "accelerator.register_for_checkpointing(lr_scheduler)\n",
    "\n",
    "weight_dtype = torch.float32\n",
    "if accelerator.mixed_precision == \"fp16\":\n",
    "    weight_dtype = torch.float16\n",
    "elif accelerator.mixed_precision == \"bf16\":\n",
    "    weight_dtype = torch.bfloat16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 4,
     "status": "ok",
     "timestamp": 1721844462283,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "kDSJ8EaTQOad"
   },
   "outputs": [],
   "source": [
    "# We need to recalculate our total training steps as the size of the training dataloader may have changed.\n",
    "num_update_steps_per_epoch = math.ceil(len(train_dataloader) / grad_accum)\n",
    "if overrode_max_train_steps:\n",
    "    max_train_steps = num_train_epochs * num_update_steps_per_epoch\n",
    "# Afterwards we recalculate our number of training epochs\n",
    "num_train_epochs = math.ceil(max_train_steps / num_update_steps_per_epoch)\n",
    "\n",
    "# We need to initialize the trackers we use, and also store our configuration.\n",
    "# The trackers initializes automatically on the main process.\n",
    "if accelerator.is_main_process:\n",
    "    accelerator.init_trackers(\"bark_semantic\", config={})\n",
    "\n",
    "# Train!\n",
    "total_batch_size = train_batch_size * accelerator.num_processes * grad_accum\n",
    "logger.info(\"***** Running training *****\")\n",
    "logger.info(f\"  Num examples = {len(train_dataset)}\")\n",
    "logger.info(f\"  Num batches each epoch = {len(train_dataloader)}\")\n",
    "logger.info(f\"  Num Epochs = {num_train_epochs}\")\n",
    "logger.info(f\"  Instantaneous batch size per device = {train_batch_size}\")\n",
    "logger.info(f\"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}\")\n",
    "logger.info(f\"  Gradient Accumulation steps = {grad_accum}\")\n",
    "logger.info(f\"  Total optimization steps = {max_train_steps}\")\n",
    "global_step = 0\n",
    "first_epoch = 0\n",
    "\n",
    "if resume_from_checkpoint:\n",
    "    if resume_from_checkpoint != \"latest\":\n",
    "        path = os.path.basename(resume_from_checkpoint)\n",
    "    else:\n",
    "        # Get the most recent checkpoint\n",
    "        dirs = os.listdir(output_dir)\n",
    "        dirs = [d for d in dirs if d.startswith(\"checkpoint\")]\n",
    "        dirs = sorted(dirs, key=lambda x: int(x.split(\"-\")[1]))\n",
    "        path = dirs[-1]\n",
    "    accelerator.print(f\"Resuming from checkpoint {path}\")\n",
    "    accelerator.load_state(os.path.join(output_dir, path))\n",
    "    global_step = int(path.split(\"-\")[1])\n",
    "\n",
    "    resume_global_step = global_step * grad_accum\n",
    "    first_epoch = resume_global_step // num_update_steps_per_epoch\n",
    "    resume_step = resume_global_step % num_update_steps_per_epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 3,
     "status": "ok",
     "timestamp": 1721844462283,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "tYEKL6dc4fJl",
    "outputId": "bc701c3b-6d69-4ed0-fad1-14ab1a9abe64"
   },
   "outputs": [],
   "source": [
    "print(num_train_epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fhuzo5PqXCvN"
   },
   "outputs": [],
   "source": [
    "# !pip3 uninstall torch torchvision torchaudio\n",
    "# !pip3 install torch torchvision torchaudio\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kx1xa9_IZ-Lg"
   },
   "outputs": [],
   "source": [
    "# os.environ['CUDA_VISIBLE_DEVICES'] ='0'\n",
    "torch.cuda.empty_cache()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 41533,
     "status": "ok",
     "timestamp": 1721844503814,
     "user": {
      "displayName": "CreeperPower storing",
      "userId": "18072225667475959743"
     },
     "user_tz": -120
    },
    "id": "0fzjLU7OQQ6d",
    "outputId": "acee7290-a354-4217-bdb8-2e52a74cb929"
   },
   "outputs": [],
   "source": [
    "if accelerator.is_main_process:\n",
    "    model.eval()\n",
    "    validation_loss = 0.0\n",
    "    num_batches = 0\n",
    "    num_samples = 0\n",
    "    with torch.no_grad():\n",
    "        for val_step, val_batch in enumerate(validation_dataloader):\n",
    "            # Similar to training, process the validation batch\n",
    "            val_targets = val_batch['semantic_tokens'][:, 1:].contiguous()\n",
    "            val_semantic_inputs = val_batch['semantic_tokens'][:, :-1]\n",
    "            val_inputs = torch.cat([val_batch['input_ids'], val_semantic_inputs], dim=1)\n",
    "\n",
    "            # Forward pass for validation\n",
    "            val_logits = model(val_inputs, training=True)\n",
    "            val_semantic_logits = val_logits[:, val_batch['input_ids'].size(1):].contiguous()\n",
    "\n",
    "            # Calculate the validation loss\n",
    "            val_loss = criterion(val_semantic_logits.view(-1, model.config.output_vocab_size), val_targets.view(-1))\n",
    "            validation_loss += val_loss.item()\n",
    "            num_batches += 1\n",
    "            num_samples += val_batch['input_ids'].size(0)\n",
    "\n",
    "    average_validation_loss = validation_loss / num_batches\n",
    "    logger.info(f\"Validation Loss: {average_validation_loss} over {num_samples} samples and {num_batches} batches.\")\n",
    "    print(f\"Validation Loss: {average_validation_loss} over {num_samples} samples and {num_batches} batches.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true,
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "49726b3cf7ea4f8ba9591f5ccf29a99a",
      "4bbf650808d641368dac047675063822",
      "032f02914c44453991ac30183b1e9575",
      "ef5b9993556b478990a7f68b3b4cb06c",
      "9a24afcbc4e54166974629ad969a5a55",
      "55758331d575454d99fd43cf923f4a21",
      "1acaab97017b47fc8533cbc5e8e56e6c",
      "f3b79041862a4609865050451a18d110",
      "c298686a944a4db6912f818b21538172",
      "ec673bba5a794916bc461a2ce0ab77c7",
      "54c63aa200014b2c9dba18393f22be6c"
     ]
    },
    "id": "_QXzh_cuQR_V",
    "outputId": "d17ae7cb-38d1-4c12-ac4e-930f3c3480e8"
   },
   "outputs": [],
   "source": [
    "# Only show the progress bar once on each machine.\n",
    "progress_bar = tqdm(range(global_step, max_train_steps), disable=not accelerator.is_local_main_process)\n",
    "progress_bar.set_description(\"Steps\")\n",
    "\n",
    "for epoch in range(first_epoch, num_train_epochs):\n",
    "    model.train()\n",
    "    for step, batch in enumerate(train_dataloader):\n",
    "        # Skip steps until we reach the resumed step\n",
    "        if resume_from_checkpoint and epoch == first_epoch and step < resume_step:\n",
    "            if step % grad_accum == 0:\n",
    "                progress_bar.update(1)\n",
    "            continue\n",
    "\n",
    "        with accelerator.accumulate(model):\n",
    "            targets = batch['semantic_tokens'][:, 1:].contiguous()\n",
    "\n",
    "            # Remove the last semantic token from the inputs since there is no target for it.\n",
    "            semantic_inputs = batch['semantic_tokens'][:, :-1]\n",
    "\n",
    "            # Combine the text and semantic tokens and feed them into the model.\n",
    "            inputs = torch.cat([batch['input_ids'], semantic_inputs], dim=1)\n",
    "            logits = model(inputs, training=True)\n",
    "\n",
    "            # We're only interested in the logits for the semantic tokens, so we ignore the logits for the input text tokens.\n",
    "            semantic_logits = logits[:, batch['input_ids'].size(1):].contiguous()\n",
    "\n",
    "            # Compute the loss.\n",
    "            loss = criterion(semantic_logits.view(-1, model.config.output_vocab_size), targets.view(-1))\n",
    "\n",
    "            accelerator.backward(loss)\n",
    "            if accelerator.sync_gradients:\n",
    "                params_to_clip = (\n",
    "                    param for param in model.parameters() if param.requires_grad\n",
    "                )\n",
    "                accelerator.clip_grad_norm_(params_to_clip, max_grad_norm)\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "        # Checks if the accelerator has performed an optimization step behind the scenes\n",
    "        if accelerator.sync_gradients:\n",
    "            progress_bar.update(1)\n",
    "            global_step += 1\n",
    "\n",
    "            if global_step % 500 == 0:\n",
    "                if accelerator.is_main_process:\n",
    "                    save_path = os.path.join(output_dir, f\"checkpoint-{global_step}\")\n",
    "                    accelerator.save_state(save_path)\n",
    "                    logger.info(f\"Saved state to {save_path}\")\n",
    "                    # also save the model to google drive\n",
    "                    accelerator.save(model.state_dict(), os.path.join(\"epoch115\", f\"checkpoint-text-{global_step}.pt\"))\n",
    "                    config = model.config.__dict__\n",
    "                    # save config\n",
    "                    with open(os.path.join(\"epoch115\", f\"config-text-{global_step}.json\"), \"w\") as f:\n",
    "                        json.dump(config, f, indent=2)\n",
    "\n",
    "        logs = {\"loss\": loss.detach().item(), \"lr\": lr_scheduler.get_last_lr()[0]}\n",
    "        progress_bar.set_postfix(**logs)\n",
    "        accelerator.log(logs, step=global_step)\n",
    "\n",
    "        if global_step >= max_train_steps:\n",
    "            break\n",
    "\n",
    "    accelerator.wait_for_everyone()\n",
    "\n",
    "if accelerator.is_main_process:\n",
    "    if lora_dim > 0:\n",
    "        model = convert_lora_to_linear_layer(model)\n",
    "    # save model\n",
    "    accelerator.save(model.state_dict(), os.path.join(output_dir, \"pytorch_model.bin\"))\n",
    "\n",
    "    config = model.config.__dict__\n",
    "    # save config\n",
    "    with open(os.path.join(output_dir, \"config.json\"), \"w\") as f:\n",
    "        json.dump(config, f, indent=2)\n",
    "\n",
    "accelerator.end_training()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TQM2RQOWQVbn"
   },
   "outputs": [],
   "source": [
    "if accelerator.is_main_process:\n",
    "    model.eval()\n",
    "    validation_loss = 0.0\n",
    "    num_batches = 0\n",
    "    num_samples = 0\n",
    "    with torch.no_grad():\n",
    "        for val_step, val_batch in enumerate(validation_dataloader):\n",
    "            # Similar to training, process the validation batch\n",
    "            val_targets = val_batch['semantic_tokens'][:, 1:].contiguous()\n",
    "            val_semantic_inputs = val_batch['semantic_tokens'][:, :-1]\n",
    "            val_inputs = torch.cat([val_batch['input_ids'], val_semantic_inputs], dim=1)\n",
    "\n",
    "            # Forward pass for validation\n",
    "            val_logits = model(val_inputs, training=True)\n",
    "            val_semantic_logits = val_logits[:, val_batch['input_ids'].size(1):].contiguous()\n",
    "\n",
    "            # Calculate the validation loss\n",
    "            val_loss = criterion(val_semantic_logits.view(-1, model.config.output_vocab_size), val_targets.view(-1))\n",
    "            validation_loss += val_loss.item()\n",
    "            num_batches += 1\n",
    "            num_samples += val_batch['input_ids'].size(0)\n",
    "\n",
    "    average_validation_loss = validation_loss / num_batches\n",
    "    logger.info(f\"Validation Loss: {average_validation_loss} over {num_samples} samples and {num_batches} batches.\")\n",
    "    print(f\"Validation Loss: {average_validation_loss} over {num_samples} samples and {num_batches} batches.\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "0283835fe96f4b5ba0400cd64ea29b75": {
     "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
     }
    },
    "02a335585b324d2499ce59979d773703": {
     "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": ""
     }
    },
    "02e03aa708cc4f6fae60796c40b8c237": {
     "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": ""
     }
    },
    "032f02914c44453991ac30183b1e9575": {
     "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": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_f3b79041862a4609865050451a18d110",
      "max": 10750,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_c298686a944a4db6912f818b21538172",
      "value": 0
     }
    },
    "03cc91e0c3da4e978951506b78633661": {
     "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
     }
    },
    "08639571c1224db687cfaefb7a80fa6f": {
     "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
     }
    },
    "12692df6cea04d0db9be202c2a296790": {
     "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
     }
    },
    "1acaab97017b47fc8533cbc5e8e56e6c": {
     "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": ""
     }
    },
    "1d0e586a34f74cac967844ef2bce71f1": {
     "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
     }
    },
    "1f1bd137a3cf430c8ca78fc0b3a3021b": {
     "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_eabcea8b31b04a88a63fd996cf90c14b",
      "placeholder": "​",
      "style": "IPY_MODEL_2e3a3d758f6f4ebd9c8c9907b1a1f8d7",
      "value": "text_2.pt: 100%"
     }
    },
    "2040f19e32cd48289fb3912754f521f2": {
     "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
     }
    },
    "226b40ea6ce446bf88f9292e8216803d": {
     "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": ""
     }
    },
    "22abab758f7e4e0a826b20a83c408916": {
     "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": ""
     }
    },
    "26ffdcdba6544f0cb815c83e7466932c": {
     "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": ""
     }
    },
    "2e3a3d758f6f4ebd9c8c9907b1a1f8d7": {
     "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": ""
     }
    },
    "32919d67aed742809f66ae64c5079acd": {
     "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_12692df6cea04d0db9be202c2a296790",
      "max": 5353258741,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_fc1c06851cc34694a223b1320f98af5b",
      "value": 5353258741
     }
    },
    "352b805f204e4c6d83b21219e6038f29": {
     "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_6c2af273159547a392089c5834c9cd8e",
       "IPY_MODEL_5ca2e215df8e4158b2c8d9efa10765c2",
       "IPY_MODEL_b430967240c24db39e4d92da89bec23a"
      ],
      "layout": "IPY_MODEL_7f616d3fcd2c42539d9d7787213e2c57"
     }
    },
    "3f22d50b47e8447a89fd4333cf870582": {
     "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
     }
    },
    "4067cbb94b4a4726aeb6ed1c3be9e851": {
     "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
     }
    },
    "4106fe16851c4d9d8b84a005ba7bde46": {
     "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_0283835fe96f4b5ba0400cd64ea29b75",
      "placeholder": "​",
      "style": "IPY_MODEL_acbc0736f23147e19a05e63a1f1ed31d",
      "value": " 996k/996k [00:00&lt;00:00, 15.9MB/s]"
     }
    },
    "42541f3e57ac4d8cafe6218c7c476ee5": {
     "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
     }
    },
    "49726b3cf7ea4f8ba9591f5ccf29a99a": {
     "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_4bbf650808d641368dac047675063822",
       "IPY_MODEL_032f02914c44453991ac30183b1e9575",
       "IPY_MODEL_ef5b9993556b478990a7f68b3b4cb06c"
      ],
      "layout": "IPY_MODEL_9a24afcbc4e54166974629ad969a5a55"
     }
    },
    "4bbf650808d641368dac047675063822": {
     "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_55758331d575454d99fd43cf923f4a21",
      "placeholder": "​",
      "style": "IPY_MODEL_1acaab97017b47fc8533cbc5e8e56e6c",
      "value": "Steps:   0%"
     }
    },
    "4c56b7a8dbac40319087ad029ae72a1c": {
     "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": ""
     }
    },
    "54c63aa200014b2c9dba18393f22be6c": {
     "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": ""
     }
    },
    "55758331d575454d99fd43cf923f4a21": {
     "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
     }
    },
    "5ca2e215df8e4158b2c8d9efa10765c2": {
     "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_df24a18ed26d4bf3a5e4f249e8cb9c78",
      "max": 625,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_226b40ea6ce446bf88f9292e8216803d",
      "value": 625
     }
    },
    "5f71cb7313124cf8bafd6ce70a887df4": {
     "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": ""
     }
    },
    "60e2f5ec5c934f1ba33d362760a28282": {
     "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
     }
    },
    "6c2af273159547a392089c5834c9cd8e": {
     "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_93f49a957d674ccc8ef402b39e9a7d91",
      "placeholder": "​",
      "style": "IPY_MODEL_951846e55fba4a6290c769c9d35645f9",
      "value": "config.json: 100%"
     }
    },
    "734f52c9430a424683d91dffbac1347e": {
     "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
     }
    },
    "778a064b9d7c4a8ea92186810264dc08": {
     "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
     }
    },
    "7f616d3fcd2c42539d9d7787213e2c57": {
     "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
     }
    },
    "8032cb188c9f41f388588cbfe2e755b1": {
     "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_42541f3e57ac4d8cafe6218c7c476ee5",
      "placeholder": "​",
      "style": "IPY_MODEL_a4a60902c35f4c9680e1821a02d3d289",
      "value": "vocab.txt: 100%"
     }
    },
    "8167c8c4fbfb4269bfa7af77ec077408": {
     "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_1f1bd137a3cf430c8ca78fc0b3a3021b",
       "IPY_MODEL_32919d67aed742809f66ae64c5079acd",
       "IPY_MODEL_c275d3e607e64079a303039928bcf035"
      ],
      "layout": "IPY_MODEL_1d0e586a34f74cac967844ef2bce71f1"
     }
    },
    "9118bb97895a4cd79613e692c8c74b3a": {
     "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": ""
     }
    },
    "91e23b7000104e77b7d3a88ba7662818": {
     "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_8032cb188c9f41f388588cbfe2e755b1",
       "IPY_MODEL_a2060bd8e3e04cfd8a6b9b7ee3ab8822",
       "IPY_MODEL_4106fe16851c4d9d8b84a005ba7bde46"
      ],
      "layout": "IPY_MODEL_eb5b52d870ac484eb8ea1e194b4ef33b"
     }
    },
    "93f49a957d674ccc8ef402b39e9a7d91": {
     "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
     }
    },
    "951846e55fba4a6290c769c9d35645f9": {
     "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": ""
     }
    },
    "99c57c8ce08146e4b2787abb65015480": {
     "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_d997312ea272454a98774230b731b748",
       "IPY_MODEL_c3fbf88ac32f4674a4c857a74014aaa4",
       "IPY_MODEL_c416ccbfd6674c1b9e72a62ecd7532d7"
      ],
      "layout": "IPY_MODEL_4067cbb94b4a4726aeb6ed1c3be9e851"
     }
    },
    "9a24afcbc4e54166974629ad969a5a55": {
     "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
     }
    },
    "a203d336509d4f74999af5b973959fef": {
     "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
     }
    },
    "a2060bd8e3e04cfd8a6b9b7ee3ab8822": {
     "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_2040f19e32cd48289fb3912754f521f2",
      "max": 995526,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_02e03aa708cc4f6fae60796c40b8c237",
      "value": 995526
     }
    },
    "a4a60902c35f4c9680e1821a02d3d289": {
     "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": ""
     }
    },
    "acbc0736f23147e19a05e63a1f1ed31d": {
     "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": ""
     }
    },
    "ad350891065d4e29b73b5bf29e494b09": {
     "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_a203d336509d4f74999af5b973959fef",
      "placeholder": "​",
      "style": "IPY_MODEL_4c56b7a8dbac40319087ad029ae72a1c",
      "value": "tokenizer_config.json: 100%"
     }
    },
    "afaf1cb0ab0343e9a933c1a6359c05b9": {
     "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
     }
    },
    "b430967240c24db39e4d92da89bec23a": {
     "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_08639571c1224db687cfaefb7a80fa6f",
      "placeholder": "​",
      "style": "IPY_MODEL_5f71cb7313124cf8bafd6ce70a887df4",
      "value": " 625/625 [00:00&lt;00:00, 23.7kB/s]"
     }
    },
    "b84a9d3290784b57b0a2af9215992c9d": {
     "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_03cc91e0c3da4e978951506b78633661",
      "placeholder": "​",
      "style": "IPY_MODEL_26ffdcdba6544f0cb815c83e7466932c",
      "value": " 49.0/49.0 [00:00&lt;00:00, 1.91kB/s]"
     }
    },
    "c22396edfa6a45bd9ddc25da8293a5f9": {
     "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": ""
     }
    },
    "c275d3e607e64079a303039928bcf035": {
     "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_60e2f5ec5c934f1ba33d362760a28282",
      "placeholder": "​",
      "style": "IPY_MODEL_22abab758f7e4e0a826b20a83c408916",
      "value": " 5.35G/5.35G [00:46&lt;00:00, 89.4MB/s]"
     }
    },
    "c298686a944a4db6912f818b21538172": {
     "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": ""
     }
    },
    "c3fbf88ac32f4674a4c857a74014aaa4": {
     "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_ca4a2016a7ad4ba9a0fb43eadb30c90d",
      "max": 1961828,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_c22396edfa6a45bd9ddc25da8293a5f9",
      "value": 1961828
     }
    },
    "c416ccbfd6674c1b9e72a62ecd7532d7": {
     "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_734f52c9430a424683d91dffbac1347e",
      "placeholder": "​",
      "style": "IPY_MODEL_9118bb97895a4cd79613e692c8c74b3a",
      "value": " 1.96M/1.96M [00:00&lt;00:00, 29.1MB/s]"
     }
    },
    "ca4a2016a7ad4ba9a0fb43eadb30c90d": {
     "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
     }
    },
    "cce2a28fb6e0455889c04575fbb68564": {
     "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": ""
     }
    },
    "d4b0fdf3a7a945619c6905a48725a8ca": {
     "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_ad350891065d4e29b73b5bf29e494b09",
       "IPY_MODEL_df8d89932919459996b19a875b453dda",
       "IPY_MODEL_b84a9d3290784b57b0a2af9215992c9d"
      ],
      "layout": "IPY_MODEL_778a064b9d7c4a8ea92186810264dc08"
     }
    },
    "d997312ea272454a98774230b731b748": {
     "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_afaf1cb0ab0343e9a933c1a6359c05b9",
      "placeholder": "​",
      "style": "IPY_MODEL_02a335585b324d2499ce59979d773703",
      "value": "tokenizer.json: 100%"
     }
    },
    "df24a18ed26d4bf3a5e4f249e8cb9c78": {
     "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
     }
    },
    "df8d89932919459996b19a875b453dda": {
     "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_3f22d50b47e8447a89fd4333cf870582",
      "max": 49,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_cce2a28fb6e0455889c04575fbb68564",
      "value": 49
     }
    },
    "eabcea8b31b04a88a63fd996cf90c14b": {
     "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
     }
    },
    "eb5b52d870ac484eb8ea1e194b4ef33b": {
     "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
     }
    },
    "ec673bba5a794916bc461a2ce0ab77c7": {
     "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
     }
    },
    "ef5b9993556b478990a7f68b3b4cb06c": {
     "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_ec673bba5a794916bc461a2ce0ab77c7",
      "placeholder": "​",
      "style": "IPY_MODEL_54c63aa200014b2c9dba18393f22be6c",
      "value": " 0/10750 [00:00&lt;?, ?it/s, loss=1.48, lr=0]"
     }
    },
    "f3b79041862a4609865050451a18d110": {
     "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
     }
    },
    "fc1c06851cc34694a223b1320f98af5b": {
     "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": ""
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}