289 lines
12 KiB
Python
289 lines
12 KiB
Python
"""RAG 管道：Embedding（本地模型，ModelScope 下载，离线可用）"""
|
||
import sys
|
||
import threading
|
||
import math
|
||
import logging
|
||
import time
|
||
import os
|
||
from functools import lru_cache
|
||
from pathlib import Path
|
||
from ..config import MODELS_DIR, EMBEDDING_MODEL
|
||
|
||
_lock = threading.Lock()
|
||
_encoder = None
|
||
_encoder_name = None
|
||
_encoder_load_ms = 0.0
|
||
_metric_lock = threading.Lock()
|
||
_last_embedding_metric: dict = {}
|
||
_embedding_dimension = ""
|
||
logger = logging.getLogger("yawen.rag.embed")
|
||
|
||
# This is deliberately a separate value from the model name. Operators can
|
||
# bump it when replacing files in the model directory without changing the
|
||
# configured model identifier (which also invalidates query-vector caches).
|
||
EMBEDDING_MODEL_VERSION = os.environ.get("EMBEDDING_MODEL_VERSION", "1")
|
||
|
||
|
||
def _set_metric(row: dict) -> None:
|
||
"""Keep the most recent bounded embedding sample for the RAG envelope.
|
||
|
||
The RAG worker is single threaded today, but protecting this assignment
|
||
makes the helper safe when a worker is configured with concurrent requests.
|
||
No source text is retained in telemetry.
|
||
"""
|
||
global _last_embedding_metric
|
||
clean = {
|
||
"process_start_ms": 0.0,
|
||
"ipc_ms": 0.0,
|
||
"model_load_ms": 0.0,
|
||
"tokenizer_ms": 0.0,
|
||
"encode_ms": 0.0,
|
||
"postprocess_ms": 0.0,
|
||
"embedding_source": "sentence-transformers",
|
||
"embedding_count": 0,
|
||
"model": str(EMBEDDING_MODEL),
|
||
"model_version": str(EMBEDDING_MODEL_VERSION),
|
||
"pid": os.getpid(),
|
||
"cache_hit": False,
|
||
}
|
||
clean.update({key: value for key, value in row.items() if value is not None})
|
||
for key in ("process_start_ms", "ipc_ms", "model_load_ms", "tokenizer_ms",
|
||
"encode_ms", "postprocess_ms"):
|
||
try:
|
||
clean[key] = round(float(clean[key]), 2)
|
||
except (TypeError, ValueError):
|
||
clean[key] = 0.0
|
||
clean["embedding_count"] = int(clean.get("embedding_count") or 0)
|
||
clean["pid"] = int(clean.get("pid") or os.getpid())
|
||
with _metric_lock:
|
||
_last_embedding_metric = clean
|
||
|
||
|
||
def last_embedding_metric() -> dict:
|
||
"""Return phase timings for the latest encode without exposing input text."""
|
||
with _metric_lock:
|
||
return dict(_last_embedding_metric)
|
||
|
||
|
||
def query_cache_snapshot() -> dict:
|
||
"""Expose the independent query-vector LRU counters for diagnostics."""
|
||
info = _cached_query.cache_info()
|
||
total = info.hits + info.misses
|
||
return {
|
||
"size": info.currsize, "max_size": info.maxsize,
|
||
"hits": info.hits, "misses": info.misses,
|
||
"hit_rate": round(info.hits / total, 4) if total else 0.0,
|
||
"model": str(EMBEDDING_MODEL), "model_version": str(EMBEDDING_MODEL_VERSION),
|
||
"dimension": int(_embedding_dimension or 0),
|
||
"pid": os.getpid(),
|
||
}
|
||
|
||
|
||
def clear_query_cache() -> None:
|
||
"""Clear only query vectors; retrieval candidate pools remain untouched."""
|
||
_cached_query.cache_clear()
|
||
|
||
|
||
def set_transport_timings(*, process_start_ms: float = 0.0, ipc_ms: float = 0.0) -> None:
|
||
"""Attach process/IPC timings measured by the parent worker transport.
|
||
|
||
``embed.py`` cannot observe subprocess creation or pipe transfer itself;
|
||
the worker pool can call this after a request and the next envelope then
|
||
contains the complete set of phases. Keeping this optional preserves the
|
||
existing embedding API for in-process callers and tests.
|
||
"""
|
||
metric = last_embedding_metric()
|
||
metric.update(process_start_ms=process_start_ms, ipc_ms=ipc_ms)
|
||
_set_metric(metric)
|
||
|
||
|
||
def ensure_model_downloaded(model_name: str) -> Path:
|
||
"""模型目录存在即用；否则从 ModelScope 下载（HF 被墙）"""
|
||
safe = model_name.replace("/", "--")
|
||
local = MODELS_DIR / safe
|
||
if (local / "config.json").exists():
|
||
return local
|
||
try:
|
||
from modelscope import snapshot_download
|
||
print(f"[embed] 从 ModelScope 下载 {model_name} ...", file=sys.stderr)
|
||
snapshot_download(model_name, local_dir=str(local))
|
||
return local
|
||
except Exception as e:
|
||
raise RuntimeError(f"模型下载失败（ModelScope）: {e}")
|
||
|
||
|
||
def get_encoder():
|
||
"""单例加载 SentenceTransformer（本地路径）"""
|
||
global _encoder, _encoder_name, _encoder_load_ms
|
||
if _encoder is not None and _encoder_name == EMBEDDING_MODEL:
|
||
return _encoder
|
||
with _lock:
|
||
if _encoder is not None and _encoder_name == EMBEDDING_MODEL:
|
||
return _encoder
|
||
started = time.perf_counter()
|
||
from sentence_transformers import SentenceTransformer
|
||
local = ensure_model_downloaded(EMBEDDING_MODEL)
|
||
print(f"[embed] 加载 {local} ...", file=sys.stderr)
|
||
_encoder = SentenceTransformer(str(local))
|
||
_encoder_name = EMBEDDING_MODEL
|
||
_encoder_load_ms = (time.perf_counter() - started) * 1000
|
||
return _encoder
|
||
|
||
|
||
def prewarm() -> bool:
|
||
"""Load the embedding model and run one tiny encode as a best-effort warmup.
|
||
|
||
Model loading is intentionally kept out of the API startup path. The
|
||
caller should run this function from a daemon thread so a slow model
|
||
download/load never delays health checks or request handling. Returning a
|
||
boolean keeps this useful for startup diagnostics without making warmup
|
||
failures fatal to the service.
|
||
"""
|
||
started = time.perf_counter()
|
||
try:
|
||
encoder = get_encoder()
|
||
# A real encode initializes tokenizer/runtime kernels as well as the
|
||
# SentenceTransformer object. The result is discarded; query vectors
|
||
# continue to use the normal cached path.
|
||
encoder.encode(["系统启动预热"], normalize_embeddings=True,
|
||
show_progress_bar=False, batch_size=1)
|
||
logger.info("Embedding 模型预热完成：%s，耗时 %.3f 秒",
|
||
EMBEDDING_MODEL, time.perf_counter() - started)
|
||
return True
|
||
except Exception:
|
||
logger.exception("Embedding 模型预热失败：%s，服务将按需加载", EMBEDDING_MODEL)
|
||
return False
|
||
|
||
|
||
EMBEDDING_POLICY = "full-token-windows-v1"
|
||
|
||
|
||
def text_windows(text: str, encoder=None) -> list[str]:
|
||
"""Cover every source character without sending over-limit windows to BGE."""
|
||
enc = encoder or get_encoder()
|
||
tokenizer = enc.tokenizer
|
||
maximum = int(enc.max_seq_length)
|
||
def fits(value):
|
||
return len(tokenizer.encode(value, add_special_tokens=True, truncation=False)) <= maximum
|
||
if fits(text):
|
||
return [text]
|
||
# Character boundaries preserve exact original text; a binary search uses the
|
||
# actual model tokenizer rather than assuming that one character is one token.
|
||
windows, offset = [], 0
|
||
while offset < len(text):
|
||
remaining = text[offset:]
|
||
if fits(remaining):
|
||
windows.append(remaining)
|
||
break
|
||
low, high = 1, min(len(remaining), maximum * 4)
|
||
while low < high:
|
||
middle = (low + high + 1) // 2
|
||
if fits(remaining[:middle]):
|
||
low = middle
|
||
else:
|
||
high = middle - 1
|
||
if not fits(remaining[:low]):
|
||
raise ValueError("单个字符超过向量模型输入上限")
|
||
windows.append(remaining[:low])
|
||
offset += low
|
||
return windows
|
||
|
||
|
||
def embed_texts(texts: list[str]) -> list[list[float]]:
|
||
if not texts:
|
||
return []
|
||
# ``get_encoder`` includes model download and construction only on the
|
||
# first call in each resident worker. Subsequent calls report zero load
|
||
# time, which makes accidental worker restarts visible in telemetry.
|
||
enc = get_encoder()
|
||
global _encoder_load_ms
|
||
model_load_ms = _encoder_load_ms
|
||
# Consume the one-shot load measurement. A resident encoder must report
|
||
# zero on later requests; a non-zero value then reliably indicates a worker
|
||
# restart or an explicit model replacement.
|
||
_encoder_load_ms = 0.0
|
||
# At most 64 originals and 64 encoded windows in one model batch. Long
|
||
# paragraphs are pooled in the existing vector space, keeping chunk IDs and
|
||
# consumption records stable while representing their tails as well.
|
||
output = []
|
||
tokenizer_ms = encode_ms = postprocess_ms = 0.0
|
||
flat_count = 0
|
||
for start in range(0, len(texts), 64):
|
||
tokenizer_started = time.perf_counter()
|
||
groups = [text_windows(str(text), enc) for text in texts[start:start + 64]]
|
||
flat = [window for group in groups for window in group]
|
||
tokenizer_ms += (time.perf_counter() - tokenizer_started) * 1000
|
||
flat_count += len(flat)
|
||
encode_started = time.perf_counter()
|
||
vectors = enc.encode(flat, normalize_embeddings=True, show_progress_bar=False, batch_size=64)
|
||
encode_ms += (time.perf_counter() - encode_started) * 1000
|
||
postprocess_started = time.perf_counter()
|
||
cursor = 0
|
||
for group in groups:
|
||
selected = vectors[cursor:cursor + len(group)]
|
||
cursor += len(group)
|
||
if len(group) == 1:
|
||
output.append([float(v) for v in selected[0]])
|
||
continue
|
||
weights = [max(1, len(enc.tokenizer.encode(window, add_special_tokens=False, truncation=False)))
|
||
for window in group]
|
||
combined = [sum(float(vector[i]) * weight for vector, weight in zip(selected, weights))
|
||
for i in range(len(selected[0]))]
|
||
norm = math.sqrt(sum(value * value for value in combined)) or 1.0
|
||
output.append([value / norm for value in combined])
|
||
postprocess_ms += (time.perf_counter() - postprocess_started) * 1000
|
||
_set_metric({
|
||
"model_load_ms": model_load_ms,
|
||
"tokenizer_ms": tokenizer_ms,
|
||
"encode_ms": encode_ms,
|
||
"postprocess_ms": postprocess_ms,
|
||
"embedding_count": len(texts),
|
||
# ``flat_count`` is useful when long documents are split into windows;
|
||
# retain it as a diagnostic without changing the public return value.
|
||
"window_count": flat_count,
|
||
"embedding_dimension": len(output[0]) if output else 0,
|
||
"cache_hit": False,
|
||
})
|
||
return output
|
||
|
||
|
||
@lru_cache(maxsize=256)
|
||
def _cached_query(model: str, policy: str, version: str, dimension: str, text: str) -> tuple[float, ...]:
|
||
"""Cache only query vectors, keyed by model identity and policy version."""
|
||
return tuple(embed_texts([text])[0])
|
||
|
||
|
||
def embed_query(text: str) -> list[float]:
|
||
# lru_cache does not expose whether this invocation hit without comparing
|
||
# its counters. The comparison is process local and avoids retaining the
|
||
# query itself in telemetry.
|
||
before = _cached_query.cache_info()
|
||
global _embedding_dimension
|
||
# Keep the cache identity stable across the first call: the actual vector
|
||
# dimension is discovered only after encoding. Model/version already
|
||
# identify the encoder contract; ``auto`` prevents the first "unknown"
|
||
# dimension from creating a second key for the same query.
|
||
result = list(_cached_query(EMBEDDING_MODEL, EMBEDDING_POLICY,
|
||
EMBEDDING_MODEL_VERSION, "auto", str(text)))
|
||
if not _embedding_dimension and result:
|
||
_embedding_dimension = str(len(result))
|
||
after = _cached_query.cache_info()
|
||
cache_hit = after.hits > before.hits
|
||
metric = last_embedding_metric()
|
||
metric.update(
|
||
embedding_count=1,
|
||
cache_hit=cache_hit,
|
||
embedding_source="query-vector-lru",
|
||
model=EMBEDDING_MODEL,
|
||
model_version=EMBEDDING_MODEL_VERSION,
|
||
pid=os.getpid(),
|
||
embedding_dimension=len(result) if result else 0,
|
||
)
|
||
if cache_hit:
|
||
# A cache hit performs no tokenizer, model inference, or serialization;
|
||
# do not leak timings from the previous (possibly indexing) call.
|
||
metric.update(model_load_ms=0.0, tokenizer_ms=0.0,
|
||
encode_ms=0.0, postprocess_ms=0.0, window_count=0)
|
||
_set_metric(metric)
|
||
return result
|