{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "93b5a3cf",
   "metadata": {},
   "source": [
    "# EmpathRAG — FAISS Index Builder\n",
    "**Run on Colab Pro with A100 GPU**\n",
    "\n",
    "Encodes Reddit Mental Health corpus into 768-dim vectors and builds a FAISS index + SQLite sidecar.\n",
    "\n",
    "- Estimated time on A100: 30–60 minutes\n",
    "- Output: `faiss_flat.index` + `metadata.db` saved to Google Drive\n",
    "- **Do not run on local RTX 3060 laptop** — would take 6–12 hours\n",
    "\n",
    "### Before running:\n",
    "1. Upload `data/raw/reddit_mental_health/` (108 CSVs, ~3.1GB) to Google Drive at `My Drive/empathrag/data/raw/reddit_mental_health/`\n",
    "2. Set Runtime → Change runtime type → **A100 GPU**\n",
    "3. Run all cells in order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9ab42f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 1: Install dependencies\n",
    "!pip install sentence-transformers faiss-cpu tqdm transformers -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9129b2f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 2: Mount Google Drive\n",
    "from google.colab import drive\n",
    "drive.mount(\"/content/drive\")\n",
    "\n",
    "BASE       = \"/content/drive/MyDrive/empathrag\"\n",
    "REDDIT_DIR = f\"{BASE}/data/raw/reddit_mental_health\"\n",
    "INDEX_PATH = f\"{BASE}/data/indexes/faiss_flat.index\"\n",
    "DB_PATH    = f\"{BASE}/data/indexes/metadata.db\"\n",
    "CKPT_PATH  = f\"{BASE}/data/indexes/embeddings_checkpoint.npy\"\n",
    "\n",
    "import os\n",
    "os.makedirs(f\"{BASE}/data/indexes\", exist_ok=True)\n",
    "print(\"Drive mounted. Paths configured.\")\n",
    "print(f\"Reddit dir exists: {os.path.exists(REDDIT_DIR)}\")\n",
    "files = [f for f in os.listdir(REDDIT_DIR) if f.endswith('.csv')] if os.path.exists(REDDIT_DIR) else []\n",
    "print(f\"CSV files found: {len(files)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54ccc2ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 3: Helper functions\n",
    "import re\n",
    "\n",
    "def clean_text(text: str) -> str:\n",
    "    text = re.sub(r\"u/\\w+\", \"\", text)\n",
    "    text = re.sub(r\"r/\\w+\", \"\", text)\n",
    "    text = re.sub(r\"http\\S+\", \"\", text)\n",
    "    text = re.sub(r\"\\[deleted\\]|\\[removed\\]\", \"\", text)\n",
    "    text = re.sub(r\"[^\\x00-\\x7F]+\", \" \", text)\n",
    "    text = re.sub(r\"\\s+\", \" \", text).strip()\n",
    "    return text\n",
    "\n",
    "def chunk_text(text, tokenizer, chunk_size=256, stride=32, max_chunks=8):\n",
    "    tokens = tokenizer.encode(text)\n",
    "    if len(tokens) < 64:\n",
    "        return [text]\n",
    "    chunks = []\n",
    "    start = 0\n",
    "    while start < len(tokens) and len(chunks) < max_chunks:\n",
    "        end = min(start + chunk_size, len(tokens))\n",
    "        chunks.append(tokenizer.decode(tokens[start:end], skip_special_tokens=True))\n",
    "        start += chunk_size - stride\n",
    "    return chunks\n",
    "\n",
    "print(\"Helper functions defined.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68b83384",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 4: Load and chunk Reddit posts\n",
    "# This reads all 108 CSVs and chunks the post text.\n",
    "# Expected: 1-5 million chunks depending on corpus size.\n",
    "import pandas as pd\n",
    "from transformers import AutoTokenizer\n",
    "from tqdm import tqdm\n",
    "\n",
    "tok = AutoTokenizer.from_pretrained(\"roberta-base\")\n",
    "\n",
    "all_posts = []\n",
    "csv_files = [f for f in os.listdir(REDDIT_DIR) if f.endswith('.csv')]\n",
    "print(f\"Reading {len(csv_files)} CSV files...\")\n",
    "\n",
    "for fname in tqdm(csv_files, desc=\"Reading CSVs\"):\n",
    "    fpath = os.path.join(REDDIT_DIR, fname)\n",
    "    try:\n",
    "        df = pd.read_csv(fpath, on_bad_lines=\"skip\",\n",
    "                         usecols=lambda c: c in [\"post\", \"body\", \"selftext\"])\n",
    "        for col in [\"post\", \"body\", \"selftext\"]:\n",
    "            if col in df.columns:\n",
    "                all_posts.extend(df[col].dropna().tolist())\n",
    "                break\n",
    "    except Exception as e:\n",
    "        print(f\"  Skipping {fname}: {e}\")\n",
    "\n",
    "print(f\"Raw posts loaded: {len(all_posts):,}\")\n",
    "\n",
    "chunks = []\n",
    "for post in tqdm(all_posts, desc=\"Chunking\"):\n",
    "    cleaned = clean_text(str(post))\n",
    "    if not cleaned:\n",
    "        continue\n",
    "    chunks.extend(chunk_text(cleaned, tok))\n",
    "\n",
    "print(f\"Total chunks: {len(chunks):,}\")\n",
    "\n",
    "# Save chunks list to Drive as checkpoint\n",
    "import pickle\n",
    "with open(f\"{BASE}/data/indexes/chunks_checkpoint.pkl\", \"wb\") as f:\n",
    "    pickle.dump(chunks, f)\n",
    "print(\"Chunks saved to Drive checkpoint.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ee62270",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 5: Encode embeddings\n",
    "# Uses A100 GPU. Saves progress every 100K chunks so Colab disconnects are recoverable.\n",
    "import numpy as np\n",
    "from sentence_transformers import SentenceTransformer\n",
    "import torch\n",
    "\n",
    "BATCH_SIZE = 256  # Safe for A100 40GB. Reduce to 64 if OOM.\n",
    "SAVE_EVERY = 100_000  # Save checkpoint every N chunks\n",
    "\n",
    "encoder = SentenceTransformer(\"sentence-transformers/all-mpnet-base-v2\")\n",
    "print(f\"Encoding on: {encoder.device}\")\n",
    "print(f\"Total chunks to encode: {len(chunks):,}\")\n",
    "print(f\"Estimated time at batch_size={BATCH_SIZE}: ~{len(chunks) // (BATCH_SIZE * 50)} minutes\")\n",
    "\n",
    "# Check if partial checkpoint exists\n",
    "all_embeddings = []\n",
    "start_idx = 0\n",
    "\n",
    "if os.path.exists(CKPT_PATH):\n",
    "    print(f\"Resuming from checkpoint: {CKPT_PATH}\")\n",
    "    all_embeddings = list(np.load(CKPT_PATH))\n",
    "    start_idx = len(all_embeddings)\n",
    "    print(f\"Resuming from chunk {start_idx:,}\")\n",
    "\n",
    "# Encode in segments\n",
    "for seg_start in tqdm(range(start_idx, len(chunks), SAVE_EVERY), desc=\"Encoding segments\"):\n",
    "    seg_end = min(seg_start + SAVE_EVERY, len(chunks))\n",
    "    seg = chunks[seg_start:seg_end]\n",
    "    emb = encoder.encode(seg, batch_size=BATCH_SIZE, show_progress_bar=False,\n",
    "                         normalize_embeddings=True, convert_to_numpy=True)\n",
    "    all_embeddings.extend(emb)\n",
    "    # Save checkpoint\n",
    "    np.save(CKPT_PATH, np.array(all_embeddings, dtype=np.float32))\n",
    "    print(f\"  Checkpoint saved at {seg_end:,}/{len(chunks):,} chunks\")\n",
    "\n",
    "embeddings = np.array(all_embeddings, dtype=np.float32)\n",
    "print(f\"Embeddings shape: {embeddings.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a761c6e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 6: Build FAISS index and SQLite sidecar\n",
    "import faiss\n",
    "import sqlite3\n",
    "\n",
    "dim = embeddings.shape[1]  # 768\n",
    "n   = embeddings.shape[0]\n",
    "\n",
    "print(f\"Building FAISS index for {n:,} vectors of dim {dim}...\")\n",
    "\n",
    "if n > 100_000:\n",
    "    quantizer = faiss.IndexFlatL2(dim)\n",
    "    index = faiss.IndexIVFFlat(quantizer, dim, 100)\n",
    "    print(\"Training IVFFlat index (this takes a few minutes)...\")\n",
    "    index.train(embeddings)\n",
    "else:\n",
    "    index = faiss.IndexFlatL2(dim)\n",
    "\n",
    "index.add(embeddings)\n",
    "faiss.write_index(index, INDEX_PATH)\n",
    "print(f\"FAISS index saved: {index.ntotal:,} vectors → {INDEX_PATH}\")\n",
    "\n",
    "# SQLite sidecar\n",
    "conn = sqlite3.connect(DB_PATH)\n",
    "c = conn.cursor()\n",
    "c.execute(\"\"\"CREATE TABLE IF NOT EXISTS chunks (\n",
    "    id INTEGER PRIMARY KEY,\n",
    "    text TEXT,\n",
    "    emotion_label INTEGER DEFAULT -1,\n",
    "    safety_score REAL DEFAULT 0.7,\n",
    "    source TEXT\n",
    ")\"\"\")\n",
    "\n",
    "BATCH = 10_000\n",
    "for i in tqdm(range(0, len(chunks), BATCH), desc=\"Writing SQLite\"):\n",
    "    batch = chunks[i:i+BATCH]\n",
    "    c.executemany(\n",
    "        \"INSERT OR REPLACE INTO chunks VALUES (?,?,?,?,?)\",\n",
    "        [(i+j, text, -1, 0.7, \"reddit\") for j, text in enumerate(batch)]\n",
    "    )\n",
    "    conn.commit()\n",
    "conn.close()\n",
    "print(f\"SQLite DB saved: {len(chunks):,} rows → {DB_PATH}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "431f2949",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 7: Verify outputs\n",
    "import faiss, sqlite3\n",
    "\n",
    "idx = faiss.read_index(INDEX_PATH)\n",
    "print(f\"FAISS index: {idx.ntotal:,} vectors\")\n",
    "\n",
    "conn = sqlite3.connect(DB_PATH)\n",
    "row_count = conn.execute(\"SELECT COUNT(*) FROM chunks\").fetchone()[0]\n",
    "sample = conn.execute(\"SELECT text FROM chunks LIMIT 3\").fetchall()\n",
    "conn.close()\n",
    "print(f\"SQLite rows: {row_count:,}\")\n",
    "print(\"Sample chunks:\")\n",
    "for i, (t,) in enumerate(sample):\n",
    "    print(f\"  [{i}] {t[:100]}...\")\n",
    "\n",
    "print()\n",
    "print(\"=== BUILD COMPLETE ===\")\n",
    "print(f\"Download these two files from Drive:\")\n",
    "print(f\"  {INDEX_PATH}\")\n",
    "print(f\"  {DB_PATH}\")\n",
    "print(\"Place them in your local data/indexes/ folder.\")"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}