File size: 7,426 Bytes
e79d917 | 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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | """Resumable, fingerprint-validated embedding cache.
Why this file exists
--------------------
model-2 lost two multi-hour runs with nothing recoverable:
* `encode_texts` streamed into ONE `<name>.npy.tmp.npy` and only promoted it to
the final name after the very last chunk. A kill at 95% left 21 GB of garbage
and the next run started from row 0.
* The cache was accepted whenever `emb.shape[0] == len(texts)`. Change the model,
the prefix or the text format and you silently reuse the wrong vectors.
Here:
* rows are written in shards of EMB["shard_rows"], each promoted atomically;
* a kill costs at most one shard, and the next run resumes at the next shard;
* `meta.json` fingerprints model id, prefix, max_len, dim, row count and a
SHA-256 of the entity_id list, so a stale cache is REJECTED, never reused.
"""
from __future__ import annotations
import hashlib
import json
import os
import time
import numpy as np
def _sha_ids(ids) -> str:
h = hashlib.sha256()
for i in ids:
h.update(str(i).encode())
h.update(b"\n")
return h.hexdigest()
def fingerprint(hf_id, prefix, max_len, dim, ids) -> dict:
return {"model": hf_id, "prefix": prefix, "max_len": int(max_len), "dim": int(dim),
"n": len(ids), "ids_sha256": _sha_ids(ids)}
def _shard_path(d, i):
return os.path.join(d, f"shard_{i:05d}.npy")
def _meta_path(d):
return os.path.join(d, "meta.json")
def cache_state(d: str, fp: dict, shard_rows: int):
"""-> (n_shards_total, n_shards_done, reason_if_invalid)."""
n_total = (fp["n"] + shard_rows - 1) // shard_rows
mp = _meta_path(d)
if not os.path.exists(mp):
return n_total, 0, None
try:
old = json.load(open(mp))
except Exception:
return n_total, 0, "unreadable meta.json"
for k in ("model", "prefix", "max_len", "dim", "n", "ids_sha256"):
if old.get(k) != fp[k]:
return n_total, 0, f"{k} changed ({old.get(k)!r} -> {fp[k]!r})"
if old.get("shard_rows") != shard_rows:
return n_total, 0, f"shard_rows changed ({old.get('shard_rows')} -> {shard_rows})"
done = 0
for i in range(n_total):
if os.path.exists(_shard_path(d, i)):
done += 1
else:
break
return n_total, done, None
class Encoder:
"""Thin wrapper over sentence-transformers with an explicit device."""
def __init__(self, hf_id, device, max_len, batch, fp16=True):
import torch
from sentence_transformers import SentenceTransformer
self.torch = torch
kw = {}
if fp16 and str(device).startswith("cuda"):
kw["model_kwargs"] = {"torch_dtype": torch.float16}
self.model = SentenceTransformer(hf_id, device=device, **kw)
try:
self.model.max_seq_length = max_len
except AttributeError:
pass # StaticEmbedding has no sequence limit
self.device = device
self.batch = batch
self.dim = self.model.get_sentence_embedding_dimension()
def encode(self, texts, prefix):
return self.model.encode([prefix + t for t in texts], batch_size=self.batch,
normalize_embeddings=True, convert_to_numpy=True,
show_progress_bar=False).astype(np.float16)
def peak_vram_gb(self):
if not str(self.device).startswith("cuda"):
return 0.0
return self.torch.cuda.max_memory_allocated() / 1e9
def encode_to_shards(texts, ids, out_dir, hf_id, prefix, max_len, batch, device,
shard_rows=250_000, log=print):
"""Encode `texts` into resumable shards under `out_dir`. Returns the meta dict.
Safe to re-run: finished shards are skipped, a killed run resumes at the
first missing shard.
"""
assert len(texts) == len(ids), f"{len(texts)} texts vs {len(ids)} ids"
os.makedirs(out_dir, exist_ok=True)
enc = None
# dim is only known after the model loads; probe lazily but fingerprint first
# with the declared dim from the zoo when resuming an existing cache.
mp = _meta_path(out_dir)
known_dim = json.load(open(mp))["dim"] if os.path.exists(mp) else None
def build_fp(dim):
return fingerprint(hf_id, prefix, max_len, dim, ids)
if known_dim is not None:
n_total, done, bad = cache_state(out_dir, build_fp(known_dim), shard_rows)
if bad:
log(f"[emb] cache INVALID ({bad}) -> re-encoding from scratch")
for f in os.listdir(out_dir):
if f.startswith("shard_") or f == "meta.json":
os.remove(os.path.join(out_dir, f))
done = 0
n_total = (len(ids) + shard_rows - 1) // shard_rows
elif done == n_total:
log(f"[emb] cache complete and valid: {out_dir} ({n_total} shards)")
return json.load(open(mp))
else:
log(f"[emb] resuming: {done}/{n_total} shards already done")
else:
done = 0
n_total = (len(ids) + shard_rows - 1) // shard_rows
t0 = time.time()
for i in range(done, n_total):
lo, hi = i * shard_rows, min((i + 1) * shard_rows, len(texts))
if enc is None:
enc = Encoder(hf_id, device, max_len, batch)
log(f"[emb] model={hf_id} dim={enc.dim} device={device}")
E = enc.encode(texts[lo:hi], prefix)
# NOTE: np.save appends ".npy" unless the name already ends in it, so the
# temp name must end in .npy or the rename below looks for the wrong file.
tmp = _shard_path(out_dir, i) + ".tmp.npy"
np.save(tmp, E)
os.replace(tmp, _shard_path(out_dir, i)) # atomic
rate = (hi - (done * shard_rows)) / max(time.time() - t0, 1e-9)
eta = (len(texts) - hi) / max(rate, 1e-9) / 60
log(f"[emb] shard {i + 1}/{n_total} rows {hi:,}/{len(texts):,} "
f"{rate:,.0f} rec/s eta {eta:,.0f} min vram {enc.peak_vram_gb():.1f} GB")
dim = enc.dim if enc is not None else known_dim
meta = {**build_fp(dim), "shard_rows": shard_rows, "n_shards": n_total,
"seconds": round(time.time() - t0, 1)}
json.dump(meta, open(mp, "w"), indent=2)
log(f"[emb] done -> {out_dir}")
return meta
def load_matrix(out_dir, mmap=True):
"""Concatenated (n, dim) float16 matrix. Built once into a single .npy so the
search stage memory-maps one file instead of holding a list of shards."""
meta = json.load(open(_meta_path(out_dir)))
flat = os.path.join(out_dir, "matrix.npy")
if not os.path.exists(flat):
M = np.lib.format.open_memmap(flat + ".tmp", mode="w+", dtype=np.float16,
shape=(meta["n"], meta["dim"]))
off = 0
for i in range(meta["n_shards"]):
S = np.load(_shard_path(out_dir, i), mmap_mode="r")
M[off:off + len(S)] = S
off += len(S)
assert off == meta["n"], f"row mismatch {off} != {meta['n']}"
M.flush()
del M
os.replace(flat + ".tmp", flat)
return np.load(flat, mmap_mode="r" if mmap else None), meta
def read_ids(path):
with open(path) as f:
return [l.rstrip("\n") for l in f if l.strip()]
def write_ids(path, ids):
tmp = path + ".tmp"
with open(tmp, "w") as f:
f.write("\n".join(str(i) for i in ids))
os.replace(tmp, path)
|