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
|