#!/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()