shap_on_biotech.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/akshat/miniconda3/envs/interp/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import shap\n",
    "import sklearn\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import sklearn\n",
    "import sklearn.ensemble\n",
    "import sklearn.metrics\n",
    "import accelerate\n",
    "import pytorch_lightning as pl\n",
    "from transformers import pipeline, AutoTokenizer, AutoModel, DataCollatorWithPadding, EvalPrediction, TrainingArguments, Trainer, OPTForSequenceClassification, AutoConfig, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "from torch.optim import AdamW\n",
    "from torch.utils.data import TensorDataset\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import evaluate\n",
    "import tqdm.notebook as tq\n",
    "from datasets import load_dataset\n",
    "from peft import PeftModel, PeftConfig, LoraConfig, get_peft_model, TaskType\n",
    "from transformers import AutoModelForCausalLM, BitsAndBytesConfig\n",
    "import os "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cpu\n"
     ]
    }
   ],
   "source": [
    "torch.cuda.empty_cache()\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "MAX_LEN = 2048\n",
    "MODEL = \"facebook/opt-350m\"\n",
    "TRAIN_BATCH_SIZE = 4\n",
    "VALID_BATCH_SIZE = 4\n",
    "TEST_BATCH_SIZE = 4\n",
    "EPOCHS = 5\n",
    "LEARNING_RATE = 3e-05"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_checkpoint_path = './OPT-350m-events_classification_biotech'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/akshat/miniconda3/envs/interp/lib/python3.10/site-packages/datasets/load.py:1461: FutureWarning: The repository for knowledgator/events_classification_biotech contains custom code which must be executed to correctly load the dataset. You can inspect the repository content at https://hf.co/datasets/knowledgator/events_classification_biotech\n",
      "You can avoid this message in future by passing the argument `trust_remote_code=True`.\n",
      "Passing `trust_remote_code=True` will be mandatory to load this dataset from the next major release of `datasets`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "dataset = load_dataset('knowledgator/events_classification_biotech')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "classes = dataset['train'].features['label 1'].names\n",
    "classes = [c for c in classes if c]\n",
    "id2class = {i: label for i, label in enumerate(classes)}\n",
    "class2id = {label: i for i, label in enumerate(classes)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of OPTForSequenceClassification were not initialized from the model checkpoint at facebook/opt-350m and are newly initialized: ['score.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 1,588,224 || all params: 332,800,000 || trainable%: 0.47723076923076924\n"
     ]
    }
   ],
   "source": [
    "lora_config = LoraConfig(\n",
    "    r=16,\n",
    "    target_modules=[\"q_proj\", \"v_proj\"],\n",
    "    task_type=TaskType.SEQ_CLS,\n",
    "    lora_alpha=32,\n",
    "    lora_dropout=0.05,\n",
    ")\n",
    "\n",
    "config, unused_kwargs = AutoConfig.from_pretrained(\n",
    "    MODEL,\n",
    "    num_labels=len(classes),\n",
    "    id2label=id2class,\n",
    "    label2id=class2id,\n",
    "    problem_type=\"multi_label_classification\",\n",
    "    return_unused_kwargs=True,\n",
    ")\n",
    "\n",
    "\n",
    "model = OPTForSequenceClassification.from_pretrained(\n",
    "    MODEL,\n",
    "    config=config,\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, lora_config)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = PeftConfig.from_pretrained(\"./OPT-350m-events_classification_biotech/\")\n",
    "model = PeftModel.from_pretrained(model, \"./OPT-350m-events_classification_biotech/\", is_trainable=False,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 15,360 || all params: 332,800,000 || trainable%: 0.004615384615384616\n"
     ]
    }
   ],
   "source": [
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(texts):\n",
    "    encoded_input = tokenizer(texts, padding='max_length', truncation=True, max_length=MAX_LEN, return_tensors='pt')\n",
    "    with torch.no_grad():\n",
    "        output = model(**encoded_input)\n",
    "    probabilities = torch.nn.functional.softmax(output.logits, dim=-1).numpy()\n",
    "    return probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_subset = dataset['train'][:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = data_subset['content']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tokenize the texts\n",
    "tokenized_background = tokenizer(\n",
    "    texts,\n",
    "    padding='max_length',\n",
    "    truncation=True,\n",
    "    max_length=MAX_LEN,\n",
    "    return_tensors='pt'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "background_data = tokenized_background['input_ids'].detach().numpy()\n",
    "explainer = shap.KernelExplainer(predict, background_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shap_values = explainer.shap_values(background_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shap.initjs()\n",
    "shap.force_plot(explainer.expected_value[0], shap_values[0], texts)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "interp",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}