{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "intro",
   "metadata": {},
   "source": [
    "# Abuse Model Inference Test\n",
    "\n",
    "This notebook loads the saved production abuse model and its local artifacts from this folder, then runs inference on a calibrated set of example sentences. It performs inference only.\n",
    "\n",
    "Note: this checkpoint is a subtype classifier, not a generic toxicity detector. It predicts `Age`, `Ethnicity`, `Gender`, `Religion`, or `Not Cyberbullying`, and it responds most strongly when the text explicitly names the targeted group."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "load-notes",
   "metadata": {},
   "source": [
    "## 1. Load the model artifacts\n",
    "\n",
    "The code below reads the local model config, tokenizer, weights, and label mapping directly from the current folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "load-code",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "37774b8a0e1f4d959c6e6a3293e0e766",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading weights:   0%|          | 0/201 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model directory: /Users/georgijkutivadze/projects/mago-text-scoring/app/app-models/prod-model-abuse\n",
      "Device: cpu\n",
      "Architecture: BertForSequenceClassification\n",
      "Tokenizer: BertTokenizer\n",
      "Max sequence length: 256\n",
      "Artifacts: ['config.json', 'model-config.yaml', 'model.safetensors', 'test_inference.ipynb', 'tokenizer.json', 'tokenizer_config.json', 'training_args.bin']\n",
      "Labels: ['Age', 'Ethnicity', 'Gender', 'Not Cyberbullying', 'Religion']\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "import os\n",
    "\n",
    "import pandas as pd\n",
    "import torch\n",
    "import yaml\n",
    "from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer\n",
    "from transformers.utils import logging as hf_logging\n",
    "\n",
    "os.environ[\"HF_HUB_DISABLE_PROGRESS_BARS\"] = \"1\"\n",
    "hf_logging.set_verbosity_error()\n",
    "\n",
    "MODEL_DIR = Path(\".\").resolve()\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "with (MODEL_DIR / \"model-config.yaml\").open(\"r\", encoding=\"utf-8\") as handle:\n",
    "    model_config = yaml.safe_load(handle)\n",
    "\n",
    "runtime_config = model_config.get(\"runtime\", {})\n",
    "MAX_LEN = int(runtime_config.get(\"max_sequence_length\", 256))\n",
    "PADDING = runtime_config.get(\"padding\", False)\n",
    "\n",
    "hf_config = AutoConfig.from_pretrained(MODEL_DIR, local_files_only=True)\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, local_files_only=True)\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    MODEL_DIR,\n",
    "    config=hf_config,\n",
    "    local_files_only=True,\n",
    ")\n",
    "model.to(device)\n",
    "model.eval()\n",
    "\n",
    "yaml_labels = model_config.get(\"labels\", {}).get(\"classes\", [])\n",
    "label_names = [\n",
    "    item.get(\"display_name\") or item.get(\"name\") or str(item.get(\"id\"))\n",
    "    for item in sorted(yaml_labels, key=lambda item: item[\"id\"])\n",
    "]\n",
    "if not label_names:\n",
    "    label_names = [hf_config.id2label[idx] for idx in sorted(hf_config.id2label)]\n",
    "\n",
    "artifact_files = sorted(path.name for path in MODEL_DIR.iterdir() if path.is_file())\n",
    "\n",
    "print(f\"Model directory: {MODEL_DIR}\")\n",
    "print(f\"Device: {device}\")\n",
    "print(f\"Architecture: {model.__class__.__name__}\")\n",
    "print(f\"Tokenizer: {tokenizer.__class__.__name__}\")\n",
    "print(f\"Max sequence length: {MAX_LEN}\")\n",
    "print(f\"Artifacts: {artifact_files}\")\n",
    "print(f\"Labels: {label_names}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "helper-notes",
   "metadata": {},
   "source": [
    "## 2. Define an inference helper\n",
    "\n",
    "This helper mirrors the production runtime settings, tokenizes input text, runs the model in evaluation mode, applies softmax, and returns the predicted label together with confidence, top-3 labels, and the full probability distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "helper-code",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.inference_mode()\n",
    "def predict_texts(texts):\n",
    "    encoded = tokenizer(\n",
    "        texts,\n",
    "        return_tensors=\"pt\",\n",
    "        truncation=bool(runtime_config.get(\"truncation\", True)),\n",
    "        padding=\"max_length\" if PADDING else False,\n",
    "        max_length=MAX_LEN,\n",
    "    )\n",
    "    encoded = {name: tensor.to(device) for name, tensor in encoded.items()}\n",
    "    logits = model(**encoded).logits\n",
    "    probabilities = torch.softmax(logits, dim=-1).cpu().tolist()\n",
    "\n",
    "    rows = []\n",
    "    for text, input_ids, probs in zip(texts, encoded[\"input_ids\"].cpu().tolist(), probabilities):\n",
    "        ranked_indices = sorted(range(len(probs)), key=probs.__getitem__, reverse=True)\n",
    "        pred_idx = ranked_indices[0]\n",
    "        rows.append(\n",
    "            {\n",
    "                \"text\": text,\n",
    "                \"predicted_label\": label_names[pred_idx],\n",
    "                \"confidence\": round(probs[pred_idx], 4),\n",
    "                \"token_count\": sum(1 for token_id in input_ids if token_id != tokenizer.pad_token_id),\n",
    "                \"top_3\": [\n",
    "                    f\"{label_names[idx]} ({probs[idx]:.4f})\"\n",
    "                    for idx in ranked_indices[:3]\n",
    "                ],\n",
    "                \"probabilities\": {\n",
    "                    label: round(score, 4)\n",
    "                    for label, score in zip(label_names, probs)\n",
    "                },\n",
    "            }\n",
    "        )\n",
    "\n",
    "    return pd.DataFrame(rows)\n",
    "\n",
    "pd.set_option(\"display.max_colwidth\", 140)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "examples-notes",
   "metadata": {},
   "source": [
    "## 3. Run inference on calibrated example sentences\n",
    "\n",
    "The first notebook version used generic insults and indirect phrasing, which made the checkpoint look inconsistent. The examples below are stronger subtype-targeted probes and should exercise all five labels much more reliably."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "examples-code",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "      <th>predicted_label</th>\n",
       "      <th>confidence</th>\n",
       "      <th>token_count</th>\n",
       "      <th>top_3</th>\n",
       "      <th>probabilities</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>go do your homework you stupid child</td>\n",
       "      <td>Age</td>\n",
       "      <td>0.9971</td>\n",
       "      <td>9</td>\n",
       "      <td>[Age (0.9971), Not Cyberbullying (0.0014), Ethnicity (0.0009)]</td>\n",
       "      <td>{'Age': 0.9971, 'Ethnicity': 0.0009, 'Gender': 0.0003, 'Not Cyberbullying': 0.0014, 'Religion': 0.0003}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>immigrants like you ruin everything in this country</td>\n",
       "      <td>Ethnicity</td>\n",
       "      <td>0.9979</td>\n",
       "      <td>10</td>\n",
       "      <td>[Ethnicity (0.9979), Not Cyberbullying (0.0013), Religion (0.0007)]</td>\n",
       "      <td>{'Age': 0.0, 'Ethnicity': 0.9979, 'Gender': 0.0001, 'Not Cyberbullying': 0.0013, 'Religion': 0.0007}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>women are stupid and should stay quiet</td>\n",
       "      <td>Gender</td>\n",
       "      <td>0.9999</td>\n",
       "      <td>9</td>\n",
       "      <td>[Gender (0.9999), Not Cyberbullying (0.0001), Ethnicity (0.0000)]</td>\n",
       "      <td>{'Age': 0.0, 'Ethnicity': 0.0, 'Gender': 0.9999, 'Not Cyberbullying': 0.0001, 'Religion': 0.0}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>muslims like you are the problem with this country</td>\n",
       "      <td>Religion</td>\n",
       "      <td>1.0000</td>\n",
       "      <td>11</td>\n",
       "      <td>[Religion (1.0000), Not Cyberbullying (0.0000), Age (0.0000)]</td>\n",
       "      <td>{'Age': 0.0, 'Ethnicity': 0.0, 'Gender': 0.0, 'Not Cyberbullying': 0.0, 'Religion': 1.0}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>the project update was sent to the team this morning</td>\n",
       "      <td>Not Cyberbullying</td>\n",
       "      <td>0.9999</td>\n",
       "      <td>12</td>\n",
       "      <td>[Not Cyberbullying (0.9999), Gender (0.0000), Religion (0.0000)]</td>\n",
       "      <td>{'Age': 0.0, 'Ethnicity': 0.0, 'Gender': 0.0, 'Not Cyberbullying': 0.9999, 'Religion': 0.0}</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                   text    predicted_label  \\\n",
       "0                  go do your homework you stupid child                Age   \n",
       "1   immigrants like you ruin everything in this country          Ethnicity   \n",
       "2                women are stupid and should stay quiet             Gender   \n",
       "3    muslims like you are the problem with this country           Religion   \n",
       "4  the project update was sent to the team this morning  Not Cyberbullying   \n",
       "\n",
       "   confidence  token_count  \\\n",
       "0      0.9971            9   \n",
       "1      0.9979           10   \n",
       "2      0.9999            9   \n",
       "3      1.0000           11   \n",
       "4      0.9999           12   \n",
       "\n",
       "                                                                 top_3  \\\n",
       "0       [Age (0.9971), Not Cyberbullying (0.0014), Ethnicity (0.0009)]   \n",
       "1  [Ethnicity (0.9979), Not Cyberbullying (0.0013), Religion (0.0007)]   \n",
       "2    [Gender (0.9999), Not Cyberbullying (0.0001), Ethnicity (0.0000)]   \n",
       "3        [Religion (1.0000), Not Cyberbullying (0.0000), Age (0.0000)]   \n",
       "4     [Not Cyberbullying (0.9999), Gender (0.0000), Religion (0.0000)]   \n",
       "\n",
       "                                                                                             probabilities  \n",
       "0  {'Age': 0.9971, 'Ethnicity': 0.0009, 'Gender': 0.0003, 'Not Cyberbullying': 0.0014, 'Religion': 0.0003}  \n",
       "1     {'Age': 0.0, 'Ethnicity': 0.9979, 'Gender': 0.0001, 'Not Cyberbullying': 0.0013, 'Religion': 0.0007}  \n",
       "2           {'Age': 0.0, 'Ethnicity': 0.0, 'Gender': 0.9999, 'Not Cyberbullying': 0.0001, 'Religion': 0.0}  \n",
       "3                 {'Age': 0.0, 'Ethnicity': 0.0, 'Gender': 0.0, 'Not Cyberbullying': 0.0, 'Religion': 1.0}  \n",
       "4              {'Age': 0.0, 'Ethnicity': 0.0, 'Gender': 0.0, 'Not Cyberbullying': 0.9999, 'Religion': 0.0}  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example_sentences = [\n",
    "    \"go do your homework you stupid child\",\n",
    "    \"immigrants like you ruin everything in this country\",\n",
    "    \"women are stupid and should stay quiet\",\n",
    "    \"muslims like you are the problem with this country\",\n",
    "    \"the project update was sent to the team this morning\",\n",
    "]\n",
    "\n",
    "results = predict_texts(example_sentences)\n",
    "results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "diagnostic-code",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Diagnostic: sample bank produced 5 distinct labels across 5 examples.\n"
     ]
    }
   ],
   "source": [
    "distinct_labels = results[\"predicted_label\"].nunique()\n",
    "if distinct_labels == 1:\n",
    "    print(\n",
    "        \"Warning: all sample predictions collapsed to one label. \"\n",
    "        \"That usually means the sample bank is too weak for this checkpoint or the export should be re-verified.\"\n",
    "    )\n",
    "else:\n",
    "    print(f\"Diagnostic: sample bank produced {distinct_labels} distinct labels across {len(results)} examples.\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}