| |
| """Check a reader head against its published evidence, with nothing but this file. |
| |
| Two checks, both against files on the RiverRider/srt-reader-heads and |
| RiverRider/srt-browser-head-118k repositories: |
| |
| parity re-encode the eight probe captions with the encoder, apply the head, |
| and compare the embedding, the projection and the top-5 gallery rows |
| with the recorded fixture (parity_<reader>.json) |
| replay score the 5,001 held-out captions (replay_123k.json) against every |
| row of the 123,287-row gallery and report t2i R@1/5/10 and median |
| rank, the numbers in the model card and reader_swap_123k.json |
| |
| pip install sentence-transformers safetensors huggingface_hub numpy torch |
| python check_head.py --reader bge-small # parity + replay |
| python check_head.py --reader e5-base --skip-replay |
| |
| The files are fetched from the Hub on first use (the gallery is 130 MB). |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import struct |
| import time |
|
|
| import numpy as np |
| import torch |
| from huggingface_hub import hf_hub_download |
| from safetensors.torch import load_file |
| from sentence_transformers import SentenceTransformer |
|
|
| HEADS = "RiverRider/srt-reader-heads" |
| GALLERY = "RiverRider/srt-browser-head-118k" |
| READERS = { |
| "bge-small": ("BAAI/bge-small-en-v1.5", ""), |
| "minilm": ("sentence-transformers/all-MiniLM-L6-v2", ""), |
| "e5-base": ("intfloat/e5-base-v2", "query: "), |
| "gte-base": ("thenlper/gte-base", ""), |
| } |
| MAGIC_F16, MAGIC_I8 = b"SRTIDX01", b"SRTIDX02" |
|
|
|
|
| def read_srtidx(path): |
| b = open(path, "rb").read() |
| magic, (dim, n) = b[:8], struct.unpack("<II", b[8:16]) |
| off = 16 |
| if magic == MAGIC_I8: |
| scale = np.frombuffer(b[off:off + 4 * n], dtype="<f4"); off += 4 * n |
| q = np.frombuffer(b[off:off + n * dim], dtype=np.int8).reshape(n, dim); off += n * dim |
| Z = q.astype(np.float32) * scale[:, None] |
| elif magic == MAGIC_F16: |
| Z = np.frombuffer(b[off:off + 2 * n * dim], dtype="<f2").reshape(n, dim).astype(np.float32); off += 2 * n * dim |
| else: |
| raise SystemExit(f"bad magic {magic!r}") |
| keys = [] |
| for _ in range(n): |
| (L,) = struct.unpack("<I", b[off:off + 4]); off += 4 |
| keys.append(b[off:off + L].decode()); off += L |
| Z /= np.linalg.norm(Z, axis=1, keepdims=True) + 1e-8 |
| return torch.from_numpy(Z), keys |
|
|
|
|
| def project(head, e): |
| p = (e - head["mu_txt"].float()) @ head["txt.weight"].float().T + head["txt.bias"].float() |
| return torch.nn.functional.normalize(p, dim=-1) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--reader", default="bge-small", choices=list(READERS)) |
| ap.add_argument("--skip-replay", action="store_true") |
| ap.add_argument("--device", default="cpu") |
| a = ap.parse_args() |
| repo_id, prefix = READERS[a.reader] |
|
|
| head = load_file(hf_hub_download(HEADS, f"text_head_{a.reader}_v3gallery.safetensors")) |
| enc = SentenceTransformer(repo_id, device=a.device) |
| enc.max_seq_length = 64 |
| encode = lambda texts: torch.tensor(enc.encode([prefix + t for t in texts], normalize_embeddings=True, batch_size=256, convert_to_numpy=True)) |
| print(f"reader {repo_id} head {a.reader} d_txt {head['txt.weight'].shape[1]} -> {head['txt.weight'].shape[0]}") |
|
|
| Zg, keys = read_srtidx(hf_hub_download(GALLERY, "gallery_123k_v3.srtidx")) |
| print(f"gallery {len(keys):,} rows x {Zg.shape[1]}") |
|
|
| |
| try: |
| fx = json.load(open(hf_hub_download(HEADS, f"parity_{a.reader}.json"))) |
| except Exception: |
| fx = None |
| print("parity: no fixture for this reader on the Hub (fixtures exist for bge-small and e5-base)") |
| if fx: |
| texts = [p["text"] for p in fx["probes"]] |
| e = encode(texts); p = project(head, e) |
| emb_gap = max(float((e[i] - torch.tensor(fx["probes"][i]["embedding"])).abs().max()) for i in range(len(texts))) |
| proj_cos = min(float(torch.nn.functional.cosine_similarity(p[i], torch.nn.functional.normalize(torch.tensor(fx["probes"][i]["projection"]), dim=0), dim=0)) for i in range(len(texts))) |
| top5 = (p @ Zg.T).topk(5, dim=1).indices |
| same1 = sum(int(top5[i, 0]) == fx["probes"][i]["top5_rows"][0] for i in range(len(texts))) |
| same5 = sum(sorted(top5[i].tolist()) == sorted(fx["probes"][i]["top5_rows"]) for i in range(len(texts))) |
| print(f"parity: max |embedding gap| {emb_gap:.2e} min projection cosine {proj_cos:.6f} top-1 identical {same1}/{len(texts)} top-5 identical {same5}/{len(texts)}") |
|
|
| if a.skip_replay: |
| return |
| rp = json.load(open(hf_hub_download(HEADS, "replay_123k.json"))) |
| caps = [c["text"] for c in rp["captions"]]; gold = torch.tensor([c["gold"] for c in rp["captions"]]) |
| t0 = time.time() |
| zt = project(head, encode(caps)) |
| S = zt @ Zg.T |
| gold_s = S[torch.arange(len(gold)), gold] |
| rank = (S > gold_s[:, None]).sum(1) + 1 |
| print(f"replay: {len(caps):,} captions vs {len(keys):,} rows in {time.time() - t0:.0f}s " |
| f"t2i R@1 {(rank <= 1).float().mean():.4f} R@5 {(rank <= 5).float().mean():.4f} R@10 {(rank <= 10).float().mean():.4f} median rank {int(rank.float().median())}") |
| ref = rp.get("pytorch_t2i", {}) |
| if ref: |
| print(f"the Qwen3-0.6B tower on the same set: R@1 {ref['r@1']} R@5 {ref['r@5']} R@10 {ref['r@10']} median rank 36") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
