Spaces:
Running
Running
| """Embedding models behind one interface, and an on-disk store keyed by text hash (plan §5.1). | |
| emb = Embedder("Qwen/Qwen3-Embedding-0.6B", devices=["cuda:0", "cuda:1"], dtype="float16") | |
| vectors = emb.encode(texts) # L2-normalised float32 rows | |
| EmbeddingStore.save(path, ids, texts, vectors, meta) # float16 matrix + ids + sha1 per row | |
| store = EmbeddingStore.load(path); store.lookup(texts) -> rows (or None when a text was never embedded) | |
| Several devices -> one worker process per device (sentence-transformers multi-process pool). Documents are encoded | |
| without an instruction prompt (Qwen3-Embedding uses prompts on the query side only). | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import time | |
| from dataclasses import dataclass, field | |
| from pathlib import Path | |
| from typing import Any, Sequence | |
| import numpy as np | |
| def text_sha1(text: str) -> str: | |
| return hashlib.sha1(text.encode("utf-8")).hexdigest() | |
| def embedding_slug(ecfg: dict[str, Any]) -> str: | |
| """Folder name for a models.embedder config under paths.embeddings: model id, plus the truncated dimension.""" | |
| return ecfg["id"].replace("/", "__") + (f"__d{ecfg['truncate_dim']}" if ecfg.get("truncate_dim") else "") | |
| class Embedder: | |
| def __init__(self, model_id: str, *, devices: Sequence[str] = ("cpu",), dtype: str = "float16", | |
| max_length: int = 1024, batch_size: int = 32, normalize: bool = True, | |
| truncate_dim: int | None = None, revision: str | None = None): | |
| import torch | |
| from sentence_transformers import SentenceTransformer | |
| self.model_id, self.devices, self.batch_size = model_id, list(devices), batch_size | |
| self.normalize, self.truncate_dim = normalize, truncate_dim | |
| torch_dtype = getattr(torch, dtype) if self.devices[0] != "cpu" else torch.float32 | |
| self.model = SentenceTransformer(model_id, device=self.devices[0], revision=revision, | |
| model_kwargs={"dtype": torch_dtype}, | |
| processor_kwargs={"padding_side": "left"}, truncate_dim=truncate_dim) | |
| self.model.max_seq_length = max_length | |
| self.dim = self.model.get_sentence_embedding_dimension() | |
| def encode(self, texts: Sequence[str], *, show_progress: bool = False, prompt: str | None = None) -> np.ndarray: | |
| """prompt: text prepended to every input (an instruction for instruction-aware models), or None.""" | |
| if not texts: | |
| return np.zeros((0, self.dim), dtype=np.float32) | |
| if len(self.devices) > 1: | |
| pool = self.model.start_multi_process_pool(target_devices=self.devices) | |
| try: | |
| out = self.model.encode(list(texts), pool=pool, batch_size=self.batch_size, prompt=prompt, | |
| normalize_embeddings=self.normalize, show_progress_bar=show_progress, | |
| chunk_size=max(1000, self.batch_size * 20)) | |
| finally: | |
| self.model.stop_multi_process_pool(pool) | |
| else: | |
| out = self.model.encode(list(texts), batch_size=self.batch_size, prompt=prompt, normalize_embeddings=self.normalize, | |
| show_progress_bar=show_progress, convert_to_numpy=True) | |
| return np.asarray(out, dtype=np.float32) | |
| class EmbeddingStore: | |
| ids: list[Any] | |
| sha1: list[str] | |
| vectors: np.ndarray | |
| meta: dict[str, Any] = field(default_factory=dict) | |
| def save(path: str | Path, ids: list[Any], texts: Sequence[str], vectors: np.ndarray, meta: dict[str, Any]) -> None: | |
| path = Path(path) | |
| path.mkdir(parents=True, exist_ok=True) | |
| np.save(path / "vectors.f16.npy", vectors.astype(np.float16)) | |
| with open(path / "rows.jsonl", "w", encoding="utf-8") as f: | |
| for i, t in zip(ids, texts): | |
| f.write(json.dumps({"id": i, "sha1": text_sha1(t)}) + "\n") | |
| (path / "meta.json").write_text(json.dumps({**meta, "rows": len(ids), "dim": int(vectors.shape[1])}, | |
| indent=1), encoding="utf-8") | |
| def load(cls, path: str | Path) -> "EmbeddingStore": | |
| path = Path(path) | |
| rows = [json.loads(l) for l in open(path / "rows.jsonl", encoding="utf-8")] | |
| return cls(ids=[r["id"] for r in rows], sha1=[r["sha1"] for r in rows], | |
| vectors=np.load(path / "vectors.f16.npy").astype(np.float32), | |
| meta=json.loads((path / "meta.json").read_text(encoding="utf-8"))) | |
| def lookup(self, texts: Sequence[str]) -> list[np.ndarray | None]: | |
| index = {h: i for i, h in enumerate(self.sha1)} | |
| return [self.vectors[index[h]] if (h := text_sha1(t)) in index else None for t in texts] | |
| def embed_and_store(embedder: Embedder, ids: list[Any], texts: list[str], path: str | Path, | |
| extra_meta: dict[str, Any] | None = None, *, reuse: bool = False) -> dict[str, Any]: | |
| """Encode texts and save them as an EmbeddingStore at path. reuse=True first loads the store already at path and | |
| encodes only texts whose sha1 it does not hold (after a rebuild that split some sections, only the new chunks).""" | |
| t0 = time.perf_counter() | |
| old = EmbeddingStore.load(path) if reuse and (Path(path) / "rows.jsonl").exists() else None | |
| found = old.lookup(texts) if old is not None else [None] * len(texts) | |
| missing = [i for i, v in enumerate(found) if v is None] | |
| fresh = embedder.encode([texts[i] for i in missing], show_progress=True) | |
| seconds = round(time.perf_counter() - t0, 1) | |
| if old is not None: | |
| vectors = np.zeros((len(texts), fresh.shape[1] if len(missing) else old.vectors.shape[1]), dtype=np.float32) | |
| for i, v in enumerate(found): | |
| if v is not None: | |
| vectors[i] = v | |
| vectors[missing] = fresh | |
| else: | |
| vectors = fresh | |
| meta = {"model": embedder.model_id, "devices": embedder.devices, "normalized": embedder.normalize, | |
| "truncate_dim": embedder.truncate_dim, "seconds": seconds, **(extra_meta or {})} | |
| if old is not None: # provenance of the reused rows (they may come from another machine and dtype) | |
| meta.update(reused_rows=len(texts) - len(missing), new_rows=len(missing), | |
| reused_from={k: old.meta.get(k) for k in ("model", "devices", "seconds", "rows", "reused_from")}) | |
| EmbeddingStore.save(path, ids, texts, vectors, meta) | |
| return {"rows": len(ids), "dim": int(vectors.shape[1]), "seconds": seconds, | |
| "rows_per_second": round(len(missing) / max(seconds, 1e-6), 1), | |
| **({"reused_rows": len(texts) - len(missing), "new_rows": len(missing)} if old is not None else {})} | |
| def document_vector(section_vectors: np.ndarray, title_vector: np.ndarray | None = None, *, | |
| title_abstract_weight: float = 0.5) -> np.ndarray: | |
| """Document embedding (plan §0.2 #9, models.embedder.title_abstract_weight = w): (1 - w) x the normalised mean | |
| of the (normalised) section vectors + w x the normalised title/abstract vector, normalised again; sections only | |
| when there is no title. Uploads use the same function and model as the index (the index records both).""" | |
| S = np.atleast_2d(np.asarray(section_vectors, dtype=np.float64)) | |
| S = S / np.maximum(np.linalg.norm(S, axis=1, keepdims=True), 1e-12) | |
| v = S.mean(0) | |
| v = v / max(np.linalg.norm(v), 1e-12) | |
| if title_vector is not None and title_abstract_weight: | |
| t = np.asarray(title_vector, dtype=np.float64) | |
| v = (1 - title_abstract_weight) * v + title_abstract_weight * t / max(np.linalg.norm(t), 1e-12) | |
| return (v / max(np.linalg.norm(v), 1e-12)).astype(np.float32) | |
| def section_mean_vector(doc_vector: np.ndarray, title_vector: np.ndarray | None, *, | |
| title_abstract_weight: float = 0.5) -> np.ndarray: | |
| """The normalised mean of the section vectors, recovered from a document vector made by document_vector with the | |
| same weight w: v is proportional to (1 - w) m + w t with m, t unit vectors, so m = (s v - w t) / (1 - w), where s | |
| is the positive root of |s v - w t| = 1 - w. Lets a stored document vector be compared without its title part | |
| (cohort diagnosis, decision record 007) when the section vectors themselves were not stored.""" | |
| v = np.asarray(doc_vector, dtype=np.float64) | |
| v = v / max(np.linalg.norm(v), 1e-12) | |
| w = float(title_abstract_weight) | |
| if title_vector is None or not w: | |
| return v.astype(np.float32) | |
| t = np.asarray(title_vector, dtype=np.float64) | |
| t = t / max(np.linalg.norm(t), 1e-12) | |
| c = float(v @ t) | |
| s = w * c + np.sqrt(max(w * w * c * c - w * w + (1 - w) ** 2, 0.0)) | |
| m = s * v - w * t | |
| return (m / max(np.linalg.norm(m), 1e-12)).astype(np.float32) | |