Spaces:
Running
Running
File size: 8,869 Bytes
7103d1c | 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 | """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)
@dataclass
class EmbeddingStore:
ids: list[Any]
sha1: list[str]
vectors: np.ndarray
meta: dict[str, Any] = field(default_factory=dict)
@staticmethod
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")
@classmethod
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)
|