File size: 2,545 Bytes
69cb606 | 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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """Cluster cached embeddings (emb.f32) on GPU with a minimal torch k-means — works on Blackwell
where faiss-gpu has no sm_103 kernels. Reuses an existing embed step; writes assign/centroids/meta.
python scripts/recluster_torch.py --clusters-dir data/pilot_clusters --clusters 10000
"""
import argparse
import json
import numpy as np
import torch
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--clusters-dir", required=True)
ap.add_argument("--clusters", type=int, required=True)
ap.add_argument("--iters", type=int, default=20)
ap.add_argument("--chunk", type=int, default=200_000)
ap.add_argument("--seed", type=int, default=0)
a = ap.parse_args()
dev = "cuda"
n = sum(1 for _ in open(f"{a.clusters_dir}/texts.jsonl"))
d = np.memmap(f"{a.clusters_dir}/emb.f32", dtype=np.float32, mode="r").shape[0] // n
emb = np.memmap(f"{a.clusters_dir}/emb.f32", dtype=np.float32, mode="r", shape=(n, d))
print(f"kmeans: n={n} d={d} k={a.clusters}", flush=True)
g = torch.Generator().manual_seed(a.seed)
init = torch.from_numpy(np.asarray(emb[torch.randperm(n, generator=g)[: a.clusters].numpy()]))
C = init.to(dev, torch.float32) # [k, d]
Xg = torch.from_numpy(np.ascontiguousarray(emb)).to(dev, torch.float32) # [n, d] (n*d*4 bytes)
assign = torch.empty(n, dtype=torch.long, device=dev)
for it in range(a.iters):
Cn = (C * C).sum(1) # [k]
for s in range(0, n, a.chunk):
xb = Xg[s : s + a.chunk]
d2 = Cn[None] - 2 * xb @ C.T # argmin over k (||x||² constant per row)
assign[s : s + a.chunk] = d2.argmin(1)
C.zero_()
C.index_add_(0, assign, Xg) # sum members
cnt = torch.bincount(assign, minlength=a.clusters).clamp(min=1)
C /= cnt[:, None]
empty = cnt == 1 # reseed empties to random points (bincount clamp hides true-empties)
if empty.any():
C[empty] = Xg[torch.randint(0, n, (int(empty.sum()),), device=dev)]
print(f" iter {it}", flush=True)
np.save(f"{a.clusters_dir}/assign.npy", assign.cpu().numpy().astype(np.int32))
np.save(f"{a.clusters_dir}/centroids.npy", C.cpu().numpy().astype(np.float32))
json.dump({"n_docs": int(n), "d": int(d), "clusters": a.clusters},
open(f"{a.clusters_dir}/meta.json", "w"))
print(f"CLUSTERED {n} -> {a.clusters} clusters", flush=True)
if __name__ == "__main__":
main()
|