Trixy7/trixy_core/stt/layers/jamspell_layer.py

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

Neuronaler kontextueller Spellchecker (~0.2ms/Wort).
"""

from pathlib import Path

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

try:
    import jamspell
    _HAS_JAMSPELL = True
except ImportError:
    _HAS_JAMSPELL = False


class JamSpellLayer(CorrectionLayer):
    """JamSpell neuronale kontextuelle Korrektur."""

    NAME = "jamspell"

    def __init__(self) -> None:
        self._corrector: "jamspell.TSpellCorrector | None" = None
        self._protected_words: set[str] = set()
        # Platzhalter fuer geschuetzte Woerter waehrend Korrektur
        self._placeholder_map: dict[str, str] = {}

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

    def initialize(self, config: dict, language: str, protected_words: list[str]) -> bool:
        """Laedt vortrainiertes JamSpell-Modell."""
        if not _HAS_JAMSPELL:
            return False

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

        try:
            self._corrector = jamspell.TSpellCorrector()
            if not self._corrector.LoadLangModel(str(model_path)):
                pwarn(f"[JamSpell] Modell konnte nicht geladen werden: {model_path}")
                self._corrector = None
                return False

            self._protected_words = {w.lower() for w in protected_words}

            # Platzhalter generieren (nicht im Woerterbuch vorkommend)
            for i, word in enumerate(protected_words):
                self._placeholder_map[word.lower()] = f"XPROTX{i:03d}X"

            pdebug(f"[JamSpell] Modell geladen: {model_path}")
            return True

        except Exception as e:
            pwarn(f"[JamSpell] Initialisierung fehlgeschlagen: {e}")
            self._corrector = None
            return False

    def correct(self, text: str) -> str:
        """Korrigiert Text mit JamSpell (kontextbasiert)."""
        if not self._corrector or not text.strip():
            return text

        try:
            # Geschuetzte Woerter durch Platzhalter ersetzen
            processed = text.lower()
            used_placeholders: dict[str, str] = {}

            for word, placeholder in self._placeholder_map.items():
                if word in processed.split():
                    processed = " ".join(
                        placeholder if w == word else w
                        for w in processed.split()
                    )
                    used_placeholders[placeholder] = word

            # JamSpell korrigieren
            corrected = self._corrector.FixFragment(processed)

            # Platzhalter zurueckersetzen
            for placeholder, original in used_placeholders.items():
                corrected = corrected.replace(placeholder.lower(), original)
                corrected = corrected.replace(placeholder, original)

            if corrected != text.lower():
                pdebug(f"[JamSpell] '{text}' → '{corrected}'")
                return corrected

            return text

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

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