model-3/src/encoder.py

"""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)