File size: 5,373 Bytes
75ce203 | 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 | """Aurelius core — text-embedding model singleton + per-search caches.
v2 concurrency fix: v1 kept a single module-global `_emb_cache`/`_dead_ends`
pair that reset_session_caches() rebound at the start of EVERY search, while
the server allowed 4 concurrent searches — concurrent runs wiped each
other's caches mid-flight and re-embedded the same titles over and over.
The globals are gone. Each search now owns an EmbeddingCache instance
(navigator-scoped), and long-lived stores can own their own instance with
whatever lifetime they need. Nothing here is shared mutable state.
"""
from __future__ import annotations
import asyncio
import os
import time
from pathlib import Path
from typing import Optional
# Once the model has been downloaded once, sentence-transformers /
# huggingface_hub still spend 1-3s on every startup doing a network
# round-trip to check for updates. Must be set before sentence_transformers
# is imported; only skips the check when a cached model already exists
# locally (first run still goes online to fetch it).
_hf_cache = Path.home() / ".cache" / "huggingface" / "hub"
if _hf_cache.exists() and any(_hf_cache.glob("models--*")):
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
import numpy as np
from sentence_transformers import SentenceTransformer
from config import EMBED_MODEL_NAME, EMBED_DEVICE, EMBED_BATCH_SIZE
_EMBED_MODEL: Optional[SentenceTransformer] = None
# Single-flight lock around encode(): torch with num_threads=1 gains nothing
# from interleaved encodes, and serializing them keeps per-call latency
# predictable when several searches run at once.
_ENCODE_LOCK = asyncio.Lock()
async def load_model():
"""Startup handler: load the model once, in-process.
No silent fallback on failure — if the load fails this re-raises so the
server refuses connections rather than degrading silently.
"""
global _EMBED_MODEL
# Free-tier hosts give one throttled vCPU; torch's default thread pool
# spawns one thread per logical core it thinks exists, each with its own
# arena. Cap to 1: no extra cores to run them anyway.
import torch
torch.set_num_threads(1)
print(f"[Embed] Loading sentence-transformers model '{EMBED_MODEL_NAME}'...")
t0 = time.time()
_EMBED_MODEL = SentenceTransformer(EMBED_MODEL_NAME, device=EMBED_DEVICE)
print(f"[Embed] {EMBED_MODEL_NAME} ready "
f"(dim={_EMBED_MODEL.get_embedding_dimension()}, {time.time()-t0:.1f}s)")
def model_loaded() -> bool:
return _EMBED_MODEL is not None
def embedding_dim() -> int:
return _EMBED_MODEL.get_embedding_dimension() if _EMBED_MODEL else 0
def cosine_similarity(a, b) -> float:
if a is None or b is None:
return 0.0
a = np.asarray(a, dtype=np.float32)
b = np.asarray(b, dtype=np.float32)
if a.size == 0 or b.size == 0:
return 0.0
na = np.linalg.norm(a)
nb = np.linalg.norm(b)
if na == 0 or nb == 0:
return 0.0
return float(np.dot(a, b) / (na * nb))
class EmbeddingCache:
"""A key → vector cache with batched, executor-offloaded encoding.
Keys are caller-chosen (the navigator uses NodeRef.key()); `texts` is
what actually gets encoded — always the enriched "{title}. {context}"
form, never a bare title (the v1 semantic-drift lesson).
"""
def __init__(self):
self._cache: dict[str, np.ndarray] = {}
self.encode_calls = 0
def get(self, key: str) -> Optional[np.ndarray]:
return self._cache.get(key)
def put(self, key: str, emb: np.ndarray):
self._cache[key] = emb
def __contains__(self, key: str) -> bool:
return key in self._cache
async def embed(self, keys: list[str],
texts: list[str] | None = None) -> list[np.ndarray]:
"""Return embeddings for keys (parallel lists), encoding only the
uncached ones in a single batched, off-loop encode call."""
if not keys or _EMBED_MODEL is None:
return [np.array([]) for _ in keys]
if texts is None:
texts = keys
uncached_idx = [i for i, k in enumerate(keys) if k not in self._cache]
if uncached_idx:
uncached_texts = [texts[i] for i in uncached_idx]
loop = asyncio.get_running_loop()
async with _ENCODE_LOCK:
embs = await loop.run_in_executor(
None,
lambda: _EMBED_MODEL.encode(
uncached_texts, convert_to_numpy=True,
batch_size=EMBED_BATCH_SIZE, show_progress_bar=False,
),
)
self.encode_calls += 1
for orig_i, emb in zip(uncached_idx, embs):
self._cache[keys[orig_i]] = emb
return [self._cache.get(k, np.array([])) for k in keys]
def encode_texts_sync(texts: list[str]) -> np.ndarray:
"""Synchronous batch encode for offline ingestion pipelines (no event
loop, no cache). Raises if the model isn't loaded."""
if _EMBED_MODEL is None:
raise RuntimeError("Embedding model not loaded — call load_model() first")
return _EMBED_MODEL.encode(texts, convert_to_numpy=True,
batch_size=EMBED_BATCH_SIZE,
show_progress_bar=False)
|