File size: 4,416 Bytes
8505f8e 48db85f 8505f8e 48db85f 8505f8e 69cb606 8505f8e | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """Stage 1: stream corpus -> BGE embeddings (sharded memmap) -> faiss k-means K clusters.
python scripts/embed_cluster.py --clusters 10000 --n-docs 4000000
Scales to K=1e6 by bumping --clusters (and --n-docs to keep >=~40*K docs).
"""
import argparse
import json
import os
import numpy as np
import torch
from datasets import load_dataset
from mxf.config import ClusterConfig
def embed_stream(cfg, model, tok, device):
"""Yield (texts, np.float32[n,d]) batches from the streamed corpus."""
ds = load_dataset(cfg.corpus, split="en", streaming=True) # Ultra-FineWeb splits are en/zh (no "train"); field is "content"
buf, seen = [], 0
for row in ds:
t = (row.get("content") or row.get("text") or "").strip()
if not (cfg.min_chars <= len(t)):
continue
buf.append(t[: cfg.max_chars])
if len(buf) == cfg.embed_batch:
yield buf, _encode(buf, model, tok, device)
seen += len(buf); buf = []
if seen >= cfg.n_docs:
return
if buf:
yield buf, _encode(buf, model, tok, device)
@torch.no_grad()
def _encode(texts, model, tok, device):
enc = tok(texts, padding=True, truncation=True, max_length=256, return_tensors="pt").to(device)
out = model(**enc).last_hidden_state[:, 0] # BGE: CLS pooling
return torch.nn.functional.normalize(out, dim=-1).float().cpu().numpy()
def main():
from transformers import AutoModel, AutoTokenizer
cfg = ClusterConfig()
ap = argparse.ArgumentParser()
for f in ("corpus", "embed_model", "out_dir"):
ap.add_argument(f"--{f.replace('_','-')}", default=getattr(cfg, f))
ap.add_argument("--n-docs", type=int, default=cfg.n_docs)
ap.add_argument("--clusters", type=int, default=cfg.clusters)
a = ap.parse_args()
cfg.corpus, cfg.embed_model, cfg.out_dir = a.corpus, a.embed_model, a.out_dir
cfg.n_docs, cfg.clusters = a.n_docs, a.clusters
os.makedirs(cfg.out_dir, exist_ok=True)
device = "cuda"
tok = AutoTokenizer.from_pretrained(cfg.embed_model)
model = AutoModel.from_pretrained(cfg.embed_model, torch_dtype=torch.bfloat16).to(device).eval()
d = model.config.hidden_size
# sharded memmap of embeddings + a jsonl of the source texts (for centroid targets later)
emb_path = os.path.join(cfg.out_dir, "emb.f32")
emb = np.memmap(emb_path, dtype=np.float32, mode="w+", shape=(cfg.n_docs, d))
txt = open(os.path.join(cfg.out_dir, "texts.jsonl"), "w")
n = 0
for texts, vecs in embed_stream(cfg, model, tok, device):
k = min(len(texts), cfg.n_docs - n)
emb[n : n + k] = vecs[:k]
for t in texts[:k]:
txt.write(json.dumps({"t": t}) + "\n")
n += k
if n // 100_000 != (n - k) // 100_000: # batch-sized steps almost never land on exact multiples
print(f"embedded {n}/{cfg.n_docs}", flush=True)
if n >= cfg.n_docs:
break
txt.close(); emb.flush()
print(f"embedded {n} docs, d={d}", flush=True)
# faiss GPU k-means (falls back to CPU/MiniBatchKMeans if faiss-gpu absent)
X = np.asarray(emb[:n])
try:
import faiss
faiss.omp_set_num_threads(min(os.cpu_count() or 32, 128))
# gpu=False: faiss-gpu wheels have no Blackwell sm_103 kernels ("no kernel image"); CPU
# faiss is multithreaded and fine for k-means. (GPU k-means: use recluster_torch.py.)
km = faiss.Kmeans(d, cfg.clusters, niter=20, verbose=True, gpu=False, seed=cfg.seed)
km.train(X)
_, assign = km.index.search(X, 1)
assign = assign.ravel()
centroids = km.centroids.reshape(cfg.clusters, d)
except Exception as e:
print("faiss unavailable, MiniBatchKMeans:", e, flush=True)
from sklearn.cluster import MiniBatchKMeans
km = MiniBatchKMeans(cfg.clusters, batch_size=10000, n_init=3, random_state=cfg.seed).fit(X)
assign, centroids = km.labels_, km.cluster_centers_
np.save(os.path.join(cfg.out_dir, "assign.npy"), assign.astype(np.int32))
np.save(os.path.join(cfg.out_dir, "centroids.npy"), centroids.astype(np.float32))
json.dump({"n_docs": int(n), "d": int(d), "clusters": cfg.clusters},
open(os.path.join(cfg.out_dir, "meta.json"), "w"))
print(f"CLUSTERED {n} docs into {cfg.clusters} clusters -> {cfg.out_dir}", flush=True)
if __name__ == "__main__":
main()
|