medtech-hackathon/etl/normalize/matcher.py

"""Сопоставление извлечённых позиций с эталонным справочником услуг.

Каскад по убыванию надёжности:
  1. код тарификатора — точное совпадение, если клиника его указала;
  2. точное совпадение нормализованного названия;
  3. эмбеддинги (многоязычная модель) — семантическая близость названий;
  4. нечёткое сравнение (RapidFuzz) как запасной сигнал;
  5. иначе — очередь ручной разметки (unmatched).

В справочнике нет синонимов, поэтому шаги 1 и 3 несут основную нагрузку. Коды у
клиник встречаются с лишним хвостом («A02.020.000.2»), поэтому сравниваем по
канонической части кода тарификатора.
"""

from __future__ import annotations

import re
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from rapidfuzz import fuzz, process  # noqa: E402

from contracts.models import MatchResult  # noqa: E402
from etl.dictionary import Service, normalize_name  # noqa: E402

# Каноническая часть кода тарификатора: буква + три группы цифр.
_CODE_CORE_RE = re.compile(r"[A-ZА-Я]\d{2}\.\d{3}\.\d{3}")


class Matcher:
    """Готовит индексы справочника и сопоставляет с ним сырые названия услуг."""

    def __init__(
        self,
        services: list[Service],
        embedder=None,
        emb_threshold: float = 0.70,
        fuzzy_threshold: int = 88,
        dict_emb=None,
    ):
        self.services = services
        self.by_code = {s.tarificator_code: s for s in services if s.tarificator_code}
        self.by_norm: dict[str, Service] = {}
        for s in services:
            self.by_norm.setdefault(s.name_norm, s)
        self.names_norm = [s.name_norm for s in services]
        self.emb_threshold = emb_threshold
        self.fuzzy_threshold = fuzzy_threshold
        self.embedder = embedder
        # Готовые эмбеддинги справочника (dict_emb) переиспользуются между прогонами:
        # при догрузке нового прайса не нужно заново векторизовать тысячи услуг.
        self.dict_emb = dict_emb
        if embedder is not None and self.dict_emb is None:
            self.dict_emb = embedder.encode(
                [s.name_ru for s in services],
                normalize_embeddings=True,
                task_type="RETRIEVAL_DOCUMENT",
            )

    def _match_by_code(self, code: str | None) -> Service | None:
        if not code:
            return None
        m = _CODE_CORE_RE.search(code)
        return self.by_code.get(m.group(0)) if m else None

    def _match_by_fuzzy(self, name_norm: str) -> MatchResult:
        best = process.extractOne(name_norm, self.names_norm, scorer=fuzz.token_set_ratio)
        if best and best[1] >= self.fuzzy_threshold:
            return MatchResult(
                service_id=self.services[best[2]].service_id,
                method="fuzzy",
                confidence=best[1] / 100,
            )
        return MatchResult()

    def match_batch(self, names: list[str], codes: list[str | None]) -> list[MatchResult]:
        """Сопоставить пачку позиций. Эмбеддинги считаются разом — так быстрее."""
        results: list[MatchResult | None] = [None] * len(names)
        pending_idx: list[int] = []
        pending_norm: list[str] = []

        for i, (name, code) in enumerate(zip(names, codes, strict=True)):
            service = self._match_by_code(code)
            if service:
                results[i] = MatchResult(
                    service_id=service.service_id, method="code", confidence=1.0
                )
                continue
            name_norm = normalize_name(name)
            exact = self.by_norm.get(name_norm)
            if exact:
                results[i] = MatchResult(
                    service_id=exact.service_id, method="exact", confidence=1.0
                )
                continue
            pending_idx.append(i)
            pending_norm.append(name_norm)

        if pending_norm and self.embedder is not None and self.dict_emb is not None:
            # Дедуп: одинаковые названия эмбеддим один раз (экономит запросы и квоту API).
            pending_names = [names[i] for i in pending_idx]
            unique: dict[str, int] = {}
            for nm in pending_names:
                unique.setdefault(nm, len(unique))
            query = self.embedder.encode(
                list(unique), normalize_embeddings=True, task_type="RETRIEVAL_QUERY"
            )
            sims = query @ self.dict_emb.T  # косинус по нормированным векторам
            best_idx = sims.argmax(axis=1)
            best_score = sims.max(axis=1)
            for k, i in enumerate(pending_idx):
                u = unique[pending_names[k]]
                suggested = self.services[int(best_idx[u])].service_id
                score = float(best_score[u])
                if score >= self.emb_threshold:
                    results[i] = MatchResult(
                        service_id=suggested,
                        method="embedding",
                        confidence=score,
                        suggested_service_id=suggested,
                        suggested_score=score,
                    )
                else:
                    # Не дотянул до порога — в очередь, но кандидата сохраняем оператору.
                    fallback = self._match_by_fuzzy(pending_norm[k])
                    fallback.suggested_service_id = suggested
                    fallback.suggested_score = score
                    results[i] = fallback
        else:
            for k, i in enumerate(pending_idx):
                results[i] = self._match_by_fuzzy(pending_norm[k])

        return [r if r is not None else MatchResult() for r in results]