50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
"""Offline, lazy CPU runtime; one encoder call per process, with a fast busy response."""
|
|
|
|
from functools import lru_cache
|
|
from threading import Lock
|
|
|
|
from fastapi import HTTPException
|
|
|
|
from app.core.config import settings
|
|
from app.services.embedding import LOCAL_MODELS, EmbeddingSpec, LocalEmbeddingProvider
|
|
|
|
|
|
class EmbeddingRuntime:
|
|
def __init__(self, spec: EmbeddingSpec, cache_dir, batch_size: int, threads: int):
|
|
self.spec = spec
|
|
self.cache_dir, self.batch_size, self.threads = cache_dir, batch_size, threads
|
|
self._provider = None
|
|
self._lock = Lock()
|
|
|
|
def encode(self, texts: list[str], *, query: bool = False) -> list[list[float]]:
|
|
if not self._lock.acquire(blocking=False):
|
|
raise HTTPException(503, "本地编码器忙，请稍后重试")
|
|
try:
|
|
if self._provider is None:
|
|
self._provider = LocalEmbeddingProvider(
|
|
self.spec,
|
|
cache_dir=self.cache_dir,
|
|
batch_size=self.batch_size,
|
|
threads=self.threads,
|
|
offline=True,
|
|
)
|
|
return self._provider.encode(texts, query=query)
|
|
finally:
|
|
self._lock.release()
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def runtime(spec, cache_dir, batch_size, threads):
|
|
return EmbeddingRuntime(spec, cache_dir, batch_size, threads)
|
|
|
|
|
|
def get_embedding_provider():
|
|
if not settings.embedding_enabled:
|
|
raise HTTPException(503, "本地 Embedding 未启用，请先完成模型缓存和运行环境配置")
|
|
return runtime(
|
|
LOCAL_MODELS[settings.embedding_model],
|
|
settings.embedding_cache_dir,
|
|
settings.embedding_batch_size,
|
|
settings.embedding_threads,
|
|
)
|