EnBergman_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Clean Wiki"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "from datasets import load_dataset, load_from_disk\n",
    "from tqdm.autonotebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "wikipedia_dataset = load_dataset(\"wikipedia\", \"20220301.en\", split=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wikipedia_dataset = wikipedia_dataset.remove_columns(\n",
    "    [col for col in wikipedia_dataset.column_names if col != \"text\"]\n",
    ")  # only keep the 'text' column"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wikipedia_dataset = wikipedia_dataset.map(lambda x: {\"len\": len(x[\"text\"])}, num_proc=24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(wikipedia_dataset[\"len\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# wikipedia_dataset = wikipedia_dataset.select(range(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !python -m spacy download en_core_web_sm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import spacy\n",
    "nlp = spacy.load(\"en_core_web_sm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_see_also(sents):\n",
    "    start_pos = -1\n",
    "    end_pos = -1\n",
    "    for i, s in enumerate(sents):\n",
    "        if len(s) < 9 and \"See also\" in s:\n",
    "            start_pos = i\n",
    "            continue\n",
    "        if start_pos > 0:\n",
    "            if s[0] == \" \":\n",
    "                end_pos = i\n",
    "            else:\n",
    "                break\n",
    "    if start_pos < 0:\n",
    "        return sents\n",
    "    if end_pos - start_pos < 1:\n",
    "        return sents\n",
    "    sents[start_pos] = sents[start_pos] + \": \" + \", \".join(sents[start_pos+1:end_pos]) + \".\"\n",
    "    sents = sents[: start_pos + 1]\n",
    "    return sents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_references(sents):\n",
    "    start_pos = -1\n",
    "    for i, s in enumerate(sents):\n",
    "        if len(s) < 12 and 'References' in s:\n",
    "            start_pos = i\n",
    "            break\n",
    "    return sents[:start_pos]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_empty(sents, th=3):\n",
    "    return [s for s in sents if len(s) > th]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clean_wiki(examples):\n",
    "    sents_merged = []\n",
    "    for text in examples[\"text\"]:\n",
    "        paragraphs = text.split(\"\\n\")\n",
    "        sents = [sent.text for p in paragraphs for sent in nlp(p).sents]\n",
    "        sents = filter_empty(sents)\n",
    "        sents = process_see_also(sents)\n",
    "        sents = remove_references(sents)\n",
    "        sents_merged.extend(sents)\n",
    "    return {\"text\": sents_merged}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wikipedia_dataset = wikipedia_dataset.map(\n",
    "    lambda x: clean_wiki(x),\n",
    "    batched=True,\n",
    "    remove_columns=wikipedia_dataset.column_names\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wikipedia_dataset.save_to_disk(\"wikipedia_dataset_cleaned.hf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# bookcorpusopen dataset cleaning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "from datasets import load_dataset, load_from_disk\n",
    "from tqdm.autonotebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "bookcorpus_dataset = load_dataset(\"bookcorpusopen\", \"plain_text\", split=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bookcorpus_dataset = bookcorpus_dataset.remove_columns(\n",
    "    [col for col in bookcorpus_dataset.column_names if col != \"text\"]\n",
    ")  # only keep the 'text' column"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import spacy\n",
    "nlp = spacy.load(\"en_core_web_sm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_empty(sents, th=3):\n",
    "    return [s for s in sents if len(s) > th]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clean_bookcorpus(examples):\n",
    "    sents_merged = []\n",
    "    for text in examples[\"text\"]:\n",
    "        paragraphs = text.split(\"\\n\")\n",
    "        sents = [sent.text for p in paragraphs for sent in nlp(p).sents]\n",
    "        sents = filter_empty(sents)\n",
    "        sents_merged.extend(sents)\n",
    "    return {\"text\": sents_merged}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bookcorpus_dataset = bookcorpus_dataset.map(\n",
    "    clean_bookcorpus,\n",
    "    batched=True,\n",
    "    remove_columns=bookcorpus_dataset.column_names,\n",
    "    batch_size=1,\n",
    "    num_proc=24\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bookcorpus_dataset.save_to_disk(\"bookcorpus_dataset_cleaned.hf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 100
    },
    "id": "e67Ut53QYEdU",
    "outputId": "437871b8-b8ac-4eaf-c2e1-61d801c5e6b2"
   },
   "source": [
    "Notebook is based on https://huggingface.co/blog/how-to-train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oK7PPVm2XBgr"
   },
   "source": [
    "## Prepare tokenized dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "from datasets import load_dataset, load_from_disk\n",
    "from tqdm.autonotebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "wikipedia_dataset = load_from_disk(\"wikipedia_dataset_cleaned.hf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "bookcorpus_dataset = load_from_disk(\"bookcorpus_dataset_cleaned.hf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['text'],\n",
       "    num_rows: 205447996\n",
       "})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wikipedia_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['text'],\n",
       "    num_rows: 97892049\n",
       "})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bookcorpus_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# wikipedia_dataset = wikipedia_dataset.map(lambda x: {\"len\": len(x[\"text\"])}, num_proc=24)\n",
    "# sum(wikipedia_dataset[\"len\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert bookcorpus_dataset.features.type == wikipedia_dataset.features.type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_datasets = datasets.combine.concatenate_datasets([wikipedia_dataset, bookcorpus_dataset])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "beeb8266253c42c9b077b7f7ec6dbf0c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)olve/main/vocab.json:   0%|          | 0.00/899k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ac9fd7e1f5e45d78824ab0a64c70e76",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)olve/main/merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5edb03e63679402b9ca14a2930c9649c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)/main/tokenizer.json:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "70f0c8fb732644e694be5a7d4100b69d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)lve/main/config.json:   0%|          | 0.00/481 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import RobertaTokenizerFast\n",
    "tokenizer = RobertaTokenizerFast.from_pretrained(\"roberta-base\", use_fast=True, add_prefix_space=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tokenizers import normalizers\n",
    "from tokenizers.normalizers import BertNormalizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.backend_tokenizer.normalizer = normalizers.Sequence([BertNormalizer()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_training_corpus():\n",
    "    return (\n",
    "        raw_datasets[i : i + 1000][\"text\"]\n",
    "        for i in tqdm(range(0, len(raw_datasets), 1000))\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.tokenize(\"Hello world\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_corpus = get_training_corpus()\n",
    "tokenizer = tokenizer.train_new_from_iterator(training_corpus, vocab_size=32_768)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.tokenize(\"Hello world\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.save_pretrained(\"bergman-tokenizer_32k\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import RobertaTokenizerFast\n",
    "\n",
    "tokenizer = RobertaTokenizerFast.from_pretrained(\"./bergman-tokenizer_32k/\", max_len=512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=12):   0%|          | 0/303340045 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Token indices sequence length is longer than the specified maximum sequence length for this model (525 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (636 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (570 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (614 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (703 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (516 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (780 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (525 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (763 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (535 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (520 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (881 > 512). Running this sequence through the model will result in indexing errors\n"
     ]
    }
   ],
   "source": [
    "raw_datasets = raw_datasets.map(\n",
    "    lambda examples: tokenizer(examples[\"text\"], return_special_tokens_mask=True),\n",
    "    batched=True,\n",
    "    remove_columns=raw_datasets.column_names,\n",
    "    num_proc=24,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "22395bb968864c1f88bd611bedc77dd3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/81 shards):   0%|          | 0/303340045 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "raw_datasets.save_to_disk(\"raw_dataset_tokenized_32k.hf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WQpUC_CDhnWW"
   },
   "source": [
    "# Train a language model from scratch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "from datasets import load_dataset, load_from_disk\n",
    "from tqdm.autonotebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_datasets = load_from_disk(\"raw_dataset_tokenized_32k.hf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u0qQzgrBi1OX"
   },
   "source": [
    "### We'll define the following config for the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "4keFBUjQFOD1"
   },
   "outputs": [],
   "source": [
    "from transformers import RobertaTokenizerFast\n",
    "tokenizer = RobertaTokenizerFast.from_pretrained(\"./bergman-tokenizer_32k/\", max_len=512)\n",
    "\n",
    "from tokenizers import normalizers\n",
    "from tokenizers.normalizers import BertNormalizer\n",
    "tokenizer.backend_tokenizer.normalizer = normalizers.Sequence([BertNormalizer()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "LTXXutqeDzPi"
   },
   "outputs": [],
   "source": [
    "from bergman import BergmanConfig\n",
    "\n",
    "# # Bergman_Apr02_06-14-51_raven_200_000\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=4,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=4,\n",
    "#     num_matrix_heads=32,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "# )\n",
    "\n",
    "# # Apr06_23-21-44_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=1,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=8,\n",
    "#     num_matrix_heads=96,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "# )\n",
    "\n",
    "# # Apr07_03-18-33_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=2,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=8,\n",
    "#     num_matrix_heads=48,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "# )\n",
    "\n",
    "# # Apr07_15-21-03_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=1,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=4,\n",
    "#     num_matrix_heads=48,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "# )\n",
    "\n",
    "# # Apr07_16-39-44_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=2,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=4,\n",
    "#     num_matrix_heads=48,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "# )\n",
    "\n",
    "# # Apr11_20-08-57_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=6,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=4,\n",
    "#     num_matrix_heads=64,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "#     rl_lr_matrix_different=True,\n",
    "# )\n",
    "\n",
    "# # Bergman_Jun21_16-21-09/\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=3,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=4,\n",
    "#     num_matrix_heads=96,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "#     rl_lr_matrix_different=False,\n",
    "# )\n",
    "\n",
    "# # Aug13_16-17-37_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=5,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=8,\n",
    "#     num_matrix_heads=48,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "#     rl_lr_matrix_different=True,\n",
    "# )\n",
    "\n",
    "# #  Aug14_14-56-05_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=3,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=4,\n",
    "#     num_matrix_heads=96,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "#     rl_lr_matrix_different=True,\n",
    "# )\n",
    "\n",
    "# # Aug23_13-22-56_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=3,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=4,\n",
    "#     num_matrix_heads=96,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "#     rl_lr_matrix_different=False,\n",
    "#     matrix_encoder_activation=\"softmax\",\n",
    "#     matrix_encoder_hidden_size=768,\n",
    "# )\n",
    "\n",
    "# # Aug25_02-56-37_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=3,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=4,\n",
    "#     num_matrix_heads=96,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "#     rl_lr_matrix_different=False,\n",
    "#     matrix_encoder_hidden_size=32,\n",
    "#     matrix_encoder_version=2\n",
    "# )\n",
    "\n",
    "# # Aug28_04-02-42_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=3,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=8,\n",
    "#     num_matrix_heads=32,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "#     rl_lr_matrix_different=False,\n",
    "#     matrix_encoder_hidden_size=64,\n",
    "#     matrix_encoder_version=2\n",
    "# )\n",
    "\n",
    "# Aug28_04-02-42_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=3,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=4,\n",
    "#     num_matrix_heads=96,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "#     rl_lr_matrix_different=False,\n",
    "#     matrix_encoder_hidden_size=16,\n",
    "#     matrix_encoder_version=2\n",
    "# )\n",
    "\n",
    "# # Nov06_06-55-42_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=3,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=4,\n",
    "#     num_matrix_heads=96,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "#     rl_lr_matrix_different=False,\n",
    "#     matrix_encoder_hidden_size=16,\n",
    "#     matrix_encoder_version=2,\n",
    "#     matrix_encoder_v2_softdiff=True,\n",
    "# )\n",
    "\n",
    "\n",
    "# # Nov07_19-56-06_raven\n",
    "# config = BergmanConfig(\n",
    "#     vocab_size=tokenizer.vocab_size,\n",
    "#     max_position_embeddings=512,\n",
    "#     num_hidden_layers=3,\n",
    "#     type_vocab_size=1,\n",
    "#     hidden_size=768,\n",
    "#     position_embedding_type=\"none\",\n",
    "#     matrix_norm_alg=None,\n",
    "#     matrix_dim=4,\n",
    "#     num_matrix_heads=96,\n",
    "#     vector_init_direction=\"one\",\n",
    "#     use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "#     networks_for_heads=\"common\",\n",
    "#     matrix_encoder_two_layers=True,\n",
    "#     #\n",
    "#     matrix_norm_loss_type=None,\n",
    "#     matrix_norm_loss_k=0.0,\n",
    "#     matrix_unitary_loss=None,\n",
    "#     matrix_unitary_loss_k = 0.0,\n",
    "#     norm_vectors=True,\n",
    "#     complex_matrix=True,\n",
    "#     complex_matrix_abs=True,\n",
    "#     rl_lr_matrix_different=False,\n",
    "#     input_convnet_filter_size=4,\n",
    "# )\n",
    "\n",
    "# Nov08_21-32-13_raven -> Nov12_05-23-52_raven\n",
    "config = BergmanConfig(\n",
    "    vocab_size=tokenizer.vocab_size,\n",
    "    max_position_embeddings=512,\n",
    "    num_hidden_layers=2,\n",
    "    type_vocab_size=1,\n",
    "    hidden_size=768,\n",
    "    position_embedding_type=\"none\",\n",
    "    matrix_norm_alg=None,\n",
    "    matrix_dim=4,\n",
    "    num_matrix_heads=96,\n",
    "    vector_init_direction=\"one\",\n",
    "    use_for_context=[\"lr_excl\", \"rl_excl\"],\n",
    "    networks_for_heads=\"common\",\n",
    "    matrix_encoder_two_layers=True,\n",
    "    #\n",
    "    matrix_norm_loss_type=None,\n",
    "    matrix_norm_loss_k=0.0,\n",
    "    matrix_unitary_loss=None,\n",
    "    matrix_unitary_loss_k = 0.0,\n",
    "    norm_vectors=True,\n",
    "    complex_matrix=True,\n",
    "    complex_matrix_abs=True,\n",
    "    rl_lr_matrix_different=False,\n",
    "    input_convnet_filter_size=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yAwQ82JiE5pi"
   },
   "source": [
    "Now let's re-create our tokenizer in transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6yNCw-3hFv9h"
   },
   "source": [
    "Finally let's initialize our model.\n",
    "\n",
    "**Important:**\n",
    "\n",
    "As we are training from scratch, we only initialize from a config, not from an existing pretrained model or checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "BzMqR-dzF4Ro"
   },
   "outputs": [],
   "source": [
    "from bergman import BergmanForMaskedLM\n",
    "\n",
    "# model = BergmanForMaskedLM(config=config)\n",
    "\n",
    "model = BergmanForMaskedLM.from_pretrained(\"./Bergman_Dec24_04-42-55_raven_120000/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "id": "jU6JhBSTKiaM",
    "outputId": "35879a60-2915-4894-f702-2d649cfa398a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "48631808"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.num_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# for name, param in model.named_parameters():\n",
    "#     if param.requires_grad:\n",
    "#         print(name, param.size(), param.numel())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jBtUHRMliOLM"
   },
   "source": [
    "### Now let's build our training Dataset\n",
    "\n",
    "We'll build our dataset by applying our tokenizer to our text file.\n",
    "\n",
    "Here, as we only have one text file, we don't even need to customize our `Dataset`. We'll just use the `LineByLineDataset` out-of-the-box."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiprocessing\n",
    "num_proc = multiprocessing.cpu_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataset = raw_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_seq_length = tokenizer.model_max_length\n",
    "max_seq_length = 150"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "merge_texts = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def group_texts(examples, max_seq_length, merge_texts):\n",
    "    \"\"\"\n",
    "    >>> group_texts({\"a\": [list(range(5))]}, 4, True)\n",
    "    {'a': [[0, 1, 2, 4], [0, 3, 4]]}\n",
    "    >>> group_texts({\"a\": [list(range(3)), list(range(4))]}, 5, True)\n",
    "    {'a': [[0, 1, 1, 2, 3]]}\n",
    "    >>> group_texts({\"a\": [list(range(3)), list(range(4))]}, 5, False)\n",
    "    {'a': [[0, 1, 2], [0, 1, 2, 3]]}\n",
    "    >>> group_texts({\"a\": [list(range(4)), list(range(4))]}, 5, True)\n",
    "    {'a': [[0, 1, 2, 3], [0, 1, 2, 3]]}\n",
    "    \"\"\"\n",
    "    # Concatenate all texts.\n",
    "    result = {}\n",
    "    for k, v in examples.items():\n",
    "        acc = []\n",
    "        for text in v:\n",
    "            if (\n",
    "                len(acc) > 0\n",
    "                and len(acc[-1]) + len(text) - 2 <= max_seq_length\n",
    "                and merge_texts\n",
    "            ):\n",
    "                acc[-1].pop()  # remove </s>\n",
    "                acc[-1].extend(text[1:])  # remove <s>\n",
    "            else:\n",
    "                b = text[0]\n",
    "                e = text[-1]\n",
    "                content = text[1:-1]\n",
    "                for i in range((len(content)) // (max_seq_length - 2) + 1):\n",
    "                    body = content[\n",
    "                        (max_seq_length - 2) * i : (i + 1) * (max_seq_length - 2)\n",
    "                    ]\n",
    "                    if len(body) > 0:\n",
    "                        acc.append(\n",
    "                            [b]  # <s> or corresponding mask\n",
    "                            + body\n",
    "                            + [e]  # </s> or corresponding mask\n",
    "                        )\n",
    "        result[k] = acc\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/eugene/Projects/matrix_network/raw_dataset_tokenized_32k.hf/cache-7a818b2d5092987a_*_of_00024.arrow\n"
     ]
    }
   ],
   "source": [
    "dataset = dataset.map(\n",
    "    lambda x: group_texts(x, max_seq_length, merge_texts),\n",
    "    batched=True,\n",
    "    num_proc=num_proc,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.set_format(type=\"torch\", columns=[\"input_ids\"])  # , 'special_tokens_mask'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<s> anarchism is a political philosophy and movement that is sceptical of authority and rejects all involuntary, coercive forms of hierarchy. anarchism calls for the abolition of the state, which it holds to be unnecessary, undesirable, and harmful. as a historically left-wing movement, placed on the farthest left of the political spectrum, it is usually described alongside communalism and libertarian marxism as the libertarian wing (libertarian socialism) of the socialist movement, and has a strong historical association with anti-capitalism and socialism. humans lived in societies without formal hierarchies long before the establishment of formal states, realms, or empires. with the rise of organised hierarchical bodies, scepticism toward authority also rose.</s>'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.decode(dataset[0][\"input_ids\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = dataset.shuffle()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hDLs73HcIHk5"
   },
   "source": [
    "Like in the [`run_language_modeling.py`](https://github.com/huggingface/transformers/blob/master/examples/language-modeling/run_language_modeling.py) script, we need to define a data_collator.\n",
    "\n",
    "This is just a small helper that will help us batch different samples of the dataset together into an object that PyTorch knows how to perform backprop on."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hDLs73HcIHk5"
   },
   "source": [
    "Like in the [`run_language_modeling.py`](https://github.com/huggingface/transformers/blob/master/examples/language-modeling/run_language_modeling.py) script, we need to define a data_collator.\n",
    "\n",
    "This is just a small helper that will help us batch different samples of the dataset together into an object that PyTorch knows how to perform backprop on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "zTgWPa9Dipk2"
   },
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForLanguageModeling\n",
    "\n",
    "data_collator = DataCollatorForLanguageModeling(\n",
    "    tokenizer=tokenizer, mlm=True, mlm_probability=0.15\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ri2BIQKqjfHm"
   },
   "source": [
    "### Finally, we are all set to initialize our Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Trainer, TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./Bergman\",\n",
    "    overwrite_output_dir=True,\n",
    "    num_train_epochs=1,\n",
    "    per_device_train_batch_size=62,\n",
    "    save_steps=10_000,\n",
    "    save_total_limit=2,\n",
    "    prediction_loss_only=True,\n",
    "    logging_steps=500,\n",
    "    learning_rate=5E-6,\n",
    "    weight_decay=0.01,\n",
    "    disable_tqdm=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers.trainer import (\n",
    "    MODEL_FOR_CAUSAL_LM_MAPPING_NAMES,\n",
    "    is_torch_tpu_available,\n",
    ")\n",
    "import torch\n",
    "\n",
    "\n",
    "class BergmanTrainer(Trainer):\n",
    "    def compute_loss(self, model, inputs, return_outputs=False):\n",
    "        \"\"\"\n",
    "        How the loss is computed by Trainer. By default, all models return the loss in the first element.\n",
    "\n",
    "        Subclass and override for custom behavior.\n",
    "        \"\"\"\n",
    "        if self.label_smoother is not None and \"labels\" in inputs:\n",
    "            labels = inputs.pop(\"labels\")\n",
    "        else:\n",
    "            labels = None\n",
    "        outputs = model(**inputs)\n",
    "        # Save past state if it exists\n",
    "        # TODO: this needs to be fixed and made cleaner later.\n",
    "        if self.args.past_index >= 0:\n",
    "            self._past = outputs[self.args.past_index]\n",
    "\n",
    "        metrics = outputs[\"metrics\"] if isinstance(outputs, dict) else outputs[-1]\n",
    "        self.metrics = {\n",
    "            m: v if isinstance(v, float) else v.detach() for m, v in metrics.items()\n",
    "        }\n",
    "\n",
    "        if labels is not None:\n",
    "            if (\n",
    "                unwrap_model(model)._get_name()\n",
    "                in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values()\n",
    "            ):\n",
    "                loss = self.label_smoother(outputs, labels, shift_labels=True)\n",
    "            else:\n",
    "                loss = self.label_smoother(outputs, labels)\n",
    "        else:\n",
    "            if isinstance(outputs, dict) and \"loss\" not in outputs:\n",
    "                raise ValueError(\n",
    "                    \"The model did not return a loss from the inputs, only the following keys: \"\n",
    "                    f\"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}.\"\n",
    "                )\n",
    "            # We don't use .loss here since the model may return tuples instead of ModelOutput.\n",
    "            loss = outputs[\"loss\"] if isinstance(outputs, dict) else outputs[0]\n",
    "\n",
    "        return (loss, outputs) if return_outputs else loss\n",
    "\n",
    "    def _maybe_log_save_evaluate(\n",
    "        self, tr_loss, model, trial, epoch, ignore_keys_for_eval\n",
    "    ):\n",
    "        if not hasattr(self, \"metrics_acc\"):\n",
    "            self.metrics_acc: Dict[str, torch.Tensor] = {}\n",
    "\n",
    "        for m, v in self.metrics.items():\n",
    "            if v is None:\n",
    "                continue\n",
    "            if m not in self.metrics_acc:\n",
    "                self.metrics_acc[m] = torch.tensor(0.0).to(model.device)\n",
    "            self.metrics_acc[m] += v\n",
    "\n",
    "        if self.control.should_log:\n",
    "            if is_torch_tpu_available():\n",
    "                xm.mark_step()\n",
    "\n",
    "            metrics = {\n",
    "                m: self._nested_gather(v).mean().item()\n",
    "                for m, v in self.metrics_acc.items()\n",
    "            }\n",
    "            # reset counters\n",
    "            self.metrics_acc = {}\n",
    "\n",
    "            logs = {\n",
    "                m: round(\n",
    "                    v / (self.state.global_step - self._globalstep_last_logged),\n",
    "                    4,\n",
    "                )\n",
    "                for m, v in metrics.items()\n",
    "            }\n",
    "\n",
    "            # all_gather + mean() to get average loss over all processes\n",
    "            tr_loss_scalar = self._nested_gather(tr_loss).mean().item()\n",
    "\n",
    "            # reset tr_loss to zero\n",
    "            tr_loss -= tr_loss\n",
    "\n",
    "            logs[\"loss\"] = round(\n",
    "                tr_loss_scalar\n",
    "                / (self.state.global_step - self._globalstep_last_logged),\n",
    "                4,\n",
    "            )\n",
    "            logs[\"learning_rate\"] = self._get_learning_rate()\n",
    "\n",
    "            self._total_loss_scalar += tr_loss_scalar\n",
    "            self._globalstep_last_logged = self.state.global_step\n",
    "            self.store_flos()\n",
    "\n",
    "            self.log(logs)\n",
    "\n",
    "        metrics = None\n",
    "        if self.control.should_evaluate:\n",
    "            if isinstance(self.eval_dataset, dict):\n",
    "                for eval_dataset_name, eval_dataset in self.eval_dataset.items():\n",
    "                    metrics = self.evaluate(\n",
    "                        eval_dataset=eval_dataset,\n",
    "                        ignore_keys=ignore_keys_for_eval,\n",
    "                        metric_key_prefix=f\"eval_{eval_dataset_name}\",\n",
    "                    )\n",
    "            else:\n",
    "                metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)\n",
    "            self._report_to_hp_search(trial, self.state.global_step, metrics)\n",
    "\n",
    "        if self.control.should_save:\n",
    "            self._save_checkpoint(model, trial, metrics=metrics)\n",
    "            self.control = self.callback_handler.on_save(\n",
    "                self.args, self.state, self.control\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "YpvnFFmZJD-N"
   },
   "outputs": [],
   "source": [
    "trainer = BergmanTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    data_collator=data_collator,\n",
    "    train_dataset=dataset,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o6sASa36Nf-N"
   },
   "source": [
    "### Start training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 738,
     "referenced_widgets": [
      "a58a66392b644b1384661e850c077a6c",
      "a491e8caa0a048beb3b5259f14eb233f",
      "837c9ddc3d594e088891874560c646b8",
      "dbf50873d62c4ba39321faefbed0cca5",
      "40bf955ba0284e84b198da6be8654219",
      "fe20a8dae6e84628b5076d02183090f5",
      "93b3f9eae3cb4e3e859cf456e3547c6d",
      "6feb10aeb43147e6aba028d065947ae8",
      "0989d41a4da24e9ebff377e02127642c",
      "42c6061ef7e44f179db5a6e3551c0f17",
      "d295dd80550447d88da0f04ce36a22ff",
      "04e7e6d291da49d5816dc98a2904e95c",
      "e7d8c3a4fecd40778e32966b29ea65a1",
      "016d7c8318f742c1943464b08232a510",
      "8388e9da9da4492c98c19235ca5fc1b5",
      "39c23c6a972b419eb2eeeebafeaedc22"
     ]
    },
    "id": "VmaHZXzmkNtJ",
    "outputId": "a19880cb-bcc6-4885-bf24-c2c6d0f56d1e",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following columns in the training set don't have a corresponding argument in `BergmanForMaskedLM.forward` and have been ignored: special_tokens_mask. If special_tokens_mask are not expected by `BergmanForMaskedLM.forward`,  you can safely ignore this message.\n",
      "/home/eugene/Projects/matrix_network/transformers/src/transformers/optimization.py:306: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "***** Running training *****\n",
      "  Num examples = 41096777\n",
      "  Num Epochs = 1\n",
      "  Instantaneous batch size per device = 62\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 62\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 662852\n",
      "  Number of trainable parameters = 48631808\n",
      "You're using a RobertaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n",
      "/home/eugene/Projects/matrix_network/venv/lib/python3.8/site-packages/torch/nn/modules/conv.py:309: UserWarning: Using padding='same' with even kernel lengths and odd dilation may require a zero-padded copy of the input be created (Triggered internally at ../aten/src/ATen/native/Convolution.cpp:895.)\n",
      "  return F.conv1d(input, weight, bias, self.stride,\n"
     ]
    }
   ],
   "source": [
    "# %#%time\n",
    "# with torch.autograd.detect_anomaly(True):\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_ZkooHz1-_2h"
   },
   "source": [
    "#### 🎉 Save final model (+ tokenizer + config) to disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QDNgPls7_l13"
   },
   "outputs": [],
   "source": [
    "trainer.save_model(\"./Bergman\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raise Exception()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d0caceCy_p1-"
   },
   "source": [
    "## 4. Check that the LM actually trained"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iIQJ8ND_AEhl"
   },
   "source": [
    "Aside from looking at the training and eval losses going down, the easiest way to check whether our language model is learning anything interesting is via the `FillMaskPipeline`.\n",
    "\n",
    "Pipelines are simple wrappers around tokenizers and models, and the 'fill-mask' one will let you input a sequence containing a masked token (here, `<mask>`) and return a list of the most probable filled sequences, with their probabilities.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'model' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[38;5;241m.\u001b[39mfrom_pretrained(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mBergman_Dec24_04-42-55_raven_120000/\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'model' is not defined"
     ]
    }
   ],
   "source": [
    "model = model.from_pretrained(\"Bergman_Dec24_04-42-55_raven_120000/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model.to(\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "ltXgXyCbAJLY"
   },
   "outputs": [],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "fill_mask = pipeline(\n",
    "    \"fill-mask\",\n",
    "    model=model,\n",
    "    tokenizer=tokenizer\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 283
    },
    "id": "UIvgZ3S6AO0z",
    "outputId": "5f3d2f00-abdc-44a9-9c1b-75e3ec328576"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/eugene/Projects/matrix_network/venv/lib/python3.8/site-packages/torch/nn/modules/conv.py:309: UserWarning: Using padding='same' with even kernel lengths and odd dilation may require a zero-padded copy of the input be created (Triggered internally at ../aten/src/ATen/native/Convolution.cpp:895.)\n",
      "  return F.conv1d(input, weight, bias, self.stride,\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'score': 0.06418540328741074,\n",
       "  'token': 3120,\n",
       "  'token_str': ' itself',\n",
       "  'sequence': ' while the term itself has been largely synonymous with anarchism'},\n",
       " {'score': 0.023493576794862747,\n",
       "  'token': 225,\n",
       "  'token_str': ' ',\n",
       "  'sequence': ' while the term  has been largely synonymous with anarchism'},\n",
       " {'score': 0.01783517189323902,\n",
       "  'token': 641,\n",
       "  'token_str': 'ite',\n",
       "  'sequence': ' while the termite has been largely synonymous with anarchism'},\n",
       " {'score': 0.01671333983540535,\n",
       "  'token': 2996,\n",
       "  'token_str': ' term',\n",
       "  'sequence': ' while the term term has been largely synonymous with anarchism'},\n",
       " {'score': 0.015532868914306164,\n",
       "  'token': 347,\n",
       "  'token_str': ' it',\n",
       "  'sequence': ' while the term it has been largely synonymous with anarchism'}]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The sun <mask>.\n",
    "# =>\n",
    "\n",
    "fill_mask(\"while the term <mask> has been largely synonymous with anarchism\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'score': 0.04951208457350731,\n",
       "  'token': 6045,\n",
       "  'token_str': ' evolution',\n",
       "  'sequence': ' newton introduces the theory of evolution.'},\n",
       " {'score': 0.04289407283067703,\n",
       "  'token': 265,\n",
       "  'token_str': ' the',\n",
       "  'sequence': ' newton introduces the theory of the.'},\n",
       " {'score': 0.037284113466739655,\n",
       "  'token': 225,\n",
       "  'token_str': ' ',\n",
       "  'sequence': ' newton introduces the theory of .'},\n",
       " {'score': 0.0347602404654026,\n",
       "  'token': 26350,\n",
       "  'token_str': ' relativity',\n",
       "  'sequence': ' newton introduces the theory of relativity.'},\n",
       " {'score': 0.03271481767296791,\n",
       "  'token': 1360,\n",
       "  'token_str': '.',\n",
       "  'sequence': ' newton introduces the theory of..'}]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fill_mask(\"Newton introduces the theory of <mask>.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.onnx.export(model, torch.LongTensor([[0,0,0,0,0]]), 'Bergman.onnx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i0qCyyhNAWZi"
   },
   "source": [
    "Ok, simple syntax/grammar works. Let’s try a slightly more interesting prompt:\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6RsGaD1qAfLP"
   },
   "source": [
    "## 5. Share your model 🎉"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5oESe8djApQw"
   },
   "source": [
    "Finally, when you have a nice model, please think about sharing it with the community:\n",
    "\n",
    "- upload your model using the CLI: `transformers-cli upload`\n",
    "- write a README.md model card and add it to the repository under `model_cards/`. Your model card should ideally include:\n",
    "    - a model description,\n",
    "    - training params (dataset, preprocessing, hyperparameters), \n",
    "    - evaluation results,\n",
    "    - intended uses & limitations\n",
    "    - whatever else is helpful! 🤓\n",
    "\n",
    "### **TADA!**\n",
    "\n",
    "➡️ Your model has a page on http://huggingface.co/models and everyone can load it using `AutoModel.from_pretrained(\"username/model_name\")`.\n",
    "\n",
    "[![tb](https://huggingface.co/blog/assets/01_how-to-train/model_page.png)](https://huggingface.co/julien-c/EsperBERTo-small)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aw9ifsgqBI2o"
   },
   "source": [
    "If you want to take a look at models in different languages, check https://huggingface.co/models\n",
    "\n",
    "[![all models](https://huggingface.co/front/thumbnails/models.png)](https://huggingface.co/models)\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "machine_shape": "hm",
   "name": "01_how-to-train.ipynb",
   "provenance": [
    {
     "file_id": "https://github.com/huggingface/blog/blob/main/notebooks/01_how_to_train.ipynb",
     "timestamp": 1675426885377
    }
   ],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "016d7c8318f742c1943464b08232a510": {
     "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
     }
    },
    "04e7e6d291da49d5816dc98a2904e95c": {
     "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_39c23c6a972b419eb2eeeebafeaedc22",
      "placeholder": "​",
      "style": "IPY_MODEL_8388e9da9da4492c98c19235ca5fc1b5",
      "value": " 15228/15228 [2:46:46&lt;00:00,  1.52it/s]"
     }
    },
    "0989d41a4da24e9ebff377e02127642c": {
     "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_d295dd80550447d88da0f04ce36a22ff",
       "IPY_MODEL_04e7e6d291da49d5816dc98a2904e95c"
      ],
      "layout": "IPY_MODEL_42c6061ef7e44f179db5a6e3551c0f17"
     }
    },
    "39c23c6a972b419eb2eeeebafeaedc22": {
     "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
     }
    },
    "40bf955ba0284e84b198da6be8654219": {
     "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": "initial"
     }
    },
    "42c6061ef7e44f179db5a6e3551c0f17": {
     "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
     }
    },
    "6feb10aeb43147e6aba028d065947ae8": {
     "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
     }
    },
    "837c9ddc3d594e088891874560c646b8": {
     "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": "Epoch: 100%",
      "description_tooltip": null,
      "layout": "IPY_MODEL_fe20a8dae6e84628b5076d02183090f5",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_40bf955ba0284e84b198da6be8654219",
      "value": 1
     }
    },
    "8388e9da9da4492c98c19235ca5fc1b5": {
     "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": ""
     }
    },
    "93b3f9eae3cb4e3e859cf456e3547c6d": {
     "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": ""
     }
    },
    "a491e8caa0a048beb3b5259f14eb233f": {
     "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
     }
    },
    "a58a66392b644b1384661e850c077a6c": {
     "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_837c9ddc3d594e088891874560c646b8",
       "IPY_MODEL_dbf50873d62c4ba39321faefbed0cca5"
      ],
      "layout": "IPY_MODEL_a491e8caa0a048beb3b5259f14eb233f"
     }
    },
    "d295dd80550447d88da0f04ce36a22ff": {
     "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": "Iteration: 100%",
      "description_tooltip": null,
      "layout": "IPY_MODEL_016d7c8318f742c1943464b08232a510",
      "max": 15228,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_e7d8c3a4fecd40778e32966b29ea65a1",
      "value": 15228
     }
    },
    "dbf50873d62c4ba39321faefbed0cca5": {
     "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_6feb10aeb43147e6aba028d065947ae8",
      "placeholder": "​",
      "style": "IPY_MODEL_93b3f9eae3cb4e3e859cf456e3547c6d",
      "value": " 1/1 [2:46:46&lt;00:00, 10006.17s/it]"
     }
    },
    "e7d8c3a4fecd40778e32966b29ea65a1": {
     "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": "initial"
     }
    },
    "fe20a8dae6e84628b5076d02183090f5": {
     "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
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}