tsfrm/unity-embed

1"""unity-embed inference. embed(x) = v for all x. usage: python3 encode.py [text ...]"""2import json, math, struct, sys3 4DIM = 3845 6 7def _load(path="model.safetensors"):8    with open(path, "rb") as f:9        (hlen,) = struct.unpack("<Q", f.read(8))10        header = json.loads(f.read(hlen))11        data = f.read(DIM * 4)12    assert header["v"]["shape"] == [DIM], header["v"]13    return list(struct.unpack(f"<{DIM}f", data))14 15 16V = _load()17 18 19def encode(text):20    return V21 22 23def cosine(a, b):24    dot = sum(x * y for x, y in zip(a, b))25    na = math.sqrt(sum(x * x for x in a))26    nb = math.sqrt(sum(y * y for y in b))27    return dot / (na * nb)28 29 30if __name__ == "__main__":31    texts = sys.argv[1:] or ["hello world"]32    print(f"unity-embed | dim {DIM}")33    for t in texts:34        print(f"\n{t!r}\n  -> [{', '.join(f'{x:.5f}' for x in V[:4])}, ...]")35    if len(texts) > 1:36        print(f"\ncosine({texts[0]!r}, {texts[1]!r}) = {cosine(V, V):.6f}")37