69a9884db5
- Загрузка архива (§4.1): POST /ingest распаковывает ZIP, дораспознаёт и догружает прайсы в базу (append). Виджет в админке. Кэш эмбеддингов справочника (dict_emb.npy) — догрузка считает только новый файл, эмбеддинги best-effort с откатом на fuzzy при сбое Gemini. - Поиск по name_norm (lowercase): «УЗИ» и «узи» дают одинаковый результат — SQLite LIKE не сворачивает регистр для кириллицы. - Фильтр названий: код-мнемоники («МОЗО.15») больше не попадают в очередь. - /stats считается из базы — цифры на дашборде обновляются после загрузки. - Деплой через Dockerfile (зависимости в образе, код монтируется томом). - ruff: ядро и весь репозиторий проходят проверку. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
"""Эмбеддинги через Gemini API (SDK google-genai) — без локального torch.
|
||
|
||
Лёгкая зависимость вместо sentence-transformers: контейнер на маленьком VPS не
|
||
раздувается и не держит модель в RAM. Размерность 768 выбрана, чтобы вектор лез в
|
||
HNSW-индекс pgvector. Требуется переменная окружения GEMINI_API_KEY.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
|
||
import numpy as np
|
||
|
||
EMBED_DIM = 768
|
||
|
||
|
||
class GeminiEmbedder:
|
||
"""Обёртка над Gemini-эмбеддингами с интерфейсом `.encode()` (как у sentence-transformers)."""
|
||
|
||
def __init__(
|
||
self,
|
||
model: str = "gemini-embedding-001",
|
||
dim: int = EMBED_DIM,
|
||
api_key: str | None = None,
|
||
batch: int = 100,
|
||
):
|
||
from google import genai
|
||
|
||
self._client = genai.Client(api_key=api_key or os.environ["GEMINI_API_KEY"])
|
||
self.model = model
|
||
self.dim = dim
|
||
self.batch = batch
|
||
|
||
def encode(
|
||
self,
|
||
texts,
|
||
normalize_embeddings: bool = True,
|
||
task_type: str = "RETRIEVAL_DOCUMENT",
|
||
show_progress_bar: bool = False,
|
||
):
|
||
"""Векторизовать список строк. Возвращает numpy-массив (n, dim) float32.
|
||
|
||
Для справочника используем RETRIEVAL_DOCUMENT, для запросов —
|
||
RETRIEVAL_QUERY: так несвязанные названия расходятся сильнее, чем при
|
||
SEMANTIC_SIMILARITY, и порог отсечения работает надёжнее.
|
||
"""
|
||
import time
|
||
|
||
from google.genai import types
|
||
|
||
items = [t if t and t.strip() else " " for t in texts]
|
||
config = types.EmbedContentConfig(task_type=task_type, output_dimensionality=self.dim)
|
||
vectors: list[list[float]] = []
|
||
for start in range(0, len(items), self.batch):
|
||
chunk = items[start : start + self.batch]
|
||
for attempt in range(8):
|
||
try:
|
||
response = self._client.models.embed_content(
|
||
model=self.model, contents=chunk, config=config
|
||
)
|
||
vectors.extend(embedding.values for embedding in response.embeddings)
|
||
break
|
||
except Exception as exc: # пауза и повтор при превышении квоты (429)
|
||
if ("RESOURCE_EXHAUSTED" in str(exc) or "429" in str(exc)) and attempt < 7:
|
||
time.sleep(20)
|
||
continue
|
||
raise
|
||
|
||
array = np.array(vectors, dtype=np.float32)
|
||
if normalize_embeddings:
|
||
array = array / np.clip(np.linalg.norm(array, axis=1, keepdims=True), 1e-9, None)
|
||
return array
|