aim_i_train_tune_predict.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q3dO24toXoAb"
      },
      "source": [
        "# Passive suicidality in a repressive U.S. political context: Aim I"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pl9inWkyX3HI"
      },
      "source": [
        "_WIP - NOT FOR DISTRIBUTION_\n",
        "\n",
        "_Formats and augments annotated Reddit data with excerpted rationales. Trains and tests baseline and in-domain-adapted BERT, RoBERTa, and DistilBERT pretrained models, and baseline Llama 3.1 for sequence classification. Tunes hyperparameters, saves best-performing model for inference. Includes LIME instance-wise explainability, regex for additional encodings, temperature scaling for calibration._"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C3S8qRUNX577"
      },
      "source": [
        "> aim_i_train_tune_predict.ipynb<br>\n",
        "> Simone J. Skeen (01-08-2025)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oVjqU-H_ZRZT"
      },
      "source": [
        "1. [Prepare](#scrollTo=8z4T1_xNZXFb)\n",
        "2. [Write](#scrollTo=vA0d1kH7bOYg)\n",
        "3. [Preprocess](#scrollTo=woZYUo6JgN1X)\n",
        "4. [Train-Adapt-Test](#scrollTo=rxzedr7a-kuK)\n",
        "5. [Tune-Regularize](#scrollTo=G614tXUn0VWw)\n",
        "6. [Infer (prelim)](#scrollTo=xnP3hWY2bGlp)\n",
        "7. [Explain](#scrollTo=I5JVWw7spBaF)\n",
        "8. [Calibrate](#scrollTo=2nSz7vQh5xYz)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8z4T1_xNZXFb"
      },
      "source": [
        "### 1. Prepare\n",
        "Installs, imports, and downloads requisite models and packages. Organizes RAP-consistent directory structure.\n",
        "***"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gg9PRwHCZlKP"
      },
      "source": [
        "**Install**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "WKF9nw1xXeLg"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "\n",
        "%pip install accelerate\n",
        "%pip install bitsandbytes\n",
        "%pip install brokenaxes\n",
        "%pip install causalnlp\n",
        "%pip install contractions\n",
        "%pip install datasets\n",
        "%pip install evaluate\n",
        "%pip install lime\n",
        "%pip install peft\n",
        "#%pip install trl\n",
        "%pip install unidecode\n",
        "%pip install wandb\n",
        "\n",
        "#%pip uninstall -y pyarrow datasets\n",
        "#%pip install pyarrow datasets\n",
        "\n",
        "#!python -m spacy download en_core_web_lg --user"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WdMnyx7TZt44"
      },
      "source": [
        "**Import**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FiUj7-fcZs6i"
      },
      "outputs": [],
      "source": [
        "import accelerate\n",
        "import ast\n",
        "import bitsandbytes as bnb\n",
        "import contractions\n",
        "#import en_core_web_lg\n",
        "import gzip\n",
        "import huggingface_hub\n",
        "import json\n",
        "import lime\n",
        "import logging\n",
        "import matplotlib.font_manager as fm\n",
        "import matplotlib.pyplot as plt\n",
        "import nltk\n",
        "import numpy as np\n",
        "import os\n",
        "import pandas as pd\n",
        "import peft\n",
        "import random\n",
        "import re\n",
        "import seaborn as sns\n",
        "import sklearn\n",
        "import spacy\n",
        "import string\n",
        "import torch\n",
        "#import trl\n",
        "import wandb.sdk\n",
        "import warnings\n",
        "\n",
        "from brokenaxes import brokenaxes\n",
        "from causalnlp import Autocoder, CausalInferenceModel\n",
        "from google.colab import drive\n",
        "from lightgbm import LGBMClassifier\n",
        "from lime.lime_text import LimeTextExplainer\n",
        "from lime.submodular_pick import SubmodularPick\n",
        "from matplotlib.lines import Line2D\n",
        "from nltk.text import Text\n",
        "from sklearn.feature_extraction import text\n",
        "from sklearn.linear_model import LogisticRegression, LinearRegression\n",
        "from sklearn.model_selection import (\n",
        "                                     KFold,\n",
        "                                     ParameterGrid,\n",
        "                                     StratifiedKFold,\n",
        "                                     train_test_split,\n",
        "                                     )\n",
        "\n",
        "from sklearn.metrics import (\n",
        "                             accuracy_score,\n",
        "                             average_precision_score,\n",
        "                             classification_report,\n",
        "                             cohen_kappa_score,\n",
        "                             f1_score,\n",
        "                             matthews_corrcoef,\n",
        "                             )\n",
        "from sklearn.preprocessing import LabelEncoder\n",
        "from sklearn.utils import shuffle\n",
        "from sklearn.utils.multiclass import type_of_target\n",
        "from textblob import TextBlob\n",
        "from torch.optim import AdamW\n",
        "from torch.nn import CrossEntropyLoss\n",
        "from torch.utils.data import (\n",
        "                              DataLoader,\n",
        "                              Dataset,\n",
        "                              TensorDataset,\n",
        "                              )\n",
        "from tqdm import tqdm\n",
        "from transformers import (\n",
        "                          AdamW,\n",
        "                          BertForSequenceClassification,\n",
        "                          BertTokenizer,\n",
        "                          DataCollatorForLanguageModeling,\n",
        "                          DistilBertForSequenceClassification,\n",
        "                          DistilBertTokenizer,\n",
        "                          RobertaForSequenceClassification,\n",
        "                          RobertaTokenizer,\n",
        "                          Trainer,\n",
        "                          TrainingArguments,\n",
        "                          )\n",
        "from unidecode import unidecode\n",
        "\n",
        "#spacy.cli.download('en_core_web_lg')\n",
        "\n",
        "from IPython.core.interactiveshell import InteractiveShell\n",
        "InteractiveShell.ast_node_interactivity = 'all'\n",
        "\n",
        "pd.options.mode.copy_on_write = True\n",
        "\n",
        "pd.set_option(\n",
        "              'display.max_columns',\n",
        "              None,\n",
        "              )\n",
        "\n",
        "pd.set_option(\n",
        "              'display.max_rows',\n",
        "              None,\n",
        "              )\n",
        "\n",
        "warnings.simplefilter(\n",
        "                      action = 'ignore',\n",
        "                      category = FutureWarning,\n",
        "                      )\n",
        "\n",
        "#!python -m prodigy stats"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IATf4h203EtW"
      },
      "source": [
        "**Set env variables**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-tUT-L9U3EBh"
      },
      "outputs": [],
      "source": [
        "os.environ['OPENAI_API_KEY'] = ''\n",
        "os.environ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1MY3yO_FZ0qc"
      },
      "source": [
        "**Mount gdrive**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VjzlnRXkZs_F"
      },
      "outputs": [],
      "source": [
        "drive.mount(\n",
        "            '/content/drive',\n",
        "            #force_remount = True,\n",
        "            )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tPQCYAa6aHrp"
      },
      "source": [
        "**Structure directories**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vIS1Ed8wZtC9"
      },
      "outputs": [],
      "source": [
        "%cd /content/drive/My Drive/Colab/bar_policy_suicidality\n",
        "#%cd /content/drive/My Drive/#<my_project_folder>\n",
        "\n",
        "#%mkdir bar_policy_suicidality\n",
        "#%cd bar_policy_suicidality"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TGHF0HLnZtG1"
      },
      "outputs": [],
      "source": [
        "#%mkdir inputs outputs code temp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DLMtdecmZtKs"
      },
      "outputs": [],
      "source": [
        "#%cd inputs\n",
        "#%mkdir annotation archives data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J9__408paVgp"
      },
      "outputs": [],
      "source": [
        "#%cd ../outputs\n",
        "#%mkdir models tables figures"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z8If94wBM0a9"
      },
      "outputs": [],
      "source": [
        "bar_policy_suicidality/\n",
        "├── inputs/\n",
        "│   ├── archives\n",
        "│   │   └── ### archive name TKTK\n",
        "│   └── data\n",
        "│       └── d_annotated.xlsx\n",
        "├── outputs/\n",
        "│   ├── models\n",
        "│   ├── tables\n",
        "│   └── figures\n",
        "├── code/\n",
        "└── temp/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vA0d1kH7bOYg"
      },
      "source": [
        "### 2. Write\n",
        "Writes and imports requisite custom scripts in .py.\n",
        "***"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0lDM3ZhLaVlL"
      },
      "outputs": [],
      "source": [
        "%cd code"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wDa-A40SbzxS"
      },
      "source": [
        "#### preprocess.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zQ47VdBub3-6"
      },
      "source": [
        "**_augment_training_data_with_rationales_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QvnNBR4saVos"
      },
      "outputs": [],
      "source": [
        "%%writefile preprocess.py\n",
        "\n",
        "import pandas as pd\n",
        "\n",
        "def augment_training_data_with_rationales(df):\n",
        "    \"\"\"\n",
        "    Identifies all pos_1 strn, duplicates as new row below, replaces new row 'text' with appended concatenated asp-dep-val 'rtnl'.\n",
        "    \"\"\"\n",
        "    augmented_rows = []\n",
        "\n",
        "    for index, row in df.iterrows():\n",
        "        augmented_rows.append(row)\n",
        "\n",
        "        if row['strn'] == 1:\n",
        "            duplicate_row = row.copy()\n",
        "            duplicate_row['text'] = duplicate_row['rtnl']\n",
        "            augmented_rows.append(duplicate_row)\n",
        "\n",
        "    df_augmented = pd.DataFrame(augmented_rows)\n",
        "\n",
        "    return df_augmented"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YkGb9oK3cFKd"
      },
      "source": [
        "**_dummy_code_augmented_rows_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5Xl1XCTXaVsN"
      },
      "outputs": [],
      "source": [
        "%%writefile -a preprocess.py\n",
        "\n",
        "import pandas as pd\n",
        "\n",
        "def dummy_code_augmented_rows(df):\n",
        "    \"\"\"\n",
        "    Identifies all rationale-augmented rows in df, dummy codes for deletion prior to evaluation.\n",
        "    \"\"\"\n",
        "    df = df.reset_index(drop = True)\n",
        "\n",
        "    df['aug'] = 0\n",
        "\n",
        "    for i in range(1, len(df)):\n",
        "        if df.at[i, 'rtnl'] != '.' and df.at[i, 'rtnl'] == df.at[i-1, 'rtnl']:\n",
        "            df.at[i, 'aug'] = 1\n",
        "\n",
        "    return df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M4QhwLbwcwR0"
      },
      "source": [
        "**_read_and_append_jsonl_posts_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ugKscEGOctEA"
      },
      "outputs": [],
      "source": [
        "%%writefile -a preprocess.py\n",
        "\n",
        "import os\n",
        "import pandas as pd\n",
        "\n",
        "def read_and_append_jsonl_posts(directory, chunk_size = 10000):\n",
        "    \"\"\"\n",
        "    Reads and appends JSONL posts archives from Arctic Shift archives dir.\n",
        "    \"\"\"\n",
        "    d_posts = pd.DataFrame()\n",
        "\n",
        "    for filename in os.listdir(directory):\n",
        "        if filename.endswith(\"_posts.jsonl\"):\n",
        "            filepath = os.path.join(\n",
        "                                    directory,\n",
        "                                    filename,\n",
        "                                    )\n",
        "\n",
        "            for chunk in pd.read_json(\n",
        "                                      filepath,\n",
        "                                      lines = True,\n",
        "                                      chunksize = chunk_size,\n",
        "                                      ):\n",
        "\n",
        "                d_posts = pd.concat([d_posts, chunk], ignore_index = True)\n",
        "\n",
        "    return d_posts"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kcaVJDy2czen"
      },
      "source": [
        "**_read_and_append_jsonl_comments_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j-5-_gbqctQ3"
      },
      "outputs": [],
      "source": [
        "%%writefile -a preprocess.py\n",
        "\n",
        "import os\n",
        "import pandas as pd\n",
        "\n",
        "def read_and_append_jsonl_comments(directory, chunk_size = 10000):\n",
        "    \"\"\"\n",
        "    Reads and appends JSONL comments archives from Arctic Shift archives dir.\n",
        "    \"\"\"\n",
        "    d_comments = pd.DataFrame()\n",
        "\n",
        "    for filename in os.listdir(directory):\n",
        "        if filename.endswith(\"_comments.jsonl\"):\n",
        "            filepath = os.path.join(\n",
        "                                    directory,\n",
        "                                    filename,\n",
        "                                    )\n",
        "\n",
        "            for chunk in pd.read_json(\n",
        "                                      filepath,\n",
        "                                      lines = True,\n",
        "                                      chunksize = chunk_size,\n",
        "                                      ):\n",
        "\n",
        "                d_comments = pd.concat([d_comments, chunk], ignore_index = True)\n",
        "\n",
        "    return d_comments"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DgGFhVlV7sYw"
      },
      "source": [
        "#### redact.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3dUaruYL7sG_"
      },
      "source": [
        "**_ner_redact_post_texts_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Szq7MCpy7q4C"
      },
      "outputs": [],
      "source": [
        "%%writefile redact.py\n",
        "\n",
        "import spacy\n",
        "nlp = spacy.load('en_core_web_lg')\n",
        "\n",
        "def ner_redact_post_texts(p_text):\n",
        "    \"\"\"\n",
        "    Redacts all named entities recognized by spaCy EntityRecognizer, replaces with <|PII|> pseudo-word token.\n",
        "    \"\"\"\n",
        "    ne = list(\n",
        "              [\n",
        "               'PERSON',   ### people, including fictional\n",
        "               'NORP',     ### nationalities or religious or political groups\n",
        "               'FAC',      ### buildings, airports, highways, bridges, etc.\n",
        "               'ORG',      ### companies, agencies, institutions, etc.\n",
        "               #'GPE',     ### countries, cities, states\n",
        "               'LOC',      ### non-GPE locations, mountain ranges, bodies of water\n",
        "               'PRODUCT',  ### objects, vehicles, foods, etc. (not services)\n",
        "               'EVENT',    ### named hurricanes, battles, wars, sports events, etc.\n",
        "               ]\n",
        "                )\n",
        "\n",
        "    doc = nlp(p_text)\n",
        "    ne_to_remove = []\n",
        "    final_string = str(p_text)\n",
        "    for sent in doc.ents:\n",
        "        if sent.label_ in ne:\n",
        "            ne_to_remove.append(str(sent.text))\n",
        "    for n in range(len(ne_to_remove)):\n",
        "        final_string = final_string.replace(\n",
        "                                            ne_to_remove[n],\n",
        "                                            '<|PII|>',\n",
        "                                            )\n",
        "    return final_string"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SS_mrveOcVru"
      },
      "source": [
        "#### bert_train.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oixV7l2zPITQ"
      },
      "source": [
        "**_set_seed_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZqT7_DnFPBfL"
      },
      "outputs": [],
      "source": [
        "%%writefile bert_train.py\n",
        "\n",
        "import torch\n",
        "import numpy as np\n",
        "import random\n",
        "\n",
        "def set_seed(seed):\n",
        "    \"\"\"\n",
        "    Set random seeds for reproducibility in Pytorch.\n",
        "    \"\"\"\n",
        "    random.seed(seed)\n",
        "    np.random.seed(seed)\n",
        "    torch.manual_seed(seed)\n",
        "    torch.cuda.manual_seed(seed)\n",
        "    torch.cuda.manual_seed_all(seed)\n",
        "    torch.backends.cudnn.deterministic = True\n",
        "    torch.backends.cudnn.benchmark = False"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aZSdSJH1jfap"
      },
      "source": [
        "**_train_eval_save_bl_models_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3i3sJ8qhHNKN"
      },
      "outputs": [],
      "source": [
        "%%writefile -a bert_train.py\n",
        "\n",
        "import torch\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "from transformers import (\n",
        "                          BertForSequenceClassification,\n",
        "                          RobertaForSequenceClassification,\n",
        "                          DistilBertForSequenceClassification,\n",
        "                          BertTokenizer,\n",
        "                          RobertaTokenizer,\n",
        "                          DistilBertTokenizer,\n",
        "                          get_linear_schedule_with_warmup,\n",
        "                          )\n",
        "from sklearn.metrics import f1_score, matthews_corrcoef, average_precision_score\n",
        "from sklearn.model_selection import StratifiedKFold, train_test_split\n",
        "from sklearn.preprocessing import LabelEncoder\n",
        "from torch.nn import CrossEntropyLoss\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "def train_eval_save_bl_models(target_datasets, targets_and_class_weights, models, save_path, cycle, hyperparameter_grid):\n",
        "    \"\"\"\n",
        "    Fine-tune and eval pre-trained baseline LMs on multiple targets using target-specific train and test datasets.\n",
        "\n",
        "    Training sets d_train_{target} are split using 5-fold cross-validation: 4 folds for training and 1 fold for validation. Training\n",
        "    folds use augmented data; validation folds use original data. The best model is selected based on average performance across the\n",
        "    folds and evaluated on separate test sets d_test_{target}.\n",
        "\n",
        "    Parameters:\n",
        "    -----------\n",
        "    target_datasets : dict\n",
        "        A dictionary where keys are target names and values are tuples containing the train and test datasets for each target.\n",
        "        For example: {'asp': (d_train_asp, d_test_asp), 'dep': (d_train_dep, d_test_dep)}\n",
        "\n",
        "    targets_and_class_weights : dict\n",
        "        A dictionary where keys are target names and values are lists of class weights corresponding to each target.\n",
        "\n",
        "    models : dict\n",
        "        A dictionary of models to evaluate. Each key is a model name, and each value is a tuple containing:\n",
        "        - the model class,\n",
        "        - the tokenizer class,\n",
        "        - the name of the pre-trained model.\n",
        "\n",
        "    save_path : str, optional\n",
        "        The path where the best models will be saved.\n",
        "\n",
        "    cycle : str\n",
        "        The cycle identifier: 'baseline', indicating performance with prespecified default params; 'adapted', indicating performance\n",
        "        with in-domain adapted params.\n",
        "\n",
        "    hyperparameter_grid : dict\n",
        "        Dictionary containing hyperparameter values: batch_size, gradient_accumulation_steps, learning_rate, num_epochs, warmup_steps, and weight_decay.\n",
        "\n",
        "    Returns:\n",
        "    --------\n",
        "    d_{cycle}_performance : pd.DataFrame\n",
        "        A df containing performance metrics for each target and model per fold per pre-specified cycle,\n",
        "        and final evaluation on test data.\n",
        "    \"\"\"\n",
        "\n",
        "    # verify cycle\n",
        "\n",
        "    print(f\"CYCLE: {cycle}\")\n",
        "\n",
        "    # check CUDA\n",
        "\n",
        "    print(\"CUDA: \", torch.cuda.is_available())\n",
        "    use_cuda = torch.cuda.is_available()\n",
        "\n",
        "    # set seed\n",
        "\n",
        "    set_seed(56)\n",
        "\n",
        "    # unpack hyperparameters\n",
        "\n",
        "    batch_size = hyperparameter_grid.get('batch_size', 4)\n",
        "    gradient_accumulation_steps = hyperparameter_grid.get('gradient_accumulation_steps', 1)\n",
        "    learning_rate = hyperparameter_grid.get('learning_rate', 2e-5)\n",
        "    num_epochs = hyperparameter_grid.get('num_epochs', 2)\n",
        "    warmup_steps = hyperparameter_grid.get('warmup_steps', 0)\n",
        "    weight_decay = hyperparameter_grid.get('weight_decay', 0.0)\n",
        "\n",
        "    # best target x model F1 tracking\n",
        "\n",
        "    best_f1_scores = {target: {'score': 0, 'model': None, 'model_instance': None} for target in targets_and_class_weights}\n",
        "    results = []\n",
        "\n",
        "    # training loop: target x model\n",
        "\n",
        "    for target, class_weights in targets_and_class_weights.items():\n",
        "        print(\"\\n======================================================================================\")\n",
        "        print(f\"Label: {target}\")\n",
        "        print(\"======================================================================================\")\n",
        "\n",
        "        # target-specific datasets\n",
        "\n",
        "        d_train, d_test = target_datasets[target]\n",
        "\n",
        "        # split augmented v. non-augmented data\n",
        "\n",
        "        d_train_aug = d_train[d_train['aug'] == 1]\n",
        "        d_train_no_aug = d_train[d_train['aug'] == 0]\n",
        "\n",
        "        # prepare fold-wise training v. validation data\n",
        "\n",
        "        X_train_aug, y_train_aug = d_train_aug['text'].values, d_train_aug[target].values\n",
        "        X_train_no_aug, y_train_no_aug = d_train_no_aug['text'].values, d_train_no_aug[target].values\n",
        "        X_test, y_test = d_test['text'].values, d_test[target].values\n",
        "\n",
        "        # determine target type, encode (as needed)\n",
        "\n",
        "        target_type = 'binary' if len(np.unique(y_train_aug)) <= 2 else 'multiclass'\n",
        "        le = LabelEncoder() # Using separate LabelEncoder for each target data group to avoid encoding mismatch issues.\n",
        "        if target_type == 'binary':\n",
        "            #le = LabelEncoder()\n",
        "            y_train_aug = le.fit_transform(y_train_aug)\n",
        "            y_train_no_aug = le.fit_transform(y_train_no_aug) # Re-encode with new encoder\n",
        "            y_test = le.transform(y_test)\n",
        "\n",
        "        # define k folds\n",
        "\n",
        "        k_fold = StratifiedKFold(\n",
        "                                 n_splits = 5,\n",
        "                                 shuffle = True,\n",
        "                                 random_state = 56,\n",
        "                                 )\n",
        "\n",
        "        for model_name, (model_class, tokenizer_class, pretrained_model_name) in models.items():\n",
        "            print(f\"\\nFine-tuning {model_name} for {target}\")\n",
        "            print(\"--------------------------------------------------------------------------------------\")\n",
        "\n",
        "            fold_f1, fold_mcc, fold_auprc = [], [], []  ### store fold-wise performance metrics\n",
        "\n",
        "            # initialize tokenizer\n",
        "\n",
        "            tokenizer = tokenizer_class.from_pretrained(pretrained_model_name)\n",
        "\n",
        "            for fold_idx, (train_no_aug_idx, valid_idx) in enumerate(k_fold.split(X_train_no_aug, y_train_no_aug)):\n",
        "                print(f\"\\nFold {fold_idx + 1}/5\")\n",
        "\n",
        "                # create training set: combine aug = 1 (augmented) with fold-specific aug = 0 (non-augmented)\n",
        "\n",
        "                X_train_fold_aug = X_train_aug\n",
        "                y_train_fold_aug = y_train_aug\n",
        "\n",
        "                X_train_fold_no_aug, X_valid_fold = X_train_no_aug[train_no_aug_idx], X_train_no_aug[valid_idx]\n",
        "                y_train_fold_no_aug, y_valid_fold = y_train_no_aug[train_no_aug_idx], y_train_no_aug[valid_idx]\n",
        "\n",
        "                # combine augmented and non-augmented training data\n",
        "\n",
        "                X_train_fold = np.concatenate([X_train_fold_aug, X_train_fold_no_aug])\n",
        "                y_train_fold = np.concatenate([y_train_fold_aug, y_train_fold_no_aug])\n",
        "\n",
        "                # tokenize training and validation data\n",
        "\n",
        "                encoded_train = tokenizer(\n",
        "                                          X_train_fold.tolist(),\n",
        "                                          padding = True,\n",
        "                                          truncation = True,\n",
        "                                          return_tensors = 'pt',\n",
        "                                          )\n",
        "\n",
        "                encoded_valid = tokenizer(\n",
        "                                          X_valid_fold.tolist(),\n",
        "                                          padding = True,\n",
        "                                          truncation = True,\n",
        "                                          return_tensors = 'pt',\n",
        "                                          )\n",
        "\n",
        "                train_dataset = TensorDataset(\n",
        "                                              encoded_train['input_ids'],\n",
        "                                              encoded_train['attention_mask'],\n",
        "                                              torch.tensor(y_train_fold),\n",
        "                                              )\n",
        "\n",
        "                valid_dataset = TensorDataset(\n",
        "                                              encoded_valid['input_ids'],\n",
        "                                              encoded_valid['attention_mask'],\n",
        "                                              torch.tensor(y_valid_fold),\n",
        "                                              )\n",
        "\n",
        "                train_loader = DataLoader(\n",
        "                                          train_dataset,\n",
        "                                          batch_size = batch_size,\n",
        "                                          shuffle = True,\n",
        "                                          )\n",
        "\n",
        "                valid_loader = DataLoader(\n",
        "                                          valid_dataset,\n",
        "                                          batch_size = batch_size,\n",
        "                                          shuffle = False,\n",
        "                                          )\n",
        "\n",
        "                # instantiate model\n",
        "\n",
        "                model = model_class.from_pretrained(pretrained_model_name)\n",
        "\n",
        "                # migrate to CUDA\n",
        "\n",
        "                use_cuda = torch.cuda.is_available()\n",
        "                if use_cuda:\n",
        "                    model = model.cuda()\n",
        "\n",
        "                # set optimizer + scheduler\n",
        "\n",
        "                optimizer = torch.optim.AdamW(\n",
        "                                              model.parameters(),\n",
        "                                              lr = learning_rate,\n",
        "                                              weight_decay = weight_decay,\n",
        "                                              )\n",
        "\n",
        "                total_steps = len(train_loader) * num_epochs\n",
        "                #scheduler = torch.optim.lr_scheduler.LinearLR(\n",
        "                #                                              optimizer,\n",
        "                #                                              start_factor = 0.1,\n",
        "                #                                              #total_iters = warmup_steps,\n",
        "                #                                              total_iters = total_steps, # Fix: corrected from warmup_steps to total_steps\n",
        "                #                                              )\n",
        "\n",
        "                scheduler = get_linear_schedule_with_warmup(\n",
        "                                                            optimizer,\n",
        "                                                            num_warmup_steps = warmup_steps,\n",
        "                                                            num_training_steps = total_steps\n",
        "                                                            )\n",
        "\n",
        "                # fine-tune model on training folds (x4)\n",
        "\n",
        "                model.train()\n",
        "                criterion = CrossEntropyLoss(weight = torch.tensor(\n",
        "                                                                   class_weights,\n",
        "                                                                   dtype = torch.float\n",
        "                                                                   ).cuda() if use_cuda else torch.tensor(\n",
        "                                                                                                          class_weights,\n",
        "                                                                                                          dtype = torch.float\n",
        "                                                                                                          )\n",
        "                                             )\n",
        "\n",
        "                for epoch in range(num_epochs):\n",
        "                    for i, batch in enumerate(train_loader):\n",
        "                        input_ids, attention_mask, labels = batch\n",
        "                        labels = labels.long()\n",
        "\n",
        "                        if use_cuda:\n",
        "                            input_ids, attention_mask, labels = input_ids.cuda(), attention_mask.cuda(), labels.cuda()\n",
        "\n",
        "                        outputs = model(input_ids, attention_mask = attention_mask)\n",
        "                        logits = outputs.logits\n",
        "                        loss = criterion(logits, labels)\n",
        "\n",
        "                        # accumulate gradients, normalize loss\n",
        "\n",
        "                        loss = loss / gradient_accumulation_steps\n",
        "                        loss.backward()\n",
        "\n",
        "                        # update model weights post-accumulation steps\n",
        "\n",
        "                        if (i + 1) % gradient_accumulation_steps == 0:\n",
        "                            optimizer.step()\n",
        "                            optimizer.zero_grad()\n",
        "\n",
        "                        # apply learning rate scheduler\n",
        "\n",
        "                        scheduler.step()\n",
        "\n",
        "                # evaluate on validation fold (x1)\n",
        "\n",
        "                model.eval()\n",
        "                all_predictions, all_true_labels = [], []\n",
        "                with torch.no_grad():\n",
        "                    for batch in valid_loader:\n",
        "                        input_ids, attention_mask, labels = batch\n",
        "\n",
        "                        if use_cuda:\n",
        "                            input_ids, attention_mask, labels = input_ids.cuda(), attention_mask.cuda(), labels.cuda()\n",
        "\n",
        "                        outputs = model(input_ids, attention_mask = attention_mask)\n",
        "                        logits = outputs.logits\n",
        "                        predictions = torch.argmax(logits, dim = 1).tolist()\n",
        "                        all_predictions.extend(predictions)\n",
        "                        all_true_labels.extend(labels.tolist())\n",
        "\n",
        "                # performance metrics per validation fold\n",
        "\n",
        "                f1_macro = f1_score(\n",
        "                                    all_true_labels,\n",
        "                                    all_predictions,\n",
        "                                    average = 'macro',\n",
        "                                    )\n",
        "\n",
        "                mcc = matthews_corrcoef(\n",
        "                                        all_true_labels,\n",
        "                                        all_predictions,\n",
        "                                        )\n",
        "\n",
        "                auprc = average_precision_score(\n",
        "                                                all_true_labels,\n",
        "                                                all_predictions,\n",
        "                                                average = 'macro',\n",
        "                                                )\n",
        "\n",
        "                fold_f1.append(f1_macro)\n",
        "                fold_mcc.append(mcc)\n",
        "                fold_auprc.append(auprc)\n",
        "\n",
        "            # mean results over folds, track best model\n",
        "\n",
        "            mean_f1 = np.mean(fold_f1)\n",
        "            if mean_f1 > best_f1_scores[target]['score']:\n",
        "                best_f1_scores[target]['score'] = mean_f1\n",
        "                best_f1_scores[target]['model'] = model_name\n",
        "                best_f1_scores[target]['model_instance'] = model\n",
        "\n",
        "                save_model_name = f'{target}_{model_name}_best_{cycle}_model.pt'\n",
        "                torch.save(model.state_dict(), save_path + save_model_name)\n",
        "\n",
        "            # store results for each fold\n",
        "\n",
        "            for i in range(5):\n",
        "                results.append({\n",
        "                                'target': target,\n",
        "                                'model': model_name,\n",
        "                                'fold': i + 1,\n",
        "                                'f1_macro': fold_f1[i],\n",
        "                                'mcc': fold_mcc[i],\n",
        "                                'auprc': fold_auprc[i]\n",
        "                                })\n",
        "\n",
        "        # test on held-out d_test_{target} df\n",
        "\n",
        "        print(f\"\\nTest on held-out d_test_{target} using the best {best_f1_scores[target]['model']} model\")\n",
        "        print(\"--------------------------------------------------------------------------------------\")\n",
        "\n",
        "        test_model = best_f1_scores[target]['model_instance']\n",
        "        test_model.eval()\n",
        "\n",
        "        #tokenizer = tokenizer_class.from_pretrained(pretrained_model_name)\n",
        "        #tokenizer = models[best_f1_scores[target]['model']][1].from_pretrained(pretrained_model_name)\n",
        "\n",
        "        # ensure correct tokenizer for testing\n",
        "\n",
        "        best_model_name = best_f1_scores[target]['model']  # Retrieve the name of the best model\n",
        "        best_pretrained_model_name = models[best_model_name][2]  # Retrieve the correct pretrained model name\n",
        "        tokenizer = models[best_model_name][1].from_pretrained(best_pretrained_model_name)  # Use correct tokenizer class\n",
        "\n",
        "        encoded_test = tokenizer(\n",
        "                                 X_test.tolist(),\n",
        "                                 padding = True,\n",
        "                                 truncation = True,\n",
        "                                 return_tensors = 'pt',\n",
        "                                 )\n",
        "\n",
        "        test_dataset = TensorDataset(\n",
        "                                     encoded_test['input_ids'],\n",
        "                                     encoded_test['attention_mask'],\n",
        "                                     torch.tensor(y_test),\n",
        "                                     )\n",
        "\n",
        "        test_loader = DataLoader(\n",
        "                                 test_dataset,\n",
        "                                 batch_size = batch_size,\n",
        "                                 shuffle = False,\n",
        "                                 )\n",
        "\n",
        "        all_test_predictions, all_test_true_labels = [], []\n",
        "\n",
        "        with torch.no_grad():\n",
        "            for batch in test_loader:\n",
        "                input_ids, attention_mask, labels = batch\n",
        "\n",
        "                if use_cuda:\n",
        "                    input_ids, attention_mask, labels = input_ids.cuda(), attention_mask.cuda(), labels.cuda()\n",
        "\n",
        "                outputs = test_model(input_ids, attention_mask = attention_mask)\n",
        "                logits = outputs.logits\n",
        "                test_predictions = torch.argmax(logits, dim = 1).tolist()\n",
        "                all_test_predictions.extend(test_predictions)\n",
        "                all_test_true_labels.extend(labels.tolist())\n",
        "\n",
        "        # preformance metrics for held-out d_test_{target} df\n",
        "\n",
        "        test_f1_macro = f1_score(\n",
        "                                 all_test_true_labels,\n",
        "                                 all_test_predictions,\n",
        "                                 average='macro',\n",
        "                                 )\n",
        "\n",
        "        test_mcc = matthews_corrcoef(\n",
        "                                     all_test_true_labels,\n",
        "                                     all_test_predictions,\n",
        "                                     )\n",
        "\n",
        "        test_auprc = average_precision_score(\n",
        "                                             all_test_true_labels,\n",
        "                                             all_test_predictions,\n",
        "                                             average = 'macro',\n",
        "                                             )\n",
        "\n",
        "        # display\n",
        "\n",
        "        print(f\"Test F1 (macro) for {target}: {test_f1_macro}\")\n",
        "        print(f\"Test MCC for {target}: {test_mcc}\")\n",
        "        print(f\"Test AUPRC for {target}: {test_auprc}\")\n",
        "\n",
        "        # store\n",
        "\n",
        "        results.append({\n",
        "                        'target': target,\n",
        "                        'model': best_f1_scores[target]['model'],\n",
        "                        'fold': 'Test',\n",
        "                        'f1_macro': test_f1_macro,\n",
        "                        'mcc': test_mcc,\n",
        "                        'auprc': test_auprc\n",
        "                        })\n",
        "\n",
        "    # summarize + return d_{cycle}_performance df\n",
        "\n",
        "    print(\"\\n--------------------------------------------------------------------------------------\")\n",
        "    print(f\"Summary\")\n",
        "    print(\"--------------------------------------------------------------------------------------\")\n",
        "\n",
        "    for target, info in best_f1_scores.items():\n",
        "        print(f\"Best F1 (macro) for {target}: {info['score']} achieved by {info['model']}\")\n",
        "\n",
        "    d_performance = pd.DataFrame(results)\n",
        "    print(f\"\\nd_{cycle}_performance:\")\n",
        "    print(d_performance.head(5))\n",
        "    d_performance.to_excel(f'{save_path}d_{cycle}_performance.xlsx')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84AhevrqZZ_6"
      },
      "source": [
        "**_performance_scatterplot_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "muGmQSbhZZGz"
      },
      "outputs": [],
      "source": [
        "%%writefile -a bert_train.py\n",
        "\n",
        "from brokenaxes import brokenaxes\n",
        "from matplotlib.lines import Line2D\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import seaborn as sns\n",
        "\n",
        "def performance_scatterplot(df, plot_name):\n",
        "    \"\"\"\n",
        "    Creates a categorical scatterplot with custom aesthetics, markers, error bars, and a legend.\n",
        "\n",
        "    Parameters:\n",
        "    df (pd.DataFrame): Input dataframe containing columns 'target', 'f1_macro', 'model', and 'fold'.\n",
        "\n",
        "    plot_name (str): The name used for saving the output plot file (without extension).\n",
        "\n",
        "    Returns:\n",
        "    --------\n",
        "    Matplotlib Axes object containing the barplot.\n",
        "    \"\"\"\n",
        "\n",
        "    # aesthetics\n",
        "\n",
        "    model_colors = [\n",
        "                    '#87bc45',\n",
        "                    '#b33dc6',\n",
        "                    '#27aeef',\n",
        "                   ]\n",
        "\n",
        "      ### SJS 10/1: last three colors in \"Retro Metro (Default)\" https://www.heavy.ai/blog/12-color-palettes-for-telling-better-stories-with-your-data\n",
        "\n",
        "    sns.set_style(\n",
        "                  style = 'whitegrid',\n",
        "                  rc = None,\n",
        "                  )\n",
        "\n",
        "    # map target: numeric position\n",
        "\n",
        "    target_mapping = {\n",
        "                      'asp': 0,\n",
        "                      'dep': 2,\n",
        "                      'val': 4,\n",
        "                      'prg': 6,\n",
        "                      'tgd': 8,\n",
        "                      'age': 10,\n",
        "                      'race': 12,\n",
        "                      'dbty': 14\n",
        "                     }\n",
        "\n",
        "    df['target_numeric'] = df['target'].map(target_mapping)\n",
        "    df['target_numeric'] = pd.to_numeric(df['target_numeric'])\n",
        "\n",
        "    # inject noise for jitter\n",
        "\n",
        "    df['target_jitter'] = df['target_numeric'] + np.random.uniform(\n",
        "                                                                   -0.35,\n",
        "                                                                   0.35,\n",
        "                                                                   size = len(df),\n",
        "                                                                   )\n",
        "\n",
        "    # initialize fig. with broken y-axis\n",
        "\n",
        "    fig = plt.figure(figsize=(12, 5.5))\n",
        "    bax = brokenaxes(\n",
        "                     ylims = ((0, 0.1), (0.4, 1)), ### y-axis bounds\n",
        "                     hspace = 0.1, ### y-axis break space\n",
        "                     )\n",
        "\n",
        "    # define colors: held-out test set ('fold' = Test)\n",
        "\n",
        "    test_colors = {\n",
        "                   'bert': '#87bc45',\n",
        "                   'roberta': '#27aeef',\n",
        "                   'distilbert': '#b33dc6',\n",
        "                   }\n",
        "\n",
        "    # distinguish markers: fold v. held-out test set\n",
        "\n",
        "    for fold_value, marker in [('Test', 'o'), ('non-Test', '.')]:\n",
        "        if fold_value == 'Test':\n",
        "            #data_subset = d_v[d_v['fold'] == 'Test']\n",
        "            data_subset = df[df['fold'] == 'Test']\n",
        "\n",
        "            for model in data_subset['model'].unique():\n",
        "                model_data = data_subset[data_subset['model'] == model]\n",
        "                bax.scatter(\n",
        "                            model_data['target_jitter'],\n",
        "                            model_data['f1_macro'],\n",
        "                            color = test_colors[model],\n",
        "                            s = 40,\n",
        "                            alpha = 0.6,\n",
        "                            label = None,\n",
        "                            marker = marker,\n",
        "                            )\n",
        "        else:\n",
        "            #data_subset = d_v[d_v['fold'] != 'Test']\n",
        "            data_subset = df[df['fold'] != 'Test']\n",
        "            for model, color in test_colors.items():\n",
        "                model_data = data_subset[data_subset['model'] == model]\n",
        "                bax.scatter(\n",
        "                            model_data['target_jitter'],\n",
        "                            model_data['f1_macro'],\n",
        "                            color = color,\n",
        "                            s = 40,\n",
        "                            alpha = 0.6,\n",
        "                            label = None,\n",
        "                            marker = marker,\n",
        "                            )\n",
        "\n",
        "#    for fold_value, marker in [('Test', 'o'), ('non-Test', '.')]:\n",
        "#        if fold_value == 'Test':\n",
        "#            data_subset = df[df['fold'] == 'Test']\n",
        "#        else:\n",
        "#            data_subset = df[df['fold'] != 'Test']\n",
        "\n",
        "#        sns.scatterplot(\n",
        "#                        data = data_subset,\n",
        "#                        x = 'target_jitter',\n",
        "#                        y = 'f1_macro',\n",
        "#                        hue = 'model',\n",
        "#                        palette = model_colors,\n",
        "#                        s = 40,\n",
        "#                        alpha = 0.6,\n",
        "#                        marker = marker,\n",
        "#                       )\n",
        "\n",
        "    # mean and SD of f1_macro for each target x model\n",
        "\n",
        "    mean_std_df = df.groupby(['target', 'model']).agg(\n",
        "                                                      mean_f1_macro = ('f1_macro', 'mean'),\n",
        "                                                      std_f1_macro = ('f1_macro', 'std'),\n",
        "                                                      ).reset_index()\n",
        "\n",
        "    # add target_numeric values to mean_std_df for plotting means and error bars\n",
        "\n",
        "    #mean_std_df['target_numeric'] = mean_std_df['target'].map(target_mapping).astype(float)\n",
        "\n",
        "    # x-axis offsets\n",
        "\n",
        "    mean_std_df['target_numeric'] = mean_std_df['target'].map(target_mapping).astype(float)\n",
        "    mean_std_df['target_offset'] = mean_std_df['target_numeric'] + mean_std_df['model'].map(\n",
        "                                                                                            {'bert': -0.3,\n",
        "                                                                                             'roberta': 0.0,\n",
        "                                                                                             'distilbert': 0.3}\n",
        "                                                                                            )\n",
        "\n",
        "    #model_offsets = {\n",
        "    #                 'bert-base-uncased': -0.3,\n",
        "    #                 'roberta-base': 0.0,\n",
        "    #                 'distilbert-base-uncased': 0.3,\n",
        "    #                 }\n",
        "\n",
        "    #mean_std_df['target_offset'] = mean_std_df['target_numeric'] + mean_std_df['model'].map(model_offsets)\n",
        "\n",
        "    # means (SDs), error bars\n",
        "\n",
        "    for model in mean_std_df['model'].unique():\n",
        "        model_data = mean_std_df[mean_std_df['model'] == model]\n",
        "\n",
        "    # inspect for NaNs\n",
        "\n",
        "        if not model_data[['target_offset', 'mean_f1_macro', 'std_f1_macro']].isnull().any().any():\n",
        "            plt.errorbar(\n",
        "                         model_data['target_offset'],\n",
        "                         model_data['mean_f1_macro'],\n",
        "                         yerr = model_data['std_f1_macro'],\n",
        "                         fmt = 'D',\n",
        "                         markersize = 7,\n",
        "                         capsize = 0,\n",
        "                         elinewidth = 1,\n",
        "                         markeredgewidth = 1,\n",
        "                         color = test_colors[model]\n",
        "                        )\n",
        "\n",
        "    # x-tick: map to targets\n",
        "\n",
        "    bax.set_xlabel(\n",
        "                   'Target',\n",
        "                   fontsize = 12,\n",
        "                   labelpad = 30,\n",
        "                   )\n",
        "\n",
        "    bax.set_ylabel(\n",
        "                   f'$F_1$ (macro): {plot_name}',\n",
        "                   fontsize = 12,\n",
        "                   labelpad = 30,\n",
        "                   )\n",
        "\n",
        "    # x-tick: label lower axis\n",
        "\n",
        "    bax.axs[1].set_xticks(list(target_mapping.values()))\n",
        "    bax.axs[1].set_xticklabels(list(target_mapping.keys()), rotation = 45, fontsize = 10)\n",
        "\n",
        "    #sns.despine(left = True)\n",
        "    bax.grid(\n",
        "             #axis='x',\n",
        "             False,\n",
        "             )\n",
        "\n",
        "    # line at 0.8 threshold\n",
        "\n",
        "    #bax.axhline(\n",
        "    #            y = 0.8,\n",
        "    #            color = 'r',\n",
        "    #            linewidth = 0.6,\n",
        "    #            linestyle = '--',\n",
        "    #            )\n",
        "\n",
        "    #plt.xticks(\n",
        "    #           [0, 2, 4, 6, 8, 10, 12, 14],\n",
        "    #           ['asp', 'dep', 'val', 'prg', 'tgd', 'age', 'race', 'dbty']\n",
        "    #          )\n",
        "\n",
        "    # label axes\n",
        "\n",
        "    #plt.ylim(0, 1)\n",
        "    #ax = plt.gca()\n",
        "    #ax.set_ylabel(\n",
        "    #              '$F_1$ (macro)',\n",
        "    #              fontsize = 12,\n",
        "    #              labelpad = 10,\n",
        "    #              )\n",
        "\n",
        "    #ax.set_xlabel(\n",
        "    #              'Target',\n",
        "    #              fontsize = 12,\n",
        "    #              labelpad = 10,\n",
        "    #              )\n",
        "\n",
        "    #sns.despine(left = True)\n",
        "    #ax.grid(axis = 'x')\n",
        "\n",
        "    # set line at 0.9 threshold\n",
        "\n",
        "    #ax.axhline(\n",
        "    #           y = 0.9,\n",
        "    #           color = 'r',\n",
        "    #           linewidth = 0.6,\n",
        "    #           linestyle = '--',\n",
        "    #           )\n",
        "\n",
        "    # custom legend\n",
        "\n",
        "    legend_elements = [\n",
        "                       Line2D([0], [0], marker = 'o', color = 'w', label = 'bert', markersize = 8, markerfacecolor = '#87bc45', lw = 0),\n",
        "                       Line2D([0], [0], marker = 'o', color = 'w', label = 'roberta', markersize = 8, markerfacecolor = '#27aeef', lw = 0),\n",
        "                       Line2D([0], [0], marker = 'o', color = 'w', label = 'distilbert', markersize = 8, markerfacecolor = '#b33dc6', lw = 0),\n",
        "                      ]\n",
        "\n",
        "    bax.axs[0].legend(\n",
        "                      handles = legend_elements,\n",
        "                      loc = 'upper center',\n",
        "                      bbox_to_anchor = (0.5, 1.15),\n",
        "                      ncol = 4,\n",
        "                      fontsize = 9,\n",
        "                      frameon = False,\n",
        "                      )\n",
        "\n",
        "    # save\n",
        "\n",
        "    plt.savefig(f'{plot_name}_low_res_scatter.png', dpi = 100)\n",
        "    plt.savefig(f'{plot_name}_high_res_scatter.png', dpi = 300)\n",
        "\n",
        "    # display\n",
        "\n",
        "    plt.show()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Op5JVhV0ceia"
      },
      "source": [
        "**_iterative_stratified_train_test_split_with_rationales_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qwk1WTf-aVvw"
      },
      "outputs": [],
      "source": [
        "%%writefile -a bert_train.py\n",
        "\n",
        "import pandas as pd\n",
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "def iterative_stratified_train_test_split_with_rationales(df, targets, test_size, random_state):\n",
        "    \"\"\"\n",
        "    Splits df into target-stratified train and test sets for each target in targets list:\n",
        "    d_train_{target}, d_test_{target}, respectively. Partitions 'rationales' (aug = 1) to train\n",
        "    set. Returns a dict with target names as keys\n",
        "    \"\"\"\n",
        "\n",
        "    # initialize dict\n",
        "\n",
        "    target_datasets = {}\n",
        "\n",
        "    for target in targets:\n",
        "\n",
        "        # create 'targets' col for stratification\n",
        "\n",
        "        df_target = df.copy()\n",
        "        df_target['targets'] = df[target]\n",
        "\n",
        "        # split augmented vs. non-augmented rows\n",
        "\n",
        "        aug_rows = df_target[df_target['aug'] == 1]\n",
        "        non_aug_rows = df_target[df_target['aug'] != 1]\n",
        "\n",
        "        if non_aug_rows.empty:\n",
        "            print(f\"No non-augmented rows for target {target}. Skipping...\")\n",
        "            continue\n",
        "\n",
        "        # stratified train-test split on non-augmented rows only\n",
        "\n",
        "        train_non_aug, test_non_aug = train_test_split(\n",
        "                                                       non_aug_rows,\n",
        "                                                       test_size = test_size,\n",
        "                                                       stratify = non_aug_rows['targets'],\n",
        "                                                       random_state = random_state,\n",
        "                                                      )\n",
        "\n",
        "        # concat augmented rows back into train set\n",
        "\n",
        "        d_train = pd.concat([train_non_aug, aug_rows])\n",
        "\n",
        "        # shuffle + reset index: train set\n",
        "\n",
        "        d_train = d_train.sample(\n",
        "                                 frac = 1,\n",
        "                                 random_state = random_state,\n",
        "                                 ).reset_index(drop = True)\n",
        "\n",
        "        # retain 'text', 'aug', target cols\n",
        "\n",
        "        d_train = d_train[['text', 'aug', target]]\n",
        "        d_test = test_non_aug[['text', 'aug', target]]\n",
        "\n",
        "        # reset index: test set\n",
        "\n",
        "        d_test = d_test.reset_index(drop = True)\n",
        "\n",
        "        # add train and test sets as tuples to target_datasets dict\n",
        "\n",
        "        target_datasets[target] = (d_train, d_test)\n",
        "\n",
        "        # inspect\n",
        "\n",
        "        print(f\"\\nVerify: d_train_{target} 'aug' count\")\n",
        "        print(d_train['aug'].value_counts(normalize = False))\n",
        "        print(f\"\\nVerify: d_test_{target} 'aug' count\")\n",
        "        print(d_test['aug'].value_counts(normalize = False))\n",
        "\n",
        "        print(f\"\\n--------------------------------------------------------------------------------------\")\n",
        "        print(f\"d_train_{target}: Augmented training data for target '{target}'\")\n",
        "        print(f\"--------------------------------------------------------------------------------------\")\n",
        "        print(d_train.shape)\n",
        "        print(d_train[target].value_counts(normalize = True))\n",
        "        print(d_train.head(6))\n",
        "\n",
        "        print(f\"\\n--------------------------------------------------------------------------------------\")\n",
        "        print(f\"d_test_{target}: De-augmented testing data for target '{target}'\")\n",
        "        print(f\"--------------------------------------------------------------------------------------\")\n",
        "        print(d_test.shape)\n",
        "        print(d_test[target].value_counts(normalize = True))\n",
        "        print(d_test.head(6))\n",
        "\n",
        "    return target_datasets\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZH6tQ422ctK-"
      },
      "source": [
        "**_tune_and_optimize_model_hyperparams_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xJGLu_wiPUXf"
      },
      "outputs": [],
      "source": [
        "%%writefile -a bert_train.py\n",
        "\n",
        "import torch\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "from torch.nn import CrossEntropyLoss\n",
        "from transformers import get_linear_schedule_with_warmup\n",
        "from sklearn.metrics import f1_score\n",
        "from sklearn.model_selection import ParameterGrid\n",
        "from tqdm import tqdm\n",
        "import pandas as pd\n",
        "\n",
        "def tune_and_optimize_model_hyperparams(tokenizer, model_class, pretrained_model_name, d_train, d_test, target, class_weights, save_path, hyperparameter_grid):\n",
        "    \"\"\"\n",
        "    Tune and optimize model hyperparameters for a specific model-target combination.\n",
        "\n",
        "    Parameters:\n",
        "    -----------\n",
        "\n",
        "    tokenizer:\n",
        "        Pre-trained tokenizer.\n",
        "\n",
        "    model_class:\n",
        "        Pre-trained model class.\n",
        "\n",
        "    pretrained_model_name:\n",
        "        Name of the pre-trained model.\n",
        "\n",
        "    d_train : pd.DataFrame\n",
        "        Training dataset.\n",
        "\n",
        "    d_test : pd.DataFrame\n",
        "        Test dataset.\n",
        "\n",
        "    target : str\n",
        "        The target variable for classification.\n",
        "\n",
        "    class_weights : torch.tensor\n",
        "        Weights for each class.\n",
        "\n",
        "    save_path : str\n",
        "        Path to save the best model.\n",
        "\n",
        "    hyperparameter_grid : dict\n",
        "        Dictionary where keys are hyperparameter names and values are lists of possible values.\n",
        "\n",
        "    Returns:\n",
        "    --------\n",
        "    d_test : pd.DataFrame\n",
        "        Test dataset with predictions and probabilities.\n",
        "\n",
        "    d_tuned_performance : pd.DataFrame\n",
        "        DataFrame with the performance metrics for each hyperparameter configuration.\n",
        "    \"\"\"\n",
        "\n",
        "    # check CUDA\n",
        "\n",
        "    use_cuda = torch.cuda.is_available()\n",
        "    print(\"CUDA: \", use_cuda)\n",
        "\n",
        "    # set seed\n",
        "\n",
        "    set_seed(56)\n",
        "\n",
        "    print(\"======================================================================================\")\n",
        "    print(f\"Optimizing: {pretrained_model_name}\\nTarget: {target}\")\n",
        "    print(\"======================================================================================\")\n",
        "\n",
        "    # tokenize train and test sets\n",
        "\n",
        "    encoded_train = tokenizer(\n",
        "                              d_train['text'].tolist(),\n",
        "                              padding = True,\n",
        "                              truncation = True,\n",
        "                              return_tensors = 'pt',\n",
        "                              )\n",
        "\n",
        "    encoded_test = tokenizer(\n",
        "                             d_test['text'].tolist(),\n",
        "                             padding = True,\n",
        "                             truncation = True,\n",
        "                             return_tensors = 'pt',\n",
        "                             )\n",
        "\n",
        "\n",
        "    # accept dynamic target variables\n",
        "\n",
        "    train_labels = torch.tensor(d_train[target].values)\n",
        "    test_labels = torch.tensor(d_test[target].values)\n",
        "\n",
        "    # prepare datasets\n",
        "\n",
        "    train_dataset = TensorDataset(\n",
        "                                  encoded_train['input_ids'],\n",
        "                                  encoded_train['attention_mask'],\n",
        "                                  train_labels,\n",
        "                                  )\n",
        "\n",
        "    test_dataset = TensorDataset(\n",
        "                                 encoded_test['input_ids'],\n",
        "                                 encoded_test['attention_mask'],\n",
        "                                 test_labels,\n",
        "                                 )\n",
        "\n",
        "    #train_loader = DataLoader(\n",
        "    #                          train_dataset,\n",
        "    #                          batch_size = 8,  ### to be updated within grid search\n",
        "    #                          shuffle = True,\n",
        "    #                          )\n",
        "\n",
        "    #test_loader = DataLoader(\n",
        "    #                         test_dataset,\n",
        "    #                         batch_size = 8,  ### to be updated within grid search\n",
        "    #                         shuffle = False,\n",
        "    #                         )\n",
        "\n",
        "    # initialize class weights\n",
        "\n",
        "    if use_cuda:\n",
        "        class_weights = class_weights.cuda()\n",
        "\n",
        "    # initialize tracking variables\n",
        "\n",
        "    best_f1_macro = 0\n",
        "    best_params = None\n",
        "    best_model_state = None\n",
        "    best_predictions = []\n",
        "    best_probabilities = []\n",
        "\n",
        "    f1_scores = []\n",
        "    performance_data = []\n",
        "\n",
        "    # hyperparam grid search: ParameterGrid\n",
        "\n",
        "    for hyperparams in ParameterGrid(hyperparameter_grid):\n",
        "        print(f\"\\nOptimizing with hyperparameters: {hyperparams}\")\n",
        "\n",
        "        train_loader = DataLoader(\n",
        "                                  train_dataset,\n",
        "                                  batch_size = hyperparams['batch_size'],\n",
        "                                  shuffle = True\n",
        "                                  )\n",
        "        test_loader = DataLoader(\n",
        "                                 test_dataset,\n",
        "                                 batch_size = hyperparams['batch_size'],\n",
        "                                 shuffle = False\n",
        "                                 )\n",
        "\n",
        "        print(f\"\\nTotal training rows: {len(train_dataset)}\")\n",
        "        print(f\"Total evaluation rows: {len(test_dataset)}\")\n",
        "        print(f\"Training batch size: {hyperparams['batch_size']}\")\n",
        "        print(f\"Evaluation batch size: {hyperparams['batch_size']}\")\n",
        "        print(f\"Total training batches: {len(train_loader)}\")\n",
        "        print(f\"Total evaluation batches: {len(test_loader)}\")\n",
        "        print(\"\\n\")\n",
        "\n",
        "        # initialize model\n",
        "\n",
        "        model = model_class.from_pretrained(pretrained_model_name)\n",
        "        if use_cuda:\n",
        "            model.cuda()\n",
        "\n",
        "        # initialize optimizer and lr scheduler\n",
        "\n",
        "        optimizer = torch.optim.AdamW(\n",
        "                                      model.parameters(),\n",
        "                                      lr = hyperparams['learning_rate'],\n",
        "                                      weight_decay = hyperparams['weight_decay']\n",
        "                                      )\n",
        "\n",
        "        # calculate total steps\n",
        "\n",
        "        total_steps = len(train_loader) * hyperparams['num_epochs']\n",
        "\n",
        "        # add scheduler with warmup steps\n",
        "\n",
        "        scheduler = get_linear_schedule_with_warmup(\n",
        "                                                    optimizer,\n",
        "                                                    num_warmup_steps = hyperparams['warmup_steps'],\n",
        "                                                    num_training_steps=total_steps\n",
        "                                                    )\n",
        "\n",
        "        criterion = CrossEntropyLoss(weight = class_weights)\n",
        "\n",
        "        # training loop\n",
        "\n",
        "        for epoch in range(hyperparams['num_epochs']):\n",
        "            model.train()\n",
        "            optimizer.zero_grad()\n",
        "            for i, batch in enumerate(tqdm(train_loader, desc = f\"Training Epoch {epoch + 1}/{hyperparams['num_epochs']}\", leave=True)):\n",
        "                input_ids, attention_mask, labels = batch\n",
        "                if use_cuda:\n",
        "                    input_ids, attention_mask, labels = input_ids.cuda(), attention_mask.cuda(), labels.cuda()\n",
        "                outputs = model(input_ids, attention_mask=attention_mask)\n",
        "                loss = criterion(outputs.logits, labels)\n",
        "                loss = loss / hyperparams['gradient_accumulation_steps']\n",
        "                loss.backward()\n",
        "                if (i + 1) % hyperparams['gradient_accumulation_steps'] == 0:\n",
        "                    optimizer.step()\n",
        "                    scheduler.step()  ### update learning rate\n",
        "                    optimizer.zero_grad()\n",
        "\n",
        "        # eval loop\n",
        "\n",
        "        model.eval()\n",
        "        all_predictions = []\n",
        "        all_true_labels = []\n",
        "        all_probabilities = []\n",
        "        with torch.no_grad():\n",
        "            for batch in test_loader:\n",
        "                input_ids, attention_mask, labels = batch\n",
        "                if use_cuda:\n",
        "                    input_ids, attention_mask, labels = input_ids.cuda(), attention_mask.cuda(), labels.cuda()\n",
        "                outputs = model(input_ids, attention_mask=attention_mask)\n",
        "                probabilities = torch.softmax(outputs.logits, dim = 1)\n",
        "                predictions = torch.argmax(probabilities, dim = 1).cpu().tolist()\n",
        "                all_predictions.extend(predictions)\n",
        "                all_true_labels.extend(labels.cpu().tolist())\n",
        "                all_probabilities.extend(probabilities.cpu().tolist())\n",
        "\n",
        "        # calculate F1 (macro)\n",
        "\n",
        "        current_f1_macro = f1_score(all_true_labels, all_predictions, average='macro')\n",
        "        f1_scores.append(current_f1_macro)\n",
        "        print(f\"\\nCurrent F1 macro with params {hyperparams}: {current_f1_macro}\")\n",
        "\n",
        "        # append F1 and current performance data\n",
        "\n",
        "        performance_data.append({\n",
        "                                 'pretrained_model_name': pretrained_model_name,\n",
        "                                 'target': target,\n",
        "                                 'f1_score': current_f1_macro,\n",
        "                                 'batch_size': hyperparams['batch_size'],\n",
        "                                 'weight_decay': hyperparams['weight_decay'],\n",
        "                                 'learning_rate': hyperparams['learning_rate'],\n",
        "                                 'warmup_steps': hyperparams['warmup_steps'],\n",
        "                                 'num_epochs': hyperparams['num_epochs'],\n",
        "                                 'gradient_accumulation_steps': hyperparams['gradient_accumulation_steps'],\n",
        "        })\n",
        "\n",
        "        if current_f1_macro > best_f1_macro:\n",
        "            best_f1_macro = current_f1_macro\n",
        "            best_params = hyperparams\n",
        "            best_model_state = model.state_dict()\n",
        "            best_predictions = all_predictions\n",
        "            best_probabilities = all_probabilities\n",
        "\n",
        "    #if len(best_predictions) == len(d_test):\n",
        "    #    d_test['predicted_labels'] = best_predictions\n",
        "    #    d_test['predicted_probabilities'] = best_probabilities\n",
        "    #else:\n",
        "    #    print(\"Error: Length of predictions does not match length of test set\")\n",
        "\n",
        "    d_test['predicted_labels'] = best_predictions\n",
        "    d_test['predicted_probabilities'] = best_probabilities\n",
        "\n",
        "    # save d_test_{target} with pred and prob\n",
        "\n",
        "    print(\"--------------------------------------------------------------------------------------\")\n",
        "    print(f\"Summary: {target}\")\n",
        "    print(\"--------------------------------------------------------------------------------------\")\n",
        "\n",
        "    print(d_test.head(6))\n",
        "    d_test.to_excel(f'{save_path}/d_test_tuned_preds_{target}.xlsx')\n",
        "\n",
        "    if best_model_state:\n",
        "        model_path = f\"{save_path}/{target}_{pretrained_model_name}_best_tuned_model.bin\"\n",
        "        torch.save(best_model_state, model_path)\n",
        "        print(\"\\nBest model saved with F1 macro:\", best_f1_macro)\n",
        "        print(\"Best hyperparameters:\", best_params)\n",
        "\n",
        "    # display F1 scores\n",
        "\n",
        "    f1_mean = sum(f1_scores) / len(f1_scores)\n",
        "    f1_std = (sum((x - f1_mean) ** 2 for x in f1_scores) / len(f1_scores)) ** 0.5\n",
        "    print(f\"Mean F1 macro: {f1_mean}\")\n",
        "    print(f\"Standard deviation of F1 macro: {f1_std}\")\n",
        "\n",
        "    # df: target-wise\n",
        "\n",
        "    d_tuned_performance = pd.DataFrame(performance_data)\n",
        "    print(d_tuned_performance.head(10))\n",
        "\n",
        "    # save: target-wise df\n",
        "\n",
        "    d_tuned_performance.to_excel(f'{save_path}/d_tuned_performance_{target}.xlsx')\n",
        "\n",
        "    return d_test, d_tuned_performance\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EghBliavly-P"
      },
      "source": [
        "**_tune_and_optimize_model_loss_accuracy_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a1MjwRn4ltsP"
      },
      "outputs": [],
      "source": [
        "%%writefile -a bert_train.py\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import os\n",
        "import torch\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "from torch.nn import CrossEntropyLoss\n",
        "from transformers import get_linear_schedule_with_warmup\n",
        "from sklearn.metrics import f1_score\n",
        "from sklearn.model_selection import ParameterGrid\n",
        "from tqdm import tqdm\n",
        "import pandas as pd\n",
        "\n",
        "def tune_and_optimize_model_loss_accuracy(tokenizer, model_class, pretrained_model_name, d_train, d_test, target, class_weights, save_path, fixed_hyperparameters, num_epochs_range, cycle):\n",
        "    \"\"\"\n",
        "    tktk\n",
        "\n",
        "    Parameters:\n",
        "    -----------\n",
        "\n",
        "    tokenizer:\n",
        "        Pre-trained tokenizer.\n",
        "\n",
        "    model_class:\n",
        "        Pre-trained model class.\n",
        "\n",
        "    pretrained_model_name:\n",
        "        Name of the pre-trained model.\n",
        "\n",
        "    d_train : pd.DataFrame\n",
        "        Training set.\n",
        "\n",
        "    d_val : pd.DataFrame\n",
        "        Validation set.\n",
        "\n",
        "    target : str\n",
        "        The target variable for classification.\n",
        "\n",
        "    class_weights : torch.tensor\n",
        "        Weights for each class.\n",
        "\n",
        "    save_path : str\n",
        "        Path to save the best model.\n",
        "\n",
        "    tktk\n",
        "\n",
        "    Returns:\n",
        "    --------\n",
        "    d_val : pd.DataFrame\n",
        "        Validation set with predictions and probabilities.\n",
        "\n",
        "    d_epochal_performance : pd.DataFrame\n",
        "        DataFrame with the performance metrics for each hyperparameter configuration.\n",
        "    \"\"\"\n",
        "\n",
        "    # check CUDA\n",
        "\n",
        "    use_cuda = torch.cuda.is_available()\n",
        "    print(\"CUDA: \", use_cuda)\n",
        "\n",
        "    # display cycle\n",
        "\n",
        "    print(f\"CYCLE: {cycle}\")\n",
        "\n",
        "    # set seed\n",
        "\n",
        "    set_seed(56)\n",
        "\n",
        "    print(\"======================================================================================\")\n",
        "    print(f\"Optimizing: {pretrained_model_name}\\nTarget: {target}\")\n",
        "    print(\"======================================================================================\")\n",
        "\n",
        "    # tokenize train and test sets\n",
        "\n",
        "    encoded_train = tokenizer(\n",
        "                              d_train['text'].tolist(),\n",
        "                              padding = True,\n",
        "                              truncation = True,\n",
        "                              return_tensors = 'pt',\n",
        "                              )\n",
        "\n",
        "    encoded_test = tokenizer(\n",
        "                             d_test['text'].tolist(),\n",
        "                             padding = True,\n",
        "                             truncation = True,\n",
        "                             return_tensors = 'pt',\n",
        "                             )\n",
        "\n",
        "\n",
        "    # accept dynamic target variables\n",
        "\n",
        "    train_labels = torch.tensor(d_train[target].values)\n",
        "    test_labels = torch.tensor(d_test[target].values)\n",
        "\n",
        "    # prep datasets\n",
        "\n",
        "    train_dataset = TensorDataset(\n",
        "                                  encoded_train['input_ids'],\n",
        "                                  encoded_train['attention_mask'],\n",
        "                                  train_labels,\n",
        "                                  )\n",
        "\n",
        "    test_dataset = TensorDataset(\n",
        "                                 encoded_test['input_ids'],\n",
        "                                 encoded_test['attention_mask'],\n",
        "                                 test_labels,\n",
        "                                 )\n",
        "\n",
        "    # initialize class weights\n",
        "\n",
        "    if use_cuda:\n",
        "        class_weights = class_weights.cuda()\n",
        "\n",
        "    # fix target-specific hyperparams\n",
        "\n",
        "    num_epochs = max(num_epochs_range)\n",
        "    print(f\"\\nTotal epochs to run: {num_epochs}\")\n",
        "\n",
        "    # merge fixed hyperparams, num_epochs\n",
        "\n",
        "    hyperparams = fixed_hyperparameters.copy()\n",
        "    hyperparams['num_epochs'] = num_epochs\n",
        "\n",
        "    train_loader = DataLoader(\n",
        "                              train_dataset,\n",
        "                              batch_size = hyperparams['batch_size'],\n",
        "                              shuffle = True,\n",
        "                              )\n",
        "\n",
        "    test_loader = DataLoader(\n",
        "                             test_dataset,\n",
        "                             batch_size = hyperparams['batch_size'],\n",
        "                             shuffle = False,\n",
        "                             )\n",
        "\n",
        "    print(f\"\\nTotal training rows: {len(train_dataset)}\")\n",
        "    print(f\"Total evaluation rows: {len(test_dataset)}\")\n",
        "    print(f\"Training batch size: {hyperparams['batch_size']}\")\n",
        "    print(f\"Evaluation batch size: {hyperparams['batch_size']}\")\n",
        "    print(f\"Total training batches: {len(train_loader)}\")\n",
        "    print(f\"Total evaluation batches: {len(test_loader)}\")\n",
        "    print(\"\\n\")\n",
        "\n",
        "    # initialize model\n",
        "\n",
        "    model = model_class.from_pretrained(\n",
        "                                        pretrained_model_name,\n",
        "                                        attention_probs_dropout_prob = fixed_hyperparameters.get('dropout', 0.1), ### default = 0.1\n",
        "                                        hidden_dropout_prob = fixed_hyperparameters.get('dropout', 0.1), ### default = 0.1\n",
        "                                        )\n",
        "\n",
        "    if use_cuda:\n",
        "        model.cuda()\n",
        "\n",
        "    # initialize optimizer and lr scheduler\n",
        "\n",
        "    optimizer = torch.optim.AdamW(\n",
        "                                  model.parameters(),\n",
        "                                  lr = hyperparams['learning_rate'],\n",
        "                                  weight_decay = hyperparams['weight_decay'],\n",
        "                                  )\n",
        "\n",
        "    # calculate total steps\n",
        "\n",
        "    total_steps = len(train_loader) * num_epochs\n",
        "\n",
        "    # add scheduler with warmup steps\n",
        "\n",
        "    scheduler = get_linear_schedule_with_warmup(\n",
        "                                                optimizer,\n",
        "                                                num_warmup_steps = hyperparams['warmup_steps'],\n",
        "                                                num_training_steps = total_steps,\n",
        "                                                )\n",
        "\n",
        "    criterion = CrossEntropyLoss(weight = class_weights)\n",
        "\n",
        "    # set checkpoints subdirectory\n",
        "\n",
        "    optimized_model_dir = os.path.join(save_path, \"RoBERTa_checkpoints\")\n",
        "    os.makedirs(optimized_model_dir, exist_ok = True)\n",
        "\n",
        "    # initialize tracking variables\n",
        "\n",
        "    f1_scores = []\n",
        "    performance_data = []\n",
        "\n",
        "    # training loop\n",
        "\n",
        "    for epoch in range(hyperparams['num_epochs']):\n",
        "        model.train()\n",
        "        total_train_loss = 0\n",
        "        correct_train_preds = 0\n",
        "        total_train_samples = 0\n",
        "\n",
        "        for i, batch in enumerate(tqdm(train_loader, desc = f\"Training Epoch {epoch + 1}/{hyperparams['num_epochs']}\", leave = True)):\n",
        "            input_ids, attention_mask, labels = batch\n",
        "            if use_cuda:\n",
        "                input_ids, attention_mask, labels = input_ids.cuda(), attention_mask.cuda(), labels.cuda()\n",
        "            outputs = model(input_ids, attention_mask = attention_mask)\n",
        "            loss = criterion(outputs.logits, labels)\n",
        "            total_train_loss += loss.item()\n",
        "\n",
        "            predictions = torch.argmax(outputs.logits, dim = 1)\n",
        "            correct_train_preds += (predictions == labels).sum().item()\n",
        "            total_train_samples += labels.size(0)\n",
        "\n",
        "            loss = loss / hyperparams['gradient_accumulation_steps']\n",
        "            loss.backward()\n",
        "            if (i + 1) % hyperparams['gradient_accumulation_steps'] == 0:\n",
        "                optimizer.step()\n",
        "                scheduler.step()\n",
        "                optimizer.zero_grad()\n",
        "\n",
        "        # calculate avg training loss, accuracy\n",
        "\n",
        "        avg_train_loss = total_train_loss / len(train_loader)\n",
        "        training_accuracy = correct_train_preds / total_train_samples\n",
        "\n",
        "        # eval loop after each epoch\n",
        "\n",
        "        model.eval()\n",
        "        total_eval_loss = 0\n",
        "        correct_eval_preds = 0\n",
        "        total_eval_samples = 0\n",
        "        all_predictions = []\n",
        "        all_true_labels = []\n",
        "        all_probabilities = []\n",
        "\n",
        "        with torch.no_grad():\n",
        "            for batch in test_loader:\n",
        "                input_ids, attention_mask, labels = batch\n",
        "                if use_cuda:\n",
        "                    input_ids, attention_mask, labels = input_ids.cuda(), attention_mask.cuda(), labels.cuda()\n",
        "                outputs = model(input_ids, attention_mask=attention_mask)\n",
        "                loss = criterion(outputs.logits, labels)\n",
        "                total_eval_loss += loss.item()\n",
        "\n",
        "                probabilities = torch.softmax(outputs.logits, dim = 1)\n",
        "                predictions = torch.argmax(probabilities, dim = 1)\n",
        "                correct_eval_preds += (predictions == labels).sum().item()\n",
        "                total_eval_samples += labels.size(0)\n",
        "\n",
        "                all_predictions.extend(predictions.cpu().tolist())\n",
        "                all_true_labels.extend(labels.cpu().tolist())\n",
        "                all_probabilities.extend(probabilities.cpu().tolist())\n",
        "\n",
        "        # calculate accuracy, avg eval loss\n",
        "\n",
        "        avg_eval_loss = total_eval_loss / len(test_loader)\n",
        "        evaluation_accuracy = correct_eval_preds / total_eval_samples\n",
        "\n",
        "        # calculate f1\n",
        "\n",
        "        current_f1_macro = f1_score(\n",
        "                                    all_true_labels,\n",
        "                                    all_predictions,\n",
        "                                    average = 'macro',\n",
        "                                    )\n",
        "        f1_scores.append(current_f1_macro)\n",
        "\n",
        "        # display\n",
        "\n",
        "        print(f\"\\nEpoch {epoch + 1}/{hyperparams['num_epochs']} Results:\")\n",
        "        print(f\"  Training Loss: {avg_train_loss:.4f}\")\n",
        "        print(f\"  Training Accuracy: {training_accuracy:.4f}\")\n",
        "        print(f\"  Evaluation Loss: {avg_eval_loss:.4f}\")\n",
        "        print(f\"  Evaluation Accuracy: {evaluation_accuracy:.4f}\")\n",
        "        print(f\"  F1 Macro: {current_f1_macro:.4f}\")\n",
        "\n",
        "        # append performance data\n",
        "\n",
        "        performance_data.append({\n",
        "                                'pretrained_model_name': pretrained_model_name,\n",
        "                                'target': target,\n",
        "                                'epoch': epoch + 1,\n",
        "                                'training_loss': avg_train_loss,\n",
        "                                'training_accuracy': training_accuracy,\n",
        "                                'evaluation_loss': avg_eval_loss,\n",
        "                                'evaluation_accuracy': evaluation_accuracy,\n",
        "                                'f1_score': current_f1_macro,\n",
        "                                'batch_size': hyperparams['batch_size'],\n",
        "                                'weight_decay': hyperparams['weight_decay'],\n",
        "                                'learning_rate': hyperparams['learning_rate'],\n",
        "                                'warmup_steps': hyperparams['warmup_steps'],\n",
        "                                'num_epochs': hyperparams['num_epochs'],\n",
        "                                'gradient_accumulation_steps': hyperparams['gradient_accumulation_steps'],\n",
        "                                })\n",
        "\n",
        "        # save model at end of epoch\n",
        "\n",
        "        model_save_path = os.path.join(optimized_model_dir, f\"{target}_roberta_optimized_epoch_{epoch + 1}_{cycle}.bin\")\n",
        "        torch.save(model.state_dict(), model_save_path)\n",
        "        print(f\"Model checkpoint saved: {model_save_path}\")\n",
        "\n",
        "    # save d_test_{target} with pred and prob\n",
        "\n",
        "    print(\"--------------------------------------------------------------------------------------\")\n",
        "    print(f\"Summary: {target}\")\n",
        "    print(\"--------------------------------------------------------------------------------------\")\n",
        "\n",
        "    # df: target-wise\n",
        "\n",
        "    d_epochal_performance = pd.DataFrame(performance_data)\n",
        "\n",
        "    # save: target-wise df\n",
        "\n",
        "    d_epochal_performance.to_excel(f'{save_path}/d_epochal_performance_{target}_{cycle}.xlsx')\n",
        "\n",
        "    # plot epochal training and evaluation loss\n",
        "\n",
        "    plt.figure(figsize = (10, 6))\n",
        "    plt.plot(\n",
        "             d_epochal_performance['epoch'],\n",
        "             d_epochal_performance['training_loss'],\n",
        "             label = 'training_loss',\n",
        "             #marker = 'o',\n",
        "             )\n",
        "    plt.plot(\n",
        "             d_epochal_performance['epoch'],\n",
        "             d_epochal_performance['evaluation_loss'],\n",
        "             label = 'evaluation_loss',\n",
        "             #marker = 'o',\n",
        "             )\n",
        "    plt.xlabel('epoch')\n",
        "    plt.ylabel('loss')\n",
        "    plt.ylim(0, 1)\n",
        "    plt.title(f'{target}: loss')\n",
        "    plt.legend()\n",
        "    plt.grid(False)\n",
        "    loss_plot_path = os.path.join(save_path, f'{target}_loss_plot_{cycle}.png')\n",
        "    plt.savefig(loss_plot_path)\n",
        "    plt.show()\n",
        "    plt.close()\n",
        "    print(f\"Loss plot saved: {loss_plot_path}\")\n",
        "\n",
        "    # plot epochal training and evaluation accuracy\n",
        "\n",
        "    plt.figure(figsize = (10, 6))\n",
        "    plt.plot(\n",
        "             d_epochal_performance['epoch'],\n",
        "             d_epochal_performance['training_accuracy'],\n",
        "             label = 'training_accuracy',\n",
        "             #marker = 'o',\n",
        "             )\n",
        "    plt.plot(\n",
        "             d_epochal_performance['epoch'],\n",
        "             d_epochal_performance['evaluation_accuracy'],\n",
        "             label = 'evaluation_accuracy',\n",
        "             #marker = 'o',\n",
        "             )\n",
        "    plt.xlabel('epoch')\n",
        "    plt.ylabel('accuracy')\n",
        "    plt.ylim(0, 1)\n",
        "    plt.title(f'{target}: accuracy')\n",
        "    plt.legend()\n",
        "    plt.grid(False)\n",
        "    accuracy_plot_path = os.path.join(save_path, f'{target}_accuracy_plot_{cycle}.png')\n",
        "    plt.savefig(accuracy_plot_path)\n",
        "    plt.show()\n",
        "    plt.close()\n",
        "    print(f\"Accuracy plot saved: {accuracy_plot_path}\")\n",
        "\n",
        "    return d_epochal_performance\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DYuotbl5dsV8"
      },
      "source": [
        "#### bert_predict.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "um7zclzxdQtt"
      },
      "source": [
        "**_load_model_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D0X75L47cADX"
      },
      "outputs": [],
      "source": [
        "%%writefile bert_predict.py\n",
        "\n",
        "import torch\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "from transformers import (\n",
        "                          DistilBertTokenizer,\n",
        "                          DistilBertForSequenceClassification,\n",
        "                          BertTokenizer,\n",
        "                          BertForSequenceClassification,\n",
        "                          RobertaTokenizer,\n",
        "                          RobertaForSequenceClassification,\n",
        "                          )\n",
        "from tqdm import tqdm\n",
        "import pandas as pd\n",
        "\n",
        "def load_model(model_path, model_class, pretrained_model_name):\n",
        "    \"\"\"\n",
        "    Loads a pre-trained fine-tined LM from a specified path.\n",
        "    \"\"\"\n",
        "    model = model_class.from_pretrained(pretrained_model_name)\n",
        "    model.load_state_dict(torch.load(model_path))\n",
        "    model.eval()\n",
        "    return model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oPOjsBILdXMY"
      },
      "source": [
        "**_preprocess_data_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NfE5xRrVcAHl"
      },
      "outputs": [],
      "source": [
        "%%writefile -a bert_predict.py\n",
        "\n",
        "import torch\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "from transformers import (\n",
        "                          DistilBertTokenizer,\n",
        "                          DistilBertForSequenceClassification,\n",
        "                          BertTokenizer,\n",
        "                          BertForSequenceClassification,\n",
        "                          RobertaTokenizer,\n",
        "                          RobertaForSequenceClassification,\n",
        "                          )\n",
        "from tqdm import tqdm\n",
        "import pandas as pd\n",
        "\n",
        "def preprocess_data(tokenizer, texts):\n",
        "    \"\"\"\n",
        "    Tokenizes a list of texts using the specified LM-specific tokenizer.\n",
        "    \"\"\"\n",
        "    encoded_texts = tokenizer(\n",
        "        texts,\n",
        "        padding=True,\n",
        "        truncation=True,\n",
        "        return_tensors='pt'\n",
        "    )\n",
        "    return encoded_texts"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mcu7TSlBdexo"
      },
      "source": [
        "**_predict_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8wKygXEFcALG"
      },
      "outputs": [],
      "source": [
        "%%writefile -a bert_predict.py\n",
        "\n",
        "import torch\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "from transformers import (\n",
        "                          DistilBertTokenizer,\n",
        "                          DistilBertForSequenceClassification,\n",
        "                          BertTokenizer,\n",
        "                          BertForSequenceClassification,\n",
        "                          RobertaTokenizer,\n",
        "                          RobertaForSequenceClassification,\n",
        "                          )\n",
        "from tqdm import tqdm\n",
        "import pandas as pd\n",
        "\n",
        "def predict(model, tokenizer, texts, batch_size = 8, use_cuda = True):\n",
        "    \"\"\"\n",
        "    Predicts labels and probabilities for a list of texts using the specified model and tokenizer.\n",
        "    \"\"\"\n",
        "    print(f\"Total number of texts to predict: {len(texts)}\")\n",
        "    encoded_texts = preprocess_data(tokenizer, texts)\n",
        "    dataset = TensorDataset(encoded_texts['input_ids'], encoded_texts['attention_mask'])\n",
        "    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
        "\n",
        "    print(f\"Batch size: {batch_size}\")\n",
        "    print(f\"Total number of batches: {len(data_loader)}\")\n",
        "\n",
        "    if use_cuda:\n",
        "        model.cuda()\n",
        "\n",
        "    all_predictions = []\n",
        "    all_probabilities = []\n",
        "\n",
        "    with torch.no_grad():\n",
        "        progress_bar = tqdm(total = len(data_loader), desc = \"Predicting\", leave = False)\n",
        "        for batch in data_loader:\n",
        "            input_ids, attention_mask = batch\n",
        "            if use_cuda:\n",
        "                input_ids, attention_mask = input_ids.cuda(), attention_mask.cuda()\n",
        "\n",
        "            outputs = model(input_ids, attention_mask = attention_mask)\n",
        "            probabilities = torch.softmax(outputs.logits, dim = 1)\n",
        "            predictions = torch.argmax(probabilities, dim = 1).cpu().tolist()\n",
        "            all_predictions.extend(predictions)\n",
        "            all_probabilities.extend(probabilities.cpu().tolist())\n",
        "            progress_bar.update(1)\n",
        "        progress_bar.close()\n",
        "\n",
        "    return all_predictions, all_probabilities"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B_K-SDa8f-4I"
      },
      "source": [
        "#### llama_train.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qcrB7C91kkJO"
      },
      "source": [
        "**_load_llama_and_tokenizer_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JgBKojgcf9y8"
      },
      "outputs": [],
      "source": [
        "%%writefile llama_train.py\n",
        "\n",
        "from peft import(\n",
        "                 get_peft_model,\n",
        "                 LoraConfig,\n",
        "                 prepare_model_for_kbit_training,\n",
        "                 )\n",
        "\n",
        "import torch\n",
        "from transformers import(\n",
        "                         AutoTokenizer,\n",
        "                         AutoModelForSequenceClassification,\n",
        "                         BitsAndBytesConfig,\n",
        "                         )\n",
        "\n",
        "def load_llama_and_tokenizer(model_name, num_labels):\n",
        "    \"\"\"\n",
        "    Loads the Llama model and tokenizer with 4-bit quantization and LoRA (Low-Rank Adaptation) applied.\n",
        "\n",
        "    Args:\n",
        "    model_name (str): name of pretrained Llama model.\n",
        "    num_labels (int): number of labels for the classification task (binary or multiclass).\n",
        "\n",
        "    Returns:\n",
        "    model (AutoModelForSequenceClassification): Llama model configured for sequence classification.\n",
        "    tokenizer (AutoTokenizer): tokenizer associated with the Llama model.\n",
        "    \"\"\"\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_name, add_prefix_space = True)\n",
        "    tokenizer.pad_token_id = tokenizer.eos_token_id\n",
        "    tokenizer.pad_token = tokenizer.eos_token\n",
        "\n",
        "    quantization_config = BitsAndBytesConfig(\n",
        "                                             load_in_4bit = True,\n",
        "                                             bnb_4bit_quant_type = 'nf4',\n",
        "                                             bnb_4bit_use_double_quant = True,\n",
        "                                             bnb_4bit_compute_dtype = torch.bfloat16,\n",
        "                                             )\n",
        "\n",
        "    model_name = model_name\n",
        "\n",
        "    model = AutoModelForSequenceClassification.from_pretrained(\n",
        "                                                               model_name,\n",
        "                                                               quantization_config = quantization_config,\n",
        "                                                               num_labels = num_labels,\n",
        "                                                               device_map = 'auto',\n",
        "                                                               )\n",
        "\n",
        "\n",
        "    # apply LoRA\n",
        "\n",
        "    lora_config = LoraConfig(\n",
        "                             r = 16,\n",
        "                             lora_alpha = 8,\n",
        "                             target_modules = [\n",
        "                                               'q_proj',\n",
        "                                               'k_proj',\n",
        "                                               'v_proj',\n",
        "                                               'o_proj',\n",
        "                                               ],\n",
        "                             lora_dropout = 0.05,\n",
        "                             bias = 'none',\n",
        "                             task_type = 'SEQ_CLS',\n",
        "                             )\n",
        "\n",
        "    model = prepare_model_for_kbit_training(model)\n",
        "    model = get_peft_model(model, lora_config)\n",
        "\n",
        "    model.config.pad_token_id = tokenizer.pad_token_id\n",
        "    model.config.use_cache = False\n",
        "    model.config.pretraining_tp = 1\n",
        "\n",
        "    return model, tokenizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0plgyiehky2i"
      },
      "source": [
        "**_llama_tokenize_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PjZJvzVFf9rO"
      },
      "outputs": [],
      "source": [
        "%%writefile -a llama_train.py\n",
        "\n",
        "from transformers import PreTrainedTokenizer\n",
        "\n",
        "def llama_tokenize(examples, tokenizer):\n",
        "    \"\"\"\n",
        "    Tokenize the input examples using the provided tokenizer.\n",
        "\n",
        "    Args:\n",
        "    examples (dict): dictionary containing the text to be tokenized. Assumes the key 'text' contains the input text.\n",
        "    tokenizer (PreTrainedTokenizer): tokenizer to be used for tokenizing the text.\n",
        "\n",
        "    Returns:\n",
        "    dict: a dictionary with tokenized input including input_ids, attention_mask, etc.\n",
        "    \"\"\"\n",
        "\n",
        "    # tokenize 'text' col\n",
        "\n",
        "    return tokenizer(\n",
        "                     examples['text'],\n",
        "                     padding = 'max_length',\n",
        "                     truncation = True,\n",
        "                     max_length = 512,\n",
        "                     )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0jJnav3Nk49J"
      },
      "source": [
        "**_compute_llama_metrics_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UVPTm8LLhrt8"
      },
      "outputs": [],
      "source": [
        "%%writefile -a llama_train.py\n",
        "\n",
        "from sklearn.metrics import average_precision_score\n",
        "#from datasets import load_metric\n",
        "import evaluate\n",
        "\n",
        "# load metrics\n",
        "\n",
        "f1_metric = evaluate.load('f1')\n",
        "mcc_metric = evaluate.load('matthews_correlation')\n",
        "\n",
        "#f1_metric = load_metric('f1')\n",
        "#mcc_metric = load_metric('matthews_correlation')\n",
        "\n",
        "def compute_llama_metrics(eval_pred):\n",
        "    \"\"\"\n",
        "    Compute evaluation metrics for the Llama model during evaluation.\n",
        "\n",
        "    Args:\n",
        "    eval_pred (tuple): a tuple containing predictions and labels. The predictions are logits, and the labels are the ground truth.\n",
        "\n",
        "    Returns:\n",
        "    dict: a dictionary containing F1 (macro), AUPRC, and MCC scores.\n",
        "    \"\"\"\n",
        "    predictions, labels = eval_pred\n",
        "    preds = predictions.argmax(-1)\n",
        "    f1 = f1_metric.compute(predictions = preds, references = labels, average = 'macro')\n",
        "    auprc = average_precision_score(labels, predictions[:, 1]) ### use second (pos) class for binary classification\n",
        "    mcc = mcc_metric.compute(predictions=preds, references = labels)\n",
        "\n",
        "    return {\n",
        "            'f1_macro': f1,\n",
        "            'auprc': auprc,\n",
        "            'mcc': mcc,\n",
        "            }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dwW_R3P5lBot"
      },
      "source": [
        "**_train_and_evaluate_llama_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "85JfapQv6BEZ"
      },
      "outputs": [],
      "source": [
        "%%writefile -a llama_train.py\n",
        "\n",
        "from accelerate import Accelerator\n",
        "from datasets import Dataset\n",
        "from huggingface_hub import login\n",
        "import pandas as pd\n",
        "from sklearn.model_selection import ParameterGrid, StratifiedKFold\n",
        "import torch\n",
        "from transformers import (\n",
        "                          AdamW,\n",
        "                          TrainingArguments,\n",
        "                          Trainer,\n",
        "                          )\n",
        "\n",
        "def train_and_evaluate_llama(target_datasets, targets_and_class_weights, model_name, hyperparameter_grid, save_path):\n",
        "    \"\"\"\n",
        "    Trains and tests Llama for multiple targets using stratified k-fold cross-validation and a held-out test set. Handles\n",
        "    model loading, tokenization, training with Hugging Face's Trainer. Computeds performance metrics by target.\n",
        "\n",
        "    Args:\n",
        "    target_datasets (dict): dictionary where keys are target names and values are target-specific tuples of\n",
        "    (d_train_{target}, d_test_{target})\n",
        "    targets_and_class_weights (dict): dictionary of target-specific inverse-freq class weights to mitigate class imbalance.\n",
        "    model_name (str): name or path of the pretrained Llama model to load.\n",
        "    hyperparameter_grid (list): grid space of hyperparameter configurations (generated using ParameterGrid).\n",
        "    save_path (str): directory to save best-performing model\n",
        "\n",
        "    Returns:\n",
        "    Saves best-performing model by target, saves df of tabulated performance metrics.\n",
        "    \"\"\"\n",
        "\n",
        "    # initialize accelerator\n",
        "\n",
        "    accelerator = Accelerator()\n",
        "\n",
        "    # HF login\n",
        "\n",
        "    login(token = '')\n",
        "\n",
        "    # initialize performance df\n",
        "\n",
        "    d_llama_performance = pd.DataFrame(columns = [\n",
        "                                                  'target',\n",
        "                                                  'model',\n",
        "                                                  'fold',\n",
        "                                                  'f1_macro',\n",
        "                                                  'mcc',\n",
        "                                                  'auprc',\n",
        "                                                  ])\n",
        "\n",
        "    for target, (d_train, d_test) in target_datasets.items():\n",
        "        class_weights = torch.tensor(targets_and_class_weights[target]).to(accelerator.device)\n",
        "        print(\"\\n======================================================================================\")\n",
        "        print(f\"Training Llama for target: {target}\")\n",
        "        print(\"======================================================================================\")\n",
        "\n",
        "        # shuffle d_train\n",
        "\n",
        "        d_train = d_train.sample(frac = 1, random_state = 56).reset_index(drop = True)\n",
        "\n",
        "        # define stratified k-fold\n",
        "\n",
        "        skf = StratifiedKFold(n_splits = 5)\n",
        "\n",
        "        # adefine augmentation mask\n",
        "\n",
        "        aug_mask = d_train['aug'] == 1\n",
        "\n",
        "        # train-validation loop\n",
        "\n",
        "        for fold, (train_index, val_index) in enumerate(skf.split(d_train, d_train[target])):\n",
        "            print(f\"\\nFold {fold + 1}/5\")\n",
        "\n",
        "            # split train and validation sets based on aug mask\n",
        "\n",
        "            train_mask = aug_mask | d_train.index.isin(train_index)\n",
        "            val_mask = ~aug_mask & d_train.index.isin(val_index)\n",
        "\n",
        "            d_train_fold = d_train[train_mask].copy()\n",
        "            d_val_fold = d_train[val_mask].copy()\n",
        "\n",
        "            print(f\"Fold {fold + 1} Training rows: {len(d_train_fold)}\")\n",
        "            print(f\"Fold {fold + 1} Validation rows: {len(d_val_fold)}\")\n",
        "\n",
        "            # rename 'target' col to 'label' for HF Trainer\n",
        "\n",
        "            d_train_fold = d_train_fold.rename(columns = {target: 'label'})\n",
        "            d_val_fold = d_val_fold.rename(columns = {target: 'label'})\n",
        "\n",
        "            # excise 'aug' col before creating HF Dataset objects\n",
        "\n",
        "            d_train_fold = d_train_fold.drop(columns = ['aug'])\n",
        "            d_val_fold = d_val_fold.drop(columns = ['aug'])\n",
        "\n",
        "            # convert to HF Dataset\n",
        "\n",
        "            train_dataset = Dataset.from_pandas(d_train_fold)\n",
        "            val_dataset = Dataset.from_pandas(d_val_fold)\n",
        "\n",
        "           # reinitialize model and tokenizer for each fold to avoid residual information\n",
        "\n",
        "            model, tokenizer = load_llama_and_tokenizer(model_name, num_labels = 2)\n",
        "\n",
        "            # tokenize\n",
        "\n",
        "            train_dataset = train_dataset.map(lambda i: llama_tokenize(i, tokenizer), batched = True)\n",
        "            val_dataset = val_dataset.map(lambda i: llama_tokenize(i, tokenizer), batched = True)\n",
        "\n",
        "            # reformat to PyTorch tensors for HF Trainer compatibility\n",
        "\n",
        "            train_dataset.set_format(type = 'torch', columns = [\n",
        "                                                                'input_ids',\n",
        "                                                                'attention_mask',\n",
        "                                                                'label',\n",
        "                                                                ]\n",
        "                                     )\n",
        "\n",
        "            val_dataset.set_format(type = 'torch', columns = [\n",
        "                                                              'input_ids',\n",
        "                                                              'attention_mask',\n",
        "                                                              'label',\n",
        "                                                              ]\n",
        "                                   )\n",
        "\n",
        "            # display training and validation details\n",
        "\n",
        "            train_batch_size = 4\n",
        "            val_batch_size = 4\n",
        "            total_train_batches = len(train_dataset) // train_batch_size\n",
        "            total_eval_batches = len(val_dataset) // val_batch_size\n",
        "\n",
        "            print(f\"Total training rows: {len(train_dataset)}\")\n",
        "            print(f\"Total validation rows: {len(val_dataset)}\")\n",
        "            print(f\"Training batch size: {train_batch_size}\")\n",
        "            print(f\"Validation batch size: {val_batch_size}\")\n",
        "            print(f\"Total training batches: {total_train_batches}\")\n",
        "            print(f\"Total evaluation batches: {total_eval_batches}\")\n",
        "\n",
        "            # HF TrainingArguments\n",
        "\n",
        "            for h in hyperparameter_grid:\n",
        "                training_args = TrainingArguments(\n",
        "                                                  output_dir = '/content/drive/MyDrive/Colab/bar_policy_suicidality/temp/',\n",
        "                                                  learning_rate = h['learning_rate'],\n",
        "                                                  per_device_train_batch_size = 4,\n",
        "                                                  per_device_eval_batch_size = 4,\n",
        "                                                  num_train_epochs = h['num_train_epochs'],\n",
        "                                                  weight_decay = h['weight_decay'],\n",
        "                                                  gradient_accumulation_steps = h['gradient_accumulation_steps'],\n",
        "                                                  warmup_steps = h['warmup_steps'],\n",
        "                                                  evaluation_strategy = 'epoch',\n",
        "                                                  save_strategy = 'epoch',\n",
        "                                                  report_to = 'none',\n",
        "                                                  push_to_hub = False,\n",
        "                                                  remove_unused_columns = True,\n",
        "                                                  fp16 = True,\n",
        "                                                  seed = 56,\n",
        "                                                  )\n",
        "\n",
        "            # HF Trainer setup\n",
        "\n",
        "            trainer = Trainer(\n",
        "                              model = model,\n",
        "                              args = training_args,\n",
        "                              train_dataset = train_dataset,\n",
        "                              eval_dataset = val_dataset,\n",
        "                              compute_metrics = compute_llama_metrics,\n",
        "                              optimizers = (AdamW(model.parameters(), lr = training_args.learning_rate), None),\n",
        "                              )\n",
        "\n",
        "            # train\n",
        "\n",
        "            trainer.train()\n",
        "\n",
        "            # append fold metrics to performance dataframe\n",
        "\n",
        "            val_metrics = trainer.evaluate(val_dataset)\n",
        "            d_llama_performance.loc[len(d_llama_performance)] = [\n",
        "                                                                 target,\n",
        "                                                                 'llama-3.1-8b',\n",
        "                                                                 fold + 1,\n",
        "                                                                 val_metrics['eval_f1_macro'],\n",
        "                                                                 val_metrics['eval_mcc'],\n",
        "                                                                 val_metrics['eval_auprc'],\n",
        "                                                                ]\n",
        "\n",
        "        # test on held-out test set\n",
        "\n",
        "        print(\"--------------------------------------------------------------------------------------\")\n",
        "        print(f\"Testing Llama for target: {target}\")\n",
        "        print(\"--------------------------------------------------------------------------------------\")\n",
        "\n",
        "        # rename 'target' col: held-out test set\n",
        "\n",
        "        d_test = d_test.rename(columns = {target: 'label'})\n",
        "\n",
        "        # excise 'aug' col: held-out test set\n",
        "\n",
        "        d_test = d_test.drop(columns = ['aug'])\n",
        "\n",
        "        test_dataset = Dataset.from_pandas(d_test)\n",
        "        test_dataset = test_dataset.map(lambda i: llama_tokenize(i, tokenizer), batched = True)\n",
        "        test_dataset.set_format(type = 'torch', columns = [\n",
        "                                                           'input_ids',\n",
        "                                                           'attention_mask',\n",
        "                                                           'label',\n",
        "                                                           ]\n",
        "                                )\n",
        "\n",
        "        # display test set details\n",
        "\n",
        "        test_batch_size = 4\n",
        "        total_test_batches = len(test_dataset) // test_batch_size\n",
        "\n",
        "        print(f\"Total test rows: {len(test_dataset)}\")\n",
        "        print(f\"Test batch size: {test_batch_size}\")\n",
        "        print(f\"Total test batches: {total_test_batches}\")\n",
        "\n",
        "        # test\n",
        "\n",
        "        test_metrics = trainer.evaluate(test_dataset)\n",
        "        d_llama_performance.loc[len(d_llama_performance)] = [\n",
        "                                                             target,\n",
        "                                                             'llama-3.1-8b',\n",
        "                                                             'Test',\n",
        "                                                             test_metrics['eval_f1_macro'],\n",
        "                                                             test_metrics['eval_mcc'],\n",
        "                                                             test_metrics['eval_auprc'],\n",
        "                                                            ]\n",
        "\n",
        "        # save target-wise trained models\n",
        "\n",
        "        print(f\"\\nSaving baseline trained Llama for target: {target}\")\n",
        "        target_save_path = f'{save_path}/{target}_llama_baseline_model'\n",
        "        model.save_pretrained(target_save_path)\n",
        "        tokenizer.save_pretrained(target_save_path)\n",
        "\n",
        "    # extract performance scores numeric values\n",
        "\n",
        "    d_llama_performance['f1_macro'] = d_llama_performance['f1_macro'].apply(lambda i: i['f1'] if isinstance(i, dict) else i)\n",
        "    d_llama_performance['mcc'] = d_llama_performance['mcc'].apply(lambda i: i['matthews_correlation'] if isinstance(i, dict) else i)\n",
        "    d_llama_performance['auprc'] = d_llama_performance['auprc'].apply(lambda i: i if isinstance(i, float) else None)  # Ensure AUPRC is numeric\n",
        "\n",
        "    print(\"\\n--------------------------------------------------------------------------------------\")\n",
        "    print(f\"Summary: Llama performance for target: {target}\")\n",
        "    print(\"--------------------------------------------------------------------------------------\")\n",
        "\n",
        "    print(d_llama_performance.head(6))\n",
        "    d_llama_performance.to_excel('d_llama_performance.xlsx')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JD00mrxLlKQm"
      },
      "source": [
        "**_tune_and_optimize_llama_hyperparams_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WocC-ShjwnFN"
      },
      "outputs": [],
      "source": [
        "%%writefile -a llama_train.py\n",
        "\n",
        "from accelerate import Accelerator\n",
        "from datasets import Dataset\n",
        "from huggingface_hub import login\n",
        "import pandas as pd\n",
        "from sklearn.model_selection import ParameterGrid\n",
        "import torch\n",
        "from transformers import (\n",
        "                          AdamW,\n",
        "                          TrainingArguments,\n",
        "                          Trainer,\n",
        "                          )\n",
        "\n",
        "def tune_and_optimize_llama_hyperparams(target_datasets, targets_and_class_weights, model_name, hyperparameter_grid, save_path):\n",
        "    \"\"\"\n",
        "    Tune and optimize hyperparameters for a Llama model using ParameterGrid search. Trains and tests on held-out target-specific\n",
        "    d_test_{target}, adjusting model in accord with pre-specified ParameterGrid. Saves best-performing model by target.\n",
        "\n",
        "    Args:\n",
        "    target_datasets (dict): dictionary where keys are target-specific tuples of (d_train_{target}, d_test_{target}).\n",
        "    targets_and_class_weights (dict): dictionary of target-specific inverse-freq class weights.\n",
        "    model_name (str): name of pretrained Llama model to load.\n",
        "    hyperparameter_grid (list): grid space of hyperparameter configurations (generated using ParameterGrid).\n",
        "    save_path (str): directory to save best-performing model\n",
        "\n",
        "    Returns:\n",
        "    Saves best-performing model by target, saves df of tabulated performance metrics.\n",
        "    \"\"\"\n",
        "    # initialize accelerator\n",
        "\n",
        "    accelerator = Accelerator()\n",
        "\n",
        "    # HF login\n",
        "\n",
        "    login(token = '')\n",
        "\n",
        "    # initialize performance df\n",
        "\n",
        "    d_llama_performance = pd.DataFrame(columns = [\n",
        "                                                  'target',\n",
        "                                                  'model',\n",
        "                                                  'f1_macro',\n",
        "                                                  'mcc',\n",
        "                                                  'auprc',\n",
        "                                                  ]\n",
        "                                       )\n",
        "\n",
        "    for target, (d_train, d_test) in target_datasets.items():\n",
        "        class_weights = torch.tensor(targets_and_class_weights[target]).to(accelerator.device)\n",
        "        print(\"\\n======================================================================================\")\n",
        "        print(f\"Tuning Llama 3.1 for target: {target}\")\n",
        "        print(\"======================================================================================\")\n",
        "\n",
        "        best_f1_macro = 0 ### tracking var: best F1 (macro)\n",
        "        best_model_state = None ### tracking var: best-performing model x hyperparam configs\n",
        "\n",
        "        for h in hyperparameter_grid:\n",
        "            print(\"\\n\")\n",
        "            print(f\"\\nTuning with hyperparam config: {h}\")\n",
        "\n",
        "            # re/initialize model and tokenizer for each hyperparameter config\n",
        "\n",
        "            model, tokenizer = load_llama_and_tokenizer(model_name, num_labels = 2)\n",
        "\n",
        "            # rename 'target' col to 'label' for HF Trainer\n",
        "\n",
        "            d_train = d_train.rename(columns = {target: 'label'})\n",
        "            d_test = d_test.rename(columns = {target: 'label'})\n",
        "\n",
        "            # convert to HF Dataset\n",
        "\n",
        "            train_dataset = Dataset.from_pandas(d_train)\n",
        "            test_dataset = Dataset.from_pandas(d_test)\n",
        "\n",
        "            # tokenize\n",
        "\n",
        "            train_dataset = train_dataset.map(lambda i: llama_tokenize(i, tokenizer), batched = True)\n",
        "            test_dataset = test_dataset.map(lambda i: llama_tokenize(i, tokenizer), batched = True)\n",
        "\n",
        "            # reformat to PyTorch tensors for HF Trainer compatibility\n",
        "\n",
        "            train_dataset.set_format(type = 'torch', columns = [\n",
        "                                                                'input_ids',\n",
        "                                                                'attention_mask',\n",
        "                                                                'label',\n",
        "                                                                ]\n",
        "                                     )\n",
        "\n",
        "            test_dataset.set_format(type = 'torch', columns = [\n",
        "                                                               'input_ids',\n",
        "                                                               'attention_mask',\n",
        "                                                               'label',\n",
        "                                                               ]\n",
        "                                    )\n",
        "\n",
        "            # display training and testing details\n",
        "\n",
        "            train_batch_size = 4\n",
        "            test_batch_size = 4\n",
        "            total_train_batches = len(train_dataset) // train_batch_size\n",
        "            total_test_batches = len(test_dataset) // test_batch_size\n",
        "\n",
        "            print(f\"Total training rows: {len(train_dataset)}\")\n",
        "            print(f\"Total test rows: {len(test_dataset)}\")\n",
        "            print(f\"Training batch size: {train_batch_size}\")\n",
        "            print(f\"Test batch size: {test_batch_size}\")\n",
        "            print(f\"Total training batches: {total_train_batches}\")\n",
        "            print(f\"Total test batches: {total_test_batches}\")\n",
        "\n",
        "            # HF TrainingArguments w/ ParameterGrid\n",
        "\n",
        "            training_args = TrainingArguments(\n",
        "                                              output_dir = '/content/drive/MyDrive/Colab/bar_policy_suicidality/temp/',\n",
        "                                              learning_rate = h['learning_rate'],\n",
        "                                              per_device_train_batch_size = train_batch_size,\n",
        "                                              per_device_eval_batch_size = test_batch_size,\n",
        "                                              num_train_epochs = h['num_train_epochs'],\n",
        "                                              weight_decay = h['weight_decay'],\n",
        "                                              gradient_accumulation_steps = h['gradient_accumulation_steps'],\n",
        "                                              warmup_steps = h['warmup_steps'],\n",
        "                                              evaluation_strategy = 'no', ### removes midstream validation\n",
        "                                              save_strategy = 'epoch',\n",
        "                                              report_to = 'none',\n",
        "                                              push_to_hub = False,\n",
        "                                              remove_unused_columns = True, ### 'aug' dropped here\n",
        "                                              fp16 = True, ### mixed precision to mitigate memory overhead\n",
        "                                              seed = 56,\n",
        "                                              )\n",
        "\n",
        "            # HF Trainer setup\n",
        "\n",
        "            trainer = Trainer(\n",
        "                              model = model,\n",
        "                              args = training_args,\n",
        "                              train_dataset = train_dataset,\n",
        "                              eval_dataset = test_dataset, ### uses d_test directly\n",
        "                              compute_metrics = compute_llama_metrics,\n",
        "                              optimizers = (AdamW(model.parameters(), lr = training_args.learning_rate), None),\n",
        "                              )\n",
        "\n",
        "            # train\n",
        "\n",
        "            trainer.train()\n",
        "\n",
        "            # test on held-out test set\n",
        "\n",
        "            print(\"--------------------------------------------------------------------------------------\")\n",
        "            print(f\"Testing Llama for target: {target}\")\n",
        "            print(\"--------------------------------------------------------------------------------------\")\n",
        "\n",
        "            test_metrics = trainer.evaluate(test_dataset)\n",
        "            print(test_metrics)\n",
        "\n",
        "            # append fold metrics to performance dataframe\n",
        "\n",
        "            d_llama_performance.loc[len(d_llama_performance)] = [\n",
        "                                                                 target,\n",
        "                                                                 'llama-3.1-8b',\n",
        "                                                                 test_metrics['eval_f1_macro'],\n",
        "                                                                 test_metrics['eval_mcc'],\n",
        "                                                                 test_metrics['eval_auprc'],\n",
        "                                                                 ]\n",
        "\n",
        "            # save best model based on F1 (macro)\n",
        "\n",
        "            if test_metrics['eval_f1_macro']['f1'] > best_f1_macro:\n",
        "                best_f1_macro = test_metrics['eval_f1_macro']['f1']\n",
        "                print(f\"\\nUpdating best model state for target: {target} with F1 (macro): {best_f1_macro}\")\n",
        "\n",
        "        # save the best model target-wise\n",
        "\n",
        "        if best_f1_macro > 0:\n",
        "            print(f\"\\nSaving best model for target: {target} with F1 (macro): {best_f1_macro}\")\n",
        "            target_save_path = f'{save_path}/{target}_llama_best_tuned_model'\n",
        "\n",
        "            # save quantized model\n",
        "\n",
        "            model.save_pretrained(target_save_path)\n",
        "\n",
        "            # save tokenizer\n",
        "\n",
        "            tokenizer.save_pretrained(target_save_path)\n",
        "\n",
        "    # extract performance scores numeric values\n",
        "\n",
        "    d_llama_performance['f1_macro'] = d_llama_performance['f1_macro'].apply(lambda i: i['f1'] if isinstance(i, dict) else i)\n",
        "    d_llama_performance['mcc'] = d_llama_performance['mcc'].apply(lambda i: i['matthews_correlation'] if isinstance(i, dict) else i)\n",
        "    d_llama_performance['auprc'] = d_llama_performance['auprc'].apply(lambda i: i if isinstance(i, float) else None)\n",
        "\n",
        "    print(\"Llama performance summary:\")\n",
        "    d_llama_performance.to_excel('d_llama_tuned_performance.xlsx')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "55RbZ1ZlgEJn"
      },
      "source": [
        "#### llama_predict.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-khn6IuzL5eI"
      },
      "source": [
        "**_llama_load_and_predict_single_target_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nk0FIBaCgDdL"
      },
      "outputs": [],
      "source": [
        "%%writefile llama_predict.py\n",
        "\n",
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
        "import pandas as pd\n",
        "\n",
        "def llama_load_and_predict_single_target(target, df, models_path, batch_size):\n",
        "    \"\"\"\n",
        "    Function to load a single model and tokenizer for a specified target, and use them to predict\n",
        "    labels and class probabilities for the text in df['text'] in batches.\n",
        "\n",
        "    Args:\n",
        "        target (str): The name of the target (e.g., 'asp', 'dep').\n",
        "        df (pd.DataFrame): DataFrame containing a 'text' column with the input texts.\n",
        "        models_path (str): Directory where the model for the target is saved (e.g., /models/).\n",
        "        batch_size (int): The batch size for processing the data in smaller chunks.\n",
        "\n",
        "    Returns:\n",
        "        pd.DataFrame: DataFrame with additional columns '{target}_pred' and '{target}_prob'.\n",
        "    \"\"\"\n",
        "\n",
        "    # load target-specific best-performing tuned Llama\n",
        "\n",
        "    model_save_path = f'{models_path}/{target}_llama_best_tuned_model'\n",
        "\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_save_path)\n",
        "    model = AutoModelForSequenceClassification.from_pretrained(model_save_path)\n",
        "\n",
        "    # set padding token\n",
        "\n",
        "    tokenizer.pad_token = tokenizer.eos_token\n",
        "    model.config.pad_token_id = model.config.eos_token_id\n",
        "\n",
        "    # migrate to GPU\n",
        "\n",
        "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "    model.to(device)\n",
        "\n",
        "    # ensure eval mode\n",
        "\n",
        "    model.eval()\n",
        "\n",
        "    predicted_labels = []\n",
        "    class_probabilities = []\n",
        "\n",
        "    # batch processing\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for i in range(0, len(df), batch_size):\n",
        "            batch_texts = df['text'][i:i + batch_size].tolist()\n",
        "\n",
        "            # tokenize\n",
        "\n",
        "            inputs = tokenizer(\n",
        "                               batch_texts,\n",
        "                               return_tensors = 'pt',\n",
        "                               truncation = True,\n",
        "                               padding = True,\n",
        "                               max_length = 512,\n",
        "                               )\n",
        "\n",
        "            inputs = {k: v.to(device) for k, v in inputs.items()}\n",
        "\n",
        "            # get logits\n",
        "\n",
        "            outputs = model(**inputs)\n",
        "            logits = outputs.logits\n",
        "\n",
        "            # logits -> probabilities via softmax\n",
        "\n",
        "            probabilities = torch.softmax(logits, dim = -1).cpu().numpy()\n",
        "\n",
        "            # get predicted labels\n",
        "\n",
        "            predicted_labels_batch = torch.argmax(logits, dim = -1).cpu().numpy()\n",
        "\n",
        "            # append results\n",
        "\n",
        "            predicted_labels.extend(predicted_labels_batch)\n",
        "            class_probabilities.extend(probabilities)\n",
        "\n",
        "    # to df\n",
        "\n",
        "    df[f'{target}_pred'] = predicted_labels\n",
        "    df[f'{target}_prob'] = class_probabilities\n",
        "\n",
        "    return df\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kz6n-pmZL8fk"
      },
      "source": [
        "**_llama_load_and_predict_multi_target_**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KIDFpqheL7_9"
      },
      "outputs": [],
      "source": [
        "%%writefile -a llama_predict.py\n",
        "\n",
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
        "import pandas as pd\n",
        "\n",
        "def llama_load_and_predict_multi_target(targets, df, models_path, batch_size):\n",
        "    \"\"\"\n",
        "    Function to load multiple models and tokenizers for multiple targets, and use them to predict\n",
        "    labels and class probabilities for the text in df['text'] in batches.\n",
        "\n",
        "    Args:\n",
        "        targets (list): List of target names (e.g., ['asp', 'dep']).\n",
        "        df (pd.DataFrame): DataFrame containing a 'text' column with the input texts.\n",
        "        models_path (str): Directory where models for each target are saved (e.g., /models/).\n",
        "        batch_size (int): The batch size for processing the data in smaller chunks.\n",
        "\n",
        "    Returns:\n",
        "        pd.DataFrame: DataFrame with additional columns '{target}_pred' and '{target}_prob' for each target.\n",
        "    \"\"\"\n",
        "\n",
        "    for target in targets:\n",
        "\n",
        "        # load target-specific best-performing tuned Llama\n",
        "\n",
        "        model_save_path = f'{models_path}/{target}_llama_best_tuned_model'\n",
        "\n",
        "        tokenizer = AutoTokenizer.from_pretrained(model_save_path)\n",
        "        model = AutoModelForSequenceClassification.from_pretrained(model_save_path)\n",
        "\n",
        "        # set padding token\n",
        "\n",
        "        tokenizer.pad_token = tokenizer.eos_token\n",
        "        model.config.pad_token_id = model.config.eos_token_id\n",
        "\n",
        "        # migrate to GPU\n",
        "\n",
        "        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "        model.to(device)\n",
        "\n",
        "        # ensure eval mode\n",
        "\n",
        "        model.eval()\n",
        "\n",
        "        predicted_labels = []\n",
        "        class_probabilities = []\n",
        "\n",
        "        # batch processing\n",
        "\n",
        "        with torch.no_grad():\n",
        "            for i in range(0, len(df), batch_size):\n",
        "                batch_texts = df['text'][i:i + batch_size].tolist()\n",
        "\n",
        "                # tokenize\n",
        "\n",
        "                inputs = tokenizer(batch_texts, return_tensors=\"pt\", truncation=True, padding=True, max_length=512)\n",
        "                inputs = {k: v.to(device) for k, v in inputs.items()}\n",
        "\n",
        "                # get logits\n",
        "\n",
        "                outputs = model(**inputs)\n",
        "                logits = outputs.logits\n",
        "\n",
        "                # logits -> probabilities via softmax\n",
        "\n",
        "                probabilities = torch.softmax(logits, dim = -1).cpu().numpy()\n",
        "\n",
        "                # get predicted labels\n",
        "\n",
        "                predicted_labels_batch = torch.argmax(logits, dim = -1).cpu().numpy()\n",
        "\n",
        "                # append results\n",
        "\n",
        "                predicted_labels.extend(predicted_labels_batch)\n",
        "                class_probabilities.extend(probabilities)\n",
        "\n",
        "        # to df\n",
        "\n",
        "        df[f'{target}_pred'] = predicted_labels\n",
        "        df[f'{target}_prob'] = class_probabilities\n",
        "\n",
        "    return df\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hr3Dsh6zd1K6"
      },
      "source": [
        "#### Import"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_QcKdI-mcASI"
      },
      "outputs": [],
      "source": [
        "from preprocess import (\n",
        "                        augment_training_data_with_rationales,\n",
        "                        dummy_code_augmented_rows,\n",
        "                        read_and_append_jsonl_posts,\n",
        "                        read_and_append_jsonl_comments,\n",
        "                        )\n",
        "\n",
        "#from redact import (\n",
        "#                    ner_redact_post_texts,\n",
        "#                    )\n",
        "\n",
        "# BERT, RoBERTa, DistilBERT modules\n",
        "\n",
        "from bert_train import (\n",
        "                        set_seed,\n",
        "                        train_eval_save_bl_models,\n",
        "                        performance_scatterplot,\n",
        "                        iterative_stratified_train_test_split_with_rationales,\n",
        "                        tune_and_optimize_model_hyperparams,\n",
        "                        tune_and_optimize_model_loss_accuracy,\n",
        "                        )\n",
        "\n",
        "from bert_predict import (\n",
        "                          load_model,\n",
        "                          preprocess_data,\n",
        "                          predict,\n",
        "                          )\n",
        "\n",
        "# Llama 3.1 modules\n",
        "\n",
        "from llama_train import (\n",
        "                         load_llama_and_tokenizer,\n",
        "                         llama_tokenize,\n",
        "                         compute_llama_metrics,\n",
        "                         train_and_evaluate_llama,\n",
        "                         tune_and_optimize_llama_hyperparams,\n",
        "                         )\n",
        "\n",
        "from llama_predict import (\n",
        "                           llama_load_and_predict_single_target,\n",
        "                           llama_load_and_predict_multi_target,\n",
        "                           )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "woZYUo6JgN1X"
      },
      "source": [
        "### 3. Preprocess\n",
        "Takes $\\mathcal{d}$<sub>annotated</sub>, builds independent $\\mathcal{d}$<sub>calibrate</sub> and rationale-augmented $\\mathcal{d}$<sub>augmented</sub> to train. Builds $\\mathcal{V}$ corpus, $\\mathcal{d}$<sub>adapt</sub> for domain adaptation. Merges, NER-anonymizes Aim II analytic sample $\\mathcal{D}$<sub>inference</sub>.\n",
        "***"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ThJ315SGh0fx"
      },
      "source": [
        "#### Merge Wave I (purposive) and Wave II (random)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7fTcqzw9v2db"
      },
      "outputs": [],
      "source": [
        "%cd ../inputs/data\n",
        "\n",
        "    ### SJS 11/26: WIP annotation files at ../inputs/annotation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1o9be0r93OEJ"
      },
      "outputs": [],
      "source": [
        "# recover 'sbrt' subreddit var\n",
        "\n",
        "d_prp_sbrt = pd.read_excel('d_cycle999_prp_sd_single.xlsx', index_col = [0])\n",
        "d_prp_01 = pd.read_excel('d_cycle999_prp_sd_ss_gpt_agreed_01.xlsx', index_col = [0])\n",
        "d_prp_02 = pd.read_excel('d_cycle999_prp_sd_ss_gpt_agreed_02.xlsx', index_col = [0])\n",
        "\n",
        "d_prp = pd.concat([\n",
        "                   d_prp_01,\n",
        "                   d_prp_02,\n",
        "                   ])\n",
        "\n",
        "d_prp = d_prp.drop(\n",
        "                   'p_id_sd',\n",
        "                   axis = 1,\n",
        "                   )\n",
        "\n",
        "d_prp.rename(\n",
        "             columns = {\n",
        "                        'p_id_ss': 'p_id',\n",
        "                        }, inplace = True,\n",
        "            )\n",
        "\n",
        "d_prp_sbrt = d_prp_sbrt[[\n",
        "                         'p_id',\n",
        "                         'sbrt',\n",
        "                         ]].copy()\n",
        "\n",
        "#d_prp.info()\n",
        "#d_prp.head(1)\n",
        "#d_prp_sbrt.head(1)\n",
        "\n",
        "d_prp = d_prp.merge(\n",
        "                    d_prp_sbrt,\n",
        "                    on = 'p_id',\n",
        "                    how = 'inner',\n",
        "                    )\n",
        "\n",
        "d_prp.head(1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nNO_RON3q1GQ"
      },
      "outputs": [],
      "source": [
        "# Wave I IAA - post-negotiated agreement\n",
        "\n",
        "asp_kappa = cohen_kappa_score(d_prp['asp_sd'], d_prp['asp_ss'])\n",
        "val_kappa = cohen_kappa_score(d_prp['val_sd'], d_prp['val_ss'])\n",
        "\n",
        "print(\"asp Kappa:\", asp_kappa)\n",
        "print(\"val Kappa:\", val_kappa)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vb4Vz-mRi4wJ"
      },
      "outputs": [],
      "source": [
        "# import purposive - SS-SD-GPT triangulated\n",
        "\n",
        "#d_prp_01 = pd.read_excel('d_cycle999_prp_sd_ss_gpt_agreed_01.xlsx', index_col = [0])\n",
        "#d_prp_02 = pd.read_excel('d_cycle999_prp_sd_ss_gpt_agreed_02.xlsx', index_col = [0])\n",
        "\n",
        "#d_prp = pd.concat([\n",
        "#                   d_prp_01,\n",
        "#                   d_prp_02,\n",
        "#                   ])\n",
        "\n",
        "# target = 1 at SS-SD agreement\n",
        "\n",
        "strain_tuples = [\n",
        "                    ('asp_ss', 'asp_sd', 'asp'),\n",
        "                    ('dep_ss', 'dep_sd', 'dep'),\n",
        "                    ('val_ss', 'val_sd', 'val'),\n",
        "                    #('prg_ss', 'prg_sd', 'prg'),\n",
        "                    #('tgd_ss', 'tgd_sd', 'tgd'),\n",
        "                    #('age_ss', 'age_sd', 'age'),\n",
        "                    #('race_ss', 'race_sd', 'race'),\n",
        "                    #('dbty_ss', 'dbty_sd', 'dbty'),\n",
        "                    ]\n",
        "\n",
        "for ss, sd, target in strain_tuples:\n",
        "    d_prp[target] = (d_prp[ss] == 1) & (d_prp[sd] == 1)\n",
        "    d_prp[target] = d_prp[target].astype(int)\n",
        "\n",
        "# target = 1 at SS-SD-GPT agreement\n",
        "\n",
        "trait_tuples = [\n",
        "                    ('prg_ss', 'prg_sd', 'prg_gpt', 'prg'),\n",
        "                    ('tgd_ss', 'tgd_sd', 'tgd_gpt', 'tgd'),\n",
        "                    ('age_ss', 'age_sd', 'age_gpt', 'age'),\n",
        "                    ('race_ss', 'race_sd', 'race_gpt', 'race'),\n",
        "                    ('dbty_ss', 'dbty_sd', 'dbty_gpt', 'dbty'),\n",
        "                    ]\n",
        "\n",
        "for ss, sd, gpt, target in trait_tuples:\n",
        "    d_prp[target] = (d_prp[ss] == 1) & (d_prp[sd] == 1) & (d_prp[gpt] == 1)\n",
        "    d_prp[target] = d_prp[target].astype(int)\n",
        "\n",
        "\n",
        "# append rationales\n",
        "\n",
        "rationale_tuples = [\n",
        "                    ('asp_rtnl_ss', 'asp_rtnl_sd', 'asp', 'asp_rtnl'),\n",
        "                    ('dep_rtnl_ss', 'dep_rtnl_sd', 'dep', 'dep_rtnl'),\n",
        "                    ('val_rtnl_ss', 'val_rtnl_sd', 'val', 'val_rtnl'),\n",
        "                    ]\n",
        "\n",
        "for ss, sd, target, rationale in rationale_tuples:\n",
        "    d_prp[rationale] = d_prp.apply(\n",
        "        lambda row: (row[ss] or ' ') + ' ' + (row[sd] or ' ') if row[target] == 1 else None, axis = 1\n",
        "    )\n",
        "\n",
        "d_prp.rename(\n",
        "             columns = {\n",
        "                        'prg_rtnl_gpt': 'prg_rtnl',\n",
        "                        'tgd_rtnl_gpt': 'tgd_rtnl',\n",
        "                        'age_rtnl_gpt': 'age_rtnl',\n",
        "                        'race_rtnl_gpt': 'race_rtnl',\n",
        "                        'dbty_rtnl_gpt': 'dbty_rtnl',\n",
        "                        }, inplace = True,\n",
        "            )\n",
        "\n",
        "# drop negotiated agreement artifacts\n",
        "\n",
        "d_prp = d_prp[[\n",
        "                #'p_au',\n",
        "                #'p_utc',\n",
        "                #'p_date',\n",
        "                'p_id',\n",
        "                #'n_cmnt',\n",
        "                'text',\n",
        "                'sbrt',\n",
        "                'p_titl',\n",
        "                'asp',\n",
        "                'asp_rtnl',\n",
        "                'dep',\n",
        "                'dep_rtnl',\n",
        "                'val',\n",
        "                'val_rtnl',\n",
        "                'prg',\n",
        "                'prg_rtnl',\n",
        "                'tgd',\n",
        "                'tgd_rtnl',\n",
        "                'age',\n",
        "                'age_rtnl',\n",
        "                'race',\n",
        "                'race_rtnl',\n",
        "                'dbty',\n",
        "                'dbty_rtnl',\n",
        "                ]].copy()\n",
        "\n",
        "\n",
        "# inspect\n",
        "\n",
        "d_prp.info()\n",
        "d_prp.head(1)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qQIb7E-yjQN8"
      },
      "outputs": [],
      "source": [
        "# random - SS-GPT triangulated\n",
        "\n",
        "d_rnd_01 = pd.read_excel('d_cycle999_rnd_ss_gpt_agreed_01.xlsx', index_col = [0])\n",
        "d_rnd_02 = pd.read_excel('d_cycle999_rnd_ss_gpt_agreed_02.xlsx', index_col = [0])\n",
        "\n",
        "d_rnd = pd.concat([\n",
        "                   d_rnd_01,\n",
        "                   d_rnd_02,\n",
        "                   ])\n",
        "\n",
        "# target = 1 at SS-GPT agreement\n",
        "\n",
        "target_tuples = [\n",
        "                    ('asp_ss', 'asp_gpt', 'asp'),\n",
        "                    ('dep_ss', 'dep_gpt', 'dep'),\n",
        "                    ('val_ss', 'val_gpt', 'val'),\n",
        "                    ('prg_ss', 'prg_gpt', 'prg'),\n",
        "                    ('tgd_ss', 'tgd_gpt', 'tgd'),\n",
        "                    ('age_ss', 'age_gpt', 'age'),\n",
        "                    ('race_ss', 'race_gpt', 'race'),\n",
        "                    ('dbty_ss', 'dbty_gpt', 'dbty'),\n",
        "                    ]\n",
        "\n",
        "for ss, gpt, target in target_tuples:\n",
        "    d_rnd[target] = (d_rnd[ss] == 1) & (d_rnd[gpt] == 1)\n",
        "    d_rnd[target] = d_rnd[target].astype(int)\n",
        "\n",
        "# append rationales\n",
        "\n",
        "rationale_tuples = [\n",
        "                    ('asp_rtnl_ss', 'asp_rtnl_gpt', 'asp', 'asp_rtnl'),\n",
        "                    ('dep_rtnl_ss', 'dep_rtnl_gpt', 'dep', 'dep_rtnl'),\n",
        "                    ('val_rtnl_ss', 'val_rtnl_gpt', 'val', 'val_rtnl'),\n",
        "                    ]\n",
        "\n",
        "for ss, gpt, target, rationale in rationale_tuples:\n",
        "    d_rnd[rationale] = d_rnd.apply(\n",
        "        lambda row: (row[ss] or ' ') + ' ' + (row[gpt] or ' ') if row[target] == 1 else None, axis = 1\n",
        "    )\n",
        "\n",
        "d_rnd.rename(\n",
        "             columns = {\n",
        "                        'prg_rtnl_gpt': 'prg_rtnl',\n",
        "                        'tgd_rtnl_gpt': 'tgd_rtnl',\n",
        "                        'age_rtnl_gpt': 'age_rtnl',\n",
        "                        'race_rtnl_gpt': 'race_rtnl',\n",
        "                        'dbty_rtnl_gpt': 'dbty_rtnl',\n",
        "                        }, inplace = True,\n",
        "            )\n",
        "\n",
        "# drop negotiated agreement artifacts\n",
        "\n",
        "d_rnd = d_rnd[[\n",
        "                #'p_au',\n",
        "                #'p_utc',\n",
        "                #'p_date',\n",
        "                'p_id',\n",
        "                #'n_cmnt',\n",
        "                'text',\n",
        "                'sbrt',\n",
        "                'p_titl',\n",
        "                'asp',\n",
        "                'asp_rtnl',\n",
        "                'dep',\n",
        "                'dep_rtnl',\n",
        "                'val',\n",
        "                'val_rtnl',\n",
        "                'prg',\n",
        "                'prg_rtnl',\n",
        "                'tgd',\n",
        "                'tgd_rtnl',\n",
        "                'age',\n",
        "                'age_rtnl',\n",
        "                'race',\n",
        "                'race_rtnl',\n",
        "                'dbty',\n",
        "                'dbty_rtnl',\n",
        "                ]].copy()\n",
        "\n",
        "# inspect\n",
        "\n",
        "d_rnd.info()\n",
        "d_rnd.head(1)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uNXIihjSjf4A"
      },
      "outputs": [],
      "source": [
        "# supplemental - SS annotated\n",
        "\n",
        "d_bp = pd.read_excel('d_bp_suppl_ss.xlsx') ### r/Blackpeople, n = 100\n",
        "d_db = pd.read_excel('d_db_suppl_ss.xlsx') ### r/Disability, n = 100\n",
        "\n",
        "d_suppl = pd.concat([\n",
        "                     d_bp,\n",
        "                     d_db,\n",
        "                     ])\n",
        "\n",
        "d_suppl = d_suppl.drop(['date','insb'], axis = 1,)\n",
        "\n",
        "d_suppl['prg_rtnl'] = ' '\n",
        "d_suppl['tgd_rtnl'] = ' '\n",
        "d_suppl['age_rtnl'] = ' '\n",
        "\n",
        "d_suppl = d_suppl[[\n",
        "                   #'p_au',\n",
        "                   #'p_utc',\n",
        "                   #'p_date',\n",
        "                   'p_id',\n",
        "                   #'n_cmnt',\n",
        "                   'text',\n",
        "                   'sbrt',\n",
        "                   'p_titl',\n",
        "                   'asp',\n",
        "                   'asp_rtnl',\n",
        "                   'dep',\n",
        "                   'dep_rtnl',\n",
        "                   'val',\n",
        "                   'val_rtnl',\n",
        "                   'prg',\n",
        "                   'prg_rtnl',\n",
        "                   'tgd',\n",
        "                   'tgd_rtnl',\n",
        "                   'age',\n",
        "                   'age_rtnl',\n",
        "                   'race',\n",
        "                   'race_rtnl',\n",
        "                   'dbty',\n",
        "                   'dbty_rtnl',\n",
        "                   ]].copy()\n",
        "\n",
        "d_suppl.info()\n",
        "d_suppl.head(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "1WEmQf_vgM1U"
      },
      "outputs": [],
      "source": [
        "# concat, housekeeping\n",
        "\n",
        "d = pd.concat([\n",
        "               d_prp,\n",
        "               d_rnd,\n",
        "               d_suppl,\n",
        "               ])\n",
        "\n",
        "# delete index\n",
        "\n",
        "#d = d.drop(\n",
        "#           'Unnamed: 0',\n",
        "#           axis = 1,\n",
        "#           )\n",
        "\n",
        "# delete empty 'text' cells\n",
        "\n",
        "d = d[d.text != ' ']\n",
        "\n",
        "d.replace(\n",
        "          ' ',\n",
        "          0,\n",
        "          inplace = True,\n",
        "          )\n",
        "\n",
        "d.fillna(\n",
        "         0,\n",
        "         inplace = True,\n",
        "         )\n",
        "\n",
        "# to int\n",
        "\n",
        "targets = [\n",
        "           'asp',\n",
        "           'dep',\n",
        "           'val',\n",
        "           'prg',\n",
        "           'tgd',\n",
        "           'age',\n",
        "           'race',\n",
        "           'dbty',\n",
        "           ]\n",
        "\n",
        "for t in targets:\n",
        "    d[t] = pd.to_numeric(d[t], errors = 'coerce')\n",
        "    d[t] = d[t].fillna(0).astype('int64')\n",
        "\n",
        "# housekeeping: del 'text' = 0, del pseudowords\n",
        "\n",
        "d = d[d['text'] != '0']\n",
        "d = d[d['text'].astype(str) != '0']\n",
        "\n",
        "\n",
        "# shuffle, reset index\n",
        "\n",
        "d_annotated = shuffle(\n",
        "                      d,\n",
        "                      random_state = 56,\n",
        "                      )\n",
        "\n",
        "d_annotated.reset_index(\n",
        "                        drop = True,\n",
        "                        inplace = True,\n",
        "                        )\n",
        "\n",
        "# subreddit: n/%\n",
        "\n",
        "sbrt_cnt = d_annotated['sbrt'].value_counts()\n",
        "sbrt_pct = d_annotated['sbrt'].value_counts(normalize = True) * 100\n",
        "\n",
        "sbrts = pd.DataFrame({\n",
        "                      'Count': sbrt_cnt,\n",
        "                      'Percentage': sbrt_pct,\n",
        "                      })\n",
        "\n",
        "print(sbrts)\n",
        "\n",
        "# target: n\n",
        "\n",
        "d_annotated[targets].apply(pd.Series.value_counts)\n",
        "\n",
        "# target / subreddit: n/%\n",
        "\n",
        "targets_cnt = d_annotated.groupby('sbrt')[targets].sum()\n",
        "targets_pct = d_annotated.groupby('sbrt')[targets].mean() * 100\n",
        "\n",
        "targets_sbrt = targets_cnt.astype(int).add_suffix('_count').join(targets_pct.round(2).add_suffix('_percent'))\n",
        "\n",
        "print(targets_sbrt)\n",
        "\n",
        "# inspect\n",
        "\n",
        "#d.dtypes\n",
        "#d_annotated.info()\n",
        "d_annotated.head(3)\n",
        "\n",
        "# export\n",
        "\n",
        "#d_annotated.to_excel('d_annotated.xlsx')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1IPem2mmjf4V"
      },
      "source": [
        "$\\mathcal{d}$<sub>calibrate</sub> ($n$<sub>posts</sub> = 400)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZKu0OIDVgM5K"
      },
      "outputs": [],
      "source": [
        "d_calibrate = d_annotated.iloc[:400]\n",
        "\n",
        "# del pseudowords\n",
        "\n",
        "texts = ['text']\n",
        "pseudoword_tokens = ['<SPL>', '<|PII|>']\n",
        "\n",
        "for t in texts:\n",
        "    d_calibrate[t] = d_calibrate[t].replace(\n",
        "                                            pseudoword_tokens,\n",
        "                                            ' ',\n",
        "                                            regex = True,\n",
        "                                            )\n",
        "\n",
        "# inspect, export\n",
        "\n",
        "d_calibrate.shape\n",
        "d_calibrate.head(3)\n",
        "\n",
        "d_calibrate.to_excel('d_calibrate.xlsx')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vyBe4kf8j9i3"
      },
      "source": [
        "$\\mathcal{d}$<sub>augmented</sub> (unique $n$<sub>posts</sub> = 2,005)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "JYA6jiMNgNAj"
      },
      "outputs": [],
      "source": [
        "d = d_annotated.iloc[400:]\n",
        "\n",
        "d.reset_index(\n",
        "              drop = True,\n",
        "              inplace = True,\n",
        "              )\n",
        "\n",
        "# 'strn' = any pos_instance of strain\n",
        "\n",
        "d['strn'] = (d['asp'] == 1) | (d['dep'] == 1) | (d['val'] == 1)\n",
        "d['strn'] = d['strn'].astype(int)\n",
        "\n",
        "# append rationales\n",
        "\n",
        "rationales = [\n",
        "              'asp_rtnl',\n",
        "              'dep_rtnl',\n",
        "              'val_rtnl',\n",
        "              'prg_rtnl',\n",
        "              'tgd_rtnl',\n",
        "              'age_rtnl',\n",
        "              'race_rtnl',\n",
        "              'dbty_rtnl',\n",
        "               ]\n",
        "\n",
        "for r in rationales:\n",
        "    d[r] = d[r].astype(str)\n",
        "    d[r] = d[r].str.replace(\n",
        "                            r'0',\n",
        "                            '.',\n",
        "                            regex = True,\n",
        "                            )\n",
        "\n",
        "d['rtnl'] = d['asp_rtnl'] + ' ' + d['dep_rtnl'] + ' ' + d['val_rtnl'] + ' ' + d['prg_rtnl'] + ' ' + d['tgd_rtnl'] + ' ' + d['age_rtnl'] + ' ' + d['race_rtnl'] + ' ' + d['dbty_rtnl']\n",
        "\n",
        "d['rtnl'] = d['rtnl'].str.replace(\n",
        "                                  r'. . . . . . . .',\n",
        "                                  '.',\n",
        "                                  regex = False,\n",
        "                                  )\n",
        "\n",
        "# del pseudowords\n",
        "\n",
        "texts = ['text', 'rtnl']\n",
        "pseudoword_tokens = ['<SPL>', '<|PII|>']\n",
        "\n",
        "for t in texts:\n",
        "    d[t] = d[t].replace(\n",
        "                        pseudoword_tokens,\n",
        "                        ' ',\n",
        "                        regex = True,\n",
        "                        )\n",
        "\n",
        "print(\"pre-augmentation\")\n",
        "print(\"--------------------------------------------------------------------------------------\")\n",
        "d.shape\n",
        "d.head(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "0TwYX6CCgNEb"
      },
      "outputs": [],
      "source": [
        "# augment\n",
        "\n",
        "d_augmented = augment_training_data_with_rationales(d)\n",
        "\n",
        "d_augmented.reset_index(\n",
        "                        drop = True,\n",
        "                        inplace = True,\n",
        "                        )\n",
        "\n",
        "# 'aug' - flag augmented rows\n",
        "\n",
        "d_augmented = dummy_code_augmented_rows(d_augmented)\n",
        "\n",
        "\n",
        "print(\"post-augmentation\")\n",
        "print(\"--------------------------------------------------------------------------------------\")\n",
        "\n",
        "d_augmented[[\n",
        "             'asp',\n",
        "             'dep',\n",
        "             'val',\n",
        "             'prg',\n",
        "             'tgd',\n",
        "             'age',\n",
        "             'race',\n",
        "             'dbty',\n",
        "             'aug'\n",
        "              ]].apply(pd.Series.value_counts)\n",
        "\n",
        "d_augmented.shape\n",
        "d_augmented.head(3)\n",
        "\n",
        "d_augmented.to_excel('d_augmented.xlsx')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HMuLvUHmksoc"
      },
      "source": [
        "#### Append $\\mathcal{V}$ corpus, derive $\\mathcal{D}$<sub>inference</sub>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "kD5s5ruNvy95"
      },
      "outputs": [],
      "source": [
        "%cd /content/drive/My Drive/Colab/bar_policy_suicidality/inputs/data\n",
        "\n",
        "archives_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/inputs/archives/'\n",
        "\n",
        "chunk_size = 10000\n",
        "\n",
        "d_p = read_and_append_jsonl_posts(archives_path)\n",
        "\n",
        "# inspect + save\n",
        "\n",
        "d_p.info()\n",
        "d_p.head(3)\n",
        "\n",
        "d_p.to_csv(\n",
        "           'd_posts_raw.csv',\n",
        "           encoding = 'utf-8',\n",
        "           #index = False,\n",
        "           header = True,\n",
        "           )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FFzJfMf-ecjL"
      },
      "outputs": [],
      "source": [
        "%cd /content/drive/My Drive/Colab/bar_policy_suicidality/inputs/data\n",
        "\n",
        "archives_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/inputs/archives/'\n",
        "\n",
        "chunk_size = 10000\n",
        "\n",
        "d_c = read_and_append_jsonl_comments(archives_path)\n",
        "\n",
        "d_c.shape\n",
        "d_c.head(3)\n",
        "\n",
        "# inspect + save\n",
        "\n",
        "d_c.info()\n",
        "d_c.to_csv(\n",
        "           'd_comments_raw.csv',\n",
        "           encoding = 'utf-8',\n",
        "           #index = False,\n",
        "           header = True,\n",
        "           )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AXDTdWwM10Jv"
      },
      "source": [
        "**Clean, condense: posts**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "JicbL1zaxl7G"
      },
      "outputs": [],
      "source": [
        "# d_p = posts\n",
        "\n",
        "d_p = d_p.drop_duplicates(subset = 'id')\n",
        "\n",
        "d_p['date'] = pd.to_datetime(\n",
        "                             d_p.created_utc,\n",
        "                             unit = 's',\n",
        "                             )\n",
        "\n",
        "d_p.set_index(\n",
        "              'date',\n",
        "              drop = False,\n",
        "              inplace = True,\n",
        "              )\n",
        "\n",
        "d_p = d_p.loc[(d_p['date'] >= '2020-12-02') & (d_p['date'] < '2024-06-24')] ### yyyy-mm-dd = Dec 2, 2020 - Jun 24, 2024\n",
        "\n",
        "d_p = d_p[~d_p['selftext'].isin(['[deleted]', '[removed]'])]\n",
        "\n",
        "# housekeeping\n",
        "\n",
        "d_p = d_p[[\n",
        "           'author',\n",
        "           'created_utc',\n",
        "           'date',\n",
        "           'id',\n",
        "           'num_comments',\n",
        "           'selftext',\n",
        "           'subreddit',\n",
        "           'title',\n",
        "           ]].copy()\n",
        "\n",
        "d_p.rename(\n",
        "           columns = {\n",
        "                      'author': 'p_au',\n",
        "                      'created_utc': 'p_utc',\n",
        "                      'date': 'p_date',\n",
        "                      'num_comments': 'n_cmnt',\n",
        "                      'selftext': 'text',\n",
        "                      'subreddit': 'p_sbrt',\n",
        "                      'title': 'p_titl',\n",
        "                      }, inplace = True,\n",
        "            )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JOsK4F2R7ZYb"
      },
      "source": [
        "**NER anonymize: posts**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zjSPJExg7YUX"
      },
      "outputs": [],
      "source": [
        "d_p['text'] = d_p['text'].astype(str).apply(lambda i: ner_redact_post_texts(i))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A4vlHfGW19v9"
      },
      "source": [
        "**Clean, condense: comments**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9PrCoKyf19PN"
      },
      "outputs": [],
      "source": [
        "# d_c = comments\n",
        "\n",
        "d_c = d_c.drop_duplicates(subset = 'id')\n",
        "\n",
        "d_c['date'] = pd.to_datetime(\n",
        "                             d_c.created_utc,\n",
        "                             unit = 's',\n",
        "                             )\n",
        "\n",
        "d_c.set_index(\n",
        "              'date',\n",
        "              drop = False,\n",
        "              inplace = True,\n",
        "              )\n",
        "\n",
        "d_c = d_c.loc[(d_c['date'] >= '2020-12-02') & (d_c['date'] < '2024-06-24')] ### yyyy-mm-dd = Dec 2, 2020 - Jun 24, 2024\n",
        "\n",
        "d_c = d_c[~d_c['body'].isin(['[deleted]', '[removed]'])]\n",
        "\n",
        "# housekeeping\n",
        "\n",
        "d_c = d_c[[\n",
        "           'author',\n",
        "           'body',\n",
        "           'date',\n",
        "           'link_id',\n",
        "           'subreddit',\n",
        "           ]].copy()\n",
        "\n",
        "d_c.rename(\n",
        "           columns = {\n",
        "                      'author': 'c_au',\n",
        "                      'body': 'c_text',\n",
        "                      'date': 'c_date',\n",
        "                      'link_id': 'id',\n",
        "                      'subreddit': 'c_sbrt',\n",
        "                      }, inplace = True,\n",
        "            )\n",
        "\n",
        "# delete comment-level 'id' prefix for merge\n",
        "\n",
        "d_c['id'] = d_c['id'].str.replace('t3_', ' ')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4P4N-5Cc2S5g"
      },
      "outputs": [],
      "source": [
        "print(\"V-corpus posts/d_adapt\")\n",
        "print(\"--------------------------------------------------------------------------------------\")\n",
        "d_p.shape\n",
        "d_p.head(3)\n",
        "d_p.tail(3)\n",
        "\n",
        "print(\"V-corpus comments\")\n",
        "print(\"--------------------------------------------------------------------------------------\")\n",
        "d_c.shape\n",
        "d_c.head(3)\n",
        "d_c.tail(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2vM8HGGDvyuG"
      },
      "outputs": [],
      "source": [
        "# for post-labeling merge\n",
        "\n",
        "d_c.to_csv(\n",
        "           'd_comments.csv',\n",
        "           encoding = 'utf-8',\n",
        "           index = False,\n",
        "           header = True,\n",
        "           )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2_uqv-aHsDGn"
      },
      "source": [
        "$\\mathcal{d}$<sub>adapt</sub>: domain adaptation set"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UeNXevpnD7AU"
      },
      "outputs": [],
      "source": [
        "d_adapt = d_p.copy()\n",
        "\n",
        "d_adapt.to_csv(\n",
        "               'd_adapt_TEST.csv',\n",
        "               encoding = 'utf-8',\n",
        "               index = False,\n",
        "               header = True,\n",
        "               )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zLgzHX1Tts0F"
      },
      "source": [
        "$\\mathcal{D}$<sub>inference</sub>: prediction set"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2paLevfVtsLi"
      },
      "outputs": [],
      "source": [
        "d_inference = d_p.sample(\n",
        "                         n = 1000000, ### TKTK - maybe\n",
        "                         random_state = 56,\n",
        "                         )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z7iSW5X3q5tR"
      },
      "source": [
        "**GPE: encoding, extraction**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M9oLft4QqvE6"
      },
      "outputs": [],
      "source": [
        "nlp = spacy.load('en_core_web_lg')\n",
        "\n",
        "# extract, count GPEs\n",
        "\n",
        "def extract_gpe(text):\n",
        "    doc = nlp(text)\n",
        "    gpes = [ent.text for ent in doc.ents if ent.label_ == 'GPE']\n",
        "    return gpes, len(gpes)\n",
        "\n",
        "d_inference[[\n",
        "             'gpe',\n",
        "             'gpe_count',\n",
        "             ]] = d_inference['text'].apply(lambda i: pd.Series(extract_gpe(i)))\n",
        "\n",
        "total_gpe_count = d_inference['gpe_count'].sum()\n",
        "\n",
        "print(f\"Total number of GPEs recognized: {total_gpe_count}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vgUzETyeygmL"
      },
      "outputs": [],
      "source": [
        "    ### SJS 8/17: validation w/ GTP-4o TKTK: is each GPE in the US?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rqhqDda4wXa1"
      },
      "source": [
        "**GPE: concordance**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SQ6rDGwvwWBQ"
      },
      "outputs": [],
      "source": [
        "# join single string for nltk entry\n",
        "\n",
        "all_texts = ' '.join(d_inference['text'].tolist())\n",
        "\n",
        "# tokenize for nltk\n",
        "\n",
        "tokens = nltk.word_tokenize(all_texts)\n",
        "nltk_text = Text(tokens)\n",
        "\n",
        "# concordance: GPEs in context\n",
        "\n",
        "for gpes in d_inference['gpe']:\n",
        "    for gpe in gpes:  # gpes = list of GPEs\n",
        "        print(f\"\\nConcordance for '{gpe}':\")\n",
        "        nltk_text.concordance(gpe)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i07lgnrTrBny"
      },
      "source": [
        "**Explicit suicidality: encoding**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gwyFFLqyqvLi"
      },
      "outputs": [],
      "source": [
        "regex = r'\\bsuicid\\S*'\n",
        "\n",
        "d_inference['sui'] = d_inference['text'].str.contains(\n",
        "                                                      regex,\n",
        "                                                      regex = True,\n",
        "                                                      ) | (d_inference['p_sbrt'] == 'SuicideWatch')\n",
        "\n",
        "\n",
        "d_inference['sui'] = d_inference['sui'].astype(int)\n",
        "\n",
        "d_inference.head(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ClXtlLxzyqci"
      },
      "outputs": [],
      "source": [
        "    ### SJS 8/17: validation w/ GTP-4o TKTK: is sui in reference to other people?\n",
        "\n",
        "#%pwd\n",
        "%cd ../inputs/data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MqAUEz2rqvXz"
      },
      "outputs": [],
      "source": [
        "d_inference.to_csv(\n",
        "                   'd_inference_TEST.csv',\n",
        "                   encoding = 'utf-8',\n",
        "                   index = False,\n",
        "                   header = True,\n",
        "                   )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rxzedr7a-kuK"
      },
      "source": [
        "### 4. Train-Adapt-Test\n",
        "Trains baseline BERT, RoBERTa, and DistilBERT, using rationale-augmented data $\\mathcal{d}$<sub>augmented</sub>, iterating over a.) strains, b.) explicit targeting, c.) implicit vulnerabilities. Evaluates using de-augmented data. Outputs model x target _$F$_<sub>1</sub> (macro) performance scores.\n",
        "***"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8Vh-AsKSAFuA"
      },
      "outputs": [],
      "source": [
        "#%pwd\n",
        "%cd ../inputs/data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "vVEyG5UoH-uN"
      },
      "outputs": [],
      "source": [
        "d_augmented = pd.read_excel(\n",
        "                            'd_augmented.xlsx',\n",
        "                            index_col = [0],\n",
        "                            )\n",
        "\n",
        "d_augmented.info()\n",
        "d_augmented.head(3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dkd3V_SzAZQR"
      },
      "source": [
        "**Condense for model entry**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HTNkgDwhAO_2"
      },
      "outputs": [],
      "source": [
        "targets = [\n",
        "           'asp',\n",
        "           'dep',\n",
        "           'val',\n",
        "           'prg',\n",
        "           'tgd',\n",
        "           'age',\n",
        "           'race',\n",
        "           'dbty',\n",
        "           ]\n",
        "\n",
        "d_augmented = d_augmented[\n",
        "                          ['text',\n",
        "                           'aug'] +\n",
        "                           targets\n",
        "                           ].copy()\n",
        "\n",
        "d_augmented[\n",
        "            ['aug'] +\n",
        "             targets\n",
        "             ].apply(pd.Series.value_counts)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZCRDXGEICOL7"
      },
      "source": [
        "**Compute weights ($w$): inverse class ($c$) freq<br>\n",
        "$w_c = N / (2 * n_c)$**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "loj3xA_kCNdY"
      },
      "outputs": [],
      "source": [
        "class_weights = {}\n",
        "\n",
        "for t in targets:\n",
        "\n",
        "    value_counts = d_augmented[t].value_counts()\n",
        "\n",
        "    w_pos = round(len(d_augmented) / (2 * value_counts.get(1, 0)), 4)\n",
        "    w_neg = round(len(d_augmented) / (2 * value_counts.get(0, 0)), 4)\n",
        "\n",
        "    class_weights[t] = {\n",
        "                        'w_pos': w_pos if w_pos != float('inf') else 0,\n",
        "                        'w_neg': w_neg if w_neg != float('inf') else 0,\n",
        "                        }\n",
        "\n",
        "class_weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I8K7x-loGZNW"
      },
      "source": [
        "#### Train and evaluate baseline models: $k$-fold cross validate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "IvjM9N1evElY"
      },
      "outputs": [],
      "source": [
        "%cd /content/drive/My Drive/Colab/bar_policy_suicidality/temp\n",
        "\n",
        "# define target-specific aug-stratified df\n",
        "\n",
        "target_datasets = iterative_stratified_train_test_split_with_rationales(\n",
        "                                                                        d_augmented,\n",
        "                                                                        targets,\n",
        "                                                                        random_state = 56,\n",
        "                                                                        test_size = 0.2,\n",
        "                                                                        )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G39CdbnVEHua"
      },
      "outputs": [],
      "source": [
        "print(target_datasets.keys())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "AR64py6weO2w"
      },
      "outputs": [],
      "source": [
        "# define target-specific df\n",
        "\n",
        "#target_datasets = {\n",
        "#                   'asp': (d_train_asp, d_test_asp),\n",
        "#                   'dep': (d_train_dep, d_test_dep),\n",
        "#                   'val': (d_train_val, d_test_val),\n",
        "#                   'prg': (d_train_prg, d_test_prg),\n",
        "#                   'tgd': (d_train_tgd, d_test_tgd),\n",
        "#                   'age': (d_train_age, d_test_age),\n",
        "#                   'race': (d_train_race, d_test_race),\n",
        "#                   'dbty': (d_train_dbty, d_test_dbty),\n",
        "#}\n",
        "\n",
        "# define targets + class weights\n",
        "\n",
        "targets_and_class_weights = {\n",
        "                             'asp': [\n",
        "                                     0.838, ### w_neg\n",
        "                                     1.2397, ### w_pos\n",
        "                                     ],\n",
        "                             'dep': [\n",
        "                                     0.5847,\n",
        "                                     3.4522,\n",
        "                                     ],\n",
        "                             'val': [\n",
        "                                     0.708,\n",
        "                                     1.7017,\n",
        "                                     ],\n",
        "                             'prg': [\n",
        "                                     0.5425,\n",
        "                                     6.385,\n",
        "                                     ],\n",
        "                            'tgd': [\n",
        "                                     0.5619,\n",
        "                                     4.5377,\n",
        "                                     ],\n",
        "                             'age': [\n",
        "                                     0.5926,\n",
        "                                     3.1996,\n",
        "                                     ],\n",
        "                             'race': [\n",
        "                                      0.506,\n",
        "                                      42.4412,\n",
        "                                      ],\n",
        "                             'dbty': [\n",
        "                                      0.5159,\n",
        "                                      16.2135,\n",
        "                                      ],\n",
        "                              }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FVlcQJqfdG0U"
      },
      "source": [
        "**BERT, RoBERTa, DistilBERT**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KC5Bv5XydDVh"
      },
      "outputs": [],
      "source": [
        "# define models\n",
        "\n",
        "models = {\n",
        "          'bert': (\n",
        "                   BertForSequenceClassification,\n",
        "                   BertTokenizer,\n",
        "                  'bert-base-uncased',\n",
        "                   ),\n",
        "\n",
        "          'roberta': (\n",
        "                      RobertaForSequenceClassification,\n",
        "                      RobertaTokenizer,\n",
        "                      'roberta-base',\n",
        "                      ),\n",
        "\n",
        "          'distilbert': (\n",
        "                         DistilBertForSequenceClassification,\n",
        "                         DistilBertTokenizer,\n",
        "                         'distilbert-base-uncased',\n",
        "                         ),\n",
        "        }\n",
        "\n",
        "# define save path\n",
        "\n",
        "save_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models/'\n",
        "\n",
        "# define hyperparameter grid\n",
        "\n",
        "hyperparameter_grid = {\n",
        "                       'batch_size': 8,\n",
        "                       'gradient_accumulation_steps': 2,\n",
        "                       'learning_rate': 2e-5,\n",
        "                       'num_epochs': 2,\n",
        "                       'warmup_steps': 0,\n",
        "                       'weight_decay': 0.00,\n",
        "                       }\n",
        "\n",
        "# set cycle\n",
        "\n",
        "cycle = 'baseline'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "-Og-lIzTygll"
      },
      "outputs": [],
      "source": [
        "%cd ../../outputs/tables\n",
        "\n",
        "# 'baseline' cycle: train-test loop\n",
        "\n",
        "train_eval_save_bl_models(\n",
        "                          target_datasets = target_datasets,\n",
        "                          targets_and_class_weights = targets_and_class_weights,\n",
        "                          models = models,\n",
        "                          save_path = save_path,\n",
        "                          cycle = cycle,\n",
        "                          hyperparameter_grid = hyperparameter_grid,\n",
        "                          )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gJESi1LsBsSR"
      },
      "source": [
        "**Llama 3.1**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qHTONcONQUWV"
      },
      "outputs": [],
      "source": [
        "# define save path\n",
        "\n",
        "save_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models/'\n",
        "\n",
        "# define hyperparameters\n",
        "\n",
        "hyperparameters = {\n",
        "                   'gradient_accumulation_steps': [1],\n",
        "                   'learning_rate': [1e-4],\n",
        "                   'num_train_epochs': [1],\n",
        "                   'warmup_steps': [0],\n",
        "                   'weight_decay': [0.01],\n",
        "                   }\n",
        "\n",
        "hyperparameter_grid = list(ParameterGrid(hyperparameters))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yxbrp2BLQUMJ"
      },
      "outputs": [],
      "source": [
        "%cd /content/drive/MyDrive/Colab/bar_policy_suicidality/temp\n",
        "#%cd ../../outputs/tables\n",
        "\n",
        "# llama train-test loop\n",
        "\n",
        "train_and_evaluate_llama(\n",
        "                         target_datasets = target_datasets,\n",
        "                         targets_and_class_weights = targets_and_class_weights,\n",
        "                         model_name = 'meta-llama/Llama-3.1-8B',\n",
        "                         hyperparameter_grid = hyperparameter_grid,\n",
        "                         save_path = save_path\n",
        "                         )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7piusGAtPZR9"
      },
      "source": [
        "**Performance scatterplots: baseline, adapted**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UDFAqUdVIMTp"
      },
      "outputs": [],
      "source": [
        "%cd /content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/tables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "84vf2OZ6JbhA"
      },
      "outputs": [],
      "source": [
        "fm.fontManager.addfont('/content/drive/MyDrive/Colab/Arial.ttf')\n",
        "plt.rcParams['font.family'] = 'Arial'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "yQxVho3DKKJC"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "\n",
        "d_v = pd.read_excel('d_baseline_performance.xlsx')\n",
        "d_v.round({'f1_macro': 4, 'mcc': 4, 'auprc': 4})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "IgOHSMq4L96Y"
      },
      "outputs": [],
      "source": [
        "%cd ../figures\n",
        "\n",
        "    ### SJS 1/2: _NOTE_ this is shit; fix TKTK, use bar_scratchpad for viz.\n",
        "\n",
        "performance_scatterplot(\n",
        "                        df = d_v,\n",
        "                        plot_name = 'baseline',\n",
        "                        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qkBdAMlV1gjG"
      },
      "source": [
        "#### Adapt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8ButJIYk6Otd"
      },
      "outputs": [],
      "source": [
        "%cd /content/drive/MyDrive/Colab/bar_policy_suicidality/inputs/data\n",
        "#%cd inputs/data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pB2WZzlO6AHX"
      },
      "outputs": [],
      "source": [
        "d_adapt = pd.read_csv('d_adapt.csv')\n",
        "\n",
        "d_adapt.info()\n",
        "d_adapt.head(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "0pa6qlHn4UjR"
      },
      "outputs": [],
      "source": [
        "d_adapt = d_adapt[[\n",
        "                   'text',\n",
        "                   'p_sbrt',\n",
        "                   ]].copy()\n",
        "\n",
        "d_adapt.shape\n",
        "d_adapt.head(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_hASR09T40YR"
      },
      "outputs": [],
      "source": [
        "%cd ../../outputs/models\n",
        "models_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models/'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bs2bRT3LJyiu"
      },
      "source": [
        "**Domain adaptation proxy task: subreddit clr**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "uFGYYiImO0Qr"
      },
      "outputs": [],
      "source": [
        "# prep dataset\n",
        "\n",
        "class CustomDataset(Dataset):\n",
        "    def __init__(self, df, tokenizer, max_length):\n",
        "        self.df = df\n",
        "        self.tokenizer = tokenizer\n",
        "        self.max_length = max_length\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.df)\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        text = self.df.iloc[index]['text']\n",
        "        label = self.df.iloc[index]['p_sbrt']\n",
        "\n",
        "        if pd.isna(text):\n",
        "            text = ' '  ### replaces NaN w/ empty string\n",
        "\n",
        "        encoding = self.tokenizer.encode_plus(\n",
        "                                              text,\n",
        "                                              add_special_tokens = True,\n",
        "                                              max_length = self.max_length,\n",
        "                                              return_token_type_ids = False,\n",
        "                                              padding = 'max_length',\n",
        "                                              truncation = True,\n",
        "                                              return_attention_mask = True,\n",
        "                                              return_tensors = 'pt'\n",
        "                                              )\n",
        "        return {\n",
        "                'input_ids': encoding['input_ids'].flatten(),\n",
        "                'attention_mask': encoding['attention_mask'].flatten(),\n",
        "                'label': torch.tensor(label, dtype = torch.long)\n",
        "                }\n",
        "\n",
        "# hyperparams\n",
        "\n",
        "epochs = 2\n",
        "batch_size = 16\n",
        "learning_rate = 2e-5\n",
        "max_length = 512\n",
        "\n",
        "# define models\n",
        "\n",
        "models_to_train = {\n",
        "                   #'BERT': {\n",
        "                   #         'model_class': BertForSequenceClassification,\n",
        "                   #         'tokenizer_class': BertTokenizer,\n",
        "                   #         'pretrained_model_name': 'bert-base-uncased',\n",
        "                   #         },\n",
        "                  #'RoBERTa': {\n",
        "                  #            'model_class': RobertaForSequenceClassification,\n",
        "                  #            'tokenizer_class': RobertaTokenizer,\n",
        "                  #            'pretrained_model_name': 'roberta-base',\n",
        "                  #          },\n",
        "                  'DistilBERT': {\n",
        "                                 'model_class': DistilBertForSequenceClassification,\n",
        "                                 'tokenizer_class': DistilBertTokenizer,\n",
        "                                 'pretrained_model_name': 'distilbert-base-uncased'\n",
        "                                 }\n",
        "                  }\n",
        "\n",
        "# encode labels\n",
        "\n",
        "label_encoder = LabelEncoder()\n",
        "d_adapt['p_sbrt'] = label_encoder.fit_transform(d_adapt['p_sbrt'])\n",
        "\n",
        "# iterate over models\n",
        "\n",
        "for model_name, model_info in models_to_train.items():\n",
        "    print(f'\\nTraining {model_name}...')\n",
        "    print(\"--------------------------------------------------------------------------------------\")\n",
        "\n",
        "    # initialize tokenizer, model\n",
        "\n",
        "    tokenizer = model_info['tokenizer_class'].from_pretrained(model_info['pretrained_model_name'])\n",
        "    model = model_info['model_class'].from_pretrained(\n",
        "                                                      model_info['pretrained_model_name'],\n",
        "                                                      num_labels = 3, ### update for true run - maps to tensor dimensions\n",
        "                                                      )\n",
        "\n",
        "    # prep dataset\n",
        "\n",
        "    dataset = CustomDataset(\n",
        "                            d_adapt,\n",
        "                            tokenizer,\n",
        "                            max_length = max_length,\n",
        "                            )\n",
        "    dataloader = DataLoader(\n",
        "                            dataset,\n",
        "                            batch_size = batch_size,\n",
        "                            shuffle = True,\n",
        "                            )\n",
        "\n",
        "    # config optimizer\n",
        "\n",
        "    optimizer = AdamW(\n",
        "                      model.parameters(), ### ensures all model params are trained\n",
        "                      lr = learning_rate,\n",
        "                      )\n",
        "\n",
        "    # training loop\n",
        "\n",
        "    model.train()\n",
        "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "    model.to(device)\n",
        "\n",
        "    for epoch in range(epochs):\n",
        "        print(f'Epoch {epoch+1}/{epochs}')\n",
        "        for batch in tqdm(dataloader):\n",
        "            optimizer.zero_grad()\n",
        "            input_ids = batch['input_ids'].to(device)\n",
        "            attention_mask = batch['attention_mask'].to(device)\n",
        "            labels = batch['label'].to(device)\n",
        "\n",
        "            outputs = model(\n",
        "                            input_ids = input_ids,\n",
        "                            attention_mask = attention_mask,\n",
        "                            labels = labels,\n",
        "                            )\n",
        "            loss = outputs.loss\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "    # save\n",
        "\n",
        "    save_path = f'{models_path}{model_name}_adapted'\n",
        "    model.save_pretrained(save_path)\n",
        "    tokenizer.save_pretrained(save_path)\n",
        "\n",
        "    print(f'{model_name} saved to {save_path}\\n')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7xoLLpCnlkpO"
      },
      "source": [
        "#### Train and evaluate domain-adapted models: $k$-fold cross validate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7-mwIf8dxRJ_"
      },
      "outputs": [],
      "source": [
        "targets_and_class_weights = {\n",
        "                              'asp': [\n",
        "                                     0.838, ### w_neg\n",
        "                                     1.2397, ### w_pos\n",
        "                                     ],\n",
        "                             'dep': [\n",
        "                                     0.5847,\n",
        "                                     3.4522,\n",
        "                                     ],\n",
        "                             'val': [\n",
        "                                     0.708,\n",
        "                                     1.7017,\n",
        "                                     ],\n",
        "                             'prg': [\n",
        "                                     0.5425,\n",
        "                                     6.385,\n",
        "                                     ],\n",
        "                            'tgd': [\n",
        "                                     0.5619,\n",
        "                                     4.5377,\n",
        "                                     ],\n",
        "                             'age': [\n",
        "                                     0.5926,\n",
        "                                     3.1996,\n",
        "                                     ],\n",
        "                             'race': [\n",
        "                                      0.506,\n",
        "                                      42.4412,\n",
        "                                      ],\n",
        "                             'dbty': [\n",
        "                                      0.5159,\n",
        "                                      16.2135,\n",
        "                                      ],\n",
        "                              }\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q_6DAaZzO0Wq"
      },
      "outputs": [],
      "source": [
        "targets_and_class_weights = targets_and_class_weights\n",
        "\n",
        "# define models path + load in-domain-adapted models\n",
        "\n",
        "models_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models/'\n",
        "\n",
        "models = {\n",
        "          'bert': (\n",
        "                   BertForSequenceClassification.from_pretrained(f'{models_path}BERT_adapted'),\n",
        "                   BertTokenizer.from_pretrained(f'{models_path}BERT_adapted'),\n",
        "                   'bert-base-uncased',\n",
        "                   ),\n",
        "\n",
        "          'roberta': (\n",
        "                      RobertaForSequenceClassification.from_pretrained(f'{models_path}RoBERTa_adapted'),\n",
        "                      RobertaTokenizer.from_pretrained(f'{models_path}RoBERTa_adapted'),\n",
        "                      'roberta-base',\n",
        "                      ),\n",
        "\n",
        "          'distilbert': (\n",
        "                         DistilBertForSequenceClassification.from_pretrained(f'{models_path}DistilBERT_adapted'),\n",
        "                         DistilBertTokenizer.from_pretrained(f'{models_path}DistilBERT_adapted'),\n",
        "                         'distilbert-base-uncased',\n",
        "                         )\n",
        "          }\n",
        "\n",
        "# define hyperparameter grid\n",
        "\n",
        "hyperparameter_grid = {\n",
        "                       'batch_size': 8,\n",
        "                       'gradient_accumulation_steps': 2,\n",
        "                       'learning_rate': 2e-5,\n",
        "                       'num_epochs': 2,\n",
        "                       'warmup_steps': 0,\n",
        "                       'weight_decay': 0.00,\n",
        "                       }\n",
        "\n",
        "# define save path\n",
        "\n",
        "save_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models/'\n",
        "\n",
        "# set cycle\n",
        "\n",
        "cycle = 'adapted'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v0PiGvMKxdIE"
      },
      "outputs": [],
      "source": [
        "%pwd"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wzyN3sIMmrGW"
      },
      "outputs": [],
      "source": [
        "#%cd ../../outputs/tables\n",
        "%cd /content/drive/My Drive/Colab/bar_policy_suicidality/outputs/tables\n",
        "\n",
        "# 'adapted' cycle: train-test loop\n",
        "\n",
        "#results = train_eval_save_bl_models(\n",
        "#                                    d_augmented,\n",
        "#                                    targets_and_class_weights,\n",
        "#                                    models,\n",
        "#                                    save_path,\n",
        "#                                    cycle,\n",
        "#                                    )\n",
        "\n",
        "train_eval_save_bl_models(\n",
        "                          target_datasets = target_datasets,\n",
        "                          targets_and_class_weights = targets_and_class_weights,\n",
        "                          models = models,\n",
        "                          save_path = save_path,\n",
        "                          cycle = cycle,\n",
        "                          hyperparameter_grid = hyperparameter_grid,\n",
        "                          )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GV0HVwTiRmyT"
      },
      "source": [
        "**Adapted: viz**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8Ye3la_pYYlC"
      },
      "outputs": [],
      "source": [
        "%cd ../outputs/tables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U7ilAKBkq2jS"
      },
      "outputs": [],
      "source": [
        "%cd /content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/tables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OvqxmMeFRb0M"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "\n",
        "d_v = pd.read_excel('d_adapted_performance.xlsx')\n",
        "d_v.round({'f1_macro': 4, 'mcc': 4, 'auprc': 4})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "79hIkMb1Rb6d"
      },
      "outputs": [],
      "source": [
        "%cd ../figures\n",
        "\n",
        "performance_barplot(\n",
        "                    d_v,\n",
        "                    'adapted_performance',\n",
        "                    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G614tXUn0VWw"
      },
      "source": [
        "### 5. Tune-Regularize\n",
        "Builds stratified train-test sets, searches hyperparam space to optimize highest-performing target x pretrained model configs. Evaluates learning rate, dropout, and self-training accuracy gain over 20 epochs.\n",
        "***"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2nH8FoS6d4sV"
      },
      "outputs": [],
      "source": [
        "%pwd"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZPmkMPkboW7Q"
      },
      "outputs": [],
      "source": [
        "%cd /content/drive/MyDrive/Colab/bar_policy_suicidality/inputs/data\n",
        "#%cd ../inputs/data\n",
        "\n",
        "d_augmented = pd.read_excel('d_augmented.xlsx')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OBVrIz_W8du0"
      },
      "outputs": [],
      "source": [
        "targets = [\n",
        "           'asp',\n",
        "           'dep',\n",
        "           'val',\n",
        "           'prg',\n",
        "           'tgd',\n",
        "           'age',\n",
        "           'race',\n",
        "           'dbty',\n",
        "           ]\n",
        "\n",
        "d_augmented = d_augmented[\n",
        "                          ['text',\n",
        "                           'aug'] +\n",
        "                           targets\n",
        "                           ].copy()\n",
        "\n",
        "d_augmented.info()\n",
        "d_augmented.head(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TkdI-9lG0Tfc"
      },
      "outputs": [],
      "source": [
        "%cd ../../temp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FKmj9QKr2Uf8"
      },
      "source": [
        "**Target-parsed $\\mathcal{d}$<sub>train</sub>($y$): augmented | $\\mathcal{d}$<sub>test</sub>($y$): de-augmented**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "IjHgWQk-7pAZ"
      },
      "outputs": [],
      "source": [
        "target_datasets = iterative_stratified_train_test_split_with_rationales(\n",
        "                                                                        d_augmented,\n",
        "                                                                        targets,\n",
        "                                                                        random_state = 56,\n",
        "                                                                        test_size = 0.2,\n",
        "                                                                        )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TQxFeCc9w9tO"
      },
      "outputs": [],
      "source": [
        "print(target_datasets.keys())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I_MirtIWQp0K"
      },
      "source": [
        "#### BERT, RoBERTa, DistilBERT"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "GJLRYddN7pKP"
      },
      "outputs": [],
      "source": [
        "# drop 'aug' + extract target-wise train/test df\n",
        "\n",
        "    ### SJS 12/3: old/bad way; validate new _for t in targets_ loop\n",
        "\n",
        "d_train_asp, d_test_asp = target_datasets['asp']\n",
        "d_train_dep, d_test_dep = target_datasets['dep']\n",
        "d_train_val, d_test_val = target_datasets['val']\n",
        "\n",
        "d_train_prg, d_test_prg = target_datasets['prg']\n",
        "d_train_tgd, d_test_tgd = target_datasets['tgd']\n",
        "\n",
        "d_train_age, d_test_age = target_datasets['age']\n",
        "d_train_race, d_test_race = target_datasets['race']\n",
        "d_train_dbty, d_test_dbty = target_datasets['dbty']\n",
        "\n",
        "# strn\n",
        "\n",
        "d_train_asp = d_train_asp.drop('aug', axis = 1)\n",
        "d_test_asp = d_test_asp.drop('aug', axis = 1)\n",
        "\n",
        "#d_train_asp.head(3)\n",
        "#d_test_asp.head(3)\n",
        "\n",
        "d_train_dep = d_train_dep.drop('aug', axis = 1)\n",
        "d_test_dep = d_test_dep.drop('aug', axis = 1)\n",
        "\n",
        "#d_train_dep.head(3)\n",
        "#d_test_dep.head(3)\n",
        "\n",
        "d_train_val = d_train_val.drop('aug', axis = 1)\n",
        "d_test_val = d_test_val.drop('aug', axis = 1)\n",
        "\n",
        "#d_train_val.head(3)\n",
        "#d_test_val.head(3)\n",
        "\n",
        "# traits\n",
        "\n",
        "d_train_prg = d_train_prg.drop('aug', axis = 1)\n",
        "d_test_prg = d_test_prg.drop('aug', axis = 1)\n",
        "\n",
        "#d_train_prg.head(3)\n",
        "#d_test_prg.head(3)\n",
        "\n",
        "d_train_tgd = d_train_tgd.drop('aug', axis = 1)\n",
        "d_test_tgd = d_test_tgd.drop('aug', axis = 1)\n",
        "\n",
        "#d_train_tgd.head(3)\n",
        "#d_test_tgd.head(3)\n",
        "\n",
        "d_train_age = d_train_age.drop('aug', axis = 1)\n",
        "d_test_age = d_test_age.drop('aug', axis = 1)\n",
        "\n",
        "#d_train_age.head(3)\n",
        "#d_test_age.head(3)\n",
        "\n",
        "d_train_race = d_train_race.drop('aug', axis = 1)\n",
        "d_test_race = d_test_race.drop('aug', axis = 1)\n",
        "\n",
        "#d_train_race.head(3)\n",
        "#d_test_race.head(3)\n",
        "\n",
        "d_train_dbty = d_train_dbty.drop('aug', axis = 1)\n",
        "d_test_dbty = d_test_dbty.drop('aug', axis = 1)\n",
        "\n",
        "#d_train_dbty.head(3)\n",
        "#d_test_dbty.head(3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YmwkvY6KLEX6"
      },
      "source": [
        "#### Grid search: RoBERTa"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "collapsed": true,
        "id": "8AYBEaro0TrZ"
      },
      "outputs": [],
      "source": [
        "%cd ../../outputs/tables\n",
        "\n",
        "# define BERT, RoBERTa, DistilBERT hyperparam grid\n",
        "\n",
        "hyperparameter_grid = {\n",
        "                       'batch_size': [8, 16],\n",
        "                       'gradient_accumulation_steps': [1, 2],\n",
        "                       'learning_rate': [2e-5, 3e-5],\n",
        "                       'num_epochs': [2, 3],\n",
        "                       'warmup_steps': [0, 500],\n",
        "                       'weight_decay': [0.0, 0.3],\n",
        "                       }\n",
        "\n",
        "models_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models'\n",
        "\n",
        "    ### SJS 9/8: no longer using dedicated subdirectories; just do save_path = models_path when calling the Fx...\n",
        "\n",
        "# define tuning param sets\n",
        "\n",
        "params = [\n",
        "\n",
        "    # asp: Best F1 (macro) for asp: 0.77 achieved by roberta - baseline\n",
        "\n",
        "          {\n",
        "           'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "           'model_class': RobertaForSequenceClassification,\n",
        "           'pretrained_model_name': 'roberta-base',\n",
        "           'd_train': d_train_asp,\n",
        "           'd_test': d_test_asp,\n",
        "           'target': 'asp',\n",
        "           'class_weights': torch.tensor([\n",
        "                                          0.838, ### w_neg\n",
        "                                          1.2397, ### w_pos\n",
        "                                          ], dtype = torch.float),\n",
        "           'save_path': models_path,\n",
        "           },\n",
        "\n",
        "    # dep: Best F1 (macro) for dep: 0.80 achieved by roberta - baseline\n",
        "\n",
        "          {\n",
        "           'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "           'model_class': RobertaForSequenceClassification,\n",
        "           'pretrained_model_name': 'roberta-base',\n",
        "           'd_train': d_train_dep,\n",
        "           'd_test': d_test_dep,\n",
        "           'target': 'dep',\n",
        "           'class_weights': torch.tensor([\n",
        "                                          0.5847,\n",
        "                                          3.4522,\n",
        "                                          ], dtype = torch.float),\n",
        "           'save_path': models_path,\n",
        "           },\n",
        "\n",
        "    # val: Best F1 (macro) for val: 0.75 achieved by roberta - baseline\n",
        "\n",
        "          {\n",
        "           'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "           'model_class': RobertaForSequenceClassification,\n",
        "           'pretrained_model_name': 'roberta-base',\n",
        "           'd_train': d_train_val,\n",
        "           'd_test': d_test_val,\n",
        "           'target': 'val',\n",
        "           'class_weights': torch.tensor([\n",
        "                                          0.708,\n",
        "                                          1.7017,\n",
        "                                          ], dtype = torch.float),\n",
        "           'save_path': models_path,\n",
        "           },\n",
        "     # prg: Best F1 (macro) for prg: 0.70 achieved by roberta - baseline\n",
        "\n",
        "          {\n",
        "           'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "           'model_class': RobertaForSequenceClassification,\n",
        "           'pretrained_model_name': 'roberta-base',\n",
        "           'd_train': d_train_prg,\n",
        "           'd_test': d_test_prg,\n",
        "           'target': 'prg',\n",
        "           'class_weights': torch.tensor([\n",
        "                                          0.5425,\n",
        "                                          6.385,\n",
        "                                          ], dtype = torch.float),\n",
        "           'save_path': models_path,\n",
        "           },\n",
        "\n",
        "    # tgd: Best F1 (macro) for tgd: 0.76 achieved by roberta - baseline\n",
        "\n",
        "          {\n",
        "           'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "           'model_class': RobertaForSequenceClassification,\n",
        "           'pretrained_model_name': 'roberta-base',\n",
        "           'd_train': d_train_tgd,\n",
        "           'd_test': d_test_tgd,\n",
        "           'target': 'tgd',\n",
        "           'class_weights': torch.tensor([\n",
        "                                          0.5619,\n",
        "                                          4.5377,\n",
        "                                          ], dtype = torch.float),\n",
        "           'save_path': models_path,\n",
        "           },\n",
        "\n",
        "        # age: Best F1 (macro) for age: 0.76 achieved by roberta - baseline\n",
        "\n",
        "          {\n",
        "           'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "           'model_class': RobertaForSequenceClassification,\n",
        "           'pretrained_model_name': 'roberta-base',\n",
        "           'd_train': d_train_age,\n",
        "           'd_test': d_test_age,\n",
        "           'target': 'age',\n",
        "           'class_weights': torch.tensor([\n",
        "                                          0.5926,\n",
        "                                          3.1996,\n",
        "                                          ], dtype = torch.float),\n",
        "           'save_path': models_path,\n",
        "           },\n",
        "\n",
        "        # race: Best F1 (macro) for race: 0.63 achieved by roberta - baseline\n",
        "\n",
        "          {\n",
        "           'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "           'model_class': RobertaForSequenceClassification,\n",
        "           'pretrained_model_name': 'roberta-base',\n",
        "           'd_train': d_train_race,\n",
        "           'd_test': d_test_race,\n",
        "           'target': 'race',\n",
        "           'class_weights': torch.tensor([\n",
        "                                          0.506,\n",
        "                                          42.4412,\n",
        "                                          ], dtype = torch.float),\n",
        "           'save_path': models_path,\n",
        "           },\n",
        "\n",
        "        # dbty: Best F1 (macro) for dbty: 0.66 achieved by roberta - baseline\n",
        "\n",
        "          {\n",
        "           'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "           'model_class': RobertaForSequenceClassification,\n",
        "           'pretrained_model_name': 'roberta-base',\n",
        "           'd_train': d_train_dbty,\n",
        "           'd_test': d_test_dbty,\n",
        "           'target': 'dbty',\n",
        "           'class_weights': torch.tensor([\n",
        "                                          0.5159,\n",
        "                                          16.2135,\n",
        "                                          ], dtype = torch.float),\n",
        "           'save_path': models_path,\n",
        "           },\n",
        "\n",
        "]\n",
        "\n",
        "\n",
        "# hyperparam tuning loop\n",
        "\n",
        "all_tuned_performance = pd.DataFrame()\n",
        "\n",
        "for p in params:\n",
        "    #print(f\"Inspecting parameter set: {p}\")  # inspect for dict format\n",
        "    d_test, d_tuned_performance = tune_and_optimize_model_hyperparams(\n",
        "                                                                      tokenizer = p['tokenizer'],\n",
        "                                                                      model_class = p['model_class'],\n",
        "                                                                      pretrained_model_name = p['pretrained_model_name'],\n",
        "                                                                      d_train = p['d_train'],\n",
        "                                                                      d_test = p['d_test'],\n",
        "                                                                      target = p['target'],\n",
        "                                                                      class_weights = p['class_weights'],\n",
        "                                                                      save_path = p['save_path'],\n",
        "                                                                      hyperparameter_grid = hyperparameter_grid,\n",
        "                                                                      )\n",
        "\n",
        "    all_tuned_performance = pd.concat([all_tuned_performance, d_tuned_performance], ignore_index = True)\n",
        "\n",
        "print(all_tuned_performance.head())\n",
        "all_tuned_performance.to_excel('all_tuned_performance.xlsx')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pr3qgvh5BkT2"
      },
      "outputs": [],
      "source": [
        "#print(all_tuned_performance.head())\n",
        "#all_tuned_performance.to_excel('all_tuned_performance.xlsx')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Thnxs8P6ic1l"
      },
      "source": [
        "#### Accuracy scoring (_post hoc_)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EWBe_1O0ib05"
      },
      "outputs": [],
      "source": [
        "# import all d_test_tuned_preds_{target}\n",
        "\n",
        "load_path = '/content/drive/My Drive/Colab/bar_policy_suicidality/outputs/tables'\n",
        "\n",
        "targets = [\n",
        "           'asp',\n",
        "           'dep',\n",
        "           'val',\n",
        "           'prg',\n",
        "           'tgd',\n",
        "           'age',\n",
        "           'race',\n",
        "           'dbty',\n",
        "           ]\n",
        "\n",
        "# create dict for all acc scores\n",
        "\n",
        "accuracy_scores = {}\n",
        "\n",
        "for t in targets:\n",
        "    file_path = os.path.join(load_path, f'd_test_tuned_preds_{t}.xlsx')\n",
        "    if os.path.exists(file_path):\n",
        "        df = pd.read_excel(file_path)\n",
        "\n",
        "        # define y_true and y_pred\n",
        "\n",
        "        y_true = df[t]\n",
        "        y_pred = df['predicted_labels']\n",
        "\n",
        "        # calculate and tabulate acc\n",
        "\n",
        "        accuracy = accuracy_score(y_true, y_pred)\n",
        "        accuracy_scores[t] = accuracy\n",
        "        print(f\"Accuracy for {t}: {accuracy:.4f}\")\n",
        "\n",
        "    else:\n",
        "        print(f\"File not found: {file_path}\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-uxf_Ffu68D_"
      },
      "source": [
        "**Viz.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nXUtSpX376sY"
      },
      "outputs": [],
      "source": [
        "# font = Arial\n",
        "\n",
        "fm.fontManager.addfont('/content/drive/MyDrive/Colab/Arial.ttf')\n",
        "plt.rcParams['font.family'] = 'Arial'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U4UMhiCwib3s"
      },
      "outputs": [],
      "source": [
        "# computed acc\n",
        "\n",
        "acc = {\n",
        "       'asp': 0.8429,\n",
        "       'dep': 0.9352,\n",
        "       'val': 0.8778,\n",
        "       'prg': 0.9601,\n",
        "       'tgd': 0.9352,\n",
        "       'age': 0.9027,\n",
        "       'race': 0.9850,\n",
        "       'dbty': 0.9576,\n",
        "       }\n",
        "\n",
        "# define custom colors for bars\n",
        "#bar_colors = [\n",
        "#    '#27aeef',  # for 'asp'\n",
        "#    '#27aeef',  # for 'dep'\n",
        "#    '#27aeef',  # for 'val'\n",
        "#    '#27aeef',  # for 'prg'\n",
        "#    '#27aeef',  # for 'tgd'\n",
        "#    '#27aeef',  # for 'age'\n",
        "#    '#27aeef',  # for 'race'\n",
        "#    '#27aeef'   # for 'dbty'\n",
        "#]\n",
        "\n",
        "# Map target: numeric position for x-axis\n",
        "#target_mapping = {\n",
        "#    'asp': 0,\n",
        "#    'dep': 1,\n",
        "#    'val': 2,\n",
        "#    'prg': 3,\n",
        "#    'tgd': 4,\n",
        "#    'age': 5,\n",
        "#    'race': 6,\n",
        "#    'dbty': 7\n",
        "#}\n",
        "\n",
        "target_mapping = {\n",
        "                  'asp': 0,\n",
        "                  'dep': 0.6,\n",
        "                  'val': 1.2,\n",
        "                  'prg': 1.8,\n",
        "                  'tgd': 2.4,\n",
        "                  'age': 3.0,\n",
        "                  'race': 3.6,\n",
        "                  'dbty': 4.2,\n",
        "                  }\n",
        "\n",
        "# extract values for barplot\n",
        "\n",
        "x_labels = list(target_mapping.keys())\n",
        "x_positions = list(target_mapping.values())\n",
        "y_values = [acc[target] for target in x_labels]\n",
        "\n",
        "# define broken y-axis\n",
        "\n",
        "fig = plt.figure(figsize = (12, 5.5))\n",
        "bax = brokenaxes(\n",
        "                 ylims = ((0, 0.1), (0.4, 1)),\n",
        "                 hspace = 0.05,\n",
        "                 )\n",
        "\n",
        "# barplot\n",
        "\n",
        "bax.bar(\n",
        "        x_positions,\n",
        "        y_values,\n",
        "        color = '#27aeef',\n",
        "        width = 0.4,\n",
        "        #edgecolor='black',\n",
        "        linewidth = 0.6,\n",
        "        alpha = 0.7,\n",
        "        )\n",
        "\n",
        "# set labels and ticks\n",
        "\n",
        "bax.set_xlabel(\n",
        "               'Target',\n",
        "               fontsize = 12,\n",
        "               labelpad = 30,\n",
        "               )\n",
        "\n",
        "bax.set_ylabel(\n",
        "               'Accuracy: tuned',\n",
        "               fontsize = 12,\n",
        "               labelpad = 30,\n",
        "               )\n",
        "\n",
        "\n",
        "bax.axs[1].set_xticks(x_positions)\n",
        "bax.axs[1].set_xticklabels(\n",
        "                           x_labels,\n",
        "                           rotation = 45,\n",
        "                           fontsize = 10,\n",
        "                           )\n",
        "\n",
        "# customize legend\n",
        "\n",
        "legend_elements = [\n",
        "    Line2D([0], [0], marker = 's', color = 'w', label = 'BERT', markersize = 8, markerfacecolor = '#87bc45', lw = 0),\n",
        "    Line2D([0], [0], marker = 's', color = 'w', label = 'RoBERTa', markersize = 8, markerfacecolor = '#27aeef', lw = 0),\n",
        "    Line2D([0], [0], marker = 's', color = 'w', label = 'DistilBERT', markersize = 8, markerfacecolor = '#b33dc6', lw = 0),\n",
        "]\n",
        "\n",
        "bax.axs[0].legend(\n",
        "                  handles = legend_elements,\n",
        "                  loc = 'upper center',\n",
        "                  bbox_to_anchor = (0.5, 1.15),\n",
        "                  ncol = 4,\n",
        "                  fontsize = 9,\n",
        "                  frameon = False,\n",
        "                  )\n",
        "\n",
        "# set y-axis tick size\n",
        "\n",
        "for ax in bax.axs:\n",
        "    ax.tick_params(\n",
        "                   axis = 'y',\n",
        "                   labelsize = 9,\n",
        "                   )\n",
        "\n",
        "\n",
        "# save\n",
        "\n",
        "plt.savefig(\n",
        "            'tuned_high_res_bar.png',\n",
        "            dpi = 300,\n",
        "            )\n",
        "\n",
        "plt.savefig(\n",
        "            'tuned_low_res_bar.png',\n",
        "            dpi = 100,\n",
        "            )\n",
        "\n",
        "# display\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-Dz-K7r7N1zx"
      },
      "outputs": [],
      "source": [
        "%pwd"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mm0hLOX24AYx"
      },
      "source": [
        "#### Self-training augmentation: RoBERTa"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WwAWtbL337-E",
        "collapsed": true
      },
      "outputs": [],
      "source": [
        "# import d_inference: labeled by hyperparam-tuned RoBERTa\n",
        "\n",
        "#%cd ../../outputs/tables\n",
        "%cd /content/drive/My Drive/Colab/bar_policy_suicidality/outputs/tables\n",
        "\n",
        "d_inference = pd.read_csv(\n",
        "                          'd_inference_pred_pilot.csv',\n",
        "                          index_col = [0],\n",
        "                          )\n",
        "\n",
        "# harmonize column names\n",
        "\n",
        "d_inference.rename(columns = {\n",
        "                              'asp_pred': 'asp',\n",
        "                              'dep_pred': 'dep',\n",
        "                              'val_pred': 'val',\n",
        "                              'prg_pred': 'prg',\n",
        "                              'tgd_pred': 'tgd',\n",
        "                              'age_pred': 'age',\n",
        "                              'race_pred': 'race',\n",
        "                              'dbty_pred': 'dbty',\n",
        "                              }, inplace = True,\n",
        "                   )\n",
        "\n",
        "# sort by pred_proba confidence\n",
        "\n",
        "targets = [\n",
        "           'asp',\n",
        "           'dep',\n",
        "           'val',\n",
        "           'prg',\n",
        "           'tgd',\n",
        "           'age',\n",
        "           'race',\n",
        "           'dbty',\n",
        "           ]\n",
        "\n",
        "for t in targets:\n",
        "    d_inference[f'{t}_prob'] = d_inference[f'{t}_prob'].apply(lambda i: ast.literal_eval(i))\n",
        "    d_inference[f'{t}_pos'] = d_inference[f'{t}_prob'].apply(lambda i: round(i[1], 4))\n",
        "    d_inference[f'{t}_neg'] = d_inference[f'{t}_prob'].apply(lambda i: round(i[0], 4))\n",
        "\n",
        "# parse self-training columns\n",
        "\n",
        "self_train_cols = ['text'] + [\n",
        "                              f'{t}' for t in targets\n",
        "                              ] + [\n",
        "                              f'{t}_pos' for t in targets\n",
        "                              ] + [\n",
        "                              f'{t}_neg' for t in targets\n",
        "                              ]\n",
        "\n",
        "d_inference = d_inference[self_train_cols]\n",
        "\n",
        "# inspect\n",
        "\n",
        "d_inference.info()\n",
        "d_inference.head(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_NoREXqLWNLJ"
      },
      "outputs": [],
      "source": [
        "targets = [\n",
        "           'asp',\n",
        "           'dep',\n",
        "           'val',\n",
        "           'prg',\n",
        "           'tgd',\n",
        "           'age',\n",
        "           'race',\n",
        "           'dbty',\n",
        "           ]\n",
        "\n",
        "# create d_self_train dataframes\n",
        "\n",
        "d_self_train = {}\n",
        "\n",
        "# n d_inference rows to sample\n",
        "\n",
        "n = 100\n",
        "\n",
        "for t in targets:\n",
        "\n",
        "    # create target-specific df for highest-confidence positive labels\n",
        "\n",
        "    d_self_train[f'{t}_pos'] = d_inference.sort_values(\n",
        "                                                       by = f'{t}_pos',\n",
        "                                                       ascending = False,\n",
        "                                                       ).head(n)[[\n",
        "                                                                  'text',\n",
        "                                                                  t,\n",
        "                                                                  f'{t}_pos',\n",
        "                                                                  ]].copy()\n",
        "\n",
        "    # create target-specific df for highest-confidence negative labels\n",
        "\n",
        "    d_self_train[f'{t}_neg'] = d_inference.sort_values(\n",
        "                                                       by = f'{t}_neg',\n",
        "                                                       ascending = False,\n",
        "                                                       ).head(n)[[\n",
        "                                                                  'text',\n",
        "                                                                  t,\n",
        "                                                                  f'{t}_neg',\n",
        "                                                                  ]].copy()\n",
        "\n",
        "# access target-specific df: d_self_train[f'{target}_pos'] or d_self_train[f'{target}_neg']\n",
        "\n",
        "d_self_train['asp_pos'].head() ### sense check"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aD1r72MlV-1f",
        "collapsed": true
      },
      "outputs": [],
      "source": [
        "# create interleaved target-specific df\n",
        "\n",
        "d_self_train_combined = {}\n",
        "\n",
        "for t in targets:\n",
        "\n",
        "    df_pos = d_self_train[f'{t}_pos']\n",
        "    df_neg = d_self_train[f'{t}_neg']\n",
        "\n",
        "    #df_pos = d_self_train_combined[f'{t}_pos']\n",
        "    #df_neg = d_self_train_combined[f'{t}_neg']\n",
        "\n",
        "    df_pos = df_pos.reset_index(drop = True)\n",
        "    df_neg = df_neg.reset_index(drop = True)\n",
        "\n",
        "    # interleave rows from pos and neg dataframes\n",
        "\n",
        "    min_length = min(len(df_pos), len(df_neg)) ### safeguard against IndexError\n",
        "    interleaved_rows = [\n",
        "        row for pair in zip(df_pos.iloc[:min_length].to_dict('records'), df_neg.iloc[:min_length].to_dict('records'))\n",
        "        for row in pair\n",
        "    ]\n",
        "\n",
        "    # convert interleaved rows to df\n",
        "\n",
        "    d_self_train_combined[t] = pd.DataFrame(interleaved_rows)\n",
        "\n",
        "    # append remaining rows (if lengths differ)\n",
        "\n",
        "    if len(df_pos) > min_length:\n",
        "        d_self_train_combined[t] = pd.concat([d_self_train_combined[t], df_pos.iloc[min_length:]], ignore_index = True)\n",
        "    elif len(df_neg) > min_length:\n",
        "        d_self_train_combined[t] = pd.concat([d_self_train_combined[t], df_neg.iloc[min_length:]], ignore_index = True)\n",
        "\n",
        "    ### SJS 12/12: note logits retained to this checkpoint for sense check; drop before save\n",
        "\n",
        "# access target-specific df via d_self_train_combined[target]\n",
        "\n",
        "d_self_train_combined['age'] ### age: intuitive example"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "di1gw3gl373p"
      },
      "outputs": [],
      "source": [
        "%cd ../../temp\n",
        "\n",
        "# drop logit columns, shuffle, save\n",
        "\n",
        "for target, df in d_self_train_combined.items():\n",
        "    df.drop(\n",
        "            [f'{target}_pos', f'{target}_neg'],\n",
        "            axis = 1,\n",
        "            inplace = True,\n",
        "            )\n",
        "    df = shuffle(\n",
        "                 df,\n",
        "                 random_state = 56,\n",
        "                 )\n",
        "    df.to_excel(f'd_self_train_{target}.xlsx', index = False)\n",
        "    print(f\"Saved: d_self_train_{target}.xlsx\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XMIPvK196VI5"
      },
      "outputs": [],
      "source": [
        "# inspect: pre-append\n",
        "\n",
        "print(\"d_train_{target}\")\n",
        "print(d_train_asp.shape)\n",
        "\n",
        "d_train_asp.head(3)\n",
        "d_train_asp.tail(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yzFa6Wxr46G7"
      },
      "outputs": [],
      "source": [
        "# append d_self_train_{target} self-training to d_train_{target} training data\n",
        "\n",
        "targets = [\n",
        "           'asp',\n",
        "           'dep',\n",
        "           'val',\n",
        "           'prg',\n",
        "           'tgd',\n",
        "           'age',\n",
        "           'race',\n",
        "           'dbty',\n",
        "           ]\n",
        "\n",
        "    ### SJS 1/6: add 100 for cycle H\n",
        "\n",
        "for t in targets:\n",
        "\n",
        "  # access first _n_ rows of interleaved target-specific dataframes\n",
        "\n",
        "  d_self_subset = d_self_train_combined[t].head(100) ### appending 100 pseudo-labeled rows\n",
        "\n",
        "  #print(d_self_subset.head(20))\n",
        "\n",
        "  d_train_var = f'd_train_{t}'\n",
        "  globals()[d_train_var] = pd.concat([\n",
        "                                      globals()[d_train_var],\n",
        "                                      d_self_subset\n",
        "                                      ], ignore_index = True,\n",
        "                                     )\n",
        "\n",
        "  # shuffle\n",
        "\n",
        "  globals()[d_train_var] = globals()[d_train_var].sample(\n",
        "                                                         frac = 1,\n",
        "                                                         random_state = 56,\n",
        "                                                         ).reset_index(drop = True)\n",
        "\n",
        "  print(f\"Updated and shuffled {d_train_var}\")\n",
        "\n",
        "# inspect: post-append\n",
        "\n",
        "    ### SJS 12/14: asp as ex for now...\n",
        "\n",
        "print(d_train_asp.shape)\n",
        "\n",
        "d_train_asp.head(3)\n",
        "d_train_asp.tail(3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dn3eOAiALPY1"
      },
      "source": [
        "#### Training and validation loss: RoBERTa"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nxJ7FUFfjbSn"
      },
      "outputs": [],
      "source": [
        "%pwd"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yyRTnUrGl1EW"
      },
      "outputs": [],
      "source": [
        "# set cycle\n",
        "\n",
        "cycle = 'I'\n",
        "\n",
        "# define models_path\n",
        "\n",
        "models_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models'\n",
        "\n",
        "# define target x model fixed hyperparams, num_epochs range\n",
        "\n",
        "params = [\n",
        "\n",
        "          # best 'asp' tuned f1 = 0.8094\n",
        "\n",
        "          {\n",
        "          'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "          'model_class': RobertaForSequenceClassification,\n",
        "          'pretrained_model_name': 'roberta-base',\n",
        "          'd_train': d_train_asp,\n",
        "          'd_test': d_test_asp,\n",
        "          'target': 'asp',\n",
        "          'class_weights': torch.tensor([\n",
        "                                        0.838,  ### w_neg\n",
        "                                        1.2397, ### w_pos\n",
        "                                        ], dtype = torch.float),\n",
        "          'cycle': cycle,\n",
        "          'save_path': models_path,\n",
        "          'fixed_hyperparameters': {\n",
        "                                    'batch_size': 16, ### optimized by grid search: enter d_tuned_performance values\n",
        "                                    'dropout': 0.1, ### new since grid search\n",
        "                                    'gradient_accumulation_steps': 2,\n",
        "                                    'learning_rate': 1e-5,\n",
        "                                    'warmup_steps': 0,\n",
        "                                    'weight_decay': 0.3,\n",
        "                                    },\n",
        "          'num_epochs_range': [1,20],\n",
        "         },\n",
        "\n",
        "        # best 'dep' tuned f1 = 0.8515\n",
        "\n",
        "        #{\n",
        "        #  'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "        #  'model_class': RobertaForSequenceClassification,\n",
        "        #  'pretrained_model_name': 'roberta-base',\n",
        "        #  'd_train': d_train_dep,\n",
        "        #  'd_test': d_test_dep,\n",
        "        #  'target': 'dep',\n",
        "        #  'class_weights': torch.tensor([\n",
        "        #                                0.5847,\n",
        "        #                                3.4522,\n",
        "        #                                ], dtype = torch.float),\n",
        "        #  'cycle': cycle,\n",
        "        #  'save_path': models_path,\n",
        "        #  'fixed_hyperparameters': {\n",
        "        #                            'batch_size': 8,\n",
        "        #                            'dropout': 0.1,\n",
        "        #                            'gradient_accumulation_steps': 1,\n",
        "        #                            'learning_rate': 1e-5,\n",
        "        #                            'warmup_steps': 500,\n",
        "        #                            'weight_decay': 0.3,\n",
        "        #                            },\n",
        "        #  'num_epochs_range': [1,20],\n",
        "      #},\n",
        "\n",
        "      # best 'val' tuned f1 = 0.8179\n",
        "\n",
        "        {\n",
        "          'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "          'model_class': RobertaForSequenceClassification,\n",
        "          'pretrained_model_name': 'roberta-base',\n",
        "          'd_train': d_train_val,\n",
        "          'd_test': d_test_val,\n",
        "          'target': 'val',\n",
        "          'class_weights': torch.tensor([\n",
        "                                        0.708,\n",
        "                                        1.7017,\n",
        "                                        ], dtype = torch.float),\n",
        "          'cycle': cycle,\n",
        "          'save_path': models_path,\n",
        "          'fixed_hyperparameters': {\n",
        "                                    'batch_size': 16,\n",
        "                                    'dropout': 0.1,\n",
        "                                    'gradient_accumulation_steps': 2,\n",
        "                                    'learning_rate': 1e-5,\n",
        "                                    'warmup_steps': 0,\n",
        "                                    'weight_decay': 0.3,\n",
        "                                    },\n",
        "          'num_epochs_range': [1,20],\n",
        "      },\n",
        "\n",
        "      # best 'prg' tuned f1 = 0.8355\n",
        "\n",
        "      #{\n",
        "      #    'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "      #    'model_class': RobertaForSequenceClassification,\n",
        "      #    'pretrained_model_name': 'roberta-base',\n",
        "      #    'd_train': d_train_prg,\n",
        "      #    'd_test': d_test_prg,\n",
        "      #    'target': 'prg',\n",
        "      #    'class_weights': torch.tensor([\n",
        "      #                                  0.5425,\n",
        "      #                                  6.385,\n",
        "      #                                  ], dtype = torch.float),\n",
        "      #    'cycle': cycle,\n",
        "      #    'save_path': models_path,\n",
        "      #    'fixed_hyperparameters': {\n",
        "      #                              'batch_size': 16,\n",
        "      #                              'dropout': 0.1,\n",
        "      #                              'gradient_accumulation_steps': 2,\n",
        "      #                              'learning_rate': 1e-5,\n",
        "      #                              'warmup_steps': 0,\n",
        "      #                              'weight_decay': 0,\n",
        "      #                              },\n",
        "      #    'num_epochs_range': [1,20],\n",
        "      #},\n",
        "\n",
        "      # best 'tgd' tuned f1 = 0.8433\n",
        "\n",
        "      #{\n",
        "      #    'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "      #    'model_class': RobertaForSequenceClassification,\n",
        "      #    'pretrained_model_name': 'roberta-base',\n",
        "      #    'd_train': d_train_tgd,\n",
        "      #    'd_test': d_test_tgd,\n",
        "      #    'target': 'tgd',\n",
        "      #    'class_weights': torch.tensor([\n",
        "      #                                  0.5619,\n",
        "      #                                  4.5377,\n",
        "      #                                  ], dtype = torch.float),\n",
        "      #    'cycle': cycle,\n",
        "      #    'save_path': models_path,\n",
        "      #    'fixed_hyperparameters': {\n",
        "      #                              'batch_size': 8,\n",
        "      #                              'dropout': 0.1,\n",
        "      #                              'gradient_accumulation_steps': 1,\n",
        "      #                              'learning_rate': 1e-5,\n",
        "      #                              'warmup_steps': 0,\n",
        "      #                              'weight_decay': 0.3,\n",
        "      #                              },\n",
        "      #    'num_epochs_range': [1,20],\n",
        "      #},\n",
        "\n",
        "      # best 'age' tuned f1 = 0.8303\n",
        "\n",
        "      #{\n",
        "      #    'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "      #    'model_class': RobertaForSequenceClassification,\n",
        "      #    'pretrained_model_name': 'roberta-base',\n",
        "      #    'd_train': d_train_age,\n",
        "      #    'd_test': d_test_age,\n",
        "      #    'target': 'age',\n",
        "      #    'class_weights': torch.tensor([\n",
        "      #                                  0.5926,\n",
        "      #                                  3.1996,\n",
        "      #                                  ], dtype = torch.float),\n",
        "      #    'cycle': cycle,\n",
        "      #    'save_path': models_path,\n",
        "      #    'fixed_hyperparameters': {\n",
        "      #                              'batch_size': 8,\n",
        "      #                              'dropout': 0.1,\n",
        "      #                              'gradient_accumulation_steps': 2,\n",
        "      #                              'learning_rate': 1e-5,\n",
        "      #                              'warmup_steps': 0,\n",
        "      #                              'weight_decay': 0.3,\n",
        "      #                              },\n",
        "      #    'num_epochs_range': [1,20],\n",
        "      #},\n",
        "\n",
        "      # best 'race' tuned f1 = 0.7462\n",
        "\n",
        "      #{\n",
        "      #    'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "      #    'model_class': RobertaForSequenceClassification,\n",
        "      #    'pretrained_model_name': 'roberta-base',\n",
        "      #    'd_train': d_train_race,\n",
        "      #    'd_test': d_test_race,\n",
        "      #    'target': 'race',\n",
        "      #    'class_weights': torch.tensor([\n",
        "      #                                  0.506,\n",
        "      #                                  42.4412,\n",
        "      #                                  ], dtype = torch.float),\n",
        "      #    'cycle': cycle,\n",
        "      #    'save_path': models_path,\n",
        "      #    'fixed_hyperparameters': {\n",
        "      #                              'batch_size': 8,\n",
        "      #                              'dropout': 0.1,\n",
        "      #                              'gradient_accumulation_steps': 2,\n",
        "      #                              'learning_rate': 1e-5,\n",
        "      #                              'warmup_steps': 500,\n",
        "      #                              'weight_decay': 0.3,\n",
        "      #                              },\n",
        "      #    'num_epochs_range': [1,20],\n",
        "      #},\n",
        "\n",
        "      # best 'dbty' tuned f1 = xx\n",
        "\n",
        "      #{\n",
        "      #    'tokenizer': RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "      #    'model_class': RobertaForSequenceClassification,\n",
        "      #    'pretrained_model_name': 'roberta-base',\n",
        "      #    'd_train': d_train_dbty,\n",
        "      #    'd_test': d_test_dbty,\n",
        "      #    'target': 'dbty',\n",
        "      #    'class_weights': torch.tensor([\n",
        "      #                                  0.5159,\n",
        "      #                                  16.2135,\n",
        "      #                                  ], dtype = torch.float),\n",
        "      #    'cycle': cycle,\n",
        "      #    'save_path': models_path,\n",
        "      #    'fixed_hyperparameters': {\n",
        "      #                              'batch_size': 8,\n",
        "      #                              'dropout': 0.1,\n",
        "      #                              'gradient_accumulation_steps': 2,\n",
        "      #                              'learning_rate': 1e-5,\n",
        "      #                              'warmup_steps': 0,\n",
        "      #                              'weight_decay': 0,\n",
        "      #                              },\n",
        "      #    'num_epochs_range': [1,20],\n",
        "      #},\n",
        "\n",
        "]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V-Lc3RLP-yLL"
      },
      "outputs": [],
      "source": [
        "# training/validation loss and accuracy loop\n",
        "\n",
        "for p in params:\n",
        "    d_epochal_performance = tune_and_optimize_model_loss_accuracy(\n",
        "                                                                  tokenizer = p['tokenizer'],\n",
        "                                                                  model_class = p['model_class'],\n",
        "                                                                  pretrained_model_name = p['pretrained_model_name'],\n",
        "                                                                  d_train = p['d_train'],\n",
        "                                                                  d_test = p['d_test'],\n",
        "                                                                  target = p['target'],\n",
        "                                                                  class_weights = p['class_weights'],\n",
        "                                                                  save_path = p['save_path'],\n",
        "                                                                  cycle = p['cycle'],\n",
        "                                                                  fixed_hyperparameters = p['fixed_hyperparameters'],\n",
        "                                                                  num_epochs_range = p['num_epochs_range'],\n",
        "                                                                  )\n",
        "\n",
        "    d_epochal_performance.head(6)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AZ6gjBzkQytt"
      },
      "source": [
        "#### Llama 3.1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u7y16K8KQwTH"
      },
      "outputs": [],
      "source": [
        "# define targets + class weights\n",
        "\n",
        "targets_and_class_weights = {\n",
        "#                             'asp': [\n",
        "#                                     0.838, ### w_neg\n",
        "#                                     1.2397, ### w_pos\n",
        "#                                     ],\n",
        "#                             'dep': [\n",
        "#                                     0.5847,\n",
        "#                                     3.4522,\n",
        "#                                     ],\n",
        "#                             'val': [\n",
        "#                                     0.708,\n",
        "#                                     1.7017,\n",
        "#                                     ],\n",
        "#                             'prg': [\n",
        "#                                     0.5425,\n",
        "#                                     6.385,\n",
        "#                                     ],\n",
        "#                            'tgd': [\n",
        "#                                     0.5619,\n",
        "#                                     4.5377,\n",
        "#                                     ],\n",
        "                             'age': [\n",
        "                                     0.5926,\n",
        "                                     3.1996,\n",
        "                                     ],\n",
        "#                             'race': [\n",
        "#                                      0.506,\n",
        "#                                      42.4412,\n",
        "#                                      ],\n",
        "#                             'dbty': [\n",
        "#                                      0.5159,\n",
        "#                                      16.2135,\n",
        "#                                      ],\n",
        "                              }\n",
        "\n",
        "\n",
        "# define Llama 3.1 hypereparam grid\n",
        "\n",
        "hyperparams = {\n",
        "               'gradient_accumulation_steps': [1,\n",
        "                                               2\n",
        "                                               ],\n",
        "               'learning_rate': [1e-4,\n",
        "                                 5e-5\n",
        "                                 ],\n",
        "               'num_train_epochs': [1,\n",
        "                                    2\n",
        "                                    ],\n",
        "               'warmup_steps': [0,\n",
        "                                500\n",
        "                                ],\n",
        "               'weight_decay': [0.0,\n",
        "                                0.01\n",
        "                                ],\n",
        "               }\n",
        "\n",
        "models_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models'\n",
        "\n",
        "hyperparameter_grid = list(ParameterGrid(hyperparams))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jHSd8Y4GR3kv"
      },
      "outputs": [],
      "source": [
        "tune_and_optimize_llama_hyperparams(\n",
        "                                    target_datasets = target_datasets,\n",
        "                                    targets_and_class_weights = targets_and_class_weights,\n",
        "                                    model_name = 'meta-llama/Llama-3.1-8b',\n",
        "                                    hyperparameter_grid = hyperparameter_grid,\n",
        "                                    save_path = models_path,\n",
        "                                    )\n",
        "\n",
        "    ### SJS 12/2: target-by-target grid search, manually reset 'targets' etc param and rename deposited df\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fXo9N041m7ju"
      },
      "source": [
        "**Performance scatterplots: grid search**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qCIglYqFoXb1"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "\n",
        "%cd /content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/tables\n",
        "\n",
        "d_v = pd.read_excel('all_tuned_performance.xlsx')\n",
        "\n",
        "d_v.rename(columns = {'pretrained_model_name': 'model'}, inplace = True)\n",
        "\n",
        "d_v.round({'f1_macro': 4, 'mcc': 4, 'auprc': 4})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yTSM7B2km6_Z"
      },
      "outputs": [],
      "source": [
        "# aesthetics\n",
        "\n",
        "### SJS 10/1: last three colors in \"Retro Metro (Default)\" https://www.heavy.ai/blog/12-color-palettes-for-telling-better-stories-with-your-data\n",
        "\n",
        "plot_name = 'tuned'\n",
        "\n",
        "# aesthetics\n",
        "\n",
        "sns.set_style(\n",
        "              style = 'whitegrid',\n",
        "              rc = None,\n",
        "              )\n",
        "\n",
        "\n",
        "model_colors = {\n",
        "                'bert': '#87bc45',\n",
        "                'roberta-base': '#27aeef',\n",
        "                'distilbert': '#b33dc6',\n",
        "                }\n",
        "\n",
        "target_mapping = {\n",
        "                  'asp': 0,\n",
        "                  'dep': 2,\n",
        "                  'val': 4,\n",
        "                  'prg': 6,\n",
        "                  'tgd': 8,\n",
        "                  'age': 10,\n",
        "                  'race': 12,\n",
        "                  'dbty': 14\n",
        "                  }\n",
        "\n",
        "d_v['target_numeric'] = d_v['target'].map(target_mapping)\n",
        "d_v['target_numeric'] = pd.to_numeric(d_v['target_numeric'])\n",
        "\n",
        "d_v['target_jitter'] = d_v['target_numeric'] + np.random.uniform(-0.35, 0.35, size = len(d_v))\n",
        "\n",
        "plt.figure(figsize = (12, 5.5))\n",
        "\n",
        "\n",
        "# initialize fig. with broken y-axis\n",
        "\n",
        "fig = plt.figure(figsize = (12, 5.5))\n",
        "bax = brokenaxes(\n",
        "                 ylims = ((0, 0.1), (0.4, 1)),\n",
        "                 hspace = 0.1\n",
        "                 )\n",
        "\n",
        "# scatterplot: categorical x model\n",
        "\n",
        "for model, color in model_colors.items():\n",
        "    model_data = d_v[d_v['model'] == model]\n",
        "    bax.scatter(\n",
        "                model_data['target_jitter'],\n",
        "                model_data['f1_score'],\n",
        "                marker = '.',\n",
        "                color = color,\n",
        "                s = 40,\n",
        "                alpha = 0.6,\n",
        "                label=None,\n",
        "                )\n",
        "\n",
        "# mean (SD) f1 for each target x model\n",
        "\n",
        "mean_std_df = d_v.groupby(['target', 'model']).agg(\n",
        "                                                   mean_f1_score = ('f1_score', 'mean'),\n",
        "                                                   std_f1_score = ('f1_score', 'std'),\n",
        "                                                   ).reset_index()\n",
        "\n",
        "mean_std_df['target_numeric'] = mean_std_df['target'].map(target_mapping).astype(float)\n",
        "mean_std_df['target_offset'] = mean_std_df['target_numeric'] + mean_std_df['model'].map(\n",
        "    {'bert': -0.3, 'roberta-base': 0.0, 'distilbert': 0.3}\n",
        ")\n",
        "\n",
        "\n",
        "# mean (SD) f1 - plot error bars\n",
        "\n",
        "for model in mean_std_df['model'].unique():\n",
        "    model_data = mean_std_df[mean_std_df['model'] == model]\n",
        "\n",
        "    if not model_data[['target_offset', 'mean_f1_score', 'std_f1_score']].isnull().any().any():\n",
        "        bax.errorbar(\n",
        "                     model_data['target_offset'],\n",
        "                     model_data['mean_f1_score'],\n",
        "                     yerr = model_data['std_f1_score'],\n",
        "                     fmt = 'D',\n",
        "                     markersize = 7,\n",
        "                     capsize = 0,\n",
        "                     elinewidth = 1,\n",
        "                     markeredgewidth = 1,\n",
        "                     color = model_colors[model],\n",
        "                     )\n",
        "\n",
        "bax.set_xlabel(\n",
        "               'Target',\n",
        "               fontsize = 12,\n",
        "               labelpad = 30,\n",
        "               )\n",
        "\n",
        "bax.set_ylabel(\n",
        "               f'$F_1$ (macro): {plot_name}',\n",
        "               fontsize = 12,\n",
        "               labelpad = 30,\n",
        "               )\n",
        "\n",
        "\n",
        "# x-ticks: lower subplot of broken axis\n",
        "\n",
        "bax.axs[1].set_xticks(list(target_mapping.values()))\n",
        "bax.axs[1].set_xticklabels(\n",
        "                           list(target_mapping.keys()),\n",
        "                           rotation = 45,\n",
        "                           fontsize = 10,\n",
        "                           )\n",
        "\n",
        "#sns.despine(left = True)\n",
        "bax.grid(\n",
        "         #axis='x',\n",
        "         False,\n",
        "         )\n",
        "\n",
        "# f1 = 0.8 marker\n",
        "\n",
        "bax.axhline(\n",
        "            y = 0.8,\n",
        "            color = 'r',\n",
        "            linewidth = 0.6,\n",
        "            linestyle = '--',\n",
        "            )\n",
        "\n",
        "# custom legend\n",
        "\n",
        "from matplotlib.lines import Line2D\n",
        "\n",
        "legend_elements = [\n",
        "                   Line2D([0], [0], marker = 'o', color = 'w', label = 'BERT', markersize = 8, markerfacecolor = '#87bc45', lw = 0),\n",
        "                   Line2D([0], [0], marker = 'o', color = 'w', label = 'RoBERTa', markersize = 8, markerfacecolor = '#27aeef', lw = 0),\n",
        "                   Line2D([0], [0], marker = 'o', color = 'w', label = 'DistilBERT', markersize = 8, markerfacecolor = '#b33dc6', lw = 0),\n",
        "                   ]\n",
        "\n",
        "bax.axs[0].legend(\n",
        "                  handles = legend_elements,\n",
        "                  loc = 'upper center',\n",
        "                  bbox_to_anchor = (0.5, 1.15),\n",
        "                  ncol = 4,\n",
        "                  fontsize = 9,\n",
        "                  frameon = False,\n",
        "                  )\n",
        "\n",
        "# save\n",
        "\n",
        "plt.savefig(f'{plot_name}_scatter.png')\n",
        "\n",
        "# display\n",
        "\n",
        "plt.show()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xnP3hWY2bGlp"
      },
      "source": [
        "### 6. Infer (prelim)\n",
        "Labels $n$<sub>posts</sub> = 10K subset of $\\mathcal{d}$<sub>adapt</sub> for prediction and model explainability, preliminary human and GPT-4o cross-validation.\n",
        "***"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AAeXGG-n0TuR"
      },
      "outputs": [],
      "source": [
        "### SJS 12/2: pilot inference: RoBERTa\n",
        "\n",
        "%cd /content/drive/MyDrive/Colab/bar_policy_suicidality/inputs/data\n",
        "#%cd inputs/data\n",
        "\n",
        "d_adapt = pd.read_csv('d_adapt.csv')\n",
        "\n",
        "# delete empty/NaN 'text' cells\n",
        "\n",
        "#d_adapt = d_adapt[d_adapt.text != ' ']\n",
        "\n",
        "d_adapt = d_adapt[d_adapt['text'].notnull() & (d_adapt['text'].str.strip() != '')]\n",
        "\n",
        "d_adapt.info()\n",
        "d_adapt.head(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6rJcflyfTt5S"
      },
      "outputs": [],
      "source": [
        "d_inference = d_adapt.sample(\n",
        "                             n = 10000,\n",
        "                             random_state = 56,\n",
        "                             ).reset_index(drop = True)\n",
        "\n",
        "d_inference = d_inference.dropna().reset_index(drop = True)\n",
        "\n",
        "d_inference.info()\n",
        "d_inference.head(3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "49THUJSqeL2C"
      },
      "source": [
        "#### BERT, RoBERTa, DistilBERT"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QDCM9Rrexj8G"
      },
      "outputs": [],
      "source": [
        "%cd ../../outputs/tables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "Xjfl2BodDkEM"
      },
      "outputs": [],
      "source": [
        "models_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models/'\n",
        "\n",
        "# define inference param sets\n",
        "\n",
        "params = [\n",
        "          {\n",
        "           'target': 'asp',\n",
        "           'model_class': RobertaForSequenceClassification,\n",
        "           'tokenizer_class': RobertaTokenizer,\n",
        "           'pretrained_model_name': 'roberta-base',\n",
        "           'model_path': f'{models_path}asp_roberta-base_best_tuned_model.bin',\n",
        "          },\n",
        "          {\n",
        "           'target': 'dep',\n",
        "           'model_class': RobertaForSequenceClassification,\n",
        "           'tokenizer_class': RobertaTokenizer,\n",
        "           'pretrained_model_name': 'roberta-base',\n",
        "           'model_path': f'{models_path}dep_roberta-base_best_tuned_model.bin',\n",
        "          },\n",
        "         {\n",
        "          'target': 'val',\n",
        "          'model_class': RobertaForSequenceClassification,\n",
        "          'tokenizer_class': RobertaTokenizer,\n",
        "          'pretrained_model_name': 'roberta-base',\n",
        "          'model_path': f'{models_path}val_roberta-base_best_tuned_model.bin',\n",
        "         },\n",
        "          {\n",
        "           'target': 'prg',\n",
        "           'model_class': RobertaForSequenceClassification,\n",
        "           'tokenizer_class': RobertaTokenizer,\n",
        "           'pretrained_model_name': 'roberta-base',\n",
        "           'model_path': f'{models_path}prg_roberta-base_best_tuned_model.bin',\n",
        "          },\n",
        "         {\n",
        "          'target': 'tgd',\n",
        "          'model_class': RobertaForSequenceClassification,\n",
        "          'tokenizer_class': RobertaTokenizer,\n",
        "          'pretrained_model_name': 'roberta-base',\n",
        "          'model_path': f'{models_path}tgd_roberta-base_best_tuned_model.bin',\n",
        "         },\n",
        "         {\n",
        "          'target': 'age',\n",
        "          'model_class': RobertaForSequenceClassification,\n",
        "          'tokenizer_class': RobertaTokenizer,\n",
        "          'pretrained_model_name': 'roberta-base',\n",
        "          'model_path': f'{models_path}age_roberta-base_best_tuned_model.bin',\n",
        "         },\n",
        "          {\n",
        "           'target': 'race',\n",
        "           'model_class': RobertaForSequenceClassification,\n",
        "           'tokenizer_class': RobertaTokenizer,\n",
        "           'pretrained_model_name': 'roberta-base',\n",
        "           'model_path': f'{models_path}race_roberta-base_best_tuned_model.bin',\n",
        "          },\n",
        "         {\n",
        "          'target': 'dbty',\n",
        "          'model_class': RobertaForSequenceClassification,\n",
        "          'tokenizer_class': RobertaTokenizer,\n",
        "          'pretrained_model_name': 'roberta-base',\n",
        "          'model_path': f'{models_path}dbty_roberta-base_best_tuned_model.bin',\n",
        "         },\n",
        "]\n",
        "\n",
        "# coerce to str\n",
        "\n",
        "d_inference['text'] = d_inference['text'].astype(str)\n",
        "texts = d_inference['text'].tolist()\n",
        "\n",
        "# inference loop\n",
        "\n",
        "for p in params:\n",
        "    target = p['target']\n",
        "\n",
        "    # load tokenizers, models\n",
        "\n",
        "    tokenizer = p['tokenizer_class'].from_pretrained(p['pretrained_model_name'])\n",
        "    model = load_model(\n",
        "                       p['model_path'],\n",
        "                       p['model_class'],\n",
        "                       p['pretrained_model_name'],\n",
        "                       )\n",
        "\n",
        "    # infer predictions, probabilities\n",
        "\n",
        "    predictions, probabilities = predict(\n",
        "                                         model,\n",
        "                                         tokenizer,\n",
        "                                         texts,\n",
        "                                         )\n",
        "\n",
        "    d_inference[f'{target}_pred'] = predictions\n",
        "    d_inference[f'{target}_prob'] = probabilities\n",
        "\n",
        "    # inspect\n",
        "\n",
        "d_inference[[\n",
        "             'asp_pred',\n",
        "             'dep_pred',\n",
        "             'val_pred',\n",
        "             'prg_pred',\n",
        "             'tgd_pred',\n",
        "             'age_pred',\n",
        "             'race_pred',\n",
        "             'dbty_pred',\n",
        "             ]].apply(pd.Series.value_counts)\n",
        "\n",
        "d_inference.head(3)\n",
        "\n",
        "# save\n",
        "\n",
        "d_inference.to_csv('d_inference_pred_pilot.csv')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aRMH0WABebFG"
      },
      "source": [
        "#### Llama 3.1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n7QdKBMHKQgq"
      },
      "outputs": [],
      "source": [
        "# clear cache\n",
        "\n",
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rMPJhqANOQr4"
      },
      "source": [
        "**Model-by-model (memory safe)**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "c9s2XnlFHSdd"
      },
      "outputs": [],
      "source": [
        "%cd /content/drive/MyDrive/Colab/bar_policy_suicidality/inputs/data\n",
        "#%cd inputs/data\n",
        "\n",
        "# single-target Llama inference\n",
        "\n",
        "models_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models/'\n",
        "\n",
        "target = 'asp'\n",
        "\n",
        "batch_size = 8\n",
        "\n",
        "d_inference = llama_load_and_predict_single_target(\n",
        "                                                   target,\n",
        "                                                   d_inference,\n",
        "                                                   models_path,\n",
        "                                                   batch_size,\n",
        "                                                   )\n",
        "# inspect + save\n",
        "\n",
        "d_inference.head(6)\n",
        "d_inference.to_csv('d_inference_pred_llama.csv')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xkOz_aZ8OXhy"
      },
      "source": [
        "**Multiple models**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cu6PvM7WHSjF"
      },
      "outputs": [],
      "source": [
        "#%cd inputs/data\n",
        "\n",
        "# multi-target Llama inference\n",
        "\n",
        "#targets = [\n",
        "#           'asp',\n",
        "#           'dep',\n",
        "#           'val',\n",
        "#           'prg',\n",
        "#           'tgd',\n",
        "#           'age',\n",
        "#           'race',\n",
        "#           'dbty',\n",
        "#           ]\n",
        "\n",
        "#models_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models/'\n",
        "\n",
        "#batch_size = 8\n",
        "\n",
        "#d_inference = llama_load_and_predict_multi_target(\n",
        "#                                                  targets,\n",
        "#                                                  d_inference,\n",
        "#                                                  model_save_dir,\n",
        "#                                                  batch_size,\n",
        "#                                                  )\n",
        "\n",
        "# inspect + save\n",
        "\n",
        "#d_inference.head(6)\n",
        "#d_inference.to_csv('d_inference_pred_llama.csv')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I5JVWw7spBaF"
      },
      "source": [
        "### 7. Explain\n",
        "Generates local instance-specific and global model-representative interpretations using LIME and SP-LIME.\n",
        "***"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y9zfVgzDrqeg"
      },
      "outputs": [],
      "source": [
        "%pwd"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-wTlTOzAIJ95"
      },
      "outputs": [],
      "source": [
        "#%cd ../../outputs/tables\n",
        "%cd /content/drive/My Drive/Colab/bar_policy_suicidality/outputs/tables\n",
        "\n",
        "d_inference = pd.read_csv('d_inference_pred_pilot.csv')\n",
        "d_inference.drop(\n",
        "                 'Unnamed: 0',\n",
        "                 axis = 1,\n",
        "                 inplace = True,\n",
        "                 )\n",
        "\n",
        "d_inference.info()\n",
        "d_inference.head(3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cvJwBfSdpLzZ"
      },
      "source": [
        "#### LIME: prediction explainability"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9D1_tRL8Rt3C"
      },
      "source": [
        "**Sort by confidence**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ORYlejYtNavH"
      },
      "outputs": [],
      "source": [
        "# define target\n",
        "\n",
        "target = 'dbty'\n",
        "\n",
        "# extract neg proba\n",
        "\n",
        "#d_inference[f'{target}_prob'] = d_inference[f'{target}_prob'].apply(lambda i: ast.literal_eval(i))\n",
        "#d_inference[f'{target}_neg'] = d_inference[f'{target}_prob'].apply(lambda i: round(i[0], 4))\n",
        "\n",
        "# extract pos proba\n",
        "\n",
        "d_inference[f'{target}_prob'] = d_inference[f'{target}_prob'].apply(lambda i: ast.literal_eval(i))\n",
        "d_inference[f'{target}_pos'] = d_inference[f'{target}_prob'].apply(lambda i: round(i[1], 4))\n",
        "\n",
        "#d_inference.head(3)\n",
        "\n",
        "# parse by 'f'{target}_neg' > 0.90\n",
        "\n",
        "#neg_index = d_inference[d_inference[f'{target}_neg'] > 0.90]\n",
        "\n",
        "# parse by 'f'{target}_pos' > 0.90\n",
        "\n",
        "pos_index = d_inference[d_inference[f'{target}_pos'] > 0.90]\n",
        "\n",
        "# shuffle\n",
        "\n",
        "pos_index = shuffle(\n",
        "                    pos_index,\n",
        "                    random_state = 56,\n",
        "                    )\n",
        "pos_index.head(10)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r82wKRWnPLHv"
      },
      "outputs": [],
      "source": [
        "# define instance\n",
        "\n",
        "instance = 6633"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WAYX8v9WSnyf"
      },
      "source": [
        "**LIMETextExplainer**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "gRXWSFS3NEL_"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "\n",
        "models_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models/'\n",
        "\n",
        "# load pre-trained HF tokenizer, config, architecture\n",
        "\n",
        "tokenizer = RobertaTokenizer.from_pretrained('roberta-base')\n",
        "model = RobertaForSequenceClassification.from_pretrained(\n",
        "                                                         'roberta-base',\n",
        "                                                          num_labels = 2,\n",
        "                                                          )\n",
        "\n",
        "# load fine-tuned model weights - dynamic target handling\n",
        "\n",
        "target_model_path = f'{models_path}{target}_roberta-base_best_tuned_model.bin'\n",
        "model.load_state_dict(torch.load(target_model_path, map_location=torch.device('cuda')))\n",
        "\n",
        "# coerce eval mode\n",
        "\n",
        "model.eval()\n",
        "\n",
        "# targets\n",
        "\n",
        "class_names = ['0', '1']\n",
        "\n",
        "# tokenize Fx\n",
        "\n",
        "def predict_proba(texts):\n",
        "    inputs = tokenizer(\n",
        "                       texts,\n",
        "                       padding = True,\n",
        "                       truncation = True,\n",
        "                       return_tensors = 'pt',\n",
        "                       max_length = 512,\n",
        "                       )\n",
        "    with torch.no_grad():\n",
        "        outputs = model(**inputs)\n",
        "        probs = torch.softmax(\n",
        "                              outputs.logits,\n",
        "                              dim = 1,\n",
        "                              )\n",
        "    return probs.cpu().numpy()\n",
        "\n",
        "# initialize explainer\n",
        "\n",
        "### SJS 9/27: re 'bow' and 'char_level' params: https://lime-ml.readthedocs.io/en/latest/lime.html\n",
        "\n",
        "explainer = LimeTextExplainer(\n",
        "                              class_names = class_names,\n",
        "                              random_state = 56,\n",
        "                              )\n",
        "\n",
        "# d_inference selection - dynamic\n",
        "\n",
        "text = d_inference.loc[instance, 'text']\n",
        "\n",
        "# explain\n",
        "\n",
        "exp = explainer.explain_instance(\n",
        "                                 text,\n",
        "                                 predict_proba,\n",
        "                                 num_features = 8,\n",
        "                                 num_samples = 1000,\n",
        "                                 distance_metric = 'cosine',\n",
        "                                 )\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_sezLG5rRngk"
      },
      "source": [
        "**Display explanation**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NMFUfMmVRg1_"
      },
      "outputs": [],
      "source": [
        "print(instance)\n",
        "exp.show_in_notebook(text = True)\n",
        "#exp.save_to_file('lime_explanation.html')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TKd4j2wMRg7p"
      },
      "outputs": [],
      "source": [
        "%cd ../figures\n",
        "exp.save_to_file(f'lime_explanation_{target}_{instance}_pos.html')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wUCAhTZlpT_L"
      },
      "source": [
        "#### SP-LIME: model explainability"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "qOFffbXipTXX"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "\n",
        "# define target\n",
        "\n",
        "target = 'dbty'\n",
        "\n",
        "# define models dir\n",
        "\n",
        "models_path = '/content/drive/MyDrive/Colab/bar_policy_suicidality/outputs/models/'\n",
        "\n",
        "# load pre-trained HF tokenizer, config, architecture\n",
        "\n",
        "tokenizer = RobertaTokenizer.from_pretrained('roberta-base')\n",
        "model = RobertaForSequenceClassification.from_pretrained(\n",
        "                                                         'roberta-base',\n",
        "                                                          num_labels = 2,\n",
        "                                                          )\n",
        "\n",
        "# load fine-tuned model weights - dynamic target handling\n",
        "\n",
        "target_model_path = f'{models_path}{target}_roberta-base_best_tuned_model.bin'\n",
        "model.load_state_dict(torch.load(target_model_path, map_location=torch.device('cuda')))\n",
        "\n",
        "# coerce eval mode\n",
        "\n",
        "model.eval()\n",
        "\n",
        "# targets\n",
        "\n",
        "class_names = ['0', '1']\n",
        "\n",
        "# tokenization function - incl. batch processing\n",
        "\n",
        "def predict_proba(texts, batch_size = 8): ### batch processing - memory safe\n",
        "\n",
        "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "    model.to(device) ### move model - same device as data\n",
        "\n",
        "    all_probs = []\n",
        "    for i in range(0, len(texts), batch_size):\n",
        "        batch_texts = texts[i:i + batch_size]\n",
        "        inputs = tokenizer(\n",
        "                           batch_texts,\n",
        "                           padding = True,\n",
        "                           truncation = True,\n",
        "                           return_tensors = 'pt',\n",
        "                           max_length = 512,\n",
        "                           )\n",
        "        with torch.no_grad():\n",
        "            inputs = {key: val.to(device) for key, val in inputs.items()} ### tokenized inputs - same device as model\n",
        "\n",
        "            # infer\n",
        "\n",
        "            outputs = model(**inputs)\n",
        "            probs = torch.softmax(outputs.logits, dim = 1)\n",
        "\n",
        "            all_probs.extend(probs.cpu().numpy()) ### logits to CPU\n",
        "\n",
        "    return np.array(all_probs)\n",
        "\n",
        "# initialize LIMETextExplainer\n",
        "\n",
        "explainer = LimeTextExplainer(\n",
        "                              class_names = class_names,\n",
        "                              random_state = 56,\n",
        "                              )\n",
        "\n",
        "# select text instances dynamically for SP-LIME\n",
        "\n",
        "selected_instances = d_inference.sample(\n",
        "                                        n = 100,\n",
        "                                        random_state = 56,\n",
        "                                        )\n",
        "\n",
        "# store explanations\n",
        "\n",
        "sp_explanations = []\n",
        "\n",
        "for instance_idx in selected_instances.index:\n",
        "    text = d_inference.loc[instance_idx, 'text']\n",
        "\n",
        "    # instance-wise LIME explanations\n",
        "\n",
        "    exp = explainer.explain_instance(\n",
        "                                     text,\n",
        "                                     predict_proba,\n",
        "                                     num_features = 10,\n",
        "                                     num_samples = 1000,\n",
        "                                     distance_metric = 'cosine',\n",
        "                                     )\n",
        "    # append\n",
        "\n",
        "    sp_explanations.append(exp)\n",
        "\n",
        "# define SubmodularPick logics\n",
        "\n",
        "sp = SubmodularPick(\n",
        "                    explainer,\n",
        "                    selected_instances['text'].tolist(),\n",
        "                    predict_proba,\n",
        "                    num_features = 10,\n",
        "                    num_exps_desired = 10, ### n explanations to provide\n",
        "                    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "5XbuY54Z0-kZ"
      },
      "outputs": [],
      "source": [
        "# viz. + save\n",
        "\n",
        "%cd ../figures\n",
        "\n",
        "for idx, exp in enumerate(sp.sp_explanations, start = 1):\n",
        "    exp.show_in_notebook(text = True)\n",
        "    exp.save_to_file(f'sp_lime_explanation_{target}_{idx:02}.html') ### save full range of num_exps_desired\n",
        "\n",
        "#for exp in sp.sp_explanations:\n",
        "#    exp.show_in_notebook(text = True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2nSz7vQh5xYz"
      },
      "source": [
        "### 8. Calibrate\n",
        "tk"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HF6wdotT5tFP"
      },
      "outputs": [],
      "source": [
        "# tk"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-odwkzTUrD32"
      },
      "source": [
        "> End of aim_i_train_tune_predict.ipynb"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "woZYUo6JgN1X",
        "HMuLvUHmksoc",
        "rxzedr7a-kuK",
        "qkBdAMlV1gjG",
        "7xoLLpCnlkpO",
        "YmwkvY6KLEX6",
        "Thnxs8P6ic1l",
        "Mm0hLOX24AYx",
        "AZ6gjBzkQytt"
      ],
      "machine_shape": "hm",
      "provenance": [],
      "gpuType": "A100"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}