scripts/build_calibration.py

#!/usr/bin/env python3
"""Build deterministic long-context calibration token caches."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from heapr.calibration import build_calibration_cache
from heapr.constants import DEFAULT_NUM_CHUNKS, DEFAULT_PRUNE_MODEL, DEFAULT_SEQ_LEN
from heapr.model_utils import load_tokenizer


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--tokenizer-id", default=DEFAULT_PRUNE_MODEL)
    parser.add_argument("--revision")
    parser.add_argument("--output-root", default="artifacts/data/calibration")
    parser.add_argument("--seq-len", type=int, default=DEFAULT_SEQ_LEN)
    parser.add_argument("--num-chunks", type=int, default=DEFAULT_NUM_CHUNKS)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--text-fraction", type=float, default=0.5)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    tokenizer = load_tokenizer(args.tokenizer_id, revision=args.revision)
    manifest = build_calibration_cache(
        tokenizer,
        output_root=args.output_root,
        seq_len=args.seq_len,
        num_chunks=args.num_chunks,
        seed=args.seed,
        text_fraction=args.text_fraction,
    )
    print(manifest.output_path)


if __name__ == "__main__":
    main()