| 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]) |
|
|
