diff --git a/services/worker-gpu/app/config.py b/services/worker-gpu/app/config.py index d7823c8..53bf0aa 100644 --- a/services/worker-gpu/app/config.py +++ b/services/worker-gpu/app/config.py @@ -49,13 +49,19 @@ class Settings(BaseSettings): # FAISS / ML FAISS_INDEX_PATH: str = "/data/index/faiss.index" FAISS_ID_MAP_PATH: str = "/data/index/faiss_id_map.json" - # EMBED_BACKEND: "ollama" (GGUF-модель через Vulkan, для AMD GPU без ROCm/atomics) - # или "sentence_transformers" (локальная загрузка, CUDA/CPU) + # EMBED_BACKEND: "ollama" (GGUF-модель через Vulkan, для AMD GPU без ROCm/atomics), + # "sentence_transformers" (локальная загрузка, CUDA/CPU) или "cloud" (облачный + # инференс той же модели — не нужен вообще никакой локальный GPU). EMBED_BACKEND: str = "ollama" EMBED_MODEL: str = "bge-m3" EMBED_DEVICE: str = "cuda" # используется только при EMBED_BACKEND=sentence_transformers EMBED_BATCH_SIZE: int = 64 EMBED_DIM: int = 1024 # bge-m3; было 768 у paraphrase-multilingual-mpnet-base-v2 + # DeepInfra — OpenAI-совместимый инференс BAAI/bge-m3 в облаке, та же модель, + # что и на VM109, дешёвый пей-пер-токен. Используется при EMBED_BACKEND=cloud. + CLOUD_EMBED_URL: str = "https://api.deepinfra.com/v1/openai/embeddings" + CLOUD_EMBED_MODEL: str = "BAAI/bge-m3" + CLOUD_EMBED_API_KEY: str = "" FAISS_NLIST: int = 1024 # Количество кластеров для IVFFlat FAISS_NPROBE: int = 64 # Количество кластеров для поиска diff --git a/services/worker-gpu/app/model_manager.py b/services/worker-gpu/app/model_manager.py index 2c2b7d7..2e4f091 100644 --- a/services/worker-gpu/app/model_manager.py +++ b/services/worker-gpu/app/model_manager.py @@ -3,6 +3,7 @@ Бэкенд переключается через EMBED_BACKEND: - "ollama" — HTTP-вызов к Ollama /api/embed (GGUF-модель через Vulkan, работает на AMD GPU) - "sentence_transformers" — локальная загрузка через sentence-transformers (CUDA/CPU) +- "cloud" — облачный инференс той же модели (DeepInfra, OpenAI-совместимый формат) Модель (в случае sentence_transformers) загружается один раз при первом обращении и кэшируется в памяти GPU. @@ -67,6 +68,8 @@ class ModelManager: if settings.EMBED_BACKEND == "ollama": return cls._encode_ollama(texts) + if settings.EMBED_BACKEND == "cloud": + return cls._encode_cloud(texts) model = cls.get_model() vectors = model.encode( @@ -98,6 +101,30 @@ class ModelManager: norms[norms == 0] = 1.0 return vectors / norms # Нормализация для cosine через inner product + @classmethod + def _encode_cloud(cls, texts: list[str]) -> np.ndarray: + """Закодировать тексты через облачный инференс (DeepInfra, OpenAI-формат).""" + 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( + settings.CLOUD_EMBED_URL, + json={"model": settings.CLOUD_EMBED_MODEL, "input": batch}, + headers={"Authorization": f"Bearer {settings.CLOUD_EMBED_API_KEY}"}, + timeout=120.0, + ) + response.raise_for_status() + data = response.json()["data"] + # OpenAI-формат может вернуть элементы не по порядку — сортируем по index + data.sort(key=lambda d: d["index"]) + all_vectors.extend(d["embedding"] for d in data) + + 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 def encode_single(cls, text: str) -> np.ndarray: """Закодировать один текст. Удобный метод."""