scripts/generate_instruction_embeddings.py

#!/usr/bin/env python
"""
Generate instruction embeddings for all groups in dataset.

Usage:
    python scripts/generate_instruction_embeddings.py \
        --dataset /path/to/dataset \
        --output /path/to/embeddings.pkl
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from dovla_cil.data.datasets import CILDataset
from dovla_cil.utils.language_embeddings import LanguageEmbedder


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Generate instruction embeddings for dataset"
    )

    parser.add_argument(
        "--dataset",
        type=Path,
        required=True,
        help="Path to CIL dataset directory"
    )

    parser.add_argument(
        "--output",
        type=Path,
        required=True,
        help="Output path for embeddings (.pkl)"
    )

    parser.add_argument(
        "--model",
        default="all-mpnet-base-v2",
        help="SentenceTransformer model name"
    )

    parser.add_argument(
        "--cache-dir",
        type=Path,
        default=None,
        help="Cache directory for embeddings"
    )

    args = parser.parse_args(argv)

    print("=" * 70)
    print("Instruction Embedding Generation")
    print("=" * 70)
    print(f"Dataset: {args.dataset}")
    print(f"Output: {args.output}")
    print(f"Model: {args.model}")
    print()

    # Load dataset
    print("Loading dataset...")
    dataset = CILDataset(args.dataset)
    print(f"Found {len(dataset.group_ids)} groups")
    print()

    # Initialize embedder
    print(f"Loading embedding model: {args.model}")
    embedder = LanguageEmbedder(
        model_name=args.model,
        cache_dir=args.cache_dir
    )
    print()

    # Generate embeddings
    embeddings = embedder.encode_dataset(dataset, save_path=args.output)

    print()
    print(f"✅ Generated {len(embeddings)} embeddings")
    print(f"✅ Embedding dimension: {next(iter(embeddings.values())).shape[0]}")
    print(f"✅ Saved to: {args.output}")

    # Sample output
    sample_groups = list(embeddings.keys())[:3]
    print()
    print("Sample embeddings:")
    for gid in sample_groups:
        records = dataset.get_group(gid)
        instruction = records[0].instruction if records else "N/A"
        emb_norm = float((embeddings[gid] ** 2).sum() ** 0.5)
        print(f"  {gid}: '{instruction}' (norm={emb_norm:.2f})")

    return 0


if __name__ == "__main__":
    sys.exit(main())