Files
medtech-hackathon/etl/extractors/vision.py
T
admins 14462e8c10 feat(etl): нормализация (каскад), валидация (8 правил), Vision-fallback
- etl/normalize/matcher.py — каскад код→точное→эмбеддинги→fuzzy→unmatched
- etl/validate/rules.py — правила ТЗ §4.4; проверено: 15588 done / 11 needs_review
- etl/extractors/vision.py — Gemini Vision для битых/нестандартных PDF
- замер % автонормализации — после установки модели эмбеддингов

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-26 16:46:18 +05:00

91 lines
4.1 KiB
Python

"""Запасной путь извлечения через Gemini Vision.
Когда нативный парсер не справился (битый текстовый слой, нестандартная вёрстка,
скан), страница PDF отрисовывается в картинку и отдаётся Gemini с требованием
вернуть позиции прайса строго по JSON-схеме. Требует переменную окружения
GEMINI_API_KEY; модель задаётся через GEMINI_MODEL (по умолчанию gemini-2.0-flash).
"""
from __future__ import annotations
import json
import os
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from contracts.models import RawRow # noqa: E402
_VALID_TIERS = {"resident", "nonresident", "cis", "far", "insurance", "partner"}
_PROMPT = """Ты извлекаешь позиции из прайс-листа медицинской клиники на изображении.
Верни СТРОГО JSON-массив объектов, без пояснений и текста вокруг. Поля объекта:
- "name": название услуги как в документе (строка);
- "code": код услуги или тарификатора, если есть, иначе null;
- "unit": единица измерения, если есть, иначе null;
- "prices": объект «тариф → число в тенге». Ключи тарифов:
"resident" (резидент / гражданин РК), "nonresident" (нерезидент / иностранец),
"cis" (СНГ / ближнее зарубежье), "far" (дальнее зарубежье),
"insurance" (страховые компании), "partner" (партнёр).
Если в строке одна цена — положи её в "resident".
Пропускай заголовки разделов, итоги и строки без цены. Числа — без пробелов и валюты."""
def _render_pages(path: str, max_pages: int, dpi: int) -> list[bytes]:
"""Отрисовать первые страницы PDF в PNG-картинки."""
import fitz # PyMuPDF
document = fitz.open(path)
images = []
for page in document[:max_pages]:
images.append(page.get_pixmap(dpi=dpi).tobytes("png"))
return images
def _parse_response(text: str) -> list[RawRow]:
try:
data = json.loads(text)
except (json.JSONDecodeError, ValueError):
return []
rows: list[RawRow] = []
for obj in data if isinstance(data, list) else []:
if not isinstance(obj, dict):
continue
name = str(obj.get("name") or "").strip()
prices = {
tier: float(value)
for tier, value in (obj.get("prices") or {}).items()
if tier in _VALID_TIERS and isinstance(value, (int, float)) and value > 0
}
if name and prices:
rows.append(
RawRow(
service_name_raw=name,
service_code_source=(obj.get("code") or None),
prices=prices,
unit=(obj.get("unit") or None),
)
)
return rows
def extract_with_vision(path: str, max_pages: int = 12, dpi: int = 140) -> list[RawRow]:
"""Извлечь позиции из PDF через Gemini Vision. Возвращает список RawRow."""
api_key = os.environ.get("GEMINI_API_KEY")
if not api_key:
raise RuntimeError("нет GEMINI_API_KEY в окружении")
import google.generativeai as genai
genai.configure(api_key=api_key)
model = genai.GenerativeModel(os.environ.get("GEMINI_MODEL", "gemini-2.0-flash"))
rows: list[RawRow] = []
for png in _render_pages(path, max_pages, dpi):
response = model.generate_content(
[_PROMPT, {"mime_type": "image/png", "data": png}],
generation_config={"response_mime_type": "application/json", "temperature": 0},
)
rows.extend(_parse_response(response.text))
return rows