#!/usr/bin/env python3
"""Lost in Embedding: invert GPT-2 inputs_embeds by nearest-neighbour on wte (optionally minus wpe)."""
import json, sys, os
import numpy as np

D = os.path.dirname(os.path.abspath(__file__))
cap = json.load(open(sys.argv[1] if len(sys.argv) > 1 else f"{D}/files/capture.json"))
E = np.array(cap["inputs_embeds"], dtype=np.float32)
wte = np.fromfile(f"{D}/model/wte.bin", dtype="<f4").reshape(50257, 768)
wpe = np.fromfile(f"{D}/model/wpe.bin", dtype="<f4").reshape(1024, 768)
vocab = {v: k for k, v in json.load(open(f"{D}/model/vocab.json")).items()}

# GPT-2 byte-level BPE: map unicode chars back to bytes
bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
cs = bs[:]; n = 0
for b in range(256):
    if b not in bs: bs.append(b); cs.append(256 + n); n += 1
byte_dec = {chr(c): b for b, c in zip(bs, cs)}

def decode(ids):
    return bytes(byte_dec[c] for i in ids for c in vocab[i]).decode("utf-8", "replace")

for name, X in [("raw", E), ("minus_wpe", E - wpe[: len(E)])]:
    d = ((X[:, None, :] - wte[None]) ** 2).sum(-1)
    ids = d.argmin(1)
    print(name, "min dist:", np.round(d.min(1), 4).tolist())
    print(name, "->", repr(decode(ids)))