Files
medtech-hackathon/etl/normalize/embedding.py
T
admins 69a9884db5 feat(ingest+поиск): приём ZIP через интерфейс, регистронезависимый поиск, чистка очереди
- Загрузка архива (§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>
2026-06-27 11:35:12 +05:00

73 lines
2.9 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Эмбеддинги через 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