146 lines
4.9 KiB
Python
146 lines
4.9 KiB
Python
"""
|
|
Bi-encoder semantic search over slot_modifiers definitions.
|
|
|
|
Replaces the old "always show all modifiers as [APPLIES]/[SEMANTIC]" dump
|
|
that caused the model to parrot rule-text (slot names, action keywords)
|
|
back as tag content.
|
|
|
|
Index shape: each row of slot_modifiers gets a 384-dim embedding from
|
|
BAAI/bge-small-en-v1.5 (same model bucket_search uses), CLS-pooled and
|
|
L2-normalized so cosine == dot. Index lives in memory; rowcount-based
|
|
hot-reload picks up DB schema changes without restart.
|
|
|
|
Threshold/top_k are calibrated against real prompts (see commit history).
|
|
0.65 / top_k=2 catches semantic paraphrases that miss the alias scan
|
|
("her soles aimed at the camera", "the foot is pointed at viewer")
|
|
while keeping distractors ("girl walking through forest", "holding a
|
|
sword") out — they all top out around 0.50-0.58.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
from typing import Any
|
|
|
|
from . import _embed_model
|
|
|
|
|
|
logger = logging.getLogger("promptchain.modifier_search")
|
|
_dbg = logging.getLogger("promptchain.ai.debug")
|
|
|
|
_lock = threading.RLock()
|
|
_state: dict[str, Any] = {
|
|
"embeddings": None, # torch.Tensor [N, 384]
|
|
"modifiers": [], # list[dict] aligned with embeddings
|
|
"fingerprint": None, # (max_rowid, count) — change → rebuild
|
|
}
|
|
|
|
|
|
def _ensure_model_loaded() -> bool:
|
|
"""Defers to the shared embed model loader. Returns False on hard
|
|
failures (caller handles the empty-result fallback) — transient
|
|
init-time errors leave the loader's state unset so the next call
|
|
retries naturally."""
|
|
return _embed_model.get() is not None
|
|
|
|
|
|
def _embed(texts: list[str]):
|
|
return _embed_model.embed(texts, batch_size=64)
|
|
|
|
|
|
def _fingerprint() -> tuple:
|
|
from .tag_builder import get_db
|
|
db = get_db()
|
|
try:
|
|
row = db.execute(
|
|
"SELECT COALESCE(MAX(rowid), 0) AS m, COUNT(*) AS c FROM slot_modifiers"
|
|
).fetchone()
|
|
return (row["m"], row["c"])
|
|
except Exception:
|
|
return (0, 0)
|
|
|
|
|
|
def _embed_text(mod: dict) -> str:
|
|
"""Text we feed to bge-small. Definition-only on purpose — calibration
|
|
showed that mixing aliases into the embed shifts the embedding toward
|
|
specific phrasings ("pointing her feet at viewer") and AWAY from the
|
|
underlying concept, hurting recall on paraphrases the curator didn't
|
|
write down. Aliases continue to drive the deterministic word-boundary
|
|
scan; the embedding stays focused on what the modifier *means*.
|
|
|
|
Falls back to the canonical tag name (with underscores → spaces) only
|
|
when no definition has been populated yet — better than skipping the
|
|
row entirely."""
|
|
definition = (mod.get("definition") or "").strip()
|
|
if definition:
|
|
return definition
|
|
return (mod.get("canonical_tag") or "").replace("_", " ").strip()
|
|
|
|
|
|
def _rebuild_index() -> None:
|
|
from . import ai_api
|
|
if not _ensure_model_loaded():
|
|
return
|
|
modifiers = ai_api._load_slot_modifiers()
|
|
if not modifiers:
|
|
_state["embeddings"] = None
|
|
_state["modifiers"] = []
|
|
_state["fingerprint"] = _fingerprint()
|
|
return
|
|
texts = [_embed_text(m) for m in modifiers]
|
|
embeddings = _embed(texts)
|
|
_state["embeddings"] = embeddings
|
|
_state["modifiers"] = list(modifiers)
|
|
_state["fingerprint"] = _fingerprint()
|
|
logger.info("modifier_search: indexed %d modifiers", len(modifiers))
|
|
|
|
|
|
def _ensure_index_fresh() -> None:
|
|
if _state["fingerprint"] == _fingerprint() and _state["embeddings"] is not None:
|
|
return
|
|
_rebuild_index()
|
|
|
|
|
|
def search(user_text: str, top_k: int = 2, threshold: float = 0.65) -> list[dict]:
|
|
"""Return modifiers whose definition cosine-matches user_text above
|
|
threshold, sorted by score desc, capped at top_k.
|
|
|
|
Each entry is the full modifier dict (from _load_slot_modifiers)
|
|
plus a `score` field. Empty list when the model failed to load —
|
|
caller should degrade to alias-only behavior.
|
|
"""
|
|
user_text = (user_text or "").strip()
|
|
if not user_text:
|
|
return []
|
|
with _lock:
|
|
if not _ensure_model_loaded():
|
|
return []
|
|
_ensure_index_fresh()
|
|
if _state["embeddings"] is None or not _state["modifiers"]:
|
|
return []
|
|
qv = _embed([user_text])
|
|
scores = (_state["embeddings"] @ qv.T).squeeze(-1).tolist()
|
|
paired = list(zip(_state["modifiers"], scores))
|
|
paired.sort(key=lambda x: -x[1])
|
|
out: list[dict] = []
|
|
for mod, score in paired[:top_k]:
|
|
if score < threshold:
|
|
break
|
|
entry = dict(mod)
|
|
entry["score"] = float(score)
|
|
out.append(entry)
|
|
return out
|
|
|
|
|
|
def warmup() -> None:
|
|
"""Best-effort eager load on server boot."""
|
|
import time
|
|
for attempt in range(3):
|
|
with _lock:
|
|
if _ensure_model_loaded():
|
|
_ensure_index_fresh()
|
|
return
|
|
if _embed_model.get_load_error():
|
|
return
|
|
time.sleep(2.0 * (attempt + 1))
|