scripts/reembed_all.py

"""
GovBridge India — Full Re-embedding Script
Run AFTER: Supabase migration 003 is executed
Purpose: Re-embed all documents with nomic-embed-text-v1 (768-dim)
Usage: python3 gov_backend/scripts/reembed_all.py
"""
import os
import time
from supabase import create_client
from sentence_transformers import SentenceTransformer
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import settings

# Load environment variables
SUPABASE_URL = settings.SUPABASE_URL
SUPABASE_KEY = settings.SUPABASE_KEY
BATCH_SIZE = 50

def main():
    if not SUPABASE_URL or not SUPABASE_KEY:
        print("❌ Missing SUPABASE_URL or SUPABASE_KEY")
        return

    print("⏳ Loading nomic-embed-text-v1 model...")
    model = SentenceTransformer(
        'nomic-ai/nomic-embed-text-v1',
        trust_remote_code=True
    )
    print("✅ Model loaded")

    supabase = create_client(SUPABASE_URL, SUPABASE_KEY)

    # Fetch all chunks where embedding is null
    result = supabase.table('document_chunks') \
        .select('id, chunk_text, scheme_title') \
        .is_('embedding', 'null') \
        .limit(1000) \
        .execute()

    chunks = result.data or []
    if not chunks:
        print("🎉 No chunks found needing re-embedding.")
        return

    print(f"Found {len(chunks)} chunks to embed. Starting migration...")

    for i in range(0, len(chunks), BATCH_SIZE):
        batch = chunks[i:i+BATCH_SIZE]
        # Use Nomic task prefix for documents
        texts = [
            f"search_document: {c.get('chunk_text', '') or c.get('scheme_title', '')}"
            for c in batch
        ]
        
        print(f"🔄 Processing batch {i//BATCH_SIZE + 1}...")
        embeddings = model.encode(
            texts,
            normalize_embeddings=True,
            show_progress_bar=False
        ).tolist()

        for chunk, embedding in zip(batch, embeddings):
            supabase.table('document_chunks') \
                .update({'embedding': embedding}) \
                .eq('id', chunk['id']) \
                .execute()

        print(f"✅ Batch {i//BATCH_SIZE + 1} complete.")
        time.sleep(0.5)

    print("🎉 All document chunks have been re-embedded with 768 dimensions.")

if __name__ == '__main__':
    main()