Upload scripts/dcpm_index.py with huggingface_hub

scripts/dcpm_index.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+"""Parse a DCPM metadata.bin index. Stage 1: read only, write nothing.
+
+Record layout, decoded from the bytes:
+    int64 name_len | name | int64 ndim | int64 shape[ndim]
+    int32 dtype | int32 ? | int64 ? | int32 ? | int64 file_len | file
+    int64 offset  | ... trailing fields
+We locate records by scanning for the length-prefixed names rather than
+assuming a fixed stride, so a variable-length shape list cannot desync us.
+"""
+import struct, sys, re, json
+
+def parse(path):
+    b = open(path, "rb").read()
+    assert b[:4] == b"DCPM", f"not DCPM: {b[:4]!r}"
+    recs = []
+    for m in re.finditer(rb"param_(\d{6})", b):
+        i = m.start()
+        nlen = struct.unpack_from("<q", b, i - 8)[0]
+        if nlen != m.end() - i:            # name must be length-prefixed
+            continue
+        j = m.end()
+        ndim = struct.unpack_from("<q", b, j)[0]; j += 8
+        if not (0 <= ndim <= 8): continue
+        shape = list(struct.unpack_from("<" + "q" * ndim, b, j)); j += 8 * ndim
+        dtype = struct.unpack_from("<i", b, j)[0]; j += 4
+        f1    = struct.unpack_from("<i", b, j)[0]; j += 4
+        f2    = struct.unpack_from("<q", b, j)[0]; j += 8
+        f3    = struct.unpack_from("<i", b, j)[0]; j += 4
+        flen  = struct.unpack_from("<q", b, j)[0]; j += 8
+        fname = b[j:j + flen].decode(errors="replace"); j += flen
+        off   = struct.unpack_from("<q", b, j)[0]; j += 8
+        recs.append(dict(idx=int(m.group(1)), name=m.group(0).decode(), ndim=ndim,
+                         shape=shape, dtype=dtype, f1=f1, f2=f2, f3=f3,
+                         file=fname, offset=off))
+    recs.sort(key=lambda r: r["idx"])
+    return recs
+
+if __name__ == "__main__":
+    recs = parse(sys.argv[1])
+    print(f"  parsed {len(recs)} param records")
+    dts = {}
+    for r in recs: dts.setdefault(r["dtype"], 0); dts[r["dtype"]] += 1
+    print(f"  dtype codes seen: {dts}")
+    files = {}
+    for r in recs: files.setdefault(r["file"], 0); files[r["file"]] += 1
+    print(f"  source files: {files}")
+    print(f"\n  {'idx':>4} {'shape':<20} {'dtype':>5} {'file':<16} {'offset':>14}")
+    for r in recs[:14]:
+        print(f"  {r['idx']:>4} {str(r['shape']):<20} {r['dtype']:>5} {r['file']:<16} {r['offset']:>14,}")
+    print("   ...")
+    for r in recs[-6:]:
+        print(f"  {r['idx']:>4} {str(r['shape']):<20} {r['dtype']:>5} {r['file']:<16} {r['offset']:>14,}")
+    tot = sum(int(__import__("math").prod(r["shape"])) for r in recs)
+    print(f"\n  total elements across all params: {tot:,}")
+    json.dump(recs, open("/tmp/dcpm_index.json", "w"))
+    print("  wrote /tmp/dcpm_index.json")