124 lines
5.0 KiB
Python
124 lines
5.0 KiB
Python
"""Semantisches Embedding-Clustering für die Baustein-Konsolidierung.
|
|
|
|
Mean-Pool-Embeddings eines mehrsprachigen Satz-Modells bilden über Cosine-Blocking +
|
|
Union-Find GLOBALE Kandidaten-Cluster (kein Chunk-Verlust). Sichere Paare (Ähnlichkeit
|
|
≥ HART) werden ohne LLM gemergt; Grenz-Paare im Band [BAND_LOW, HART) gibt der Aufrufer
|
|
einem LLM-Judge zur ja/nein-Entscheidung. Fehlen `transformers`/`torch` oder lädt das
|
|
Modell nicht → `embed_sims()` liefert `None`, der Aufrufer fällt auf den alten
|
|
Panel-Judge-Pfad zurück (silente Deaktivierung, wie das Lesbarkeits-Gate).
|
|
|
|
CPU genügt; der Aufrufer wrappt die blockierende Inferenz in `asyncio.to_thread`.
|
|
`numpy` ist transitiv über torch vorhanden (bewusst nicht in requirements.txt, analog 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")
|
|
|
|
_modell_cache = None # (tokenizer, model, torch) — Singleton
|
|
_ladeversuch = False # schon versucht zu laden?
|
|
|
|
EMBEDDING_BATCH = 32 # Inferenz-Batchgröße (CPU)
|
|
EMBEDDING_MAX_LEN = 128 # Titel + Kurzbeschreibung sind kurz → kleiner Truncation-Cap genügt
|
|
|
|
|
|
def _modell():
|
|
"""Lädt das Modell einmalig. None = Clustering aus (deaktiviert oder Lade-Fehler)."""
|
|
global _modell_cache, _ladeversuch
|
|
if _ladeversuch:
|
|
return _modell_cache
|
|
_ladeversuch = 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()
|
|
_modell_cache = (tok, model, torch)
|
|
log.info("Embedding-Modell geladen: %s", EMBEDDING_MODELL)
|
|
except Exception as e:
|
|
log.warning("Embedding-Clustering deaktiviert (Modell nicht ladbar): %s", e)
|
|
_modell_cache = None
|
|
return _modell_cache
|
|
|
|
|
|
def verfuegbar() -> bool:
|
|
"""True, wenn das Modell geladen werden konnte. Lädt beim ersten Aufruf (blockierend)."""
|
|
return _modell() is not None
|
|
|
|
|
|
def embed(texts: list[str]) -> "np.ndarray | None":
|
|
"""Texte → (n, d) L2-normalisierte, mean-gepoolte Embeddings. None = Modell aus."""
|
|
if _modell() is None:
|
|
return None
|
|
tok, model, torch = _modell_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 ohne Padding
|
|
vec = torch.nn.functional.normalize(vec, p=2, dim=1) # L2 → Cosine = Skalarprodukt
|
|
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) # kleinster Index = Wurzel (deterministisch)
|
|
|
|
|
|
def embed_sims(texts: list[str]):
|
|
"""Texte → (n, n) Cosine-Matrix · None = Modell nicht verfügbar (Fallback)."""
|
|
embs = embed(texts)
|
|
if embs is None:
|
|
return None
|
|
return embs @ embs.T # (n, n) Cosine, float32 (~2 MB bei n=700)
|
|
|
|
|
|
def capped_blocks(sims, floor: float | None = None, cap: int | None = None) -> list[list[int]]:
|
|
"""Grobe Ähnlichkeits-Blocks für den LLM — High-Recall, aber Größe gedeckelt.
|
|
|
|
Greedy: alle Paare mit Cosine ≥ `floor` nach Cosine absteigend; zwei Blocks werden nur
|
|
verschmolzen, wenn der resultierende Block ≤ `cap` bleibt. Verhindert den Giant-Component
|
|
(reines Threshold-Blocking verkettet sonst fast alles) und hält die LLM-Listen kurz.
|
|
→ Liste von Blocks (Index-Listen), jeder Knoten in genau einem 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]
|
|
# höchste Cosine zuerst → engste Paare bilden zuerst Blocks
|
|
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())
|