Testing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1022db2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33maaquib111\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import wandb\n",
    "wandb.login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0ba5c4ad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "wandb version 0.13.10 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.13.7"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/gs/gsfs0/home/asyed/pw/jobs/ICLR/wandb/run-20230221_181009-3lubuoz0</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href=\"https://wandb.ai/aaquib111/ICLR/runs/3lubuoz0\" target=\"_blank\">opt-350m Wikitext Test</a></strong> to <a href=\"https://wandb.ai/aaquib111/ICLR\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset wikitext (/gs/gsfs0/users/asyed/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126)\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "from utils.save_utils import load_unmasked_model\n",
    "import torch\n",
    "from torch.nn.utils import prune\n",
    "from transformers import AutoTokenizer, OPTForCausalLM\n",
    "from datasets import load_dataset\n",
    "import gc\n",
    "\n",
    "# Constants\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "model_name = \"opt-350m\"\n",
    "token_length=2048\n",
    "stride = 2048\n",
    "\n",
    "wandb.init(project=\"ICLR\", name = f'{model_name} Wikitext Test', config={'token_length': token_length,\n",
    "                                                             'model_name': model_name,\n",
    "                                                             'stride': stride})\n",
    "# Load tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(f'facebook/{model_name}', \n",
    "                                          padding_side='left', \n",
    "                                          use_fast=False)\n",
    "# Load dataset\n",
    "test_set = load_dataset('wikitext', 'wikitext-2-raw-v1', split='test')\n",
    "encodings = tokenizer(\"\\n\\n\".join(test_set['text']), return_tensors='pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "67dfdea4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 388/388 [00:00<00:00, 14131.31it/s]\n",
      "  0%|          | 0/141 [00:00<?, ?it/s]/gs/gsfs0/users/asyed/.local/lib/python3.9/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      " 99%|█████████▉| 140/141 [00:28<00:00,  4.92it/s]\n",
      "100%|██████████| 388/388 [00:00<00:00, 15975.48it/s]\n",
      " 99%|█████████▉| 140/141 [00:28<00:00,  4.90it/s]\n",
      "100%|██████████| 388/388 [00:00<00:00, 15505.74it/s]\n",
      " 99%|█████████▉| 140/141 [00:28<00:00,  4.87it/s]\n",
      "100%|██████████| 388/388 [00:00<00:00, 58837.63it/s]\n",
      " 99%|█████████▉| 140/141 [00:28<00:00,  4.85it/s]\n",
      "100%|██████████| 388/388 [00:00<00:00, 58961.27it/s]\n",
      " 99%|█████████▉| 140/141 [00:28<00:00,  4.85it/s]\n",
      "100%|██████████| 388/388 [00:00<00:00, 58377.51it/s]\n",
      " 36%|███▌      | 51/141 [00:10<00:18,  4.83it/s]\n",
      "\n",
      "KeyboardInterrupt\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "seq_len = encodings.input_ids.size(1)\n",
    "SPARSITIES = [0.1, 0.2, 0.3, 0.5, 0.7, 0.9, 1]#, 0.4, 0.6, 0.8, 1\n",
    "\n",
    "for SPARSITY in SPARSITIES:\n",
    "    loaded_model = OPTForCausalLM.from_pretrained(f'facebook/{model_name}', output_attentions=True, output_hidden_states=True).to(device=device) # type: ignore\n",
    "    loaded_model = torch.nn.DataParallel(loaded_model, device_ids=[0,1,2,3])\n",
    "    if SPARSITY != 1:\n",
    "        load_unmasked_model(loaded_model, f'pruned_models/{model_name}-{SPARSITY}.pt')\n",
    "    \n",
    "    loaded_model.eval()\n",
    "    _ = loaded_model(torch.randint(high=20, size=(1,10)))\n",
    "    \n",
    "    nlls = []\n",
    "    prev_end_loc = 0\n",
    "    for begin_loc in tqdm(range(0, seq_len, stride)):\n",
    "        end_loc = min(begin_loc + token_length, seq_len)\n",
    "        trg_len = end_loc - prev_end_loc\n",
    "        input_ids = encodings.input_ids[:, begin_loc:end_loc].to(device=device)\n",
    "        target_ids = input_ids.clone()\n",
    "        target_ids[:,:-trg_len] = -100\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            outputs = loaded_model(input_ids, labels=target_ids)\n",
    "            neg_log_likelihood = outputs.loss * trg_len\n",
    "            \n",
    "        nlls.append(neg_log_likelihood)\n",
    "        \n",
    "        prev_end_loc = end_loc\n",
    "        if end_loc == seq_len:\n",
    "            break\n",
    "            \n",
    "    ppl = torch.exp(torch.stack(nlls).sum() / end_loc)\n",
    "    wandb.log({\"perplexity\": ppl, 'density': SPARSITY})\n",
    "    \n",
    "    del loaded_model\n",
    "    gc.collect()\n",
    "    torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6697d6ae",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.10 64-bit (microsoft store)",
   "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.10.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "10a62883c8cf88488dd1d1962da2497146a937f4733084ac36b0a34d78c1145f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}