143 lines
5.1 KiB
Python
143 lines
5.1 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Генерация эмбеддинга по JSON шага 5 через Ollama (шаг 6).
|
||
Текст для эмбеддинга собирается из framework, insights, application по embed_input_spec.txt.
|
||
Вход по умолчанию: merged_with_tags.json (5). Выход по умолчанию: embedding.json (вектор).
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
DIR = Path(__file__).resolve().parent
|
||
DEFAULT_MERGED = DIR.parent / "5_мерж_анализа_и_тегов" / "merged_with_tags.json"
|
||
DEFAULT_OUTPUT = DIR / "embedding.json"
|
||
|
||
OLLAMA_URL = "http://localhost:11434"
|
||
EMBED_MODEL = "bge-m3"
|
||
|
||
|
||
def get_embedding_ollama(base_url: str, model: str, text: str) -> list[float]:
|
||
"""
|
||
Запрашивает эмбеддинг текста у Ollama API (POST /api/embed).
|
||
|
||
Args:
|
||
base_url: Базовый URL Ollama (например http://localhost:11434).
|
||
model: Имя модели эмбеддингов (например bge-m3).
|
||
text: Текст для эмбеддинга.
|
||
|
||
Returns:
|
||
Вектор эмбеддинга (список float).
|
||
|
||
Raises:
|
||
urllib.error.HTTPError: При ошибке HTTP.
|
||
ValueError: Если в ответе нет ожидаемой структуры.
|
||
"""
|
||
url = f"{base_url.rstrip('/')}/api/embed"
|
||
payload = {"model": model, "input": text}
|
||
body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||
req = urllib.request.Request(
|
||
url,
|
||
data=body,
|
||
headers={"Content-Type": "application/json"},
|
||
method="POST",
|
||
)
|
||
with urllib.request.urlopen(req, timeout=120) as resp:
|
||
data: dict[str, Any] = json.loads(resp.read().decode("utf-8"))
|
||
if "embeddings" not in data or not data["embeddings"]:
|
||
raise ValueError("В ответе Ollama нет поля embeddings")
|
||
embedding = data["embeddings"][0]
|
||
if not isinstance(embedding, list):
|
||
raise ValueError("embeddings[0] не является массивом")
|
||
return [float(x) for x in embedding]
|
||
|
||
|
||
def main() -> int:
|
||
"""Собирает текст из merged JSON, вызывает Ollama /api/embed, пишет вектор в файл."""
|
||
from embed_cli import merged_json_to_embed_text
|
||
|
||
parser = argparse.ArgumentParser(
|
||
description="Эмбеддинг по JSON шага 5 через Ollama (шаг 6). На выход — вектор (JSON).",
|
||
)
|
||
parser.add_argument(
|
||
"--merged",
|
||
type=Path,
|
||
default=DEFAULT_MERGED,
|
||
help=f"Путь к merged_with_tags.json (по умолчанию: {DEFAULT_MERGED})",
|
||
)
|
||
parser.add_argument(
|
||
"--model",
|
||
default=EMBED_MODEL,
|
||
help=f"Модель эмбеддингов в Ollama (по умолчанию: {EMBED_MODEL})",
|
||
)
|
||
parser.add_argument(
|
||
"--ollama-url",
|
||
default=OLLAMA_URL,
|
||
help=f"URL Ollama (по умолчанию: {OLLAMA_URL})",
|
||
)
|
||
parser.add_argument(
|
||
"-o",
|
||
"--output",
|
||
type=Path,
|
||
default=DEFAULT_OUTPUT,
|
||
help=f"Путь к выходному JSON с вектором (по умолчанию: {DEFAULT_OUTPUT})",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
if not args.merged.is_file():
|
||
print(f"Файл не найден: {args.merged}", file=sys.stderr)
|
||
return 1
|
||
|
||
print("Загрузка merged_with_tags.json...")
|
||
try:
|
||
with open(args.merged, encoding="utf-8") as f:
|
||
merged = json.load(f)
|
||
except json.JSONDecodeError as e:
|
||
print(f"Ошибка разбора JSON: {e}", file=sys.stderr)
|
||
return 1
|
||
|
||
text = merged_json_to_embed_text(merged)
|
||
if not text:
|
||
print(
|
||
"Ошибка: текст для эмбеддинга пуст (нет framework/insights/application).",
|
||
file=sys.stderr,
|
||
)
|
||
return 1
|
||
|
||
print(f"Вызов Ollama {args.model} — генерация эмбеддинга...")
|
||
t0 = time.monotonic()
|
||
try:
|
||
vector = get_embedding_ollama(args.ollama_url, args.model, text)
|
||
except urllib.error.HTTPError as e:
|
||
print(f"Ошибка HTTP {e.code}: {e.reason}", file=sys.stderr)
|
||
if e.fp:
|
||
try:
|
||
body = e.fp.read().decode("utf-8")
|
||
print(body[:500], file=sys.stderr)
|
||
except Exception:
|
||
pass
|
||
return 1
|
||
except urllib.error.URLError as e:
|
||
print(f"Ошибка запроса: {e.reason}", file=sys.stderr)
|
||
return 1
|
||
except ValueError as e:
|
||
print(f"Ошибка: {e}", file=sys.stderr)
|
||
return 1
|
||
elapsed = time.monotonic() - t0
|
||
print(f"Эмбеддинг получен за {elapsed:.1f} сек, размерность {len(vector)}")
|
||
|
||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||
with open(args.output, "w", encoding="utf-8") as f:
|
||
json.dump(vector, f, ensure_ascii=False)
|
||
print(f"Записано: {args.output}")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|