File size: 956 Bytes
87fe5d9 49b416c 87fe5d9 49b416c 87fe5d9 49b416c 87fe5d9 49b416c 87fe5d9 49b416c 87fe5d9 49b416c 87fe5d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | """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}")
|