83ec1aca24
- build_section_docs.py: indications, contraindications, warnings and interactions only, one representative per active-substance set - ingest: --dedupe, --retry-failed - embedding server: bounded PyTorch allocator cache and serialized encoding (fixes memory growth that crashed long ingests) - rebuild_vdb.py: recompute vector stores from the graph after corruption
88 lines
2.8 KiB
Python
88 lines
2.8 KiB
Python
"""
|
||
Локальный OpenAI-compatible embedding сервер на базе sentence-transformers.
|
||
Модель: paraphrase-multilingual-MiniLM-L12-v2 (поддерживает словацкий язык!)
|
||
|
||
Запуск:
|
||
python embedding_server.py
|
||
|
||
Тест:
|
||
curl http://localhost:8010/v1/embeddings -H "Content-Type: application/json" \
|
||
-d '{"model": "local-embed", "input": "test"}'
|
||
"""
|
||
|
||
import os
|
||
|
||
# Без этого кэш CPU-аллокатора PyTorch растёт с каждой новой длиной входа
|
||
# (за ночь сервер дорос до ~45 ГБ виртуальной памяти и упал).
|
||
os.environ.setdefault("LRU_CACHE_CAPACITY", "1")
|
||
|
||
import threading
|
||
import time
|
||
import json
|
||
from fastapi import FastAPI, Request
|
||
from fastapi.concurrency import run_in_threadpool
|
||
from fastapi.responses import JSONResponse
|
||
import uvicorn
|
||
from sentence_transformers import SentenceTransformer
|
||
|
||
MODEL_NAME = "paraphrase-multilingual-MiniLM-L12-v2"
|
||
PORT = 8010
|
||
|
||
print(f"Загрузка модели {MODEL_NAME}...")
|
||
model = SentenceTransformer(MODEL_NAME)
|
||
EMBED_DIM = model.get_sentence_embedding_dimension()
|
||
print(f"Модель загружена. Размерность: {EMBED_DIM}")
|
||
|
||
app = FastAPI(title="Local Embedding Server")
|
||
|
||
# Одновременно считается только один батч: параллельные encode() в потоках
|
||
# не ускоряют CPU-инференс, но умножают потребление памяти.
|
||
_encode_lock = threading.Lock()
|
||
|
||
|
||
def _encode(texts):
|
||
with _encode_lock:
|
||
return model.encode(texts, normalize_embeddings=True)
|
||
|
||
|
||
@app.get("/health")
|
||
def health():
|
||
return {"status": "ok", "model": MODEL_NAME, "dim": EMBED_DIM}
|
||
|
||
|
||
async def _handle_embeddings(request: Request):
|
||
body = await request.json()
|
||
inp = body.get("input", "")
|
||
if isinstance(inp, str):
|
||
texts = [inp]
|
||
else:
|
||
texts = inp
|
||
|
||
# encode() блокирует; в пуле потоков сервер продолжает принимать запросы
|
||
vecs = (await run_in_threadpool(_encode, texts)).tolist()
|
||
|
||
data = [
|
||
{"object": "embedding", "index": i, "embedding": vec}
|
||
for i, vec in enumerate(vecs)
|
||
]
|
||
return JSONResponse({
|
||
"object": "list",
|
||
"data": data,
|
||
"model": MODEL_NAME,
|
||
"usage": {"prompt_tokens": sum(len(t.split()) for t in texts), "total_tokens": sum(len(t.split()) for t in texts)}
|
||
})
|
||
|
||
|
||
@app.post("/v1/embeddings")
|
||
async def embeddings_v1(request: Request):
|
||
return await _handle_embeddings(request)
|
||
|
||
|
||
@app.post("/embeddings")
|
||
async def embeddings_root(request: Request):
|
||
return await _handle_embeddings(request)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="warning")
|