Day 3: proper Colab ipynb notebook for FAISS build

notebooks/colab_build_faiss_index.ipynb

@@ -0,0 +1,274 @@
+{
+ "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
+}

notebooks/colab_build_faiss_index.py

@@ -1,164 +0,0 @@
-# EmpathRAG — FAISS Index Builder
-# Run on Google Colab Pro (A100).
-# Estimated time: 30-60 minutes for full Reddit Mental Health corpus.
-#
-# SETUP INSTRUCTIONS:
-# 1. Upload the entire data/raw/reddit_mental_health/ folder to Google Drive
-#    at: My Drive/empathrag/data/raw/reddit_mental_health/
-# 2. Set Colab runtime to A100 GPU
-# 3. Run all cells in order
-# 4. Download faiss_flat.index and metadata.db from Drive when done
-
-# ── Cell 1: Install ──────────────────────────────────────────────────────────
-# !pip install sentence-transformers faiss-cpu tqdm -q
-# !pip install transformers -q
-
-# ── Cell 2: Mount Drive ──────────────────────────────────────────────────────
-# from google.colab import drive
-# drive.mount("/content/drive")
-
-# BASE = "/content/drive/MyDrive/empathrag"
-# REDDIT_DIR = f"{BASE}/data/raw/reddit_mental_health"
-# INDEX_PATH = f"{BASE}/data/indexes/faiss_flat.index"
-# DB_PATH    = f"{BASE}/data/indexes/metadata.db"
-# import os
-# os.makedirs(f"{BASE}/data/indexes", exist_ok=True)
-
-# ── Cell 3: Build index ──────────────────────────────────────────────────────
-import os
-import re
-import sqlite3
-import faiss
-import numpy as np
-import pandas as pd
-from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer
-from tqdm import tqdm
-
-# ── Config ───────────────────────────────────────────────────────────────────
-# When running locally for testing, override these paths:
-REDDIT_DIR = "data/raw/reddit_mental_health"
-INDEX_PATH = "data/indexes/faiss_flat.index"
-DB_PATH    = "data/indexes/metadata.db"
-
-MODEL_NAME  = "sentence-transformers/all-mpnet-base-v2"
-CHUNK_SIZE  = 256
-STRIDE      = 32
-MAX_CHUNKS  = 8
-# A100: batch_size=128 is safe. RTX 3060 6GB laptop: use 16.
-BATCH_SIZE  = 128
-
-
-# ── Text cleaning ─────────────────────────────────────────────────────────────
-def clean_text(text: str) -> str:
-    text = re.sub(r"u/\w+", "", text)
-    text = re.sub(r"r/\w+", "", text)
-    text = re.sub(r"http\S+", "", text)
-    text = re.sub(r"\[deleted\]|\[removed\]", "", text)
-    text = re.sub(r"[^\x00-\x7F]+", " ", text)
-    text = re.sub(r"\s+", " ", text).strip()
-    return text
-
-
-# ── Chunking ──────────────────────────────────────────────────────────────────
-def chunk_text(text, tokenizer, chunk_size=CHUNK_SIZE, stride=STRIDE, max_chunks=MAX_CHUNKS):
-    tokens = tokenizer.encode(text)
-    if len(tokens) < 64:
-        return [text]
-    chunks = []
-    start = 0
-    while start < len(tokens) and len(chunks) < max_chunks:
-        end = min(start + chunk_size, len(tokens))
-        chunks.append(tokenizer.decode(tokens[start:end], skip_special_tokens=True))
-        start += chunk_size - stride
-    return chunks
-
-
-# ── Load posts ────────────────────────────────────────────────────────────────
-def load_reddit_posts(data_dir):
-    all_posts = []
-    files = [f for f in os.listdir(data_dir) if f.endswith(".csv")]
-    print(f"Loading from {len(files)} CSV files...")
-    for fname in tqdm(files, desc="Reading CSVs"):
-        fpath = os.path.join(data_dir, fname)
-        try:
-            df = pd.read_csv(fpath, on_bad_lines="skip", usecols=lambda c: c in ["post", "body", "selftext"])
-            for col in ["post", "body", "selftext"]:
-                if col in df.columns:
-                    all_posts.extend(df[col].dropna().tolist())
-                    break
-        except Exception as e:
-            print(f"  Skipping {fname}: {e}")
-    print(f"Total raw posts loaded: {len(all_posts):,}")
-    return all_posts
-
-
-# ── Main ──────────────────────────────────────────────────────────────────────
-def build_index():
-    os.makedirs(os.path.dirname(INDEX_PATH), exist_ok=True)
-
-    # Load and chunk
-    tok = AutoTokenizer.from_pretrained("roberta-base")
-    all_posts = load_reddit_posts(REDDIT_DIR)
-
-    chunks = []
-    for post in tqdm(all_posts, desc="Chunking"):
-        cleaned = clean_text(str(post))
-        if not cleaned:
-            continue
-        chunks.extend(chunk_text(cleaned, tok))
-
-    print(f"Total chunks to encode: {len(chunks):,}")
-
-    # Encode — use GPU automatically if available
-    encoder = SentenceTransformer(MODEL_NAME)
-    print(f"Encoding on: {encoder.device}")
-    embeddings = encoder.encode(
-        chunks,
-        batch_size=BATCH_SIZE,
-        show_progress_bar=True,
-        normalize_embeddings=True,
-        convert_to_numpy=True,
-    )
-    embeddings = np.array(embeddings, dtype=np.float32)
-    print(f"Embeddings shape: {embeddings.shape}")
-
-    # Build FAISS index
-    dim = embeddings.shape[1]  # 768
-    if len(chunks) > 100_000:
-        quantizer = faiss.IndexFlatL2(dim)
-        index = faiss.IndexIVFFlat(quantizer, dim, 100)
-        print("Training IVFFlat index...")
-        index.train(embeddings)
-    else:
-        index = faiss.IndexFlatL2(dim)
-    index.add(embeddings)
-    faiss.write_index(index, INDEX_PATH)
-    print(f"FAISS index saved: {index.ntotal:,} vectors → {INDEX_PATH}")
-
-    # SQLite sidecar
-    conn = sqlite3.connect(DB_PATH)
-    c = conn.cursor()
-    c.execute("""CREATE TABLE IF NOT EXISTS chunks (
-        id INTEGER PRIMARY KEY,
-        text TEXT,
-        emotion_label INTEGER DEFAULT -1,
-        safety_score REAL DEFAULT 0.7,
-        source TEXT
-    )""")
-    # Insert in batches to avoid memory spike
-    ISERT_BATCH = 10_000
-    for i in range(0, len(chunks), ISERT_BATCH):
-        batch = chunks[i:i+ISERT_BATCH]
-        c.executemany(
-            "INSERT OR REPLACE INTO chunks VALUES (?,?,?,?,?)",
-            [(i+j, text, -1, 0.7, "reddit") for j, text in enumerate(batch)]
-        )
-        conn.commit()
-    conn.close()
-    print(f"SQLite DB saved: {len(chunks):,} rows → {DB_PATH}")
-    print("Done. Download faiss_flat.index and metadata.db from Drive.")
-
-
-if __name__ == "__main__":
-    build_index()