Trixy7/trixy_core/stt/layers/kenlm_layer.py

# -*- coding: utf-8 -*-
"""
KenLM-Korrekturschicht.

Nutzt N-Gram Sprachmodell fuer kontextuelle Kandidatenauswahl.
Sehr schnell (<1ms/Query).
"""

from pathlib import Path

from trixy_core.stt.layers.base import CorrectionLayer
from trixy_core.utils.debug import pdebug, pwarn

try:
    import kenlm
    _HAS_KENLM = True
except ImportError:
    _HAS_KENLM = False


class KenLMLayer(CorrectionLayer):
    """KenLM N-Gram basierte kontextuelle Korrektur."""

    NAME = "kenlm"

    def __init__(self) -> None:
        self._model: "kenlm.Model | None" = None
        self._protected_words: set[str] = set()

    def is_available(self) -> bool:
        """Prueft ob kenlm installiert ist."""
        return _HAS_KENLM

    def initialize(self, config: dict, language: str, protected_words: list[str]) -> bool:
        """Laedt N-Gram Modell aus konfiguriertem Pfad."""
        if not _HAS_KENLM:
            return False

        model_path = config.get("kenlm_model_path", "")
        if not model_path or not Path(model_path).exists():
            pdebug(f"[KenLM] Modell nicht gefunden: {model_path}")
            return False

        try:
            self._model = kenlm.Model(str(model_path))
            self._protected_words = {w.lower() for w in protected_words}
            pdebug(f"[KenLM] Modell geladen: {model_path} (order={self._model.order})")
            return True

        except Exception as e:
            pwarn(f"[KenLM] Modell-Laden fehlgeschlagen: {e}")
            self._model = None
            return False

    def correct(self, text: str) -> str:
        """Bewertet Text und waehlt beste Kandidaten per N-Gram Score."""
        if not self._model or not text.strip():
            return text

        try:
            words = text.lower().split()
            if len(words) < 2:
                return text

            # Score des Original-Satzes
            original_score = self._model.score(text.lower(), bos=True, eos=True)

            # Kandidaten generieren: pro Wort Edit-Distanz-1 Varianten
            best_text = text.lower()
            best_score = original_score

            for i, word in enumerate(words):
                if word in self._protected_words:
                    continue

                candidates = self._generate_candidates(word)
                for candidate in candidates:
                    test_words = words.copy()
                    test_words[i] = candidate
                    test_text = " ".join(test_words)
                    score = self._model.score(test_text, bos=True, eos=True)

                    if score > best_score:
                        best_score = score
                        best_text = test_text

            if best_text != text.lower():
                pdebug(f"[KenLM] '{text}' → '{best_text}' (score: {original_score:.2f} → {best_score:.2f})")
                return best_text

            return text

        except Exception as e:
            pdebug(f"[KenLM] Fehler bei Korrektur: {e}")
            return text

    def _generate_candidates(self, word: str) -> list[str]:
        """Generiert Edit-Distanz-1 Kandidaten fuer ein Wort."""
        candidates = set()
        letters = "abcdefghijklmnopqrstuvwxyzäöüß"

        # Loeschungen
        for i in range(len(word)):
            candidates.add(word[:i] + word[i + 1:])

        # Transpositionen
        for i in range(len(word) - 1):
            candidates.add(word[:i] + word[i + 1] + word[i] + word[i + 2:])

        # Ersetzungen
        for i in range(len(word)):
            for c in letters:
                if c != word[i]:
                    candidates.add(word[:i] + c + word[i + 1:])

        # Einfuegungen
        for i in range(len(word) + 1):
            for c in letters:
                candidates.add(word[:i] + c + word[i:])

        # Leere und zu kurze entfernen
        candidates.discard("")
        candidates.discard(word)

        return list(candidates)

    def shutdown(self) -> None:
        """Gibt Ressourcen frei."""
        self._model = None