backend/scripts/generate_embeddings.py

"""
Script to generate and store embeddings for Procedure, Fine, Office, Advisory models.
"""
import argparse
import os
import sys
from pathlib import Path
from typing import List, Tuple
import numpy as np

ROOT_DIR = Path(__file__).resolve().parents[2]
BACKEND_DIR = ROOT_DIR / "backend"
HUE_PORTAL_DIR = BACKEND_DIR / "hue_portal"

# Add backend directory to sys.path so Django can find hue_portal package
# Django needs to import hue_portal.hue_portal.settings, so backend/ must be in path
# IMPORTANT: Only add BACKEND_DIR, not HUE_PORTAL_DIR, because Django needs to find
# the hue_portal package (which is in backend/hue_portal), not the hue_portal directory itself
if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))

# Add root for other imports if needed (but not HUE_PORTAL_DIR as it breaks Django imports)
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "hue_portal.hue_portal.settings")

import django
django.setup()

from hue_portal.core.models import Procedure, Fine, Office, Advisory, LegalSection
from hue_portal.core.embeddings import (
    get_embedding_model,
    generate_embeddings_batch,
    get_embedding_dimension
)


def prepare_text_for_embedding(obj) -> str:
    """
    Prepare text from model instance for embedding.
    """
    if isinstance(obj, Procedure):
        fields = [obj.title, obj.domain, obj.level, obj.conditions, obj.dossier]
    elif isinstance(obj, Fine):
        fields = [obj.name, obj.code, obj.article, obj.decree, obj.remedial]
    elif isinstance(obj, Office):
        fields = [obj.unit_name, obj.address, obj.district, obj.service_scope]
    elif isinstance(obj, Advisory):
        fields = [obj.title, obj.summary]
    elif isinstance(obj, LegalSection):
        fields = [obj.section_code, obj.section_title, obj.content, getattr(obj.document, "title", "")]
    else:
        return ""
    
    # Combine non-empty fields
    text = " ".join(str(f) for f in fields if f and str(f).strip())
    return text.strip()


def generate_embeddings_for_model(model_class, model_name: str, batch_size: int = 32, dry_run: bool = False):
    """
    Generate embeddings for all instances of a model.
    
    Args:
        model_class: Django model class.
        model_name: Name of the model (for display).
        batch_size: Batch size for processing.
        dry_run: If True, only show what would be done without saving.
    """
    print(f"\n{'='*60}")
    print(f"Processing {model_name}")
    print(f"{'='*60}")
    
    # Get all instances
    instances = list(model_class.objects.all())
    total = len(instances)
    
    if total == 0:
        print(f"No {model_name} instances found. Skipping.")
        return 0, 0
    
    print(f"Found {total} {model_name} instances")
    
    # Prepare texts
    texts = []
    valid_indices = []
    for idx, instance in enumerate(instances):
        text = prepare_text_for_embedding(instance)
        if text:
            texts.append(text)
            valid_indices.append(idx)
        else:
            print(f"⚠️ Skipping {model_name} ID {instance.id}: empty text")
    
    if not texts:
        print(f"No valid texts found for {model_name}. Skipping.")
        return 0, 0
    
    print(f"Generating embeddings for {len(texts)} valid instances...")
    
    # Load model
    model = get_embedding_model()
    if model is None:
        print(f"❌ Cannot load embedding model. Skipping {model_name}.")
        return 0, 0
    
    # Generate embeddings
    embeddings = generate_embeddings_batch(texts, model=model, batch_size=batch_size)
    
    # Save embeddings (if not dry run)
    saved = 0
    failed = 0
    
    for idx, embedding in zip(valid_indices, embeddings):
        instance = instances[idx]
        
        if embedding is None:
            print(f"⚠️ Failed to generate embedding for {model_name} ID {instance.id}")
            failed += 1
            continue
        
        if not dry_run:
            # Convert numpy array to binary for storage
            try:
                import pickle
                embedding_binary = pickle.dumps(embedding)
                instance.embedding = embedding_binary
                instance.save(update_fields=['embedding'])
                print(f"✅ Generated and saved embedding for {model_name} ID {instance.id} (dim={len(embedding)})")
                saved += 1
            except Exception as e:
                print(f"❌ Error saving embedding for {model_name} ID {instance.id}: {e}")
                failed += 1
        else:
            print(f"[DRY RUN] Would save embedding for {model_name} ID {instance.id} (dim={len(embedding)})")
            saved += 1
    
    print(f"\n{model_name} Summary: {saved} saved, {failed} failed")
    return saved, failed


def main():
    parser = argparse.ArgumentParser(description="Generate embeddings for all models")
    parser.add_argument("--model", choices=["procedure", "fine", "office", "advisory", "legal", "all"], 
                       default="all", help="Which model to process")
    parser.add_argument("--batch-size", type=int, default=32, help="Batch size for embedding generation")
    parser.add_argument("--dry-run", action="store_true", help="Simulate without saving")
    parser.add_argument("--model-name", type=str, help="Override embedding model name")
    args = parser.parse_args()
    
    print("="*60)
    print("Embedding Generation Script")
    print("="*60)
    
    if args.dry_run:
        print("⚠️ DRY RUN MODE - No changes will be saved")
    
    if args.model_name:
        print(f"Using model: {args.model_name}")
        get_embedding_model(model_name=args.model_name, force_reload=True)
    else:
        print(f"Using default model: keepitreal/vietnamese-sbert-v2")
    
    # Check model dimension
    dim = get_embedding_dimension()
    if dim > 0:
        print(f"Embedding dimension: {dim}")
    else:
        print("⚠️ Could not determine embedding dimension")
    
    total_saved = 0
    total_failed = 0
    
    models_to_process = []
    if args.model == "all":
        models_to_process = [
            (Procedure, "Procedure"),
            (Fine, "Fine"),
            (Office, "Office"),
            (Advisory, "Advisory"),
            (LegalSection, "LegalSection"),
        ]
    else:
        model_map = {
            "procedure": (Procedure, "Procedure"),
            "fine": (Fine, "Fine"),
            "office": (Office, "Office"),
            "advisory": (Advisory, "Advisory"),
            "legal": (LegalSection, "LegalSection"),
        }
        if args.model in model_map:
            models_to_process = [model_map[args.model]]
    
    for model_class, model_name in models_to_process:
        saved, failed = generate_embeddings_for_model(
            model_class, model_name, 
            batch_size=args.batch_size,
            dry_run=args.dry_run
        )
        total_saved += saved
        total_failed += failed
    
    print("\n" + "="*60)
    print("Final Summary")
    print("="*60)
    print(f"Total saved: {total_saved}")
    print(f"Total failed: {total_failed}")
    print("="*60)


if __name__ == "__main__":
    main()