scripts/make_golden.py

"""Record native (PyTorch) GLiNER2.5-Decide answers for tests/cases.py into tests/golden.json.

Downloads the pinned source checkpoint (~1.7 GB) on first run. The tests compare the Core ML
router against this file, so they need neither PyTorch weights nor a network connection.

    uv run python scripts/make_golden.py
"""
import json
import sys
import warnings
from pathlib import Path

warnings.filterwarnings("ignore")
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "tests"))
sys.path.insert(0, str(ROOT / "scripts" / "build"))

from gliner2 import AutoExtractor  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402

from cases import CASES  # noqa: E402
from convert_names import MODEL_ID, MODEL_REVISION  # noqa: E402


def main():
    source = snapshot_download(
        MODEL_ID, revision=MODEL_REVISION,
        allow_patterns=["config.json", "encoder_config/*", "model.safetensors", "tokenizer.json",
                        "tokenizer_config.json", "special_tokens_map.json"],
    )
    native = AutoExtractor.from_pretrained(source, map_location="cpu").eval()
    golden = {}
    for case in CASES:
        golden[case["id"]] = native.classify_text(case["text"], case["tasks"], include_confidence=True)
        print(case["id"], json.dumps(golden[case["id"]]))
    out = ROOT / "tests" / "golden.json"
    out.write_text(json.dumps({"source_model": MODEL_ID, "source_revision": MODEL_REVISION, "cases": golden},
                              indent=2) + "\n")
    print("wrote", out)


if __name__ == "__main__":
    main()