nlp/BERTLM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Masked Language Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The cache for model files in Transformers v4.22.0 has been updated. Migrating your old cache. This is a one-time only operation. You can interrupt this and resume the migration later on by calling `transformers.utils.move_cache()`.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8fe56304a1aa495caec7ab8e40d80536",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "522e1e36c77041239e2cbd1093f4536d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)lve/main/config.json:   0%|          | 0.00/483 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\lkk68\\.conda\\envs\\mycondapy39\\lib\\site-packages\\huggingface_hub\\file_download.py:133: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\lkk68\\.cache\\huggingface\\hub. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n",
      "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to see activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n",
      "  warnings.warn(message)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ad87e798cb0c4171a071192ddf6beb58",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading model.safetensors:   0%|          | 0.00/268M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import AutoModelForMaskedLM\n",
    "\n",
    "model_checkpoint = \"distilbert-base-uncased\"\n",
    "model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'>>> DistilBERT number of parameters: 67M'\n",
      "'>>> BERT number of parameters: 110M'\n"
     ]
    }
   ],
   "source": [
    "distilbert_num_parameters = model.num_parameters() / 1_000_000\n",
    "print(f\"'>>> DistilBERT number of parameters: {round(distilbert_num_parameters)}M'\")\n",
    "print(f\"'>>> BERT number of parameters: 110M'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"This is a great [MASK].\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c597a27bd2134188bc9455afaf4af33f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)okenizer_config.json:   0%|          | 0.00/28.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "62d9f4fbabd042f3ac87e9c8e74dec4f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)solve/main/vocab.txt:   0%|          | 0.00/232k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "40d6b2b3cf8345cbb4bab4ac9e6169f2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)/main/tokenizer.json:   0%|          | 0.00/466k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'>>> This is a great deal.'\n",
      "'>>> This is a great success.'\n",
      "'>>> This is a great adventure.'\n",
      "'>>> This is a great idea.'\n",
      "'>>> This is a great feat.'\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "token_logits = model(**inputs).logits\n",
    "# Find the location of [MASK] and extract its logits\n",
    "mask_token_index = torch.where(inputs[\"input_ids\"] == tokenizer.mask_token_id)[1]\n",
    "mask_token_logits = token_logits[0, mask_token_index, :]\n",
    "# Pick the [MASK] candidates with the highest logits\n",
    "top_5_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()\n",
    "\n",
    "for token in top_5_tokens:\n",
    "    print(f\"'>>> {text.replace(tokenizer.mask_token, tokenizer.decode([token]))}'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "248485d1dc4241dbb1e591d14a91dec7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading builder script:   0%|          | 0.00/4.31k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4547b0b9dcd24fddb6ca646fedb7ce42",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading metadata:   0%|          | 0.00/2.17k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a290e2ef85cb4c11a0c85e3730c7eaa3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading readme:   0%|          | 0.00/7.59k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and preparing dataset imdb/plain_text to C:/Users/lkk68/.cache/huggingface/datasets/imdb/plain_text/1.0.0/d613c88cf8fa3bab83b4ded3713f1f74830d1100e171db75bbddb80b3345c9c0...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5e98f31f9cff4cd39eec0fa27b022f8e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/84.1M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ef6f3b92cf384ae3b1ca3842a05762ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split:   0%|          | 0/25000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "48040ed3e76a4af484e35d0a0ff03e45",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating test split:   0%|          | 0/25000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3d70de22a7ce4dbc8687d5c014238a5d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating unsupervised split:   0%|          | 0/50000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset imdb downloaded and prepared to C:/Users/lkk68/.cache/huggingface/datasets/imdb/plain_text/1.0.0/d613c88cf8fa3bab83b4ded3713f1f74830d1100e171db75bbddb80b3345c9c0. Subsequent calls will reuse this data.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0d15ce0742f74a00acd54e2a2650d70d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['text', 'label'],\n",
       "        num_rows: 25000\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['text', 'label'],\n",
       "        num_rows: 25000\n",
       "    })\n",
       "    unsupervised: Dataset({\n",
       "        features: ['text', 'label'],\n",
       "        num_rows: 50000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "imdb_dataset = load_dataset(\"imdb\")\n",
    "imdb_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "'>>> Review: There is no relation at all between Fortier and Profiler but the fact that both are police series about violent crimes. Profiler looks crispy, Fortier looks classic. Profiler plots are quite simple. Fortier's plot are far more complicated... Fortier looks more like Prime Suspect, if we have to spot similarities... The main character is weak and weirdo, but have \"clairvoyance\". People like to compare, to judge, to evaluate. How about just enjoying? Funny thing too, people writing Fortier looks American but, on the other hand, arguing they prefer American series (!!!). Maybe it's the language, or the spirit, but I think this series is more English than American. By the way, the actors are really good and funny. The acting is not superficial at all...'\n",
      "'>>> Label: 1'\n",
      "\n",
      "'>>> Review: This movie is a great. The plot is very true to the book which is a classic written by Mark Twain. The movie starts of with a scene where Hank sings a song with a bunch of kids called \"when you stub your toe on the moon\" It reminds me of Sinatra's song High Hopes, it is fun and inspirational. The Music is great throughout and my favorite song is sung by the King, Hank (bing Crosby) and Sir \"Saggy\" Sagamore. OVerall a great family movie or even a great Date movie. This is a movie you can watch over and over again. The princess played by Rhonda Fleming is gorgeous. I love this movie!! If you liked Danny Kaye in the Court Jester then you will definitely like this movie.'\n",
      "'>>> Label: 1'\n",
      "\n",
      "'>>> Review: George P. Cosmatos' \"Rambo: First Blood Part II\" is pure wish-fulfillment. The United States clearly didn't win the war in Vietnam. They caused damage to this country beyond the imaginable and this movie continues the fairy story of the oh-so innocent soldiers. The only bad guys were the leaders of the nation, who made this war happen. The character of Rambo is perfect to notice this. He is extremely patriotic, bemoans that US-Americans didn't appreciate and celebrate the achievements of the single soldier, but has nothing but distrust for leading officers and politicians. Like every film that defends the war (e.g. \"We Were Soldiers\") also this one avoids the need to give a comprehensible reason for the engagement in South Asia. And for that matter also the reason for every single US-American soldier that was there. Instead, Rambo gets to take revenge for the wounds of a whole nation. It would have been better to work on how to deal with the memories, rather than suppressing them. \"Do we get to win this time?\" Yes, you do.'\n",
      "'>>> Label: 0'\n"
     ]
    }
   ],
   "source": [
    "sample = imdb_dataset[\"train\"].shuffle(seed=42).select(range(3))\n",
    "\n",
    "for row in sample:\n",
    "    print(f\"\\n'>>> Review: {row['text']}'\")\n",
    "    print(f\"'>>> Label: {row['label']}'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_function(examples):\n",
    "    result = tokenizer(examples[\"text\"])\n",
    "    if tokenizer.is_fast:\n",
    "        result[\"word_ids\"] = [result.word_ids(i) for i in range(len(result[\"input_ids\"]))]\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c3d92d6aa5ca480290999b97c874f7ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/25000 [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 (720 > 512). Running this sequence through the model will result in indexing errors\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "de16125c9efe46a1ac0fbc040ffcae70",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/25000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b456afcb26fd4c8da1728940fea657f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/50000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'word_ids'],\n",
       "        num_rows: 25000\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'word_ids'],\n",
       "        num_rows: 25000\n",
       "    })\n",
       "    unsupervised: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'word_ids'],\n",
       "        num_rows: 50000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Use batched=True to activate fast multithreading!\n",
    "tokenized_datasets = imdb_dataset.map(\n",
    "    tokenize_function, batched=True, remove_columns=[\"text\", \"label\"]\n",
    ")\n",
    "tokenized_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "512"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.model_max_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "chunk_size = 128 #chunk the token to save the GPU memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': [101,\n",
       "  1045,\n",
       "  12524,\n",
       "  1045,\n",
       "  2572,\n",
       "  8025,\n",
       "  1011,\n",
       "  3756,\n",
       "  2013,\n",
       "  2026,\n",
       "  2678,\n",
       "  3573,\n",
       "  2138,\n",
       "  1997,\n",
       "  2035,\n",
       "  1996,\n",
       "  6704,\n",
       "  2008,\n",
       "  5129,\n",
       "  2009,\n",
       "  2043,\n",
       "  2009,\n",
       "  2001,\n",
       "  2034,\n",
       "  2207,\n",
       "  1999,\n",
       "  3476,\n",
       "  1012,\n",
       "  1045,\n",
       "  2036,\n",
       "  2657,\n",
       "  2008,\n",
       "  2012,\n",
       "  2034,\n",
       "  2009,\n",
       "  2001,\n",
       "  8243,\n",
       "  2011,\n",
       "  1057,\n",
       "  1012,\n",
       "  1055,\n",
       "  1012,\n",
       "  8205,\n",
       "  2065,\n",
       "  2009,\n",
       "  2412,\n",
       "  2699,\n",
       "  2000,\n",
       "  4607,\n",
       "  2023,\n",
       "  2406,\n",
       "  1010,\n",
       "  3568,\n",
       "  2108,\n",
       "  1037,\n",
       "  5470,\n",
       "  1997,\n",
       "  3152,\n",
       "  2641,\n",
       "  1000,\n",
       "  6801,\n",
       "  1000,\n",
       "  1045,\n",
       "  2428,\n",
       "  2018,\n",
       "  2000,\n",
       "  2156,\n",
       "  2023,\n",
       "  2005,\n",
       "  2870,\n",
       "  1012,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1996,\n",
       "  5436,\n",
       "  2003,\n",
       "  8857,\n",
       "  2105,\n",
       "  1037,\n",
       "  2402,\n",
       "  4467,\n",
       "  3689,\n",
       "  3076,\n",
       "  2315,\n",
       "  14229,\n",
       "  2040,\n",
       "  4122,\n",
       "  2000,\n",
       "  4553,\n",
       "  2673,\n",
       "  2016,\n",
       "  2064,\n",
       "  2055,\n",
       "  2166,\n",
       "  1012,\n",
       "  1999,\n",
       "  3327,\n",
       "  2016,\n",
       "  4122,\n",
       "  2000,\n",
       "  3579,\n",
       "  2014,\n",
       "  3086,\n",
       "  2015,\n",
       "  2000,\n",
       "  2437,\n",
       "  2070,\n",
       "  4066,\n",
       "  1997,\n",
       "  4516,\n",
       "  2006,\n",
       "  2054,\n",
       "  1996,\n",
       "  2779,\n",
       "  25430,\n",
       "  14728,\n",
       "  2245,\n",
       "  2055,\n",
       "  3056,\n",
       "  2576,\n",
       "  3314,\n",
       "  2107,\n",
       "  2004,\n",
       "  1996,\n",
       "  5148,\n",
       "  2162,\n",
       "  1998,\n",
       "  2679,\n",
       "  3314,\n",
       "  1999,\n",
       "  1996,\n",
       "  2142,\n",
       "  2163,\n",
       "  1012,\n",
       "  1999,\n",
       "  2090,\n",
       "  4851,\n",
       "  8801,\n",
       "  1998,\n",
       "  6623,\n",
       "  7939,\n",
       "  4697,\n",
       "  3619,\n",
       "  1997,\n",
       "  8947,\n",
       "  2055,\n",
       "  2037,\n",
       "  10740,\n",
       "  2006,\n",
       "  4331,\n",
       "  1010,\n",
       "  2016,\n",
       "  2038,\n",
       "  3348,\n",
       "  2007,\n",
       "  2014,\n",
       "  3689,\n",
       "  3836,\n",
       "  1010,\n",
       "  19846,\n",
       "  1010,\n",
       "  1998,\n",
       "  2496,\n",
       "  2273,\n",
       "  1012,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  2054,\n",
       "  8563,\n",
       "  2033,\n",
       "  2055,\n",
       "  1045,\n",
       "  2572,\n",
       "  8025,\n",
       "  1011,\n",
       "  3756,\n",
       "  2003,\n",
       "  2008,\n",
       "  2871,\n",
       "  2086,\n",
       "  3283,\n",
       "  1010,\n",
       "  2023,\n",
       "  2001,\n",
       "  2641,\n",
       "  26932,\n",
       "  1012,\n",
       "  2428,\n",
       "  1010,\n",
       "  1996,\n",
       "  3348,\n",
       "  1998,\n",
       "  16371,\n",
       "  25469,\n",
       "  5019,\n",
       "  2024,\n",
       "  2261,\n",
       "  1998,\n",
       "  2521,\n",
       "  2090,\n",
       "  1010,\n",
       "  2130,\n",
       "  2059,\n",
       "  2009,\n",
       "  1005,\n",
       "  1055,\n",
       "  2025,\n",
       "  2915,\n",
       "  2066,\n",
       "  2070,\n",
       "  10036,\n",
       "  2135,\n",
       "  2081,\n",
       "  22555,\n",
       "  2080,\n",
       "  1012,\n",
       "  2096,\n",
       "  2026,\n",
       "  2406,\n",
       "  3549,\n",
       "  2568,\n",
       "  2424,\n",
       "  2009,\n",
       "  16880,\n",
       "  1010,\n",
       "  1999,\n",
       "  4507,\n",
       "  3348,\n",
       "  1998,\n",
       "  16371,\n",
       "  25469,\n",
       "  2024,\n",
       "  1037,\n",
       "  2350,\n",
       "  18785,\n",
       "  1999,\n",
       "  4467,\n",
       "  5988,\n",
       "  1012,\n",
       "  2130,\n",
       "  13749,\n",
       "  7849,\n",
       "  24544,\n",
       "  1010,\n",
       "  15835,\n",
       "  2037,\n",
       "  3437,\n",
       "  2000,\n",
       "  2204,\n",
       "  2214,\n",
       "  2879,\n",
       "  2198,\n",
       "  4811,\n",
       "  1010,\n",
       "  2018,\n",
       "  3348,\n",
       "  5019,\n",
       "  1999,\n",
       "  2010,\n",
       "  3152,\n",
       "  1012,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1045,\n",
       "  2079,\n",
       "  4012,\n",
       "  3549,\n",
       "  2094,\n",
       "  1996,\n",
       "  16587,\n",
       "  2005,\n",
       "  1996,\n",
       "  2755,\n",
       "  2008,\n",
       "  2151,\n",
       "  3348,\n",
       "  3491,\n",
       "  1999,\n",
       "  1996,\n",
       "  2143,\n",
       "  2003,\n",
       "  3491,\n",
       "  2005,\n",
       "  6018,\n",
       "  5682,\n",
       "  2738,\n",
       "  2084,\n",
       "  2074,\n",
       "  2000,\n",
       "  5213,\n",
       "  2111,\n",
       "  1998,\n",
       "  2191,\n",
       "  2769,\n",
       "  2000,\n",
       "  2022,\n",
       "  3491,\n",
       "  1999,\n",
       "  26932,\n",
       "  12370,\n",
       "  1999,\n",
       "  2637,\n",
       "  1012,\n",
       "  1045,\n",
       "  2572,\n",
       "  8025,\n",
       "  1011,\n",
       "  3756,\n",
       "  2003,\n",
       "  1037,\n",
       "  2204,\n",
       "  2143,\n",
       "  2005,\n",
       "  3087,\n",
       "  5782,\n",
       "  2000,\n",
       "  2817,\n",
       "  1996,\n",
       "  6240,\n",
       "  1998,\n",
       "  14629,\n",
       "  1006,\n",
       "  2053,\n",
       "  26136,\n",
       "  3832,\n",
       "  1007,\n",
       "  1997,\n",
       "  4467,\n",
       "  5988,\n",
       "  1012,\n",
       "  2021,\n",
       "  2428,\n",
       "  1010,\n",
       "  2023,\n",
       "  2143,\n",
       "  2987,\n",
       "  1005,\n",
       "  1056,\n",
       "  2031,\n",
       "  2172,\n",
       "  1997,\n",
       "  1037,\n",
       "  5436,\n",
       "  1012,\n",
       "  102],\n",
       " 'attention_mask': [1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1],\n",
       " 'word_ids': [None,\n",
       "  0,\n",
       "  1,\n",
       "  2,\n",
       "  3,\n",
       "  4,\n",
       "  5,\n",
       "  6,\n",
       "  7,\n",
       "  8,\n",
       "  9,\n",
       "  10,\n",
       "  11,\n",
       "  12,\n",
       "  13,\n",
       "  14,\n",
       "  15,\n",
       "  16,\n",
       "  17,\n",
       "  18,\n",
       "  19,\n",
       "  20,\n",
       "  21,\n",
       "  22,\n",
       "  23,\n",
       "  24,\n",
       "  25,\n",
       "  26,\n",
       "  27,\n",
       "  28,\n",
       "  29,\n",
       "  30,\n",
       "  31,\n",
       "  32,\n",
       "  33,\n",
       "  34,\n",
       "  35,\n",
       "  36,\n",
       "  37,\n",
       "  38,\n",
       "  39,\n",
       "  40,\n",
       "  41,\n",
       "  42,\n",
       "  43,\n",
       "  44,\n",
       "  45,\n",
       "  46,\n",
       "  47,\n",
       "  48,\n",
       "  49,\n",
       "  50,\n",
       "  51,\n",
       "  52,\n",
       "  53,\n",
       "  54,\n",
       "  55,\n",
       "  56,\n",
       "  57,\n",
       "  58,\n",
       "  59,\n",
       "  60,\n",
       "  61,\n",
       "  62,\n",
       "  63,\n",
       "  64,\n",
       "  65,\n",
       "  66,\n",
       "  67,\n",
       "  68,\n",
       "  69,\n",
       "  70,\n",
       "  71,\n",
       "  72,\n",
       "  73,\n",
       "  74,\n",
       "  75,\n",
       "  76,\n",
       "  77,\n",
       "  78,\n",
       "  79,\n",
       "  80,\n",
       "  81,\n",
       "  82,\n",
       "  83,\n",
       "  84,\n",
       "  85,\n",
       "  86,\n",
       "  87,\n",
       "  88,\n",
       "  89,\n",
       "  90,\n",
       "  91,\n",
       "  92,\n",
       "  93,\n",
       "  94,\n",
       "  95,\n",
       "  96,\n",
       "  97,\n",
       "  98,\n",
       "  99,\n",
       "  100,\n",
       "  101,\n",
       "  102,\n",
       "  103,\n",
       "  104,\n",
       "  105,\n",
       "  106,\n",
       "  107,\n",
       "  107,\n",
       "  108,\n",
       "  109,\n",
       "  110,\n",
       "  111,\n",
       "  112,\n",
       "  113,\n",
       "  114,\n",
       "  115,\n",
       "  116,\n",
       "  117,\n",
       "  118,\n",
       "  118,\n",
       "  119,\n",
       "  120,\n",
       "  121,\n",
       "  122,\n",
       "  123,\n",
       "  124,\n",
       "  125,\n",
       "  126,\n",
       "  127,\n",
       "  128,\n",
       "  129,\n",
       "  130,\n",
       "  131,\n",
       "  132,\n",
       "  133,\n",
       "  134,\n",
       "  135,\n",
       "  136,\n",
       "  137,\n",
       "  138,\n",
       "  139,\n",
       "  140,\n",
       "  141,\n",
       "  142,\n",
       "  143,\n",
       "  143,\n",
       "  143,\n",
       "  144,\n",
       "  145,\n",
       "  146,\n",
       "  147,\n",
       "  148,\n",
       "  149,\n",
       "  150,\n",
       "  151,\n",
       "  152,\n",
       "  153,\n",
       "  154,\n",
       "  155,\n",
       "  156,\n",
       "  157,\n",
       "  158,\n",
       "  159,\n",
       "  160,\n",
       "  161,\n",
       "  162,\n",
       "  163,\n",
       "  164,\n",
       "  165,\n",
       "  166,\n",
       "  167,\n",
       "  168,\n",
       "  169,\n",
       "  170,\n",
       "  171,\n",
       "  172,\n",
       "  173,\n",
       "  174,\n",
       "  175,\n",
       "  176,\n",
       "  177,\n",
       "  178,\n",
       "  179,\n",
       "  180,\n",
       "  181,\n",
       "  182,\n",
       "  183,\n",
       "  184,\n",
       "  185,\n",
       "  186,\n",
       "  187,\n",
       "  188,\n",
       "  189,\n",
       "  190,\n",
       "  191,\n",
       "  192,\n",
       "  193,\n",
       "  194,\n",
       "  195,\n",
       "  196,\n",
       "  197,\n",
       "  198,\n",
       "  199,\n",
       "  199,\n",
       "  200,\n",
       "  201,\n",
       "  202,\n",
       "  203,\n",
       "  204,\n",
       "  205,\n",
       "  206,\n",
       "  207,\n",
       "  208,\n",
       "  209,\n",
       "  210,\n",
       "  211,\n",
       "  212,\n",
       "  213,\n",
       "  214,\n",
       "  215,\n",
       "  216,\n",
       "  216,\n",
       "  217,\n",
       "  218,\n",
       "  218,\n",
       "  219,\n",
       "  220,\n",
       "  221,\n",
       "  222,\n",
       "  222,\n",
       "  223,\n",
       "  224,\n",
       "  225,\n",
       "  226,\n",
       "  227,\n",
       "  228,\n",
       "  229,\n",
       "  230,\n",
       "  231,\n",
       "  232,\n",
       "  232,\n",
       "  233,\n",
       "  234,\n",
       "  235,\n",
       "  236,\n",
       "  237,\n",
       "  238,\n",
       "  239,\n",
       "  240,\n",
       "  241,\n",
       "  242,\n",
       "  242,\n",
       "  243,\n",
       "  244,\n",
       "  245,\n",
       "  246,\n",
       "  247,\n",
       "  248,\n",
       "  249,\n",
       "  250,\n",
       "  251,\n",
       "  252,\n",
       "  253,\n",
       "  254,\n",
       "  255,\n",
       "  256,\n",
       "  257,\n",
       "  258,\n",
       "  259,\n",
       "  260,\n",
       "  261,\n",
       "  262,\n",
       "  263,\n",
       "  264,\n",
       "  265,\n",
       "  266,\n",
       "  267,\n",
       "  268,\n",
       "  269,\n",
       "  270,\n",
       "  271,\n",
       "  272,\n",
       "  272,\n",
       "  272,\n",
       "  273,\n",
       "  274,\n",
       "  275,\n",
       "  276,\n",
       "  277,\n",
       "  278,\n",
       "  279,\n",
       "  280,\n",
       "  281,\n",
       "  282,\n",
       "  283,\n",
       "  284,\n",
       "  285,\n",
       "  286,\n",
       "  287,\n",
       "  288,\n",
       "  289,\n",
       "  290,\n",
       "  291,\n",
       "  292,\n",
       "  293,\n",
       "  294,\n",
       "  295,\n",
       "  296,\n",
       "  297,\n",
       "  298,\n",
       "  299,\n",
       "  300,\n",
       "  301,\n",
       "  302,\n",
       "  303,\n",
       "  304,\n",
       "  305,\n",
       "  306,\n",
       "  307,\n",
       "  308,\n",
       "  309,\n",
       "  310,\n",
       "  311,\n",
       "  312,\n",
       "  313,\n",
       "  314,\n",
       "  315,\n",
       "  316,\n",
       "  317,\n",
       "  318,\n",
       "  319,\n",
       "  320,\n",
       "  321,\n",
       "  322,\n",
       "  323,\n",
       "  324,\n",
       "  325,\n",
       "  326,\n",
       "  327,\n",
       "  328,\n",
       "  329,\n",
       "  330,\n",
       "  331,\n",
       "  332,\n",
       "  333,\n",
       "  334,\n",
       "  335,\n",
       "  336,\n",
       "  337,\n",
       "  338,\n",
       "  339,\n",
       "  340,\n",
       "  341,\n",
       "  342,\n",
       "  343,\n",
       "  344,\n",
       "  345,\n",
       "  346,\n",
       "  347,\n",
       "  348,\n",
       "  None]}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized_datasets[\"train\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'>>> Review 0 length: 363'\n",
      "'>>> Review 1 length: 304'\n",
      "'>>> Review 2 length: 133'\n"
     ]
    }
   ],
   "source": [
    "# Slicing produces a list of lists for each feature\n",
    "tokenized_samples = tokenized_datasets[\"train\"][:3]\n",
    "\n",
    "for idx, sample in enumerate(tokenized_samples[\"input_ids\"]):\n",
    "    print(f\"'>>> Review {idx} length: {len(sample)}'\") #length of the input_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'>>> Concatenated reviews length: 800'\n"
     ]
    }
   ],
   "source": [
    "concatenated_examples = {\n",
    "    k: sum(tokenized_samples[k], []) for k in tokenized_samples.keys()\n",
    "}\n",
    "total_length = len(concatenated_examples[\"input_ids\"])\n",
    "print(f\"'>>> Concatenated reviews length: {total_length}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "800 is 363+304+133, concate the three list (three reviews) together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['input_ids', 'attention_mask', 'word_ids'])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized_samples.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[101,\n",
       "  1045,\n",
       "  12524,\n",
       "  1045,\n",
       "  2572,\n",
       "  8025,\n",
       "  1011,\n",
       "  3756,\n",
       "  2013,\n",
       "  2026,\n",
       "  2678,\n",
       "  3573,\n",
       "  2138,\n",
       "  1997,\n",
       "  2035,\n",
       "  1996,\n",
       "  6704,\n",
       "  2008,\n",
       "  5129,\n",
       "  2009,\n",
       "  2043,\n",
       "  2009,\n",
       "  2001,\n",
       "  2034,\n",
       "  2207,\n",
       "  1999,\n",
       "  3476,\n",
       "  1012,\n",
       "  1045,\n",
       "  2036,\n",
       "  2657,\n",
       "  2008,\n",
       "  2012,\n",
       "  2034,\n",
       "  2009,\n",
       "  2001,\n",
       "  8243,\n",
       "  2011,\n",
       "  1057,\n",
       "  1012,\n",
       "  1055,\n",
       "  1012,\n",
       "  8205,\n",
       "  2065,\n",
       "  2009,\n",
       "  2412,\n",
       "  2699,\n",
       "  2000,\n",
       "  4607,\n",
       "  2023,\n",
       "  2406,\n",
       "  1010,\n",
       "  3568,\n",
       "  2108,\n",
       "  1037,\n",
       "  5470,\n",
       "  1997,\n",
       "  3152,\n",
       "  2641,\n",
       "  1000,\n",
       "  6801,\n",
       "  1000,\n",
       "  1045,\n",
       "  2428,\n",
       "  2018,\n",
       "  2000,\n",
       "  2156,\n",
       "  2023,\n",
       "  2005,\n",
       "  2870,\n",
       "  1012,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1996,\n",
       "  5436,\n",
       "  2003,\n",
       "  8857,\n",
       "  2105,\n",
       "  1037,\n",
       "  2402,\n",
       "  4467,\n",
       "  3689,\n",
       "  3076,\n",
       "  2315,\n",
       "  14229,\n",
       "  2040,\n",
       "  4122,\n",
       "  2000,\n",
       "  4553,\n",
       "  2673,\n",
       "  2016,\n",
       "  2064,\n",
       "  2055,\n",
       "  2166,\n",
       "  1012,\n",
       "  1999,\n",
       "  3327,\n",
       "  2016,\n",
       "  4122,\n",
       "  2000,\n",
       "  3579,\n",
       "  2014,\n",
       "  3086,\n",
       "  2015,\n",
       "  2000,\n",
       "  2437,\n",
       "  2070,\n",
       "  4066,\n",
       "  1997,\n",
       "  4516,\n",
       "  2006,\n",
       "  2054,\n",
       "  1996,\n",
       "  2779,\n",
       "  25430,\n",
       "  14728,\n",
       "  2245,\n",
       "  2055,\n",
       "  3056,\n",
       "  2576,\n",
       "  3314,\n",
       "  2107,\n",
       "  2004,\n",
       "  1996,\n",
       "  5148,\n",
       "  2162,\n",
       "  1998,\n",
       "  2679,\n",
       "  3314,\n",
       "  1999,\n",
       "  1996,\n",
       "  2142,\n",
       "  2163,\n",
       "  1012,\n",
       "  1999,\n",
       "  2090,\n",
       "  4851,\n",
       "  8801,\n",
       "  1998,\n",
       "  6623,\n",
       "  7939,\n",
       "  4697,\n",
       "  3619,\n",
       "  1997,\n",
       "  8947,\n",
       "  2055,\n",
       "  2037,\n",
       "  10740,\n",
       "  2006,\n",
       "  4331,\n",
       "  1010,\n",
       "  2016,\n",
       "  2038,\n",
       "  3348,\n",
       "  2007,\n",
       "  2014,\n",
       "  3689,\n",
       "  3836,\n",
       "  1010,\n",
       "  19846,\n",
       "  1010,\n",
       "  1998,\n",
       "  2496,\n",
       "  2273,\n",
       "  1012,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  2054,\n",
       "  8563,\n",
       "  2033,\n",
       "  2055,\n",
       "  1045,\n",
       "  2572,\n",
       "  8025,\n",
       "  1011,\n",
       "  3756,\n",
       "  2003,\n",
       "  2008,\n",
       "  2871,\n",
       "  2086,\n",
       "  3283,\n",
       "  1010,\n",
       "  2023,\n",
       "  2001,\n",
       "  2641,\n",
       "  26932,\n",
       "  1012,\n",
       "  2428,\n",
       "  1010,\n",
       "  1996,\n",
       "  3348,\n",
       "  1998,\n",
       "  16371,\n",
       "  25469,\n",
       "  5019,\n",
       "  2024,\n",
       "  2261,\n",
       "  1998,\n",
       "  2521,\n",
       "  2090,\n",
       "  1010,\n",
       "  2130,\n",
       "  2059,\n",
       "  2009,\n",
       "  1005,\n",
       "  1055,\n",
       "  2025,\n",
       "  2915,\n",
       "  2066,\n",
       "  2070,\n",
       "  10036,\n",
       "  2135,\n",
       "  2081,\n",
       "  22555,\n",
       "  2080,\n",
       "  1012,\n",
       "  2096,\n",
       "  2026,\n",
       "  2406,\n",
       "  3549,\n",
       "  2568,\n",
       "  2424,\n",
       "  2009,\n",
       "  16880,\n",
       "  1010,\n",
       "  1999,\n",
       "  4507,\n",
       "  3348,\n",
       "  1998,\n",
       "  16371,\n",
       "  25469,\n",
       "  2024,\n",
       "  1037,\n",
       "  2350,\n",
       "  18785,\n",
       "  1999,\n",
       "  4467,\n",
       "  5988,\n",
       "  1012,\n",
       "  2130,\n",
       "  13749,\n",
       "  7849,\n",
       "  24544,\n",
       "  1010,\n",
       "  15835,\n",
       "  2037,\n",
       "  3437,\n",
       "  2000,\n",
       "  2204,\n",
       "  2214,\n",
       "  2879,\n",
       "  2198,\n",
       "  4811,\n",
       "  1010,\n",
       "  2018,\n",
       "  3348,\n",
       "  5019,\n",
       "  1999,\n",
       "  2010,\n",
       "  3152,\n",
       "  1012,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1045,\n",
       "  2079,\n",
       "  4012,\n",
       "  3549,\n",
       "  2094,\n",
       "  1996,\n",
       "  16587,\n",
       "  2005,\n",
       "  1996,\n",
       "  2755,\n",
       "  2008,\n",
       "  2151,\n",
       "  3348,\n",
       "  3491,\n",
       "  1999,\n",
       "  1996,\n",
       "  2143,\n",
       "  2003,\n",
       "  3491,\n",
       "  2005,\n",
       "  6018,\n",
       "  5682,\n",
       "  2738,\n",
       "  2084,\n",
       "  2074,\n",
       "  2000,\n",
       "  5213,\n",
       "  2111,\n",
       "  1998,\n",
       "  2191,\n",
       "  2769,\n",
       "  2000,\n",
       "  2022,\n",
       "  3491,\n",
       "  1999,\n",
       "  26932,\n",
       "  12370,\n",
       "  1999,\n",
       "  2637,\n",
       "  1012,\n",
       "  1045,\n",
       "  2572,\n",
       "  8025,\n",
       "  1011,\n",
       "  3756,\n",
       "  2003,\n",
       "  1037,\n",
       "  2204,\n",
       "  2143,\n",
       "  2005,\n",
       "  3087,\n",
       "  5782,\n",
       "  2000,\n",
       "  2817,\n",
       "  1996,\n",
       "  6240,\n",
       "  1998,\n",
       "  14629,\n",
       "  1006,\n",
       "  2053,\n",
       "  26136,\n",
       "  3832,\n",
       "  1007,\n",
       "  1997,\n",
       "  4467,\n",
       "  5988,\n",
       "  1012,\n",
       "  2021,\n",
       "  2428,\n",
       "  1010,\n",
       "  2023,\n",
       "  2143,\n",
       "  2987,\n",
       "  1005,\n",
       "  1056,\n",
       "  2031,\n",
       "  2172,\n",
       "  1997,\n",
       "  1037,\n",
       "  5436,\n",
       "  1012,\n",
       "  102],\n",
       " [101,\n",
       "  1000,\n",
       "  1045,\n",
       "  2572,\n",
       "  8025,\n",
       "  1024,\n",
       "  3756,\n",
       "  1000,\n",
       "  2003,\n",
       "  1037,\n",
       "  15544,\n",
       "  19307,\n",
       "  1998,\n",
       "  3653,\n",
       "  6528,\n",
       "  20771,\n",
       "  19986,\n",
       "  8632,\n",
       "  1012,\n",
       "  2009,\n",
       "  2987,\n",
       "  1005,\n",
       "  1056,\n",
       "  3043,\n",
       "  2054,\n",
       "  2028,\n",
       "  1005,\n",
       "  1055,\n",
       "  2576,\n",
       "  5328,\n",
       "  2024,\n",
       "  2138,\n",
       "  2023,\n",
       "  2143,\n",
       "  2064,\n",
       "  6684,\n",
       "  2022,\n",
       "  2579,\n",
       "  5667,\n",
       "  2006,\n",
       "  2151,\n",
       "  2504,\n",
       "  1012,\n",
       "  2004,\n",
       "  2005,\n",
       "  1996,\n",
       "  4366,\n",
       "  2008,\n",
       "  19124,\n",
       "  3287,\n",
       "  16371,\n",
       "  25469,\n",
       "  2003,\n",
       "  2019,\n",
       "  6882,\n",
       "  13316,\n",
       "  1011,\n",
       "  2459,\n",
       "  1010,\n",
       "  2008,\n",
       "  3475,\n",
       "  1005,\n",
       "  1056,\n",
       "  2995,\n",
       "  1012,\n",
       "  1045,\n",
       "  1005,\n",
       "  2310,\n",
       "  2464,\n",
       "  1054,\n",
       "  1011,\n",
       "  6758,\n",
       "  3152,\n",
       "  2007,\n",
       "  3287,\n",
       "  16371,\n",
       "  25469,\n",
       "  1012,\n",
       "  4379,\n",
       "  1010,\n",
       "  2027,\n",
       "  2069,\n",
       "  3749,\n",
       "  2070,\n",
       "  25085,\n",
       "  5328,\n",
       "  1010,\n",
       "  2021,\n",
       "  2073,\n",
       "  2024,\n",
       "  1996,\n",
       "  1054,\n",
       "  1011,\n",
       "  6758,\n",
       "  3152,\n",
       "  2007,\n",
       "  21226,\n",
       "  24728,\n",
       "  22144,\n",
       "  2015,\n",
       "  1998,\n",
       "  20916,\n",
       "  4691,\n",
       "  6845,\n",
       "  2401,\n",
       "  1029,\n",
       "  7880,\n",
       "  1010,\n",
       "  2138,\n",
       "  2027,\n",
       "  2123,\n",
       "  1005,\n",
       "  1056,\n",
       "  4839,\n",
       "  1012,\n",
       "  1996,\n",
       "  2168,\n",
       "  3632,\n",
       "  2005,\n",
       "  2216,\n",
       "  10231,\n",
       "  7685,\n",
       "  5830,\n",
       "  3065,\n",
       "  1024,\n",
       "  8040,\n",
       "  7317,\n",
       "  5063,\n",
       "  2015,\n",
       "  11820,\n",
       "  1999,\n",
       "  1996,\n",
       "  9478,\n",
       "  2021,\n",
       "  2025,\n",
       "  1037,\n",
       "  17962,\n",
       "  21239,\n",
       "  1999,\n",
       "  4356,\n",
       "  1012,\n",
       "  1998,\n",
       "  2216,\n",
       "  3653,\n",
       "  6528,\n",
       "  20771,\n",
       "  10271,\n",
       "  5691,\n",
       "  2066,\n",
       "  1996,\n",
       "  2829,\n",
       "  16291,\n",
       "  1010,\n",
       "  1999,\n",
       "  2029,\n",
       "  2057,\n",
       "  1005,\n",
       "  2128,\n",
       "  5845,\n",
       "  2000,\n",
       "  1996,\n",
       "  2609,\n",
       "  1997,\n",
       "  6320,\n",
       "  25624,\n",
       "  1005,\n",
       "  1055,\n",
       "  17061,\n",
       "  3779,\n",
       "  1010,\n",
       "  2021,\n",
       "  2025,\n",
       "  1037,\n",
       "  7637,\n",
       "  1997,\n",
       "  5061,\n",
       "  5710,\n",
       "  2006,\n",
       "  9318,\n",
       "  7367,\n",
       "  5737,\n",
       "  19393,\n",
       "  1012,\n",
       "  2077,\n",
       "  6933,\n",
       "  1006,\n",
       "  2030,\n",
       "  20242,\n",
       "  1007,\n",
       "  1000,\n",
       "  3313,\n",
       "  1011,\n",
       "  3115,\n",
       "  1000,\n",
       "  1999,\n",
       "  5609,\n",
       "  1997,\n",
       "  16371,\n",
       "  25469,\n",
       "  1010,\n",
       "  1996,\n",
       "  10597,\n",
       "  27885,\n",
       "  5809,\n",
       "  2063,\n",
       "  2323,\n",
       "  2202,\n",
       "  2046,\n",
       "  4070,\n",
       "  2028,\n",
       "  14477,\n",
       "  6767,\n",
       "  8524,\n",
       "  6321,\n",
       "  5793,\n",
       "  28141,\n",
       "  4489,\n",
       "  2090,\n",
       "  2273,\n",
       "  1998,\n",
       "  2308,\n",
       "  1024,\n",
       "  2045,\n",
       "  2024,\n",
       "  2053,\n",
       "  8991,\n",
       "  18400,\n",
       "  2015,\n",
       "  2006,\n",
       "  4653,\n",
       "  2043,\n",
       "  19910,\n",
       "  3544,\n",
       "  15287,\n",
       "  1010,\n",
       "  1998,\n",
       "  1996,\n",
       "  2168,\n",
       "  3685,\n",
       "  2022,\n",
       "  2056,\n",
       "  2005,\n",
       "  1037,\n",
       "  2158,\n",
       "  1012,\n",
       "  1999,\n",
       "  2755,\n",
       "  1010,\n",
       "  2017,\n",
       "  3227,\n",
       "  2180,\n",
       "  1005,\n",
       "  1056,\n",
       "  2156,\n",
       "  2931,\n",
       "  8991,\n",
       "  18400,\n",
       "  2015,\n",
       "  1999,\n",
       "  2019,\n",
       "  2137,\n",
       "  2143,\n",
       "  1999,\n",
       "  2505,\n",
       "  2460,\n",
       "  1997,\n",
       "  22555,\n",
       "  2030,\n",
       "  13216,\n",
       "  14253,\n",
       "  2050,\n",
       "  1012,\n",
       "  2023,\n",
       "  6884,\n",
       "  3313,\n",
       "  1011,\n",
       "  3115,\n",
       "  2003,\n",
       "  2625,\n",
       "  1037,\n",
       "  3313,\n",
       "  3115,\n",
       "  2084,\n",
       "  2019,\n",
       "  4914,\n",
       "  2135,\n",
       "  2139,\n",
       "  24128,\n",
       "  3754,\n",
       "  2000,\n",
       "  2272,\n",
       "  2000,\n",
       "  3408,\n",
       "  20547,\n",
       "  2007,\n",
       "  1996,\n",
       "  19008,\n",
       "  1997,\n",
       "  2308,\n",
       "  1005,\n",
       "  1055,\n",
       "  4230,\n",
       "  1012,\n",
       "  102],\n",
       " [101,\n",
       "  2065,\n",
       "  2069,\n",
       "  2000,\n",
       "  4468,\n",
       "  2437,\n",
       "  2023,\n",
       "  2828,\n",
       "  1997,\n",
       "  2143,\n",
       "  1999,\n",
       "  1996,\n",
       "  2925,\n",
       "  1012,\n",
       "  2023,\n",
       "  2143,\n",
       "  2003,\n",
       "  5875,\n",
       "  2004,\n",
       "  2019,\n",
       "  7551,\n",
       "  2021,\n",
       "  4136,\n",
       "  2053,\n",
       "  2522,\n",
       "  11461,\n",
       "  2466,\n",
       "  1012,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  2028,\n",
       "  2453,\n",
       "  2514,\n",
       "  6819,\n",
       "  5339,\n",
       "  8918,\n",
       "  2005,\n",
       "  3564,\n",
       "  27046,\n",
       "  2009,\n",
       "  2138,\n",
       "  2009,\n",
       "  12817,\n",
       "  2006,\n",
       "  2061,\n",
       "  2116,\n",
       "  2590,\n",
       "  3314,\n",
       "  2021,\n",
       "  2009,\n",
       "  2515,\n",
       "  2061,\n",
       "  2302,\n",
       "  2151,\n",
       "  5860,\n",
       "  11795,\n",
       "  3085,\n",
       "  15793,\n",
       "  1012,\n",
       "  1996,\n",
       "  13972,\n",
       "  3310,\n",
       "  2185,\n",
       "  2007,\n",
       "  2053,\n",
       "  2047,\n",
       "  15251,\n",
       "  1006,\n",
       "  4983,\n",
       "  2028,\n",
       "  3310,\n",
       "  2039,\n",
       "  2007,\n",
       "  2028,\n",
       "  2096,\n",
       "  2028,\n",
       "  1005,\n",
       "  1055,\n",
       "  2568,\n",
       "  17677,\n",
       "  2015,\n",
       "  1010,\n",
       "  2004,\n",
       "  2009,\n",
       "  2097,\n",
       "  26597,\n",
       "  2079,\n",
       "  2076,\n",
       "  2023,\n",
       "  23100,\n",
       "  2143,\n",
       "  1007,\n",
       "  1012,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  2028,\n",
       "  2453,\n",
       "  2488,\n",
       "  5247,\n",
       "  2028,\n",
       "  1005,\n",
       "  1055,\n",
       "  2051,\n",
       "  4582,\n",
       "  2041,\n",
       "  1037,\n",
       "  3332,\n",
       "  2012,\n",
       "  1037,\n",
       "  3392,\n",
       "  3652,\n",
       "  1012,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  1026,\n",
       "  7987,\n",
       "  1013,\n",
       "  1028,\n",
       "  102]]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized_samples['input_ids']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[101,\n",
       " 1045,\n",
       " 12524,\n",
       " 1045,\n",
       " 2572,\n",
       " 8025,\n",
       " 1011,\n",
       " 3756,\n",
       " 2013,\n",
       " 2026,\n",
       " 2678,\n",
       " 3573,\n",
       " 2138,\n",
       " 1997,\n",
       " 2035,\n",
       " 1996,\n",
       " 6704,\n",
       " 2008,\n",
       " 5129,\n",
       " 2009,\n",
       " 2043,\n",
       " 2009,\n",
       " 2001,\n",
       " 2034,\n",
       " 2207,\n",
       " 1999,\n",
       " 3476,\n",
       " 1012,\n",
       " 1045,\n",
       " 2036,\n",
       " 2657,\n",
       " 2008,\n",
       " 2012,\n",
       " 2034,\n",
       " 2009,\n",
       " 2001,\n",
       " 8243,\n",
       " 2011,\n",
       " 1057,\n",
       " 1012,\n",
       " 1055,\n",
       " 1012,\n",
       " 8205,\n",
       " 2065,\n",
       " 2009,\n",
       " 2412,\n",
       " 2699,\n",
       " 2000,\n",
       " 4607,\n",
       " 2023,\n",
       " 2406,\n",
       " 1010,\n",
       " 3568,\n",
       " 2108,\n",
       " 1037,\n",
       " 5470,\n",
       " 1997,\n",
       " 3152,\n",
       " 2641,\n",
       " 1000,\n",
       " 6801,\n",
       " 1000,\n",
       " 1045,\n",
       " 2428,\n",
       " 2018,\n",
       " 2000,\n",
       " 2156,\n",
       " 2023,\n",
       " 2005,\n",
       " 2870,\n",
       " 1012,\n",
       " 1026,\n",
       " 7987,\n",
       " 1013,\n",
       " 1028,\n",
       " 1026,\n",
       " 7987,\n",
       " 1013,\n",
       " 1028,\n",
       " 1996,\n",
       " 5436,\n",
       " 2003,\n",
       " 8857,\n",
       " 2105,\n",
       " 1037,\n",
       " 2402,\n",
       " 4467,\n",
       " 3689,\n",
       " 3076,\n",
       " 2315,\n",
       " 14229,\n",
       " 2040,\n",
       " 4122,\n",
       " 2000,\n",
       " 4553,\n",
       " 2673,\n",
       " 2016,\n",
       " 2064,\n",
       " 2055,\n",
       " 2166,\n",
       " 1012,\n",
       " 1999,\n",
       " 3327,\n",
       " 2016,\n",
       " 4122,\n",
       " 2000,\n",
       " 3579,\n",
       " 2014,\n",
       " 3086,\n",
       " 2015,\n",
       " 2000,\n",
       " 2437,\n",
       " 2070,\n",
       " 4066,\n",
       " 1997,\n",
       " 4516,\n",
       " 2006,\n",
       " 2054,\n",
       " 1996,\n",
       " 2779,\n",
       " 25430,\n",
       " 14728,\n",
       " 2245,\n",
       " 2055,\n",
       " 3056,\n",
       " 2576,\n",
       " 3314,\n",
       " 2107,\n",
       " 2004,\n",
       " 1996,\n",
       " 5148,\n",
       " 2162,\n",
       " 1998,\n",
       " 2679,\n",
       " 3314,\n",
       " 1999,\n",
       " 1996,\n",
       " 2142,\n",
       " 2163,\n",
       " 1012,\n",
       " 1999,\n",
       " 2090,\n",
       " 4851,\n",
       " 8801,\n",
       " 1998,\n",
       " 6623,\n",
       " 7939,\n",
       " 4697,\n",
       " 3619,\n",
       " 1997,\n",
       " 8947,\n",
       " 2055,\n",
       " 2037,\n",
       " 10740,\n",
       " 2006,\n",
       " 4331,\n",
       " 1010,\n",
       " 2016,\n",
       " 2038,\n",
       " 3348,\n",
       " 2007,\n",
       " 2014,\n",
       " 3689,\n",
       " 3836,\n",
       " 1010,\n",
       " 19846,\n",
       " 1010,\n",
       " 1998,\n",
       " 2496,\n",
       " 2273,\n",
       " 1012,\n",
       " 1026,\n",
       " 7987,\n",
       " 1013,\n",
       " 1028,\n",
       " 1026,\n",
       " 7987,\n",
       " 1013,\n",
       " 1028,\n",
       " 2054,\n",
       " 8563,\n",
       " 2033,\n",
       " 2055,\n",
       " 1045,\n",
       " 2572,\n",
       " 8025,\n",
       " 1011,\n",
       " 3756,\n",
       " 2003,\n",
       " 2008,\n",
       " 2871,\n",
       " 2086,\n",
       " 3283,\n",
       " 1010,\n",
       " 2023,\n",
       " 2001,\n",
       " 2641,\n",
       " 26932,\n",
       " 1012,\n",
       " 2428,\n",
       " 1010,\n",
       " 1996,\n",
       " 3348,\n",
       " 1998,\n",
       " 16371,\n",
       " 25469,\n",
       " 5019,\n",
       " 2024,\n",
       " 2261,\n",
       " 1998,\n",
       " 2521,\n",
       " 2090,\n",
       " 1010,\n",
       " 2130,\n",
       " 2059,\n",
       " 2009,\n",
       " 1005,\n",
       " 1055,\n",
       " 2025,\n",
       " 2915,\n",
       " 2066,\n",
       " 2070,\n",
       " 10036,\n",
       " 2135,\n",
       " 2081,\n",
       " 22555,\n",
       " 2080,\n",
       " 1012,\n",
       " 2096,\n",
       " 2026,\n",
       " 2406,\n",
       " 3549,\n",
       " 2568,\n",
       " 2424,\n",
       " 2009,\n",
       " 16880,\n",
       " 1010,\n",
       " 1999,\n",
       " 4507,\n",
       " 3348,\n",
       " 1998,\n",
       " 16371,\n",
       " 25469,\n",
       " 2024,\n",
       " 1037,\n",
       " 2350,\n",
       " 18785,\n",
       " 1999,\n",
       " 4467,\n",
       " 5988,\n",
       " 1012,\n",
       " 2130,\n",
       " 13749,\n",
       " 7849,\n",
       " 24544,\n",
       " 1010,\n",
       " 15835,\n",
       " 2037,\n",
       " 3437,\n",
       " 2000,\n",
       " 2204,\n",
       " 2214,\n",
       " 2879,\n",
       " 2198,\n",
       " 4811,\n",
       " 1010,\n",
       " 2018,\n",
       " 3348,\n",
       " 5019,\n",
       " 1999,\n",
       " 2010,\n",
       " 3152,\n",
       " 1012,\n",
       " 1026,\n",
       " 7987,\n",
       " 1013,\n",
       " 1028,\n",
       " 1026,\n",
       " 7987,\n",
       " 1013,\n",
       " 1028,\n",
       " 1045,\n",
       " 2079,\n",
       " 4012,\n",
       " 3549,\n",
       " 2094,\n",
       " 1996,\n",
       " 16587,\n",
       " 2005,\n",
       " 1996,\n",
       " 2755,\n",
       " 2008,\n",
       " 2151,\n",
       " 3348,\n",
       " 3491,\n",
       " 1999,\n",
       " 1996,\n",
       " 2143,\n",
       " 2003,\n",
       " 3491,\n",
       " 2005,\n",
       " 6018,\n",
       " 5682,\n",
       " 2738,\n",
       " 2084,\n",
       " 2074,\n",
       " 2000,\n",
       " 5213,\n",
       " 2111,\n",
       " 1998,\n",
       " 2191,\n",
       " 2769,\n",
       " 2000,\n",
       " 2022,\n",
       " 3491,\n",
       " 1999,\n",
       " 26932,\n",
       " 12370,\n",
       " 1999,\n",
       " 2637,\n",
       " 1012,\n",
       " 1045,\n",
       " 2572,\n",
       " 8025,\n",
       " 1011,\n",
       " 3756,\n",
       " 2003,\n",
       " 1037,\n",
       " 2204,\n",
       " 2143,\n",
       " 2005,\n",
       " 3087,\n",
       " 5782,\n",
       " 2000,\n",
       " 2817,\n",
       " 1996,\n",
       " 6240,\n",
       " 1998,\n",
       " 14629,\n",
       " 1006,\n",
       " 2053,\n",
       " 26136,\n",
       " 3832,\n",
       " 1007,\n",
       " 1997,\n",
       " 4467,\n",
       " 5988,\n",
       " 1012,\n",
       " 2021,\n",
       " 2428,\n",
       " 1010,\n",
       " 2023,\n",
       " 2143,\n",
       " 2987,\n",
       " 1005,\n",
       " 1056,\n",
       " 2031,\n",
       " 2172,\n",
       " 1997,\n",
       " 1037,\n",
       " 5436,\n",
       " 1012,\n",
       " 102,\n",
       " 101,\n",
       " 1000,\n",
       " 1045,\n",
       " 2572,\n",
       " 8025,\n",
       " 1024,\n",
       " 3756,\n",
       " 1000,\n",
       " 2003,\n",
       " 1037,\n",
       " 15544,\n",
       " 19307,\n",
       " 1998,\n",
       " 3653,\n",
       " 6528,\n",
       " 20771,\n",
       " 19986,\n",
       " 8632,\n",
       " 1012,\n",
       " 2009,\n",
       " 2987,\n",
       " 1005,\n",
       " 1056,\n",
       " 3043,\n",
       " 2054,\n",
       " 2028,\n",
       " 1005,\n",
       " 1055,\n",
       " 2576,\n",
       " 5328,\n",
       " 2024,\n",
       " 2138,\n",
       " 2023,\n",
       " 2143,\n",
       " 2064,\n",
       " 6684,\n",
       " 2022,\n",
       " 2579,\n",
       " 5667,\n",
       " 2006,\n",
       " 2151,\n",
       " 2504,\n",
       " 1012,\n",
       " 2004,\n",
       " 2005,\n",
       " 1996,\n",
       " 4366,\n",
       " 2008,\n",
       " 19124,\n",
       " 3287,\n",
       " 16371,\n",
       " 25469,\n",
       " 2003,\n",
       " 2019,\n",
       " 6882,\n",
       " 13316,\n",
       " 1011,\n",
       " 2459,\n",
       " 1010,\n",
       " 2008,\n",
       " 3475,\n",
       " 1005,\n",
       " 1056,\n",
       " 2995,\n",
       " 1012,\n",
       " 1045,\n",
       " 1005,\n",
       " 2310,\n",
       " 2464,\n",
       " 1054,\n",
       " 1011,\n",
       " 6758,\n",
       " 3152,\n",
       " 2007,\n",
       " 3287,\n",
       " 16371,\n",
       " 25469,\n",
       " 1012,\n",
       " 4379,\n",
       " 1010,\n",
       " 2027,\n",
       " 2069,\n",
       " 3749,\n",
       " 2070,\n",
       " 25085,\n",
       " 5328,\n",
       " 1010,\n",
       " 2021,\n",
       " 2073,\n",
       " 2024,\n",
       " 1996,\n",
       " 1054,\n",
       " 1011,\n",
       " 6758,\n",
       " 3152,\n",
       " 2007,\n",
       " 21226,\n",
       " 24728,\n",
       " 22144,\n",
       " 2015,\n",
       " 1998,\n",
       " 20916,\n",
       " 4691,\n",
       " 6845,\n",
       " 2401,\n",
       " 1029,\n",
       " 7880,\n",
       " 1010,\n",
       " 2138,\n",
       " 2027,\n",
       " 2123,\n",
       " 1005,\n",
       " 1056,\n",
       " 4839,\n",
       " 1012,\n",
       " 1996,\n",
       " 2168,\n",
       " 3632,\n",
       " 2005,\n",
       " 2216,\n",
       " 10231,\n",
       " 7685,\n",
       " 5830,\n",
       " 3065,\n",
       " 1024,\n",
       " 8040,\n",
       " 7317,\n",
       " 5063,\n",
       " 2015,\n",
       " 11820,\n",
       " 1999,\n",
       " 1996,\n",
       " 9478,\n",
       " 2021,\n",
       " 2025,\n",
       " 1037,\n",
       " 17962,\n",
       " 21239,\n",
       " 1999,\n",
       " 4356,\n",
       " 1012,\n",
       " 1998,\n",
       " 2216,\n",
       " 3653,\n",
       " 6528,\n",
       " 20771,\n",
       " 10271,\n",
       " 5691,\n",
       " 2066,\n",
       " 1996,\n",
       " 2829,\n",
       " 16291,\n",
       " 1010,\n",
       " 1999,\n",
       " 2029,\n",
       " 2057,\n",
       " 1005,\n",
       " 2128,\n",
       " 5845,\n",
       " 2000,\n",
       " 1996,\n",
       " 2609,\n",
       " 1997,\n",
       " 6320,\n",
       " 25624,\n",
       " 1005,\n",
       " 1055,\n",
       " 17061,\n",
       " 3779,\n",
       " 1010,\n",
       " 2021,\n",
       " 2025,\n",
       " 1037,\n",
       " 7637,\n",
       " 1997,\n",
       " 5061,\n",
       " 5710,\n",
       " 2006,\n",
       " 9318,\n",
       " 7367,\n",
       " 5737,\n",
       " 19393,\n",
       " 1012,\n",
       " 2077,\n",
       " 6933,\n",
       " 1006,\n",
       " 2030,\n",
       " 20242,\n",
       " 1007,\n",
       " 1000,\n",
       " 3313,\n",
       " 1011,\n",
       " 3115,\n",
       " 1000,\n",
       " 1999,\n",
       " 5609,\n",
       " 1997,\n",
       " 16371,\n",
       " 25469,\n",
       " 1010,\n",
       " 1996,\n",
       " 10597,\n",
       " 27885,\n",
       " 5809,\n",
       " 2063,\n",
       " 2323,\n",
       " 2202,\n",
       " 2046,\n",
       " 4070,\n",
       " 2028,\n",
       " 14477,\n",
       " 6767,\n",
       " 8524,\n",
       " 6321,\n",
       " 5793,\n",
       " 28141,\n",
       " 4489,\n",
       " 2090,\n",
       " 2273,\n",
       " 1998,\n",
       " 2308,\n",
       " 1024,\n",
       " 2045,\n",
       " 2024,\n",
       " 2053,\n",
       " 8991,\n",
       " 18400,\n",
       " 2015,\n",
       " 2006,\n",
       " 4653,\n",
       " 2043,\n",
       " 19910,\n",
       " 3544,\n",
       " 15287,\n",
       " 1010,\n",
       " 1998,\n",
       " 1996,\n",
       " 2168,\n",
       " 3685,\n",
       " 2022,\n",
       " 2056,\n",
       " 2005,\n",
       " 1037,\n",
       " 2158,\n",
       " 1012,\n",
       " 1999,\n",
       " 2755,\n",
       " 1010,\n",
       " 2017,\n",
       " 3227,\n",
       " 2180,\n",
       " 1005,\n",
       " 1056,\n",
       " 2156,\n",
       " 2931,\n",
       " 8991,\n",
       " 18400,\n",
       " 2015,\n",
       " 1999,\n",
       " 2019,\n",
       " 2137,\n",
       " 2143,\n",
       " 1999,\n",
       " 2505,\n",
       " 2460,\n",
       " 1997,\n",
       " 22555,\n",
       " 2030,\n",
       " 13216,\n",
       " 14253,\n",
       " 2050,\n",
       " 1012,\n",
       " 2023,\n",
       " 6884,\n",
       " 3313,\n",
       " 1011,\n",
       " 3115,\n",
       " 2003,\n",
       " 2625,\n",
       " 1037,\n",
       " 3313,\n",
       " 3115,\n",
       " 2084,\n",
       " 2019,\n",
       " 4914,\n",
       " 2135,\n",
       " 2139,\n",
       " 24128,\n",
       " 3754,\n",
       " 2000,\n",
       " 2272,\n",
       " 2000,\n",
       " 3408,\n",
       " 20547,\n",
       " 2007,\n",
       " 1996,\n",
       " 19008,\n",
       " 1997,\n",
       " 2308,\n",
       " 1005,\n",
       " 1055,\n",
       " 4230,\n",
       " 1012,\n",
       " 102,\n",
       " 101,\n",
       " 2065,\n",
       " 2069,\n",
       " 2000,\n",
       " 4468,\n",
       " 2437,\n",
       " 2023,\n",
       " 2828,\n",
       " 1997,\n",
       " 2143,\n",
       " 1999,\n",
       " 1996,\n",
       " 2925,\n",
       " 1012,\n",
       " 2023,\n",
       " 2143,\n",
       " 2003,\n",
       " 5875,\n",
       " 2004,\n",
       " 2019,\n",
       " 7551,\n",
       " 2021,\n",
       " 4136,\n",
       " 2053,\n",
       " 2522,\n",
       " 11461,\n",
       " 2466,\n",
       " 1012,\n",
       " 1026,\n",
       " 7987,\n",
       " 1013,\n",
       " 1028,\n",
       " 1026,\n",
       " 7987,\n",
       " 1013,\n",
       " 1028,\n",
       " 2028,\n",
       " 2453,\n",
       " 2514,\n",
       " 6819,\n",
       " 5339,\n",
       " 8918,\n",
       " 2005,\n",
       " 3564,\n",
       " 27046,\n",
       " 2009,\n",
       " 2138,\n",
       " 2009,\n",
       " 12817,\n",
       " 2006,\n",
       " 2061,\n",
       " 2116,\n",
       " 2590,\n",
       " 3314,\n",
       " 2021,\n",
       " 2009,\n",
       " 2515,\n",
       " 2061,\n",
       " 2302,\n",
       " 2151,\n",
       " 5860,\n",
       " 11795,\n",
       " 3085,\n",
       " 15793,\n",
       " 1012,\n",
       " 1996,\n",
       " 13972,\n",
       " 3310,\n",
       " 2185,\n",
       " 2007,\n",
       " 2053,\n",
       " 2047,\n",
       " 15251,\n",
       " 1006,\n",
       " 4983,\n",
       " 2028,\n",
       " 3310,\n",
       " 2039,\n",
       " 2007,\n",
       " 2028,\n",
       " 2096,\n",
       " 2028,\n",
       " 1005,\n",
       " 1055,\n",
       " 2568,\n",
       " 17677,\n",
       " 2015,\n",
       " 1010,\n",
       " 2004,\n",
       " 2009,\n",
       " 2097,\n",
       " 26597,\n",
       " 2079,\n",
       " 2076,\n",
       " 2023,\n",
       " 23100,\n",
       " 2143,\n",
       " 1007,\n",
       " 1012,\n",
       " 1026,\n",
       " 7987,\n",
       " 1013,\n",
       " 1028,\n",
       " 1026,\n",
       " 7987,\n",
       " 1013,\n",
       " 1028,\n",
       " 2028,\n",
       " 2453,\n",
       " 2488,\n",
       " 5247,\n",
       " 2028,\n",
       " 1005,\n",
       " 1055,\n",
       " 2051,\n",
       " 4582,\n",
       " 2041,\n",
       " 1037,\n",
       " 3332,\n",
       " 2012,\n",
       " 1037,\n",
       " 3392,\n",
       " 3652,\n",
       " 1012,\n",
       " 1026,\n",
       " 7987,\n",
       " 1013,\n",
       " 1028,\n",
       " 1026,\n",
       " 7987,\n",
       " 1013,\n",
       " 1028,\n",
       " 102]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum(tokenized_samples['input_ids'], [])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cut the concatenated_examples into equal size chunk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'>>> Chunk length: 128'\n",
      "'>>> Chunk length: 128'\n",
      "'>>> Chunk length: 128'\n",
      "'>>> Chunk length: 128'\n",
      "'>>> Chunk length: 128'\n",
      "'>>> Chunk length: 128'\n",
      "'>>> Chunk length: 32'\n"
     ]
    }
   ],
   "source": [
    "chunks = {\n",
    "    k: [t[i : i + chunk_size] for i in range(0, total_length, chunk_size)]\n",
    "    for k, t in concatenated_examples.items()\n",
    "}\n",
    "\n",
    "for chunk in chunks[\"input_ids\"]:\n",
    "    print(f\"'>>> Chunk length: {len(chunk)}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Group all these into one function \"group_texts\", drop the last chunk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def group_texts(examples):\n",
    "    # Concatenate all texts\n",
    "    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}\n",
    "    # Compute length of concatenated texts\n",
    "    total_length = len(concatenated_examples[list(examples.keys())[0]])\n",
    "    # We drop the last chunk if it's smaller than chunk_size\n",
    "    total_length = (total_length // chunk_size) * chunk_size\n",
    "    # Split by chunks of max_len\n",
    "    result = {\n",
    "        k: [t[i : i + chunk_size] for i in range(0, total_length, chunk_size)]\n",
    "        for k, t in concatenated_examples.items()\n",
    "    }\n",
    "    # Create a new labels column\n",
    "    result[\"labels\"] = result[\"input_ids\"].copy()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1130311bad364ad1985f1b7eb15bd09d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/25000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f357355b45724ce6a2e6eec0319ab788",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/25000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b4225f5279f844ee9e16d4e583e48a19",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/50000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'word_ids', 'labels'],\n",
       "        num_rows: 61291\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'word_ids', 'labels'],\n",
       "        num_rows: 59904\n",
       "    })\n",
       "    unsupervised: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'word_ids', 'labels'],\n",
       "        num_rows: 122957\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm_datasets = tokenized_datasets.map(group_texts, batched=True)\n",
    "lm_datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Token Mask Data Collator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForLanguageModeling\n",
    "#mlm_probability means the percentage of [MASK]\n",
    "data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = [lm_datasets[\"train\"][i] for i in range(2)]\n",
    "for sample in samples:\n",
    "    _ = sample.pop(\"word_ids\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a DistilBertTokenizerFast 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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "'>>> [CLS] i [MASK] i am curious - yellow from my video store because [MASK] all [MASK] controversy [MASK] surrounded it when it was first released in [MASK]. i also heard that at first it was seized by [MASK]. s. customs [MASK] [MASK] ever tried [MASK] enter this country, therefore being a fan of films considered [MASK] controversial \" i really edison to see this for myself. < br [MASK] > < br / > [MASK] plot is centered around a [MASK] swedish drama student named lena who wants to learn [MASK] she can about life [MASK] in [MASK] she wants to focus her attention [MASK] to making [MASK] sort of [MASK] on what the [MASK] swede thought about certain political issues such'\n",
      "\n",
      "'>>> as the vietnam war and race issues in the united states. in between asking politicians and ordinary denizens of stockholm about [MASK] opinions on politics, she has sex with her drama [MASK], [MASK], [MASK] married men. < br [MASK] > < br / > what kills me about i am [MASK] - yellow is that 40 years ago, this [MASK] considered pornographic [MASK] really, the sex and nudity [MASK] are [MASK] and far between, even then [MASK] [MASK] s not shot like some cheaply [MASK] porno [MASK] while my countrymen mind findtler shocking, in reality [MASK] and nudity are [MASK] major staple in swedish [MASK]. even ingmar bergman belinda'\n"
     ]
    }
   ],
   "source": [
    "for chunk in data_collator(samples)[\"input_ids\"]:\n",
    "    print(f\"\\n'>>> {tokenizer.decode(chunk)}'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "'>>> ['[CLS]', 'i', 'rented', '[MASK]', 'am', 'curious', '-', '[MASK]', 'from', 'my', 'video', 'store', '[MASK]', 'of', 'all', 'the', 'controversy', 'that', '[MASK]', 'it', 'when', '[MASK]', 'was', '[MASK]', 'released', 'in', '[MASK]', '[MASK]', 'i', 'also', '[MASK]', 'that', 'at', 'first', 'it', 'was', 'seized', 'by', 'u', '.', 's', '.', 'customs', 'if', 'it', '[MASK]', 'tried', '[MASK]', 'enter', 'this', 'country', ',', 'therefore', 'being', 'a', 'fan', 'of', 'films', 'considered', '\"', 'controversial', '\"', 'i', 'really', 'had', 'to', 'see', 'this', 'for', 'myself', '.', '<', 'br', '[MASK]', '>', '<', 'br', '/', '>', 'the', 'plot', 'is', '[MASK]', 'around', 'a', 'young', '[MASK]', 'drama', 'student', 'named', 'lena', 'who', '[MASK]', 'to', 'learn', 'everything', '[MASK]', 'can', 'about', '[MASK]', '.', 'in', 'particular', 'she', 'wants', 'to', 'focus', 'her', 'attention', '##s', 'to', 'making', 'some', 'sort', 'of', 'documentary', 'on', 'what', 'the', '[MASK]', 'sw', '##ede', 'thought', 'about', 'certain', 'political', 'issues', 'such']'\n",
      "\n",
      "'>>> ['as', 'the', 'vietnam', 'war', 'and', '[MASK]', 'issues', 'in', 'the', 'united', 'states', '[MASK]', 'in', 'between', '[MASK]', 'politicians', 'and', 'ordinary', 'den', '[MASK]', '##ns', 'of', 'stockholm', 'about', 'their', 'opinions', 'on', 'politics', ',', '[MASK]', '[MASK]', 'sex', 'with', 'her', '[MASK]', 'teacher', '[MASK]', 'classmates', ',', 'and', 'married', '[MASK]', '.', '<', 'br', '/', '>', '<', 'br', '[MASK]', '>', 'what', 'kills', 'me', 'about', 'i', 'am', 'curious', '-', 'yellow', 'is', 'that', '40', 'years', 'ago', ',', 'this', 'was', 'considered', 'pornographic', '[MASK]', 'really', '[MASK]', 'the', '[MASK]', 'and', 'nu', '[MASK]', 'scenes', 'are', 'few', 'and', 'far', 'between', '[MASK]', 'even', 'then', 'it', \"'\", 's', 'not', 'shot', 'like', 'some', 'cheap', '[MASK]', 'made', 'porn', '[MASK]', '[MASK]', 'while', 'my', 'country', '##men', 'mind', 'find', 'it', '[MASK]', ',', 'in', 'reality', 'sex', 'and', 'nu', '##dity', 'are', 'a', 'major', '[MASK]', 'in', '[MASK]', 'cinema', '.', '[MASK]', 'ing', '##mar', 'bergman', '[MASK]']'\n"
     ]
    }
   ],
   "source": [
    "for chunk in data_collator(samples)[\"input_ids\"]:\n",
    "    print(f\"\\n'>>> {tokenizer.convert_ids_to_tokens(chunk)}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[MASK] token are randomly inserted into the text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'word_ids', 'labels'],\n",
       "        num_rows: 10000\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'word_ids', 'labels'],\n",
       "        num_rows: 1000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_size = 10_000\n",
    "test_size = int(0.1 * train_size)\n",
    "\n",
    "downsampled_dataset = lm_datasets[\"train\"].train_test_split(\n",
    "    train_size=train_size, test_size=test_size, seed=42\n",
    ")\n",
    "downsampled_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments\n",
    "\n",
    "batch_size = 64\n",
    "# Show the training loss with every epoch\n",
    "logging_steps = len(downsampled_dataset[\"train\"]) // batch_size\n",
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=f\"./output/{model_name}-finetuned-imdb\",\n",
    "    overwrite_output_dir=True,\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    weight_decay=0.01,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    push_to_hub=False,\n",
    "    fp16=True,\n",
    "    logging_steps=logging_steps,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Trainer\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=downsampled_dataset[\"train\"],\n",
    "    eval_dataset=downsampled_dataset[\"test\"],\n",
    "    data_collator=data_collator,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\lkk68\\.conda\\envs\\mycondapy39\\lib\\site-packages\\transformers\\optimization.py:411: 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"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "07c73b1e5aa74d0c9c3e7e4fcec83012",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/471 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 2.6977, 'learning_rate': 1.3503184713375796e-05, 'epoch': 0.99}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b6f54c4790042db969a72e26b4bda1b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 2.5231194496154785, 'eval_runtime': 0.6522, 'eval_samples_per_second': 1533.243, 'eval_steps_per_second': 24.532, 'epoch': 1.0}\n",
      "{'loss': 2.5618, 'learning_rate': 6.878980891719745e-06, 'epoch': 1.99}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "be323bd9018442c4b24d5a8cc9f6a3f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 2.474677085876465, 'eval_runtime': 0.6521, 'eval_samples_per_second': 1533.41, 'eval_steps_per_second': 24.535, 'epoch': 2.0}\n",
      "{'loss': 2.5276, 'learning_rate': 2.547770700636943e-07, 'epoch': 2.98}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "14af5556ecf24d1ab0f744dc6ac716be",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 2.4369473457336426, 'eval_runtime': 0.6596, 'eval_samples_per_second': 1516.04, 'eval_steps_per_second': 24.257, 'epoch': 3.0}\n",
      "{'train_runtime': 68.397, 'train_samples_per_second': 438.616, 'train_steps_per_second': 6.886, 'train_loss': 2.59651494481761, 'epoch': 3.0}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=471, training_loss=2.59651494481761, metrics={'train_runtime': 68.397, 'train_samples_per_second': 438.616, 'train_steps_per_second': 6.886, 'train_loss': 2.59651494481761, 'epoch': 3.0})"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "df6fdcec33014e4584f291977f1aee24",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">>> Perplexity: 11.41\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "\n",
    "eval_results = trainer.evaluate()\n",
    "print(f\">>> Perplexity: {math.exp(eval_results['eval_loss']):.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perplexity lower is better"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use Whole Word Masking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "import numpy as np\n",
    "\n",
    "from transformers import default_data_collator\n",
    "\n",
    "wwm_probability = 0.2\n",
    "\n",
    "\n",
    "def whole_word_masking_data_collator(features):\n",
    "    for feature in features:\n",
    "        word_ids = feature.pop(\"word_ids\")\n",
    "\n",
    "        # Create a map between words and corresponding token indices\n",
    "        mapping = collections.defaultdict(list)\n",
    "        current_word_index = -1\n",
    "        current_word = None\n",
    "        for idx, word_id in enumerate(word_ids):\n",
    "            if word_id is not None:\n",
    "                if word_id != current_word:\n",
    "                    current_word = word_id\n",
    "                    current_word_index += 1\n",
    "                mapping[current_word_index].append(idx)\n",
    "\n",
    "        # Randomly mask words\n",
    "        mask = np.random.binomial(1, wwm_probability, (len(mapping),))\n",
    "        input_ids = feature[\"input_ids\"]\n",
    "        labels = feature[\"labels\"]\n",
    "        new_labels = [-100] * len(labels)\n",
    "        for word_id in np.where(mask)[0]:\n",
    "            word_id = word_id.item()\n",
    "            for idx in mapping[word_id]:\n",
    "                new_labels[idx] = labels[idx]\n",
    "                input_ids[idx] = tokenizer.mask_token_id\n",
    "\n",
    "    return default_data_collator(features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "'>>> [CLS] i [MASK] i am curious - yellow [MASK] my video store [MASK] [MASK] all the controversy that surrounded it when it was first released in 1967. i also heard that at first [MASK] [MASK] [MASK] by u. s. customs if [MASK] ever tried to enter this country [MASK] therefore [MASK] [MASK] fan [MASK] films considered \" controversial \" i [MASK] had to see this [MASK] myself [MASK] < br / [MASK] < br / > the plot [MASK] centered around a young swedish drama student [MASK] lena who wants to learn everything she can about life. in particular she wants [MASK] focus her attentions to making some sort [MASK] [MASK] on [MASK] [MASK] average [MASK] [MASK] thought [MASK] [MASK] political issues such'\n",
      "\n",
      "'>>> as [MASK] [MASK] war and race issues in the [MASK] states. in between [MASK] [MASK] and ordinary denizens of stockholm about their opinions on [MASK], she has [MASK] with her drama teacher, classmates, and married men [MASK] < br / > < br / > what kills me about [MASK] am curious - [MASK] [MASK] that 40 [MASK] ago, this was [MASK] [MASK]. really, the [MASK] [MASK] nudity scenes are few [MASK] far between, even then it's not shot like some cheaply made porno [MASK] while my countrymen mind find it [MASK], [MASK] reality sex and nudity are a major staple in swedish cinema [MASK] [MASK] [MASK] [MASK] bergman,'\n"
     ]
    }
   ],
   "source": [
    "samples = [lm_datasets[\"train\"][i] for i in range(2)]\n",
    "batch = whole_word_masking_data_collator(samples)\n",
    "\n",
    "for chunk in batch[\"input_ids\"]:\n",
    "    print(f\"\\n'>>> {tokenizer.decode(chunk)}'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached split indices for dataset at C:\\Users\\lkk68\\.cache\\huggingface\\datasets\\imdb\\plain_text\\1.0.0\\d613c88cf8fa3bab83b4ded3713f1f74830d1100e171db75bbddb80b3345c9c0\\cache-1de8ceb78de95865.arrow and C:\\Users\\lkk68\\.cache\\huggingface\\datasets\\imdb\\plain_text\\1.0.0\\d613c88cf8fa3bab83b4ded3713f1f74830d1100e171db75bbddb80b3345c9c0\\cache-7a495c3653c6aa44.arrow\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'word_ids', 'labels'],\n",
       "        num_rows: 10000\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'word_ids', 'labels'],\n",
       "        num_rows: 1000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_size = 10_000\n",
    "test_size = int(0.1 * train_size)\n",
    "\n",
    "downsampled_dataset = lm_datasets[\"train\"].train_test_split(\n",
    "    train_size=train_size, test_size=test_size, seed=42\n",
    ")\n",
    "downsampled_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments\n",
    "\n",
    "batch_size = 64\n",
    "# Show the training loss with every epoch\n",
    "logging_steps = len(downsampled_dataset[\"train\"]) // batch_size\n",
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=f\"./output/{model_name}-finetuned-imdb\",\n",
    "    overwrite_output_dir=True,\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    weight_decay=0.01,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    push_to_hub=False,\n",
    "    fp16=True,\n",
    "    logging_steps=logging_steps,\n",
    "    remove_unused_columns=False, #do not remove word_ids, needed for Whole word masking\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Trainer\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=downsampled_dataset[\"train\"],\n",
    "    eval_dataset=downsampled_dataset[\"test\"],\n",
    "    data_collator=whole_word_masking_data_collator, #use whole word maksing data collator\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c86470ab6adf487796178fc5060162a9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/471 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7367, 'learning_rate': 1.337579617834395e-05, 'epoch': 0.99}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eafbc7a51a0d427bb34a8230afee309e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.6717138290405273, 'eval_runtime': 0.6619, 'eval_samples_per_second': 1510.791, 'eval_steps_per_second': 24.173, 'epoch': 1.0}\n",
      "{'loss': 0.6734, 'learning_rate': 6.751592356687898e-06, 'epoch': 1.99}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "85d82243f0c14531ae9786e7f000df5d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.6662740707397461, 'eval_runtime': 0.6607, 'eval_samples_per_second': 1513.63, 'eval_steps_per_second': 24.218, 'epoch': 2.0}\n",
      "{'loss': 0.6689, 'learning_rate': 1.2738853503184715e-07, 'epoch': 2.98}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "134dcb93528742eba348660c23075757",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.6505001783370972, 'eval_runtime': 0.6659, 'eval_samples_per_second': 1501.657, 'eval_steps_per_second': 24.027, 'epoch': 3.0}\n",
      "{'train_runtime': 65.5165, 'train_samples_per_second': 457.9, 'train_steps_per_second': 7.189, 'train_loss': 0.6928917070862594, 'epoch': 3.0}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=471, training_loss=0.6928917070862594, metrics={'train_runtime': 65.5165, 'train_samples_per_second': 457.9, 'train_steps_per_second': 7.189, 'train_loss': 0.6928917070862594, 'epoch': 3.0})"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d316a319b6984f4ebacdbc52c6f7b8c2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">>> Perplexity: 1.93\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "\n",
    "eval_results = trainer.evaluate()\n",
    "print(f\">>> Perplexity: {math.exp(eval_results['eval_loss']):.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"This is a great [MASK].\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "inputs = inputs.to('cuda')\n",
    "token_logits = model(**inputs).logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'>>> This is a great film.'\n",
      "'>>> This is a great movie.'\n",
      "'>>> This is a great idea.'\n",
      "'>>> This is a great adventure.'\n",
      "'>>> This is a great one.'\n"
     ]
    }
   ],
   "source": [
    "# Find the location of [MASK] and extract its logits\n",
    "mask_token_index = torch.where(inputs[\"input_ids\"] == tokenizer.mask_token_id)[1]\n",
    "mask_token_logits = token_logits[0, mask_token_index, :]\n",
    "# Pick the [MASK] candidates with the highest logits\n",
    "top_5_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()\n",
    "\n",
    "for token in top_5_tokens:\n",
    "    print(f\"'>>> {text.replace(tokenizer.mask_token, tokenizer.decode([token]))}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Change data_collator for data evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForLanguageModeling\n",
    "\n",
    "data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "#insert MASK into the original dataset for evaluator\n",
    "def insert_random_mask(batch):\n",
    "    features = [dict(zip(batch, t)) for t in zip(*batch.values())]\n",
    "    masked_inputs = data_collator(features)\n",
    "    # Create a new \"masked\" column for each column in the dataset\n",
    "    return {\"masked_\" + k: v.numpy() for k, v in masked_inputs.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached split indices for dataset at C:\\Users\\lkk68\\.cache\\huggingface\\datasets\\imdb\\plain_text\\1.0.0\\d613c88cf8fa3bab83b4ded3713f1f74830d1100e171db75bbddb80b3345c9c0\\cache-1de8ceb78de95865.arrow and C:\\Users\\lkk68\\.cache\\huggingface\\datasets\\imdb\\plain_text\\1.0.0\\d613c88cf8fa3bab83b4ded3713f1f74830d1100e171db75bbddb80b3345c9c0\\cache-7a495c3653c6aa44.arrow\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'word_ids', 'labels'],\n",
       "        num_rows: 10000\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['input_ids', 'attention_mask', 'word_ids', 'labels'],\n",
       "        num_rows: 1000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_size = 10_000\n",
    "test_size = int(0.1 * train_size)\n",
    "\n",
    "downsampled_dataset = lm_datasets[\"train\"].train_test_split(\n",
    "    train_size=train_size, test_size=test_size, seed=42\n",
    ")\n",
    "downsampled_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "53660b3d528e4842b52a654658fa1943",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "downsampled_dataset = downsampled_dataset.remove_columns([\"word_ids\"])\n",
    "eval_dataset = downsampled_dataset[\"test\"].map(\n",
    "    insert_random_mask,\n",
    "    batched=True,\n",
    "    remove_columns=downsampled_dataset[\"test\"].column_names,\n",
    ")\n",
    "eval_dataset = eval_dataset.rename_columns(\n",
    "    {\n",
    "        \"masked_input_ids\": \"input_ids\",\n",
    "        \"masked_attention_mask\": \"attention_mask\",\n",
    "        \"masked_labels\": \"labels\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['input_ids', 'attention_mask', 'labels'],\n",
       "    num_rows: 1000\n",
       "})"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from transformers import default_data_collator\n",
    "\n",
    "batch_size = 64\n",
    "train_dataloader = DataLoader(\n",
    "    downsampled_dataset[\"train\"],\n",
    "    shuffle=True,\n",
    "    batch_size=batch_size,\n",
    "    collate_fn=data_collator,\n",
    ")\n",
    "eval_dataloader = DataLoader(\n",
    "    eval_dataset, batch_size=batch_size, collate_fn=default_data_collator\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.optim import AdamW\n",
    "optimizer = AdamW(model.parameters(), lr=5e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "from accelerate import Accelerator\n",
    "accelerator = Accelerator()\n",
    "model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(\n",
    "    model, optimizer, train_dataloader, eval_dataloader\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import get_scheduler\n",
    "num_train_epochs = 8\n",
    "num_update_steps_per_epoch = len(train_dataloader)\n",
    "num_training_steps = num_train_epochs * num_update_steps_per_epoch\n",
    "\n",
    "lr_scheduler = get_scheduler(\n",
    "    \"linear\",\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=0,\n",
    "    num_training_steps=num_training_steps,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir='./output'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "510a06ec65514757ac52922aca0062bd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1256 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">>> Epoch 0: Perplexity: 11.221171807718862\n",
      ">>> Epoch 1: Perplexity: 10.768010645279727\n",
      ">>> Epoch 2: Perplexity: 10.55645597665742\n",
      ">>> Epoch 3: Perplexity: 10.342456958531864\n",
      ">>> Epoch 4: Perplexity: 10.151993669983593\n",
      ">>> Epoch 5: Perplexity: 10.03863898173546\n",
      ">>> Epoch 6: Perplexity: 9.959069167523797\n",
      ">>> Epoch 7: Perplexity: 9.911696030015712\n"
     ]
    }
   ],
   "source": [
    "from tqdm.auto import tqdm\n",
    "import torch\n",
    "import math\n",
    "\n",
    "progress_bar = tqdm(range(num_training_steps))\n",
    "\n",
    "for epoch in range(num_train_epochs):\n",
    "    # Training\n",
    "    model.train()\n",
    "    for batch in train_dataloader:\n",
    "        outputs = model(**batch)\n",
    "        loss = outputs.loss\n",
    "        accelerator.backward(loss)\n",
    "\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "        optimizer.zero_grad()\n",
    "        progress_bar.update(1)\n",
    "\n",
    "    # Evaluation\n",
    "    model.eval()\n",
    "    losses = []\n",
    "    for step, batch in enumerate(eval_dataloader):\n",
    "        with torch.no_grad():\n",
    "            outputs = model(**batch)\n",
    "\n",
    "        loss = outputs.loss\n",
    "        losses.append(accelerator.gather(loss.repeat(batch_size)))\n",
    "\n",
    "    losses = torch.cat(losses)\n",
    "    losses = losses[: len(eval_dataset)]\n",
    "    try:\n",
    "        perplexity = math.exp(torch.mean(losses))\n",
    "    except OverflowError:\n",
    "        perplexity = float(\"inf\")\n",
    "\n",
    "    print(f\">>> Epoch {epoch}: Perplexity: {perplexity}\")\n",
    "\n",
    "    # Save and upload\n",
    "    accelerator.wait_for_everyone()\n",
    "    unwrapped_model = accelerator.unwrap_model(model)\n",
    "    unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)\n",
    "    # if accelerator.is_main_process:\n",
    "    #     tokenizer.save_pretrained(output_dir)\n",
    "    #     repo.push_to_hub(\n",
    "    #         commit_message=f\"Training in progress epoch {epoch}\", blocking=False\n",
    "    #     )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"This is a great [MASK]\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([5])"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "mask_token_index = torch.where(inputs[\"input_ids\"] == tokenizer.mask_token_id)[1]\n",
    "mask_token_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs=inputs.to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "logits = model(**inputs).logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 7, 30522])"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logits.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-11.4297, -10.2422, -10.4297,  ...,  -8.8047, -11.2344,  -3.4688]],\n",
       "       device='cuda:0', grad_fn=<IndexBackward0>)"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mask_token_logits = logits[0, mask_token_index, :]\n",
    "mask_token_logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 30522])"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mask_token_logits.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_3_tokens = torch.topk(mask_token_logits, 3, dim=1).indices[0].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[999, 1012, 2143]"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_3_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This is a great !\n",
      "This is a great .\n",
      "This is a great film\n"
     ]
    }
   ],
   "source": [
    "for token in top_3_tokens:\n",
    "    print(text.replace(tokenizer.mask_token, tokenizer.decode([token])))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mycondapy39",
   "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.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}