scripts/generate_embeddings.py

#!/usr/bin/env python3
"""
generate_embeddings.py
======================
Generates or regenerates semantic embeddings for all chunks in fiqh.db.

Usage:
    # Normal run (resumes from where it left off):
    python scripts/generate_embeddings.py

    # Switch to a NEW model (wipes old embeddings first, then rebuilds all):
    python scripts/generate_embeddings.py --model CAMeL-Lab/bert-base-arabic-camelbert-ca --reset

    # Check which model was used for current embeddings:
    python scripts/generate_embeddings.py --info
"""
from __future__ import annotations

import sys
import time
import argparse
import sqlite3
from pathlib import Path

CURRENT_FILE = Path(__file__).resolve()
API_DIR = CURRENT_FILE.parents[1]
sys.path.insert(0, str(API_DIR))

from app.config import DB_PATH  # noqa: E402

# ── Default model ──────────────────────────────────────────────────────────────
# CAMeL-BERT is trained on Classical + Modern Arabic, much better for fiqh texts.
# asafaya/bert-base-arabic = general Arabic (old default)
# CAMeL-Lab/bert-base-arabic-camelbert-ca = Classical Arabic (recommended for fiqh)
DEFAULT_MODEL = "CAMeL-Lab/bert-base-arabic-camelbert-ca"


def get_current_model(conn: sqlite3.Connection) -> str | None:
    """Read which model was used to generate the stored embeddings."""
    try:
        row = conn.execute(
            "SELECT value FROM embedding_meta WHERE key = 'model_name'"
        ).fetchone()
        return row[0] if row else None
    except sqlite3.OperationalError:
        return None


