notebooks/train_mlm.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "module_path = os.path.abspath(os.path.join('..'))\n",
    "if module_path not in sys.path:\n",
    "    sys.path.append(module_path)\n",
    "\n",
    "from transformers import AutoTokenizer, AutoModel, Trainer, TrainingArguments\n",
    "from collections import OrderedDict\n",
    "\n",
    "import src.domlm as model \n",
    "import src.dataset as dataset\n",
    "from src.data_collator import DataCollatorForDOMNodeMask"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup model and load roberta weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"roberta-base\")\n",
    "roberta = AutoModel.from_pretrained(\"roberta-base\")\n",
    "roberta_config = roberta.config\n",
    "\n",
    "roberta_config_dict = roberta_config.to_dict()\n",
    "roberta_config_dict[\"_name_or_path\"] = \"domlm\"\n",
    "roberta_config_dict[\"architectures\"] = [\"DOMLMForMaskedLM\"]\n",
    "domlm_config = model.DOMLMConfig.from_dict(roberta_config_dict)\n",
    "# domlm_config.save_pretrained(\"../domlm-config/\")\n",
    "domlm = model.DOMLMForMaskedLM(domlm_config)\n",
    "\n",
    "state_dict = OrderedDict((f\"domlm.{k}\",v) for k,v in roberta.state_dict().items())\n",
    "domlm.load_state_dict(state_dict,strict=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare SWDE data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = \"../data/swde_preprocessed\"\n",
    "print(f\"Loading datasets from {dataset_path}...\")\n",
    "train_ds = dataset.SWDEDataset(dataset_path)\n",
    "test_ds = dataset.SWDEDataset(dataset_path,split=\"test\")\n",
    "\n",
    "# tokenizer.pad_token = tokenizer.eos_token # why do we need this?\n",
    "data_collator = DataCollatorForDOMNodeMask(tokenizer=tokenizer, mlm_probability=0.15)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train in Masked LM fashion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO: add evaluation metrics (ppl, etc.)\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./results\",\n",
    "    evaluation_strategy=\"steps\",\n",
    "    # optim=\"adamw_apex_fused\", # only with apex installed\n",
    "    weight_decay=0.01,\n",
    "    num_train_epochs=5,\n",
    "    warmup_ratio=0.1,\n",
    "    learning_rate=1e-4,\n",
    "    per_device_train_batch_size=16,\n",
    "    per_device_eval_batch_size=16,\n",
    "    gradient_accumulation_steps=4,\n",
    "    # gradient_checkpointing=True, # vram is enough without checkpointing\n",
    "    fp16 = True, # If Ampere: bf16 = True\n",
    "    # tf32 = True, # Ampere Only\n",
    "    dataloader_num_workers=8,\n",
    "    dataloader_pin_memory=True\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=domlm,\n",
    "    args=training_args,\n",
    "    train_dataset=train_ds,\n",
    "    eval_dataset=test_ds,\n",
    "    data_collator=data_collator,\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tg",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "5e78104479dc0910b09b3b35473d6210b102fd81376b73f3c3b60ba92738a5ea"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}