feat(embeddings): переключить эмбеддинги на Ollama/bge-m3 (GPU через Vulkan)
Раньше эмбеддинг-модель (уровень 3) гоняла на CPU внутри worker-gpu —
GPU CT108 использовался только под LLM-парафраз (уровень 4). Теперь
эмбеддинги идут через Ollama /api/embed на отдельной VM с RX 580
(Vulkan-бэкенд, без возни с ROCm/HIP для этой карты). EMBED_BACKEND
переключаемый ("ollama" | "sentence_transformers"), дефолт — ollama.
Модель сменилась на bge-m3 (1024-мерный вектор вместо 768 у
paraphrase-multilingual-mpnet-base-v2) — несовместимо с уже посчитанным
FAISS-индексом, нужна полная переиндексация корпуса после деплоя.
Заодно докстринг delete_task в tasks.py — снятое раньше ограничение
"нельзя удалить processing" оставляло враньё в докстринге.
This commit is contained in:
@@ -107,7 +107,7 @@ async def delete_task(
|
|||||||
current_user: User = Depends(get_current_user),
|
current_user: User = Depends(get_current_user),
|
||||||
db: AsyncSession = Depends(get_db),
|
db: AsyncSession = Depends(get_db),
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Удалить задачу. Нельзя удалить задачу в статусе 'processing'."""
|
"""Удалить задачу."""
|
||||||
result = await db.execute(
|
result = await db.execute(
|
||||||
select(Task).where(
|
select(Task).where(
|
||||||
Task.public_id == public_id,
|
Task.public_id == public_id,
|
||||||
|
|||||||
@@ -41,10 +41,13 @@ class Settings(BaseSettings):
|
|||||||
# FAISS / ML
|
# FAISS / ML
|
||||||
FAISS_INDEX_PATH: str = "/data/index/faiss.index"
|
FAISS_INDEX_PATH: str = "/data/index/faiss.index"
|
||||||
FAISS_ID_MAP_PATH: str = "/data/index/faiss_id_map.json"
|
FAISS_ID_MAP_PATH: str = "/data/index/faiss_id_map.json"
|
||||||
EMBED_MODEL: str = "paraphrase-multilingual-mpnet-base-v2"
|
# EMBED_BACKEND: "ollama" (GGUF-модель через Vulkan, для AMD GPU без ROCm/atomics)
|
||||||
EMBED_DEVICE: str = "cuda"
|
# или "sentence_transformers" (локальная загрузка, CUDA/CPU)
|
||||||
|
EMBED_BACKEND: str = "ollama"
|
||||||
|
EMBED_MODEL: str = "bge-m3"
|
||||||
|
EMBED_DEVICE: str = "cuda" # используется только при EMBED_BACKEND=sentence_transformers
|
||||||
EMBED_BATCH_SIZE: int = 64
|
EMBED_BATCH_SIZE: int = 64
|
||||||
EMBED_DIM: int = 768 # Размерность вектора paraphrase-multilingual-mpnet-base-v2
|
EMBED_DIM: int = 1024 # bge-m3; было 768 у paraphrase-multilingual-mpnet-base-v2
|
||||||
FAISS_NLIST: int = 1024 # Количество кластеров для IVFFlat
|
FAISS_NLIST: int = 1024 # Количество кластеров для IVFFlat
|
||||||
FAISS_NPROBE: int = 64 # Количество кластеров для поиска
|
FAISS_NPROBE: int = 64 # Количество кластеров для поиска
|
||||||
|
|
||||||
|
|||||||
@@ -1,10 +1,16 @@
|
|||||||
"""Singleton менеджер sentence-transformers модели.
|
"""Singleton менеджер эмбеддинг-модели.
|
||||||
|
|
||||||
Модель загружается один раз при первом обращении и кэшируется в памяти GPU.
|
Бэкенд переключается через EMBED_BACKEND:
|
||||||
|
- "ollama" — HTTP-вызов к Ollama /api/embed (GGUF-модель через Vulkan, работает на AMD GPU)
|
||||||
|
- "sentence_transformers" — локальная загрузка через sentence-transformers (CUDA/CPU)
|
||||||
|
|
||||||
|
Модель (в случае sentence_transformers) загружается один раз при первом обращении
|
||||||
|
и кэшируется в памяти GPU.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
|
|
||||||
|
import httpx
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
from app.config import settings
|
from app.config import settings
|
||||||
@@ -54,11 +60,14 @@ class ModelManager:
|
|||||||
texts: Список текстов для кодирования
|
texts: Список текстов для кодирования
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
numpy массив формы (len(texts), 768), нормализованный для cosine similarity
|
numpy массив формы (len(texts), EMBED_DIM), нормализованный для cosine similarity
|
||||||
"""
|
"""
|
||||||
if not texts:
|
if not texts:
|
||||||
return np.array([]).reshape(0, settings.EMBED_DIM)
|
return np.array([]).reshape(0, settings.EMBED_DIM)
|
||||||
|
|
||||||
|
if settings.EMBED_BACKEND == "ollama":
|
||||||
|
return cls._encode_ollama(texts)
|
||||||
|
|
||||||
model = cls.get_model()
|
model = cls.get_model()
|
||||||
vectors = model.encode(
|
vectors = model.encode(
|
||||||
texts,
|
texts,
|
||||||
@@ -69,6 +78,26 @@ class ModelManager:
|
|||||||
)
|
)
|
||||||
return vectors.astype(np.float32)
|
return vectors.astype(np.float32)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _encode_ollama(cls, texts: list[str]) -> np.ndarray:
|
||||||
|
"""Закодировать тексты через Ollama /api/embed (пакетами по EMBED_BATCH_SIZE)."""
|
||||||
|
all_vectors: list[list[float]] = []
|
||||||
|
batch_size = settings.EMBED_BATCH_SIZE
|
||||||
|
for i in range(0, len(texts), batch_size):
|
||||||
|
batch = texts[i : i + batch_size]
|
||||||
|
response = httpx.post(
|
||||||
|
f"{settings.OLLAMA_URL}/api/embed",
|
||||||
|
json={"model": settings.EMBED_MODEL, "input": batch},
|
||||||
|
timeout=120.0,
|
||||||
|
)
|
||||||
|
response.raise_for_status()
|
||||||
|
all_vectors.extend(response.json()["embeddings"])
|
||||||
|
|
||||||
|
vectors = np.array(all_vectors, dtype=np.float32)
|
||||||
|
norms = np.linalg.norm(vectors, axis=1, keepdims=True)
|
||||||
|
norms[norms == 0] = 1.0
|
||||||
|
return vectors / norms # Нормализация для cosine через inner product
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def encode_single(cls, text: str) -> np.ndarray:
|
def encode_single(cls, text: str) -> np.ndarray:
|
||||||
"""Закодировать один текст. Удобный метод."""
|
"""Закодировать один текст. Удобный метод."""
|
||||||
|
|||||||
Reference in New Issue
Block a user