src/postprocess.py

"""Постобработка распознанного текста словарём якутского языка.

Прямое исправление текста спеллчекером вредит: на 5400 реальных строках оно
подняло долю ошибок с 0.24% до 0.89%, испортив 838 строк против 44 улучшенных.
Спеллчекер не видит изображение и «чинит» имена, русские вставки и обрывки
слов на переносах.

Поэтому словарь здесь только предлагает кандидатов, а решение принимает сама
сеть: замена принимается, если выигрыш языковой модели перевешивает потерю
по CTC. При весе 1.0 такая постобработка не испортила ни одной строки ни на
одной из проверочных выборок, снимая при этом 5–8% ошибок там, где
распознаватель слаб.

Пакет sakhaspell необязателен: без него обработка просто не применяется.
"""
import os

ALPHA = float(os.environ.get("SAKHA_LM_ALPHA", "1.0"))

_state = {"ready": False, "rescore": None}


def _load():
    if _state["ready"]:
        return _state["rescore"]
    _state["ready"] = True
    try:
        import sakhaspell  # noqa: F401
    except ImportError:
        return None
    try:
        from lm_rescore import Rescorer
    except ImportError:
        return None
    _state["rescore"] = Rescorer
    return Rescorer


class Corrector:
    """Ленивая обёртка: словарь и языковая модель грузятся при первом вызове."""

    def __init__(self, charset, alpha=ALPHA):
        self.charset = charset
        self.alpha = alpha
        self._r = None
        self._tried = False

    def available(self):
        if not self._tried:
            self._tried = True
            cls = _load()
            if cls is not None:
                try:
                    self._r = cls(self.charset, alpha=self.alpha)
                except Exception:
                    self._r = None
        return self._r is not None

    def __call__(self, texts, logprobs):
        """Строки идут в порядке чтения, поэтому перенос виден по предыдущей."""
        if not self.available():
            return texts
        out, prev_hyphen = [], False
        for t, lp in zip(texts, logprobs):
            if not t or lp is None:
                out.append(t)
                prev_hyphen = bool(t) and t.rstrip().endswith(("-", "\u2010", "\u2011"))
                continue
            try:
                out.append(self._r(lp, t, prev_hyphen=prev_hyphen))
            except Exception:
                out.append(t)
            prev_hyphen = t.rstrip().endswith(("-", "\u2010", "\u2011"))
        return out