{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "[notice] A new release of pip is available: 24.0 -> 26.1.1\n",
      "[notice] To update, run: python.exe -m pip install --upgrade pip\n"
     ]
    }
   ],
   "source": [
    "%pip install -q transformers datasets accelerate evaluate seqeval ftfy nltk pandas numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import re\n",
    "import json\n",
    "import glob\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import nltk\n",
    "import ftfy\n",
    "from nltk.tokenize import sent_tokenize\n",
    "from transformers import AutoTokenizer, AutoModelForTokenClassification, logging as tf_logging\n",
    "\n",
    "tf_logging.set_verbosity_error()\n",
    "tf_logging.disable_progress_bar()\n",
    "\n",
    "from datasets import logging as ds_logging\n",
    "ds_logging.set_verbosity_error()\n",
    "ds_logging.disable_progress_bar()\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using device: {device}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected model: electra-small\n"
     ]
    }
   ],
   "source": [
    "model_choice = \"electra-small\"\n",
    "print(f\"Selected model: {model_choice}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading locally trained model from: ./gotcha-extractor-model/electra-small\n",
      "Model loaded successfully!\n"
     ]
    }
   ],
   "source": [
    "\n",
    "fallback_map = {\n",
    "    \"electra-small\": \"google/electra-small-discriminator\",\n",
    "    \"bert-tiny\": \"prajjwal1/bert-tiny\",\n",
    "    \"tinybert\": \"huawei-noah/TinyBERT_General_4L_312D\",\n",
    "    \"legal-bert\": \"nlpaueb/legal-bert-base-uncased\",\n",
    "    \"bert-mini\": \"prajjwal1/bert-mini\",\n",
    "    \"albert-base\": \"albert-base-v2\",\n",
    "}\n",
    "\n",
    "id2label = {0: 'O', 1: 'B-RISK', 2: 'I-RISK'}\n",
    "label2id = {'O': 0, 'B-RISK': 1, 'I-RISK': 2}\n",
    "\n",
    "local_path = f\"./gotcha-extractor-model/{model_choice}-optuna-search\"\n",
    "if not os.path.exists(local_path) or not os.path.exists(os.path.join(local_path, \"config.json\")):\n",
    "    local_path = f\"./gotcha-extractor-model/{model_choice}\"\n",
    "has_local = os.path.exists(local_path) and os.path.exists(os.path.join(local_path, \"config.json\"))\n",
    "\n",
    "if has_local:\n",
    "    model_path = local_path\n",
    "    print(f\"Loading locally trained model from: {model_path}\")\n",
    "else:\n",
    "    model_path = fallback_map[model_choice]\n",
    "    print(f\"Local trained model not found. Falling back to base model: {model_path}\")\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
    "prev_verbosity = tf_logging.get_verbosity()\n",
    "tf_logging.set_verbosity_error()\n",
    "tf_logging.disable_progress_bar()\n",
    "model = AutoModelForTokenClassification.from_pretrained(\n",
    "    model_path,\n",
    "    num_labels=len(label2id),\n",
    "    id2label=id2label,\n",
    "    label2id=label2id,\n",
    "    ignore_mismatched_sizes=True\n",
    ").to(device)\n",
    "tf_logging.enable_progress_bar()\n",
    "tf_logging.set_verbosity(prev_verbosity)\n",
    "\n",
    "model.eval()\n",
    "print(f\"Model loaded successfully!\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "KEYWORDS_HIGH = [\n",
    "    r\"arbitrat\", r\"class\\s+action\", r\"waiver\", r\"dispute\",\n",
    "    r\"reserve\\s+the\\s+right\\s+to\", r\"modify\", r\"revise\", r\"update\", r\"without\\s+notice\",\n",
    "    r\"sell\", r\"market\", r\"advertis\", r\"third\\s+part\",\n",
    "    r\"cannot\\s+(ensure|warrant|guarantee)\", r\"no\\s+warranty\", r\"indemni\"\n",
    "]\n",
    "\n",
    "BOILERPLATE_PATTERNS = [\n",
    "    r\"this\\s+privacy\\s+policy\\s+(\\([^)]+\\)\\s+)?describes\\s+the\\s+practices\",\n",
    "    r\"this\\s+privacy\\s+policy\\s+applies\\s+only\\s+to\",\n",
    "    r\"summary\\s+the\\s+notifications\\s+provided\\s+by\\s+this\\s+privacy\\s+policy\\s+include\",\n",
    "    r\"^[a-zA-Z\\s]+is\\s+data\\s+that\\s+can\\s+be\\s+used\\s+to\\s+identify\",\n",
    "    r\"^[a-zA-Z\\s]+\\s+means\\s+any\\s+information\",\n",
    "    r\"legal\\s+grounds\\s+for\\s+processing\\s+personal\\s+data\",\n",
    "    r\"we\\s+restrict\\s+access\\s+to\\s+personal\\s+information\\s+collected.*to\\s+our\\s+employees\",\n",
    "    r\"please\\s+note\\s+that\\s+we\\s+have\\s+a\\s+separate\\s+privacy\\s+disclosure\\s+statement\\s+to\\s+address\\s+our\\s+protocols.*located\\s+here\",\n",
    "    r\"children\\s+under\\s+13\", r\"younger\\s+than\\s+13\", r\"receive\\s+parental\\s+consent\",\n",
    "    r\"privacy\\s+policy\\s+effective\\s+date\"\n",
    "]\n",
    "\n",
    "def clean_boilerplate_header(sentence):\n",
    "    sentence_clean = sentence.strip()\n",
    "    sentence_lower = sentence_clean.lower()\n",
    "\n",
    "    if re.match(r\"^[A-Z\\s\\d/_:,,\\'\\\"]{3,50}$\", sentence_clean):\n",
    "        return True\n",
    "\n",
    "    for pattern in BOILERPLATE_PATTERNS:\n",
    "        if re.search(pattern, sentence_lower):\n",
    "            return True\n",
    "\n",
    "    return False\n",
    "\n",
    "def determine_risk_level(sentence, risk_tokens, has_high_keyword):\n",
    "    if not risk_tokens:\n",
    "        return None\n",
    "\n",
    "    probs = [t[\"prob\"] for t in risk_tokens]\n",
    "    max_prob = max(probs)\n",
    "\n",
    "    if max_prob >= 0.80 or (has_high_keyword and max_prob >= 0.68):\n",
    "        return \"HIGH RISK\"\n",
    "    elif has_high_keyword or max_prob >= 0.62:\n",
    "        return \"MEDIUM RISK\"\n",
    "    else:\n",
    "        return \"LOW RISK\"\n",
    "\n",
    "KEYWORDS_PRO_USER = [\n",
    "    r\"you\\s+may\\s+(access|correct|request\\s+deletion|delete|port|object)\",\n",
    "    r\"request\\s+that\\s+we\\s+stop\\s+(any\\s+)?processing\",\n",
    "    r\"freely\\s+visit\\s+our\\s+(website|platform)\\s+anonymously\",\n",
    "    r\"without\\s+being\\s+required\\s+to\\s+provide\\s+us\\s+with\\s+any\\s+personal\\s+information\",\n",
    "    r\"rights\\s+related\\s+to\\s+the\\s+european\\s+union\",\n",
    "    r\"rights\\s+related\\s+to\\s+gdpr\",\n",
    "    r\"your\\s+right\\s+to\\s+(access|delete|rectify|restrict)\",\n",
    "    r\"opt[- ]out\\s+of\\s+receiving\\s+(marketing|promotional|newsletter)\",\n",
    "    r\"under\\s+the\\s+general\\s+data\\s+protection\\s+regulation\",\n",
    "    r\"right\\s+to\\s+request\\s+that\\s+we\\s+disclose\",\n",
    "    r\"right\\s+to\\s+know\\s+what\\s+personal\\s+information\",\n",
    "]\n",
    "\n",
    "def check_pro_user_override(sentence):\n",
    "    sentence_clean = sentence.strip()\n",
    "    sentence_lower = sentence_clean.lower()\n",
    "\n",
    "    for pattern in KEYWORDS_PRO_USER:\n",
    "        if re.search(pattern, sentence_lower):\n",
    "            return True\n",
    "\n",
    "    if re.search(r\"\\b(right(s)?\\s+to|you\\s+have\\s+the\\s+right\\s+to)\\s+.*\\b(access|correct|delete|erase|rectify|update|portability|restrict)\\b\", sentence_lower):\n",
    "        return True\n",
    "\n",
    "    if re.search(r\"\\b(visit|browse)\\b.*\\banonymously\\b\", sentence_lower) and not re.search(r\"\\b(cannot|unable|restrict)\\b\", sentence_lower):\n",
    "        return True\n",
    "\n",
    "    if re.search(r\"\\brights\\s+related\\s+to\\b.*\\b(gdpr|ccpa|california\\s+consumer|protection\\s+regulation)\\b\", sentence_lower):\n",
    "        return True\n",
    "\n",
    "    return False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def extract_gotchas(raw_text, min_risk_tokens=3):\n",
    "    if not raw_text or not raw_text.strip():\n",
    "        return []\n",
    "\n",
    "    cleaned_text = clean_text_pipeline(raw_text)\n",
    "    sentences = sent_tokenize(cleaned_text)\n",
    "    all_extracted = []\n",
    "\n",
    "    for sentence in sentences:\n",
    "        if not sentence.strip():\n",
    "            continue\n",
    "\n",
    "        if clean_boilerplate_header(sentence):\n",
    "            continue\n",
    "\n",
    "        if check_pro_user_override(sentence):\n",
    "            continue\n",
    "\n",
    "        inputs = tokenizer(\n",
    "            sentence, \n",
    "            return_tensors=\"pt\", \n",
    "            truncation=True, \n",
    "            max_length=512\n",
    "        )\n",
    "        tokens = tokenizer.convert_ids_to_tokens(inputs[\"input_ids\"][0])\n",
    "        inputs = {k: v.to(device) for k, v in inputs.items()}\n",
    "\n",
    "        with torch.no_grad():\n",
    "            outputs = model(**inputs)\n",
    "\n",
    "        if isinstance(outputs, dict):\n",
    "            logits = outputs['logits'][0]\n",
    "        else:\n",
    "            logits = outputs.logits[0]\n",
    "\n",
    "        probs = torch.softmax(logits, dim=-1)\n",
    "        predictions = torch.argmax(logits, dim=-1)\n",
    "\n",
    "        risk_tokens = []\n",
    "        for t_idx, pred in enumerate(predictions):\n",
    "            label = id2label[pred.item()]\n",
    "            token_str = tokens[t_idx]\n",
    "            if token_str in ('[CLS]', '[SEP]', '[PAD]'):\n",
    "                continue\n",
    "            prob = probs[t_idx][pred.item()].item()\n",
    "            if label in ('B-RISK', 'I-RISK'):\n",
    "                risk_tokens.append({\"token\": token_str, \"prob\": prob})\n",
    "\n",
    "        if len(risk_tokens) >= min_risk_tokens:\n",
    "            max_prob = max(t[\"prob\"] for t in risk_tokens)\n",
    "\n",
    "            has_high_keyword = False\n",
    "            sentence_lower = sentence.lower()\n",
    "            for pattern in KEYWORDS_HIGH:\n",
    "                if re.search(pattern, sentence_lower):\n",
    "                    has_high_keyword = True\n",
    "                    break\n",
    "\n",
    "            keep = False\n",
    "            if has_high_keyword:\n",
    "                if max_prob >= 0.55:\n",
    "                    keep = True\n",
    "            else:\n",
    "                if max_prob >= 0.70:\n",
    "                    keep = True\n",
    "\n",
    "            if keep:\n",
    "                level = determine_risk_level(sentence, risk_tokens, has_high_keyword)\n",
    "                all_extracted.append((sentence, level))\n",
    "\n",
    "    return all_extracted\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "--- Test 1: Simple Inline Text ---\n",
      "Raw Text: Welcome to the platform. By continuing, you agree to forced arbitration in the event of a dispute. We also reserve the right to sell your location data and usage habits to unverified third parties.\n",
      "\n",
      "Extracted Gotchas:\n",
      "1. [HIGH RISK] By continuing, you agree to forced arbitration in the event of a dispute.\n",
      "2. [HIGH RISK] We also reserve the right to sell your location data and usage habits to unverified third parties.\n"
     ]
    }
   ],
   "source": [
    "sample_legalese = (\n",
    "    \"Welcome to the platform. By continuing, you agree to forced arbitration in the event of a dispute. \"\n",
    "    \"We also reserve the right to sell your location data and usage habits to unverified third parties.\"\n",
    ")\n",
    "\n",
    "print(\"\\n--- Test 1: Simple Inline Text ---\")\n",
    "print(\"Raw Text:\", sample_legalese)\n",
    "print(\"\\nExtracted Gotchas:\")\n",
    "for idx, (clause, level) in enumerate(extract_gotchas(sample_legalese)):\n",
    "    try:\n",
    "        print(f\"{idx + 1}. [{level}] {clause}\")\n",
    "    except UnicodeEncodeError:\n",
    "        print(f\"{idx + 1}. [{level}] {clause.encode('ascii', errors='replace').decode('ascii')}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}