notebooks/221_train_language_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ccd0eb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f374256",
   "metadata": {},
   "source": [
    "# Train Roberta masked language model + tokenizer\n",
    "\n",
    "Using the names from 220, train a Roberta masked language model and tokenizer\n",
    "\n",
    "- train tokenizer first (ByteLevelBPETokenizer)\n",
    "  - this does not have to be the same tokenizer we use when training the bi-encoder, but it could be\n",
    "- then train a Roberta masked language model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2525d3d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "\n",
    "from datasets import load_dataset, Dataset\n",
    "import torch\n",
    "from tokenizers import ByteLevelBPETokenizer\n",
    "from tokenizers.processors import BertProcessing\n",
    "from transformers import RobertaTokenizerFast, RobertaForMaskedLM, RobertaConfig, \\\n",
    "                         PreTrainedTokenizer, DataCollatorForLanguageModeling, \\\n",
    "                         Trainer, TrainingArguments, \\\n",
    "                         pipeline\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd5420dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "given_surname = 'surname'\n",
    "\n",
    "tokenizer_vocab_size = 265\n",
    "tokenizer_min_frequency = 2\n",
    "tokenizer_max_length = 32\n",
    "roberta_attention_heads = 12\n",
    "roberta_hidden_layers = 6\n",
    "\n",
    "name_paths = [\n",
    "#     f\"../data/processed/all-tree-preferred-{given_surname}-sample-1m.txt\",\n",
    "    f\"../data/processed/all-tree-hr-{given_surname}-sample-10m.txt\",  # 10m line sample \n",
    "#     f\"../data/processed/all-tree-preferred-{given_surname}.txt\",\n",
    "#     f\"../data/processed/all-tree-hr-{given_surname}.txt\",\n",
    "]\n",
    "\n",
    "roberta_dir = f\"../data/models/roberta-{given_surname}-10m-{tokenizer_vocab_size}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30554d29",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(roberta_dir):\n",
    "    os.makedirs(roberta_dir) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da1ac45b",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()\n",
    "print(torch.cuda.is_available())\n",
    "print(\"cuda total\", torch.cuda.get_device_properties(0).total_memory)\n",
    "print(\"cuda reserved\", torch.cuda.memory_reserved(0))\n",
    "print(\"cuda allocated\", torch.cuda.memory_allocated(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97c53c89",
   "metadata": {},
   "source": [
    "## Train tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de3b6475",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time \n",
    "\n",
    "# Initialize a tokenizer\n",
    "tokenizer = ByteLevelBPETokenizer()\n",
    "\n",
    "# Customize training\n",
    "tokenizer.train(files=name_paths, \n",
    "                vocab_size=tokenizer_vocab_size, \n",
    "                min_frequency=tokenizer_min_frequency, \n",
    "                special_tokens=[\n",
    "    \"<s>\",\n",
    "    \"<pad>\",\n",
    "    \"</s>\",\n",
    "    \"<unk>\",\n",
    "    \"<mask>\",\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ea427d1",
   "metadata": {},
   "source": [
    "### Save tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4eeea464",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.save_model(roberta_dir)\n",
    "\n",
    "# need to manually create a json config file for the pipeline function to work\n",
    "# copying https://zablo.net/blog/post/training-roberta-from-scratch-the-missing-guide-polish-language-model/\n",
    "tokenizer_config = {\n",
    "    \"architectures\": [\"RobertaForMaskedLM\"], \n",
    "    \"max_position_embeddings\": tokenizer_max_length+2, \n",
    "    \"vocab_size\": tokenizer_vocab_size,\n",
    "}\n",
    "with open(os.path.join(roberta_dir, \"config.json\"), \"w\") as f:\n",
    "    json.dump(tokenizer_config, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f60ac201",
   "metadata": {},
   "source": [
    "### Test tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04a27833",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_tokenizer(roberta_dir, tokenizer_max_length):\n",
    "    tokenizer = ByteLevelBPETokenizer(\n",
    "        os.path.join(roberta_dir, \"vocab.json\"),\n",
    "        os.path.join(roberta_dir, \"merges.txt\"),\n",
    "    )\n",
    "    tokenizer._tokenizer.post_processor = BertProcessing(\n",
    "        (\"</s>\", tokenizer.token_to_id(\"</s>\")),\n",
    "        (\"<s>\", tokenizer.token_to_id(\"<s>\")),\n",
    "    )\n",
    "    tokenizer.enable_truncation(max_length=tokenizer_max_length) \n",
    "    return tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4111d517",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load tokenizer\n",
    "tokenizer = load_tokenizer(roberta_dir, tokenizer_max_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dbcc79a",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.encode(\"richard\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca16a688",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.encode(\"richard\").tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9672bcfc",
   "metadata": {},
   "source": [
    "## Train Roberta Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa314c33",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = RobertaTokenizerFast.from_pretrained(roberta_dir, max_len=tokenizer_max_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "918d37e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.encode(\"richard\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b7d26e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = RobertaConfig(\n",
    "    vocab_size=tokenizer_vocab_size,\n",
    "    max_position_embeddings=tokenizer_max_length+2,\n",
    "    num_attention_heads=roberta_attention_heads,\n",
    "    num_hidden_layers=roberta_hidden_layers,\n",
    "    type_vocab_size=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "361d41cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = RobertaForMaskedLM(config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e45ee39",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.num_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60bc228e",
   "metadata": {},
   "source": [
    "### Create dataset and collator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "725c730c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[\"text\"], truncation=True)\n",
    "\n",
    "dataset = load_dataset(\"text\", data_files=name_paths).map(tokenize_function, batched=True)[\"train\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb6df52e",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "661b4918",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9d5f438",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_collator = DataCollatorForLanguageModeling(\n",
    "    tokenizer=tokenizer, mlm=True, mlm_probability=0.15, \n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "537c5cbb",
   "metadata": {},
   "source": [
    "### Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dab0ac3",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir=roberta_dir,\n",
    "    overwrite_output_dir=True,\n",
    "    num_train_epochs=1,\n",
    "    per_device_train_batch_size=64,\n",
    "    save_steps=len(dataset) / 20,\n",
    "    save_total_limit=20,\n",
    "    prediction_loss_only=True,\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    data_collator=data_collator,\n",
    "    train_dataset=dataset,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec7cbdf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "070283c0",
   "metadata": {},
   "source": [
    "### Save model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81cb43e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_model(roberta_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3ca2118",
   "metadata": {},
   "source": [
    "### Test model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b36f14c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fill_mask = pipeline(\n",
    "    \"fill-mask\",\n",
    "    model=roberta_dir,\n",
    "    tokenizer=roberta_dir\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3ca4f0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fill_mask(\"mari<mask>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9308b73",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nama",
   "language": "python",
   "name": "nama"
  },
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}