import sys, time, numpy as np, coremltools as ct
from transformers import AutoTokenizer
path = sys.argv[1]
tok = AutoTokenizer.from_pretrained("google/embeddinggemma-2")
texts = ["task: search result | query: quando finisce la posizione dei QA?",
         "title: none | text: Marco: allora il contratto del tester scade il primo gennaio 2027, poi vediamo.",
         "title: none | text: Giulia porta il vino alla cena di domenica, Maria Grazia il dolce.",
         "task: search result | query: chi porta il dolce?"]
ref = np.load("ref.npy")
for units in [ct.ComputeUnit.CPU_AND_NE]:
    m = ct.models.MLModel(path, compute_units=units)
    for L in [int(x) for x in sys.argv[2].split(",")]:
        outs = []
        for t in texts:
            ids = np.zeros((1, L), dtype=np.int32); enc = tok(t)["input_ids"][:L]; ids[0, :len(enc)] = enc
            outs.append(m.predict({"input_ids": ids})["embedding"][0])
        outs = np.array(outs)
        start = time.time()
        for _ in range(20): m.predict({"input_ids": ids})
        ms = (time.time() - start) / 20 * 1000
        print(f"L={L}: coseno con ufficiale {np.round((outs*ref).sum(1),4)} · {ms:.1f} ms per pezzo")
    print("similarità QA:", np.round(outs @ outs.T, 3)[0], "dolce:", np.round(outs @ outs.T, 3)[3])