124 lines
4.9 KiB
Python
124 lines
4.9 KiB
Python
"""Semantic embedding clustering for block consolidation.
|
|
|
|
Mean-pool embeddings of a multilingual sentence model build GLOBAL candidate
|
|
clusters via cosine blocking + union-find (no chunk loss). Safe pairs (similarity
|
|
≥ HARD) are merged without an LLM; borderline pairs in the band [BAND_LOW, HARD) are
|
|
handed by the caller to an LLM judge for a yes/no decision. If `transformers`/`torch`
|
|
are missing or the model won't load → `embed_sims()` returns `None`, and the caller
|
|
falls back to the old panel-judge path (silent deactivation, like the readability gate).
|
|
|
|
CPU is enough; the caller wraps the blocking inference in `asyncio.to_thread`.
|
|
`numpy` comes in transitively via torch (deliberately not in requirements.txt, like torch).
|
|
"""
|
|
|
|
import logging
|
|
|
|
import numpy as np
|
|
|
|
from config import EMBEDDING_AKTIV, EMBEDDING_MODELL, EMBEDDING_BLOCK_FLOOR, EMBEDDING_BLOCK_CAP
|
|
|
|
log = logging.getLogger("creator.embedding")
|
|
|
|
_model_cache = None # (tokenizer, model, torch) — singleton
|
|
_load_attempt = False # already tried to load?
|
|
|
|
EMBEDDING_BATCH = 32 # inference batch size (CPU)
|
|
EMBEDDING_MAX_LEN = 128 # title + short description are short → a small truncation cap suffices
|
|
|
|
|
|
def _model():
|
|
"""Load the model once. None = clustering off (disabled or load error)."""
|
|
global _model_cache, _load_attempt
|
|
if _load_attempt:
|
|
return _model_cache
|
|
_load_attempt = True
|
|
if not EMBEDDING_AKTIV:
|
|
return None
|
|
try:
|
|
import torch
|
|
from transformers import AutoModel, AutoTokenizer
|
|
tok = AutoTokenizer.from_pretrained(EMBEDDING_MODELL)
|
|
model = AutoModel.from_pretrained(EMBEDDING_MODELL)
|
|
model.eval()
|
|
_model_cache = (tok, model, torch)
|
|
log.info("embedding model loaded: %s", EMBEDDING_MODELL)
|
|
except Exception as e:
|
|
log.warning("embedding clustering disabled (model not loadable): %s", e)
|
|
_model_cache = None
|
|
return _model_cache
|
|
|
|
|
|
def available() -> bool:
|
|
"""True if the model could be loaded. Loads on the first call (blocking)."""
|
|
return _model() is not None
|
|
|
|
|
|
def embed(texts: list[str]) -> "np.ndarray | None":
|
|
"""Texts → (n, d) L2-normalized, mean-pooled embeddings. None = model off."""
|
|
if _model() is None:
|
|
return None
|
|
tok, model, torch = _model_cache
|
|
out = []
|
|
for i in range(0, len(texts), EMBEDDING_BATCH):
|
|
batch = texts[i:i + EMBEDDING_BATCH]
|
|
enc = tok(batch, return_tensors="pt", truncation=True, max_length=EMBEDDING_MAX_LEN, padding=True)
|
|
with torch.no_grad():
|
|
hidden = model(**enc).last_hidden_state # (b, t, d)
|
|
mask = enc["attention_mask"].unsqueeze(-1).type_as(hidden)
|
|
vec = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9) # mean-pool without padding
|
|
vec = torch.nn.functional.normalize(vec, p=2, dim=1) # L2 → cosine = dot product
|
|
out.append(vec.cpu().numpy())
|
|
return np.vstack(out).astype(np.float32)
|
|
|
|
|
|
def _find(parent: list[int], x: int) -> int:
|
|
while parent[x] != x:
|
|
parent[x] = parent[parent[x]] # Pfad-Kompression
|
|
x = parent[x]
|
|
return x
|
|
|
|
|
|
def _union(parent: list[int], a: int, b: int) -> None:
|
|
ra, rb = _find(parent, a), _find(parent, b)
|
|
if ra != rb:
|
|
parent[max(ra, rb)] = min(ra, rb) # smallest index = root (deterministic)
|
|
|
|
|
|
def embed_sims(texts: list[str]):
|
|
"""Texts → (n, n) cosine matrix · None = model not available (fallback)."""
|
|
embs = embed(texts)
|
|
if embs is None:
|
|
return None
|
|
return embs @ embs.T # (n, n) cosine, float32 (~2 MB at n=700)
|
|
|
|
|
|
def capped_blocks(sims, floor: float | None = None, cap: int | None = None) -> list[list[int]]:
|
|
"""Coarse similarity blocks for the LLM — high recall, but size-capped.
|
|
|
|
Greedy: all pairs with cosine ≥ `floor` in descending cosine order; two blocks are merged
|
|
only if the resulting block stays ≤ `cap`. Prevents the giant component (pure threshold
|
|
blocking would otherwise chain almost everything together) and keeps the LLM lists short.
|
|
→ list of blocks (index lists), each node in exactly one block.
|
|
"""
|
|
fl = EMBEDDING_BLOCK_FLOOR if floor is None else floor
|
|
cp = EMBEDDING_BLOCK_CAP if cap is None else cap
|
|
n = len(sims)
|
|
parent = list(range(n))
|
|
size = [1] * n
|
|
if n >= 2:
|
|
iu = np.triu_indices(n, k=1)
|
|
s = sims[iu]
|
|
kept = np.where(s >= fl)[0]
|
|
# highest cosine first → the tightest pairs form blocks first
|
|
for k in kept[np.argsort(-s[kept])]:
|
|
i, j = int(iu[0][k]), int(iu[1][k])
|
|
ri, rj = _find(parent, i), _find(parent, j)
|
|
if ri != rj and size[ri] + size[rj] <= cp:
|
|
_union(parent, i, j)
|
|
r = _find(parent, i)
|
|
size[r] = size[ri] + size[rj]
|
|
blocks: dict[int, list[int]] = {}
|
|
for i in range(n):
|
|
blocks.setdefault(_find(parent, i), []).append(i)
|
|
return list(blocks.values())
|