Trixy7/trixy_core/stt/layers/symspell_layer.py

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

Nutzt symspellpy fuer Wort-Segmentierung und offensichtliche Nicht-Wort-Fehler.
Sehr schnell (~0.03ms/Wort).
"""

from pathlib import Path

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

try:
    from symspellpy import SymSpell, Verbosity
    _HAS_SYMSPELL = True
except ImportError:
    _HAS_SYMSPELL = False


class SymSpellLayer(CorrectionLayer):
    """SymSpell-basierte Wortkorrektur und Segmentierung."""

    NAME = "symspell"

    def __init__(self) -> None:
        self._sym_spell: "SymSpell | None" = None
        self._max_edit_distance: int = 2
        self._protected_words: set[str] = set()

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

    def initialize(self, config: dict, language: str, protected_words: list[str]) -> bool:
        """Initialisiert SymSpell mit Frequency-Dictionary."""
        if not _HAS_SYMSPELL:
            return False

        self._max_edit_distance = config.get("symspell_max_edit_distance", 2)
        self._protected_words = {w.lower() for w in protected_words}

        try:
            self._sym_spell = SymSpell(
                max_dictionary_edit_distance=self._max_edit_distance,
                prefix_length=7,
            )

            # Dictionary laden — Suchprioriaet:
            # 1. Explizit konfigurierter Pfad
            # 2. Bundled Dictionary in models/symspell/
            # 3. symspellpy-Paket (Fallback, hat nur Englisch)
            dict_path = config.get("symspell_dictionary_path", "")
            dict_loaded = False

            if dict_path and Path(dict_path).exists():
                self._sym_spell.load_dictionary(
                    dict_path, term_index=0, count_index=1,
                )
                pdebug(f"[SymSpell] Dictionary geladen: {dict_path}")
                dict_loaded = True

            if not dict_loaded:
                # Bundled Dictionary in models/symspell/
                bundled = Path("models") / "symspell" / f"frequency_dictionary_{language}.txt"
                if bundled.exists():
                    self._sym_spell.load_dictionary(
                        str(bundled), term_index=0, count_index=1,
                    )
                    pdebug(f"[SymSpell] Bundled Dictionary geladen: {bundled}")
                    dict_loaded = True

            if not dict_loaded:
                # Fallback: symspellpy-Paket
                import pkg_resources
                pkg_base = Path(pkg_resources.resource_filename("symspellpy", ""))
                for pattern in [
                    f"frequency_dictionary_{language}_*.txt",
                    f"frequency_dictionary_{language}*.txt",
                ]:
                    found = list(pkg_base.glob(pattern))
                    if found:
                        self._sym_spell.load_dictionary(
                            str(found[0]), term_index=0, count_index=1,
                        )
                        pdebug(f"[SymSpell] Paket-Dictionary geladen: {found[0]}")
                        dict_loaded = True
                        break

            if not dict_loaded:
                pwarn(f"[SymSpell] Kein {language}-Dictionary gefunden")
                return False

            # Geschuetzte Woerter mit hoher Frequenz hinzufuegen
            # (verhindert Korrektur von z.B. "Eminem" zu "einem")
            for word in protected_words:
                self._sym_spell.create_dictionary_entry(word.lower(), 1_000_000)

            pdebug(f"[SymSpell] Initialisiert (max_edit_distance={self._max_edit_distance})")
            return True

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

    def correct(self, text: str) -> str:
        """Korrigiert Text mit SymSpell lookup_compound und word_segmentation."""
        if not self._sym_spell or not text.strip():
            return text

        try:
            # lookup_compound fuer Multi-Wort-Korrektur
            suggestions = self._sym_spell.lookup_compound(
                text,
                max_edit_distance=self._max_edit_distance,
            )

            if suggestions:
                corrected = suggestions[0].term

                # Geschuetzte Woerter zuruecksetzen falls faelschlich korrigiert
                corrected = self._restore_protected(text, corrected)

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

            return text

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

    def _restore_protected(self, original: str, corrected: str) -> str:
        """Stellt geschuetzte Woerter wieder her falls korrigiert."""
        orig_words = original.lower().split()
        corr_words = corrected.split()

        # Nur wenn gleiche Wortanzahl
        if len(orig_words) != len(corr_words):
            return corrected

        result = []
        for ow, cw in zip(orig_words, corr_words):
            if ow in self._protected_words and ow != cw:
                result.append(ow)
            else:
                result.append(cw)

        return " ".join(result)

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