#!/usr/bin/env python3
# Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE).
"""Regenerate `smoke-reference.json`: fp32 reference embeddings for three fictional sentences, computed with
sentence-transformers on the ORIGINAL google/embeddinggemma-2 checkpoint (text only, CPU, float32).

    python3 make_smoke_reference.py --model google/embeddinggemma-2 --out smoke-reference.json

Needs: sentence-transformers >= 6.1, transformers >= 5.19, torch, torchvision and pillow (the checkpoint's processor
is multimodal and imports them even for text). The shipped file was produced with sentence-transformers 6.1.0 /
transformers 5.19.0 / torch 2.14.1; torchvision 0.29.1 is the matching version (see ../python/requirements.lock).
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

CASES = [
    {"kind": "search_query", "text": "When does the warehouse lease in Brindlemoor expire?"},
    {"kind": "document", "text": "The warehouse lease in Brindlemoor runs until the last day of March; the landlord agreed to fix the window before then."},
    {"kind": "document", "text": "Lista della spesa: latte d'avena, due limoni, riso basmati, coriandolo e croccantini per il cane."},
]
PROMPT_NAME = {"search_query": "query", "document": "document"}


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--model", default="google/embeddinggemma-2", help="HF id or local checkpoint folder")
    ap.add_argument("--out", default="smoke-reference.json")
    ap.add_argument("--threshold", type=float, default=0.999)
    args = ap.parse_args()
    import sentence_transformers
    import torch
    import transformers
    from sentence_transformers import SentenceTransformer

    model = SentenceTransformer(args.model, device="cpu", model_kwargs={"torch_dtype": torch.float32})
    out = {
        "reference": "sentence-transformers fp32 on CPU, original checkpoint, text only",
        "versions": {"sentence_transformers": sentence_transformers.__version__, "transformers": transformers.__version__,
                     "torch": torch.__version__},
        "threshold_cosine": args.threshold,
        "cases": [],
    }
    for c in CASES:
        v = model.encode(c["text"], prompt_name=PROMPT_NAME[c["kind"]], normalize_embeddings=True, convert_to_numpy=True)
        out["cases"].append({**c, "embedding": [round(float(x), 7) for x in v.tolist()]})
        print(c["kind"], len(v), c["text"][:50])
    Path(args.out).write_text(json.dumps(out, ensure_ascii=False, indent=1), encoding="utf-8")
    print("wrote", args.out)


if __name__ == "__main__":
    main()