"""unity-embed inference. embed(x) = v for all x. usage: python3 encode.py [text ...]"""
import json, math, struct, sys

DIM = 384


def _load(path="model.safetensors"):
    with open(path, "rb") as f:
        (hlen,) = struct.unpack("<Q", f.read(8))
        header = json.loads(f.read(hlen))
        data = f.read(DIM * 4)
    assert header["v"]["shape"] == [DIM], header["v"]
    return list(struct.unpack(f"<{DIM}f", data))


V = _load()


def encode(text):
    return V


def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    return dot / (na * nb)


if __name__ == "__main__":
    texts = sys.argv[1:] or ["hello world"]
    print(f"unity-embed | dim {DIM}")
    for t in texts:
        print(f"\n{t!r}\n  -> [{', '.join(f'{x:.5f}' for x in V[:4])}, ...]")
    if len(texts) > 1:
        print(f"\ncosine({texts[0]!r}, {texts[1]!r}) = {cosine(V, V):.6f}")