diff --git a/etl/extractors/__init__.py b/etl/extractors/__init__.py index 0043313..8a7696c 100644 --- a/etl/extractors/__init__.py +++ b/etl/extractors/__init__.py @@ -6,6 +6,7 @@ PDF с битым текстовым слоем помечается как scan """ from __future__ import annotations +import os import sys from datetime import date from pathlib import Path @@ -53,6 +54,19 @@ def extract(path: str) -> ParsedDocument: except Exception as exc: # noqa: BLE001 doc.parse_log.append(f"ошибка разбора таблицы: {type(exc).__name__}: {exc}") - if not doc.rows and doc.file_format != "scan_pdf": - doc.parse_log.append("позиции не извлечены — возможно, нужен Vision или иной парсер") + # Трудные PDF (скан, битый слой, нулевой разбор) добираем через Gemini Vision. + needs_vision = fmt == "pdf" and (doc.file_format == "scan_pdf" or not doc.rows) + if needs_vision and os.environ.get("GEMINI_API_KEY"): + try: + from etl.extractors.vision import extract_with_vision + + vision_rows = extract_with_vision(path) + doc.rows.extend(vision_rows) + doc.file_format = "scan_pdf" + doc.parse_log.append(f"Vision добрал {len(vision_rows)} позиций") + except Exception as exc: # noqa: BLE001 + doc.parse_log.append(f"Vision не сработал: {type(exc).__name__}: {exc}") + + if not doc.rows: + doc.parse_log.append("позиции не извлечены") return doc diff --git a/etl/extractors/common.py b/etl/extractors/common.py index 72cba14..4c490d6 100644 --- a/etl/extractors/common.py +++ b/etl/extractors/common.py @@ -76,6 +76,18 @@ def looks_like_code(text) -> bool: return bool(s) and len(s) <= 16 and bool(_CODE_RE.match(s)) +def is_real_name(text) -> bool: + """Похоже ли значение на настоящее название услуги, а не на код или число. + + Защищает нормализацию от мусора: коды («E04.000.012»), номера и цены, по ошибке + попавшие в колонку названия, не должны уходить в сопоставление со справочником. + Требуем минимум три кириллические буквы и не код-подобную форму. + """ + s = clean(text) + cyrillic = sum(1 for ch in s if "а" <= ch.lower() <= "я" or ch.lower() == "ё") + return cyrillic >= 3 and not looks_like_code(s) + + def classify_price_tier(header_text) -> str | None: """По тексту заголовка ценовой колонки определить тип тарифа.""" h = clean(header_text).lower() diff --git a/etl/extractors/grid.py b/etl/extractors/grid.py index 3497a06..03f1e95 100644 --- a/etl/extractors/grid.py +++ b/etl/extractors/grid.py @@ -25,6 +25,7 @@ from etl.extractors.common import ( # noqa: E402 classify_price_tier, clean, find_header_row, + is_real_name, looks_like_code, parse_price, ) @@ -157,8 +158,8 @@ def extract_rows_from_grid(grid: Grid) -> list[RawRow]: continue name = r[name_col] if name_col < len(r) else "" - if not name or not prices: - continue # без названия или без цены это не позиция прайса + if not prices or not is_real_name(name): + continue # не позиция: нет цены либо в «названии» код/число, а не услуга rows.append( RawRow( diff --git a/etl/extractors/vision.py b/etl/extractors/vision.py index ef8428f..03b947e 100644 --- a/etl/extractors/vision.py +++ b/etl/extractors/vision.py @@ -33,9 +33,9 @@ _PROMPT = """Ты извлекаешь позиции из прайс-листа def _render_pages(path: str, max_pages: int, dpi: int) -> list[bytes]: """Отрисовать первые страницы PDF в PNG-картинки.""" - import fitz # PyMuPDF + import pymupdf # PyMuPDF - document = fitz.open(path) + document = pymupdf.open(path) images = [] for page in document[:max_pages]: images.append(page.get_pixmap(dpi=dpi).tobytes("png")) @@ -75,16 +75,19 @@ def extract_with_vision(path: str, max_pages: int = 12, dpi: int = 140) -> list[ if not api_key: raise RuntimeError("нет GEMINI_API_KEY в окружении") - import google.generativeai as genai + from google import genai + from google.genai import types - genai.configure(api_key=api_key) - model = genai.GenerativeModel(os.environ.get("GEMINI_MODEL", "gemini-2.0-flash")) + client = genai.Client(api_key=api_key) + model = os.environ.get("GEMINI_MODEL", "gemini-2.5-flash") + config = types.GenerateContentConfig(response_mime_type="application/json", temperature=0) 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}, + response = client.models.generate_content( + model=model, + contents=[types.Part.from_bytes(data=png, mime_type="image/png"), _PROMPT], + config=config, ) rows.extend(_parse_response(response.text)) return rows diff --git a/etl/normalize/embedding.py b/etl/normalize/embedding.py new file mode 100644 index 0000000..cdd3334 --- /dev/null +++ b/etl/normalize/embedding.py @@ -0,0 +1,59 @@ +"""Эмбеддинги через 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, и порог отсечения работает надёжнее. + """ + 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] + response = self._client.models.embed_content(model=self.model, contents=chunk, config=config) + vectors.extend(embedding.values for embedding in response.embeddings) + + 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 diff --git a/etl/normalize/matcher.py b/etl/normalize/matcher.py index a6352e3..28cd4aa 100644 --- a/etl/normalize/matcher.py +++ b/etl/normalize/matcher.py @@ -35,7 +35,7 @@ class Matcher: self, services: list[Service], embedder=None, - emb_threshold: float = 0.60, + emb_threshold: float = 0.70, fuzzy_threshold: int = 88, ): self.services = services @@ -50,7 +50,9 @@ class Matcher: self.dict_emb = None if embedder is not None: self.dict_emb = embedder.encode( - [s.name_ru for s in services], normalize_embeddings=True, show_progress_bar=False + [s.name_ru for s in services], + normalize_embeddings=True, + task_type="RETRIEVAL_DOCUMENT", ) def _match_by_code(self, code: str | None) -> Service | None: @@ -88,7 +90,9 @@ class Matcher: if pending_norm and self.embedder is not None and self.dict_emb is not None: query = self.embedder.encode( - [names[i] for i in pending_idx], normalize_embeddings=True, show_progress_bar=False + [names[i] for i in pending_idx], + normalize_embeddings=True, + task_type="RETRIEVAL_QUERY", ) sims = query @ self.dict_emb.T # косинус по нормированным векторам best_idx = sims.argmax(axis=1) diff --git a/pyproject.toml b/pyproject.toml index 961cb3d..6485f12 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,8 +15,8 @@ dependencies = [ "pdfplumber>=0.11", # PDF с текстовым слоем "pymupdf>=1.24", # рендер страниц для Vision "rapidfuzz>=3.10", # нечёткое сопоставление - "sentence-transformers>=3.3", # многоязычные эмбеддинги - "google-generativeai>=0.8", # Gemini Vision + "numpy>=2.0", # векторные операции при сопоставлении + "google-genai>=1.0", # Gemini Vision + эмбеддинги (API, без локального torch) ] [tool.ruff] diff --git a/scripts/proof/normalize_all.py b/scripts/proof/normalize_all.py deleted file mode 100644 index 639aba5..0000000 --- a/scripts/proof/normalize_all.py +++ /dev/null @@ -1,59 +0,0 @@ -"""Замер автонормализации: извлечь весь архив, сопоставить со справочником, дать %.""" -import sys -from collections import Counter -from pathlib import Path - -REPO = Path(__file__).resolve().parents[2] -sys.path.insert(0, str(REPO)) - -from sentence_transformers import SentenceTransformer - -from etl.dictionary import load_dictionary -from etl.extractors import extract -from etl.normalize import Matcher - -# --- собрать все извлечённые позиции --- -rows: list[tuple[str, str | None]] = [] -for path in sorted((REPO / "data/raw").glob("*")): - for row in extract(str(path)).rows: - rows.append((row.service_name_raw, row.service_code_source)) -print(f"позиций всего: {len(rows)}") - -# дедуп по (имя, код) — одинаковые строки сопоставляем один раз -occurrences: Counter[tuple[str, str | None]] = Counter(rows) -keys = list(occurrences) -print(f"уникальных (имя, код): {len(keys)}") - -services = load_dictionary(REPO / "data/reference/dictionary.xlsx") -print("загружаю модель эмбеддингов (первый раз — скачивание)…") -embedder = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2") -matcher = Matcher(services, embedder=embedder) - -print("сопоставляю…") -results = matcher.match_batch([k[0] for k in keys], [k[1] for k in keys]) - -by_method: Counter[str | None] = Counter() -total = matched = 0 -samples: dict[str | None, list] = {} -canon_by_id = {s.service_id: s.name_ru for s in services} -for key, res in zip(keys, results, strict=True): - count = occurrences[key] - total += count - by_method[res.method] += count - if res.service_id: - matched += count - bucket = samples.setdefault(res.method, []) - if len(bucket) < 5: - bucket.append((key[0], canon_by_id.get(res.service_id, "—"), round(res.confidence, 2))) - -print(f"\nАВТОНОРМАЛИЗАЦИЯ: {matched}/{total} = {100 * matched / total:.0f}% (цель ТЗ ≥70%)") -for method in ("code", "exact", "embedding", "fuzzy", None): - print(f" {method or 'unmatched':10}: {by_method[method]} строк") - -print("\nпримеры сопоставлений:") -for method in ("code", "exact", "embedding", "fuzzy"): - for raw, canon, score in samples.get(method, [])[:3]: - print(f" [{method:9}] «{raw[:36]}» → «{canon[:36]}» ({score})") -print("\nпримеры unmatched:") -for raw, _canon, _score in samples.get(None, [])[:6]: - print(f" «{raw[:52]}»") diff --git a/scripts/proof/normalize_tune.py b/scripts/proof/normalize_tune.py new file mode 100644 index 0000000..d294cbc --- /dev/null +++ b/scripts/proof/normalize_tune.py @@ -0,0 +1,79 @@ +"""Подбор порога нормализации: выборка позиций → каскад → покрытие при порогах.""" +import random +import sys +from pathlib import Path + +import numpy as np + +REPO = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO)) + +from rapidfuzz import fuzz, process + +from etl.dictionary import load_dictionary, normalize_name +from etl.extractors import extract +from etl.normalize.embedding import GeminiEmbedder +from etl.normalize.matcher import _CODE_CORE_RE + +rows = [] +for path in sorted((REPO / "data/raw").glob("*")): + for row in extract(str(path)).rows: + rows.append((row.service_name_raw, row.service_code_source)) +random.seed(0) +sample = random.sample(rows, min(2000, len(rows))) +print(f"всего {len(rows)}, выборка {len(sample)}") + +services = load_dictionary(REPO / "data/reference/dictionary.xlsx") +by_code = {s.tarificator_code: s for s in services if s.tarificator_code} +by_norm: dict[str, object] = {} +for s in services: + by_norm.setdefault(s.name_norm, s) +names_norm = [s.name_norm for s in services] + +embedder = GeminiEmbedder() +print("эмбеддинг справочника (1281)…") +dict_emb = embedder.encode([s.name_ru for s in services], task_type="RETRIEVAL_DOCUMENT") + +code_n = exact_n = 0 +pending: list[tuple[str, str]] = [] +for name, code in sample: + m = _CODE_CORE_RE.search(code) if code else None + if m and m.group(0) in by_code: + code_n += 1 + continue + norm = normalize_name(name) + if norm in by_norm: + exact_n += 1 + continue + pending.append((name, norm)) + +print(f"код: {code_n} | точное: {exact_n} | дальше эмбеддинги/fuzzy: {len(pending)}") +print("эмбеддинг запросов…") +query_emb = embedder.encode([p[0] for p in pending], task_type="RETRIEVAL_QUERY") +sims = query_emb @ dict_emb.T +top_idx = sims.argmax(axis=1) +top_score = sims.max(axis=1) +fuzzy_score = np.array( + [process.extractOne(nn, names_norm, scorer=fuzz.token_set_ratio)[1] for _, nn in pending] +) + +base = code_n + exact_n +print(f"\nБаза (код+точное): {base}/{len(sample)} = {100 * base / len(sample):.0f}%") +for thr in (0.58, 0.60, 0.62, 0.64, 0.66, 0.70): + emb_matched = int((top_score >= thr).sum()) + fuzzy_matched = int(((fuzzy_score >= 88) & (top_score < thr)).sum()) + coverage = base + emb_matched + fuzzy_matched + print(f" порог {thr}: эмб {emb_matched} + fuzzy {fuzzy_matched} → покрытие {100 * coverage / len(sample):.0f}%") + +print("\nпримеры эмбеддинг-совпадений (порог 0.62):") +shown = 0 +for k, (name, _) in enumerate(pending): + if top_score[k] >= 0.62 and shown < 7: + print(f" «{name[:34]}» → «{services[int(top_idx[k])].name_ru[:34]}» ({top_score[k]:.2f})") + shown += 1 +print("примеры unmatched (top<0.62 и fuzzy<88):") +shown = 0 +for k, (name, _) in enumerate(pending): + if top_score[k] < 0.62 and fuzzy_score[k] < 88 and shown < 6: + print(f" «{name[:48]}» (лучший {services[int(top_idx[k])].name_ru[:22]} {top_score[k]:.2f})") + shown += 1