def set_current_model(conn: sqlite3.Connection, model_name: str) -> None:
    """Persist the model name into embedding_meta table."""
    conn.execute("""
        CREATE TABLE IF NOT EXISTS embedding_meta (
            key TEXT PRIMARY KEY,
            value TEXT NOT NULL
        )
    """)
    conn.execute(
        "INSERT OR REPLACE INTO embedding_meta (key, value) VALUES ('model_name', ?)",
        (model_name,)
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate semantic embeddings for fiqh.db chunks.")
    parser.add_argument(
        "--model",
        default=DEFAULT_MODEL,
        help=f"HuggingFace model name (default: {DEFAULT_MODEL})"
    )
    parser.add_argument(
        "--reset",
        action="store_true",
        help="Wipe all existing embeddings and regenerate from scratch (required when switching models)"
    )
    parser.add_argument(
        "--info",
        action="store_true",
        help="Print info about current embeddings and exit"
    )
    parser.add_argument(
        "--threads",
        type=int,
        default=4,
        help="PyTorch CPU thread count (default: 4)"
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=96,
        help="Encoding batch size (default: 96)"
    )
    args = parser.parse_args()

    if not DB_PATH.exists():
        print(f"❌ Error: Database not found at {DB_PATH}. Run ingest first.")
        sys.exit(1)

    conn = sqlite3.connect(DB_PATH)
    conn.execute("PRAGMA journal_mode = WAL")
    conn.execute("PRAGMA synchronous = OFF")
    conn.execute("PRAGMA temp_store = MEMORY")
    conn.execute("PRAGMA cache_size = -128000")

    # ── Info mode ──────────────────────────────────────────────────────────────
    if args.info:
        current_model = get_current_model(conn)
        try:
            count = conn.execute("SELECT COUNT(*) FROM chunk_embeddings").fetchone()[0]
            total = conn.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
            print(f"Embedding model : {current_model or 'unknown (old index)'}")
            print(f"Indexed         : {count:,} / {total:,} chunks ({count/total*100:.1f}%)")
        except Exception as e:
            print(f"No embeddings found: {e}")
        conn.close()
        return

    # ── Detect model mismatch ──────────────────────────────────────────────────
    current_model = get_current_model(conn)
    if current_model and current_model != args.model and not args.reset:
        print(f"⚠️  WARNING: Existing embeddings were generated with: {current_model}")
        print(f"   You are trying to add embeddings with          : {args.model}")
        print(f"   These are INCOMPATIBLE vector spaces!")
        print(f"   Run with --reset to wipe and rebuild from scratch:")
        print(f"   python scripts/generate_embeddings.py --model {args.model} --reset")
        conn.close()
        sys.exit(1)

    # ── Initialize / Reset embedding table ────────────────────────────────────
    conn.execute("""
        CREATE TABLE IF NOT EXISTS chunk_embeddings (
            chunk_id INTEGER PRIMARY KEY REFERENCES chunks(id),
            embedding BLOB NOT NULL
        )
    """)
    conn.commit()

    if args.reset:
        print(f"🗑️  --reset: Dropping all existing embeddings...")
        conn.execute("DELETE FROM chunk_embeddings")
        conn.commit()
        print(f"   Wiped. Starting fresh with model: {args.model}")

    # ── Fetch chunks without embeddings ───────────────────────────────────────
    rows = conn.execute("""
        SELECT id, text_normalized FROM chunks
        WHERE id NOT IN (SELECT chunk_id FROM chunk_embeddings)
        ORDER BY id
    """).fetchall()

    if not rows:
        print(f"✅ All chunks already have embeddings (model: {current_model or args.model})")
        conn.close()
        return

    total_chunks = len(rows)
    print(f"Found {total_chunks:,} chunk(s) needing semantic embeddings.")
    print(f"Model : {args.model}")
    print(f"Device: Optimized CPU ({args.threads} threads)")
    print(f"Batch : {args.batch_size}")

    # ── Load model ─────────────────────────────────────────────────────────────
    import torch
    torch.set_num_threads(args.threads)
    from sentence_transformers import SentenceTransformer

    device = "cpu"
    if torch.cuda.is_available():
        device = "cuda"
    elif torch.backends.mps.is_available():
        device = "mps"

    print(f"\nLoading model '{args.model}' on device: {device}...")
    t_load = time.time()
    model = SentenceTransformer(args.model, device=device)
    print(f"Model loaded in {time.time() - t_load:.1f}s\n")

    # Persist model name before we start writing vectors
    set_current_model(conn, args.model)
    conn.commit()

    # ── Encode in batches ──────────────────────────────────────────────────────
    t_start = time.time()
    committed = 0
    COMMIT_EVERY = 40  # batches

    for batch_idx, i in enumerate(range(0, total_chunks, args.batch_size)):
        batch = rows[i : i + args.batch_size]
        texts = [r[1] for r in batch]
        ids   = [r[0] for r in batch]

        encode_kwargs = {
            "show_progress_bar": False,
            "batch_size": args.batch_size,
            "convert_to_numpy": True,
        }

        if device == "cuda":
            with torch.cuda.amp.autocast():
                vectors = model.encode(texts, **encode_kwargs)
        else:
            vectors = model.encode(texts, **encode_kwargs)

        conn.executemany(
            "INSERT OR REPLACE INTO chunk_embeddings (chunk_id, embedding) VALUES (?, ?)",
            [
                (chunk_id, vector.astype("float32").tobytes())
                for chunk_id, vector in zip(ids, vectors)
            ]
        )

        if batch_idx % COMMIT_EVERY == 0:
            conn.commit()
            committed = i + len(batch)

        # Progress + ETA
        done = i + len(batch)
        elapsed = time.time() - t_start
        rate = done / elapsed if elapsed > 0 else 1
        remaining = (total_chunks - done) / rate if rate > 0 else 0
        eta_min = int(remaining // 60)
        eta_sec = int(remaining % 60)
        print(
            f"Progress: {done:>7,}/{total_chunks:,} ({done/total_chunks*100:5.1f}%)  "
            f"| {rate:,.0f} chunks/s  "
            f"| ETA {eta_min}m {eta_sec:02d}s"
        )

    conn.commit()
    conn.close()

    total_time = time.time() - t_start
    print(f"\n★ Done! Indexed {total_chunks:,} chunks with '{args.model}'")
    print(f"  Total time: {int(total_time//60)}m {int(total_time%60):02d}s")
    print(f"\nNext step: re-deploy or restart the API server.")


if __name__ == "__main__":
    main()