2_visualize_layers.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import seaborn as sns\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "import json\n",
    "import time\n",
    "import wandb\n",
    "import torch\n",
    "from torch import cuda\n",
    "from torch.utils.data import DataLoader\n",
    "from sklearn.manifold import TSNE\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from utils.classes import (\n",
    "    get_AGNews_datasets,\n",
    "    train, test, accuracy,\n",
    "    dynamic_masking,\n",
    "    RobertaMLM_with_classifier,\n",
    "    visualize_layers\n",
    ")\n",
    "%env WANDB_PROJECT=TAPT_roberta\n",
    "%env WANDB_LOG_MODEL='end'\n",
    "\n",
    "# Import relevant models, tokenizers, and related libs\n",
    "from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline\n",
    "import datasets\n",
    "\n",
    "# Statics\n",
    "DEVICE = 'cuda' if cuda.is_available() else 'cpu'\n",
    "# DEVICE = 'cpu'\n",
    "SEED = 69\n",
    "SEEDED_GEN = torch.Generator().manual_seed(SEED)\n",
    "\n",
    "# Confirm device type, should say CUDA if you have a GPU\n",
    "print(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hyperparams_TAPT = {\n",
    "    \"EPOCHS\" : 20,\n",
    "    \"MASK_PROB\" : 0.1,\n",
    "    'TRAINING_BATCH_SIZE' : 32,\n",
    "    \"MAX_LEN\" : 77\n",
    "}\n",
    "\n",
    "hyperparams = {\n",
    "    \"TRAIN_PCT\" : 0.9,\n",
    "    \"TRAIN_BATCH_SIZE\" : 200,\n",
    "    \"VALID_BATCH_SIZE\" : 200,\n",
    "    \"TEST_BATCH_SIZE\" : 200,\n",
    "    \"MAX_LEN\" : 77,\n",
    "    \"EPOCHS\" : 25,\n",
    "    \"LR\" : 0.005,\n",
    "    \"L2_REG\" : 0.000000,\n",
    "    \"ADAM_BETAS\" : (0.87, 0.98),\n",
    "    \"ADAM_EPS\" : 1e-6,\n",
    "    \"FC_HIDDEN\" : 768,\n",
    "    \"FC_DROPOUT\" : 0.05,\n",
    "    \"SCH_ENDFACTOR\" : 0.1,\n",
    "    \"RUN_SUFFIX\" : \"6\"\n",
    "}\n",
    "\n",
    "# Choose either 1) fine-tuned or 2) pre-trained MLM Model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1) Load Fine-Tuned model pytorch saved\n",
    "### Load model and data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_type = \"distilroberta-base\"\n",
    "PATH = f\"models/distilroberta-base_base_finetuned_1682486148.pt\"\n",
    "MLM_layers = AutoModelForMaskedLM.from_pretrained(model_type).roberta\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_type)\n",
    "lazarus_model = RobertaMLM_with_classifier(MLM_layers, fc_hidden=hyperparams['FC_HIDDEN'], fc_dropout=hyperparams['FC_DROPOUT'])\n",
    "lazarus_model.load_state_dict(torch.load(PATH))\n",
    "\n",
    "lazarus_model.to(DEVICE)\n",
    "for param in lazarus_model.parameters():\n",
    "    param.requires_grad = False\n",
    "lazarus_model.eval()\n",
    "\n",
    "_, _, test_dataset = get_AGNews_datasets(\n",
    "    tokenizer,\n",
    "    DEVICE,\n",
    "    max_length=hyperparams['MAX_LEN'],\n",
    "    train_pct=hyperparams['TRAIN_PCT'],\n",
    "    generator=SEEDED_GEN\n",
    ")\n",
    "test_dataloader = DataLoader(dataset=test_dataset, batch_size=hyperparams['TEST_BATCH_SIZE'], shuffle=True)\n",
    "t_sne = TSNE(2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vised_layers = visualize_layers(\n",
    "    lazarus_model.mlm, \n",
    "    test_dataloader, \n",
    "    hyperparams['MAX_LEN'], \n",
    "    DEVICE, \n",
    "    layers=range(0,7), \n",
    "    saved_model_name=\"baseline\"\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2) Load pretrained model\n",
    "### Load model and data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_type = \"distilroberta-base\"\n",
    "huggingmodel = \"checkpoints/TAPT_Roberta_DAPT_TAPT/checkpoint-70000\"\n",
    "lazarus_model = AutoModelForMaskedLM.from_pretrained(huggingmodel).roberta\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_type)\n",
    "\n",
    "lazarus_model.to(DEVICE)\n",
    "for param in lazarus_model.parameters():\n",
    "    param.requires_grad = False\n",
    "lazarus_model.eval()\n",
    "\n",
    "_, _, test_dataset = get_AGNews_datasets(\n",
    "    tokenizer,\n",
    "    DEVICE,\n",
    "    max_length=hyperparams['MAX_LEN'],\n",
    "    train_pct=hyperparams['TRAIN_PCT'],\n",
    "    generator=SEEDED_GEN\n",
    ")\n",
    "test_dataloader = DataLoader(dataset=test_dataset, batch_size=hyperparams['TEST_BATCH_SIZE'], shuffle=True)\n",
    "t_sne = TSNE(2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vised_layers = visualize_layers(\n",
    "    lazarus_model, \n",
    "    test_dataloader,\n",
    "    hyperparams['MAX_LEN'], \n",
    "    DEVICE, \n",
    "    layers=range(0,7), \n",
    "    saved_model_name=\"TAPT_Roberta_DAPT_TAPT_ckpt_7\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}