"""Which token embeddings did fine-tuning move most? Streams only the embedding tensors."""
import sys, glob, os, json
import torch
from huggingface_hub import snapshot_download
from safetensors import safe_open
from transformers import AutoTokenizer

def emb(rid):
    p = snapshot_download(rid, allow_patterns=["*.safetensors", "*.json"])
    for f in sorted(glob.glob(os.path.join(p, "*.safetensors"))):
        with safe_open(f, framework="pt") as fh:
            for k in fh.keys():
                if k.endswith("embed_tokens.weight") or k.endswith("wte.weight"):
                    return fh.get_tensor(k).float()

m, p = sys.argv[1], sys.argv[2]
watch = sys.argv[3:]
tok = AutoTokenizer.from_pretrained(m)
Em, Ep = emb(m), emb(p)
n = min(Em.shape[0], Ep.shape[0], len(tok))
rel = (Em[:n] - Ep[:n]).norm(dim=1) / (Ep[:n].norm(dim=1) + 1e-6)
order = rel.argsort(descending=True)
rank = torch.empty_like(order); rank[order] = torch.arange(n)
print("median rel change", rel.median().item(), "n tokens", n)
print("top 40:", [(tok.convert_tokens_to_string([tok.convert_ids_to_tokens(int(i))]), round(rel[i].item(), 4)) for i in order[:40]])
for w in watch:
    for v in (w, " " + w, " " + w.capitalize(), w.capitalize()):
        ids = tok(v, add_special_tokens=False)["input_ids"]
        print(repr(v), ids, [(int(rank[i]), round(rel[i].item(), 4)) for i in ids])