28 lines
903 B
Python
28 lines
903 B
Python
"""EmbeddingGemma 2 через Ollama (dim=768)."""
|
|
from __future__ import annotations
|
|
import httpx
|
|
from . import config
|
|
|
|
_client = httpx.Client(timeout=120)
|
|
|
|
def embed(text: str) -> list[float]:
|
|
"""Вернуть вектор для одного текста."""
|
|
r = _client.post(f"{config.OLLAMA_URL}/api/embed", json={
|
|
"model": config.EMBED_MODEL,
|
|
"input": text,
|
|
})
|
|
r.raise_for_status()
|
|
data = r.json()
|
|
vec = data["embeddings"][0]
|
|
assert len(vec) == config.EMBED_DIM, f"unexpected dim {len(vec)}"
|
|
return vec
|
|
|
|
def embed_batch(texts: list[str]) -> list[list[float]]:
|
|
"""Вернуть векторы для пачки текстов (один запрос)."""
|
|
r = _client.post(f"{config.OLLAMA_URL}/api/embed", json={
|
|
"model": config.EMBED_MODEL,
|
|
"input": texts,
|
|
})
|
|
r.raise_for_status()
|
|
return r.json()["embeddings"]
|