| """npu-forge-ear: train a tiny voice-verifier HEAD on EmbeddingGemma-300m |
| embeddings. At runtime the embeddings come from FLM's NPU (embed-gemma:300m, |
| unquantized) and the head is ~20KB of JSON applied in plain JS — a fully |
| NPU-accelerated 'trained ear' with zero extra runtimes. |
| |
| modal run train_ear_head.py |
| """ |
| import modal |
|
|
| app = modal.App("npu-forge-ear") |
| image = ( |
| modal.Image.debian_slim(python_version="3.11") |
| .pip_install("torch==2.5.1", "transformers==4.56.2", "sentence-transformers==5.1.0", "numpy", "huggingface_hub") |
| .add_local_dir("C:/Users/Forgemind/Desktop/voice-harness/scratch/verifier-data", remote_path="/data") |
| ) |
| vol = modal.Volume.from_name("npu-forge-out", create_if_missing=True) |
|
|
|
|
| @app.function(image=image, gpu="T4", timeout=2400, volumes={"/out": vol}, secrets=[modal.Secret.from_name("huggingface-token")]) |
| def train(): |
| import json, numpy as np, torch |
| from sentence_transformers import SentenceTransformer |
|
|
| def load(p): |
| return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()] |
|
|
| labels = json.load(open("/data/labels.json")) |
| lab2id = {l: i for i, l in enumerate(labels)} |
| train_rows, eval_rows, probe_rows = load("/data/train.jsonl"), load("/data/eval.jsonl"), load("/data/probe.jsonl") |
|
|
| candidates = ["google/embeddinggemma-300m", "unsloth/embeddinggemma-300m", "onnx-community/embeddinggemma-300m-ONNX"] |
| model = None |
| for cid in candidates: |
| try: |
| model = SentenceTransformer(cid, device="cuda") |
| print("[embed model] " + cid) |
| break |
| except Exception as e: |
| print(f"[skip] {cid}: {str(e)[:120]}") |
| if model is None: |
| raise RuntimeError("no embeddinggemma variant loadable") |
|
|
| def embed(rows): |
| return np.asarray(model.encode([r["text"] for r in rows], batch_size=64, show_progress_bar=False, normalize_embeddings=True)) |
|
|
| Xtr, Xev, Xpr = embed(train_rows), embed(eval_rows), embed(probe_rows) |
| ytr = np.array([lab2id[r["label"]] for r in train_rows]) |
|
|
| |
| d, k = Xtr.shape[1], len(labels) |
| W = torch.zeros(d, k, requires_grad=True) |
| b = torch.zeros(k, requires_grad=True) |
| Xt = torch.tensor(Xtr, dtype=torch.float32) |
| yt = torch.tensor(ytr) |
| opt = torch.optim.Adam([W, b], lr=0.05) |
| for i in range(400): |
| loss = torch.nn.functional.cross_entropy(Xt @ W + b, yt) + 1e-4 * W.pow(2).sum() |
| opt.zero_grad(); loss.backward(); opt.step() |
| print(f"[head] final loss {loss.item():.4f}, dim={d}") |
|
|
| def acc(X, rows): |
| pred = (torch.tensor(X, dtype=torch.float32) @ W + b).argmax(-1).numpy() |
| per = {} |
| for r, p in zip(rows, pred): |
| s = per.setdefault(r["label"], [0, 0]); s[1] += 1 |
| if labels[p] == r["label"]: s[0] += 1 |
| overall = sum(v[0] for v in per.values()) / max(1, len(rows)) |
| return overall, {l: f"{v[0]}/{v[1]}" for l, v in per.items()} |
|
|
| ev_overall, ev_per = acc(Xev, eval_rows) |
| pr_overall, pr_per = acc(Xpr, probe_rows) |
|
|
| head = {"labels": labels, "dim": d, "normalize": True, |
| "W": W.detach().numpy().round(5).tolist(), "b": b.detach().numpy().round(5).tolist()} |
| import os |
| os.makedirs("/out/ear-head", exist_ok=True) |
| with open("/out/ear-head/head.json", "w") as f: |
| json.dump(head, f) |
| vol.commit() |
| return {"eval_overall": round(ev_overall, 4), "eval_per_class": ev_per, |
| "probe_recognized": pr_per, "probe_overall": round(pr_overall, 4), |
| "head_kb": round(len(json.dumps(head)) / 1024)} |
|
|
|
|
| @app.local_entrypoint() |
| def main(): |
| import json |
| print(json.dumps(train.remote(), indent=2)) |
|
|
