File size: 2,611 Bytes
adc02fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | #!/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())
|