scripts/switch_model.py

#!/usr/bin/env python3
"""
switch_model.py
===============
Convenience script to switch the embedding model and rebuild all vectors.

Usage:
    # Switch to CAMeL classical Arabic (recommended for fiqh):
    python scripts/switch_model.py camel

    # Switch back to the general Arabic model:
    python scripts/switch_model.py asafaya

    # Use any custom HuggingFace model ID:
    python scripts/switch_model.py custom --model "aubmindlab/bert-base-arabertv2"

    # Dry-run: just show what would happen
    python scripts/switch_model.py camel --dry-run
"""
from __future__ import annotations

import sys
import argparse
import subprocess
from pathlib import Path

MODELS = {
    "camel":   "CAMeL-Lab/bert-base-arabic-camelbert-ca",
    "asafaya": "asafaya/bert-base-arabic",
    "arabert": "aubmindlab/bert-base-arabertv2",
    "arbert":  "UBC-NLP/ARBERT",
}

SCRIPT = Path(__file__).parent / "generate_embeddings.py"


def main() -> None:
    parser = argparse.ArgumentParser(description="Switch embedding model and rebuild vectors.")
    parser.add_argument(
        "preset",
        choices=[*MODELS.keys(), "custom"],
        help="Model preset name, or 'custom' to pass --model directly",
    )
    parser.add_argument(
        "--model",
        help="HuggingFace model ID (required when preset=custom)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print the command that would be run without executing it",
    )
    parser.add_argument(
        "--threads",
        type=int,
        default=4,
        help="CPU thread count for encoding (default: 4)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=96,
        help="Encoding batch size (default: 96)",
    )
    args = parser.parse_args()

    if args.preset == "custom":
        if not args.model:
            print("❌  --model is required when preset=custom")
            sys.exit(1)
        model_id = args.model
    else:
        model_id = MODELS[args.preset]

    cmd = [
        sys.executable,
        str(SCRIPT),
        "--model", model_id,
        "--reset",
        "--threads", str(args.threads),
        "--batch-size", str(args.batch_size),
    ]

    print("=" * 60)
    print(f"Switching embedding model to: {model_id}")
    print(f"This will WIPE all existing embeddings and rebuild from scratch.")
    print(f"Estimated time: 1-2 hours on Mac CPU, ~20 min on a GPU.")
    print("=" * 60)

    if args.dry_run:
        print("\n[DRY RUN] Would execute:")
        print("  " + " ".join(cmd))
        return

    confirm = input("\nProceed? [y/N] ").strip().lower()
    if confirm != "y":
        print("Aborted.")
        sys.exit(0)

    subprocess.run(cmd, check=True)


if __name__ == "__main__":
    main()