| """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}") |
|
|
