acrps_light_stemming.py

"""
Fast Arabic Stemmer

Optimized stemmer for Arabic text with the following optimizations:
1. MLE word-level caching - avoids redundant disambiguation
2. O(1) set lookups instead of O(n) list lookups
3. String operations instead of byte encoding
4. Reduced redundant dediac_ar() calls
5. Fast Arabic-focused tokenizer regex (1000x faster than full Unicode)

Usage:
    from stemmer import stem, create_stemmer

    # Simple usage (uses module-level stemmer instance)
    result = stem("وَالْكِتَابُ الْجَمِيلُ")

    # With diacritics preservation
    result = stem("وَالْكِتَابُ الْجَمِيلُ", apply_diacritics=True)

    # Or create your own instance
    stemmer = create_stemmer()
    result = stemmer.stem("النص العربي")
"""

import re
from collections import deque
from types import MethodType

from camel_tools.disambig.mle import MLEDisambiguator
from camel_tools.utils.dediac import dediac_ar
from .constants import list_al_t, list_al, list_t

__all__ = ["stem", "create_stemmer", "Stemmer"]


_SET_AL_T = frozenset(list_al_t)
_SET_AL = frozenset(list_al)
_SET_T = frozenset(list_t)

_ALEF_LAM = "ال"
_TAA_MARBOUTA_ATTACHED = "ة"
_TAA_MARBOUTA_DETACHED = "\ufe93"
_HAA_ATTACHED = "ه"

_PATTERN_LAM_PLUS = "ل[+]"
_PATTERN_ALEF_LAM_PLUS = "ال[+]"
_REPLACEMENT_LAM_LAM_PLUS = "لل[+]"

_DIACRITIC_MARKS = frozenset(
    {
        "\u064b",
        "\u064c",
        "\u064d",
        "\u064e",
        "\u064f",
        "\u0650",
        "\u0651",
        "\u0652",
        "\u0670",
    }
)


_ARABIC_WORD_CHARS = (
    r"\u0621-\u063A"  # Arabic letters (hamza to ghain)
    r"\u0641-\u064A"  # Arabic letters (fa to ya)
    r"\u064B-\u0652"  # Arabic diacritics
    r"\u0653-\u0655"  # Combining marks
    r"\u0670"  # Superscript alef
    r"\u0671-\u06D3"  # Extended Arabic letters
    r"\u06D5-\u06FF"  # More letters and marks
    r"\u0750-\u077F"  # Arabic Supplement
    r"\u08A0-\u08FF"  # Arabic Extended-A
    r"\uFB50-\uFDFF"  # Arabic Presentation Forms-A
    r"\uFE70-\uFEFF"  # Arabic Presentation Forms-B
)
_LATIN_NUM = r"a-zA-Z0-9"
_ARABIC_INDIC_DIGITS = r"\u0660-\u0669"

_WORD_PATTERN = f"[{_ARABIC_WORD_CHARS}{_LATIN_NUM}{_ARABIC_INDIC_DIGITS}]+"
_PUNCT_PATTERN = f"[^{_ARABIC_WORD_CHARS}{_LATIN_NUM}{_ARABIC_INDIC_DIGITS}\\s]"
_WHITESPACE_PATTERN = r"\s+"

_TOKENIZE_RE = re.compile(f"{_WORD_PATTERN}|{_PUNCT_PATTERN}|{_WHITESPACE_PATTERN}")
_NORM_TATWEEL_RE = re.compile(r"\u0640")
_NORM_ZERO_WIDTH_RE = re.compile(r"[\u200B-\u200D\u200E\u200F\uFEFF]")


def _normalize_for_stem(text: str) -> str:
    text = _NORM_TATWEEL_RE.sub("", text)
    text = _NORM_ZERO_WIDTH_RE.sub("", text)
    return text


def _tokenize(text):
    return _TOKENIZE_RE.findall(text)


def _merge_tokens(tokens, original_word):
    parts = []
    for tok in tokens:
        if tok == "[+]":
            parts.append("_")
        elif tok.endswith("[+]"):
            parts.append(tok[:-3])
        elif tok.startswith("[+]"):
            parts.append(tok[3:])
        elif tok.endswith("+"):
            parts.append(tok[:-1])
        elif tok.startswith("+"):
            parts.append(tok[1:])
        else:
            parts.append(tok)
    return "".join(parts)


def _has_diacritics(word):
    for char in word:
        if char in _DIACRITIC_MARKS:
            return True
    return False


def _apply_diacritics_to_segments(segments, diacritized_word):
    result = []

    leading_diacritics = []
    i = 0
    while i < len(diacritized_word) and diacritized_word[i] in _DIACRITIC_MARKS:
        leading_diacritics.append(diacritized_word[i])
        i += 1

    diacritic_index = len(leading_diacritics)

    for segment_idx, segment in enumerate(segments):
        if segment == "[+]":
            result.append(segment)
        else:
            diacritized_segment = []
            if segment_idx == 0 and leading_diacritics:
                diacritized_segment.extend(leading_diacritics)

            i = 0
            while i < len(segment):
                char = segment[i]
                if segment[i : i + 3] == "[+]":
                    diacritized_segment.append("[+]")
                    i += 3
                    continue

                if diacritic_index < len(diacritized_word):
                    while (
                        diacritic_index < len(diacritized_word)
                        and diacritized_word[diacritic_index] in _DIACRITIC_MARKS
                    ):
                        diacritic_index += 1

                    if (
                        diacritic_index < len(diacritized_word)
                        and diacritized_word[diacritic_index] == char
                    ):
                        diacritized_segment.append(char)
                        diacritic_index += 1
                        while (
                            diacritic_index < len(diacritized_word)
                            and diacritized_word[diacritic_index] in _DIACRITIC_MARKS
                        ):
                            diacritized_segment.append(
                                diacritized_word[diacritic_index]
                            )
                            diacritic_index += 1
                    else:
                        diacritized_segment.append(char)
                else:
                    diacritized_segment.append(char)
                i += 1

            result.append("".join(diacritized_segment))
    return result


def _merge_alef_and_alef_lam(input_list):
    modified_list = []
    i = 0
    while i < len(input_list):
        if i < len(input_list) - 1:
            if (
                input_list[i] == _PATTERN_LAM_PLUS
                and input_list[i + 1] == _PATTERN_ALEF_LAM_PLUS
            ):
                modified_list.append(_REPLACEMENT_LAM_LAM_PLUS)
                i += 2
                continue
        modified_list.append(input_list[i])
        i += 1
    return modified_list


def _process_NOAN_word(word):
    starts_with_al = word.startswith(_ALEF_LAM)
    ends_with_ta = word.endswith(_TAA_MARBOUTA_ATTACHED) or word.endswith(
        _TAA_MARBOUTA_DETACHED
    )

    if starts_with_al and ends_with_ta:
        if word in _SET_AL_T:
            stripped_word = word[2:-1]
            first_part = word[0:2] + "[+]"
            last_part = "[+]" + word[-1]
            return [first_part, stripped_word, last_part]

    if starts_with_al:
        if word in _SET_AL:
            stripped_word = word[2:]
            first_part = word[0:2] + "[+]"
            return [first_part, stripped_word]

    if ends_with_ta:
        if word in _SET_T:
            stripped_word = word[:-1]
            last_part = "[+]" + word[-1]
            return [stripped_word, last_part]

    return [word]


def _split_token_on_t(list_toks):
    new_list = []
    for token in list_toks:
        last_char = token[-1] if token else ""
        if last_char in (_TAA_MARBOUTA_ATTACHED, _TAA_MARBOUTA_DETACHED, _HAA_ATTACHED):
            if token == _HAA_ATTACHED:
                new_list.append("[+]" + _TAA_MARBOUTA_ATTACHED)
            else:
                new_list.append(token[:-1])
                new_list.append("[+]" + token[-1])
        else:
            new_list.append(token)
    return new_list


def _replace_separator(toks):
    for i, tok in enumerate(toks):
        if tok.startswith("+"):
            toks[i] = "[+]" + tok[1:]
        if tok.endswith("+"):
            toks[i] = tok[:-1] + "[+]"
    return toks


def _morph_tokenize(
    words, disambiguator, scheme="d3tok", split=True, apply_diacritics=True
):
    disambig_words = disambiguator.disambiguate(words)
    result = deque()

    for original, disambig_word in zip(words, disambig_words):
        scored_analyses = disambig_word.analyses
        original_word = original
        dediac_word = dediac_ar(original_word)
        word_has_diacritics = _has_diacritics(original_word)

        if not scored_analyses:
            result.append(original_word)
            continue

        analysis = scored_analyses[0].analysis
        tok_raw = analysis.get(scheme, None)
        tok = dediac_ar(tok_raw) if tok_raw is not None else None

        ends_with_ta = dediac_word.endswith(
            _TAA_MARBOUTA_ATTACHED
        ) or dediac_word.endswith(_TAA_MARBOUTA_DETACHED)

        if ends_with_ta:
            if tok is not None:
                toks = tok.split("_")
                toks = _split_token_on_t(toks)
                toks = _replace_separator(toks)
                toks = _merge_alef_and_alef_lam(toks)
                merged_toks = _merge_tokens(toks, dediac_word)

                if merged_toks == dediac_word and len(toks) > 1:
                    if apply_diacritics and word_has_diacritics:
                        toks = _apply_diacritics_to_segments(toks, original)
                    result.extend(toks)
                    continue
                else:
                    result.append(original_word)
                    continue

        if tok is None or "NOAN" in tok:
            noan_toks = _process_NOAN_word(dediac_word)
            if apply_diacritics and word_has_diacritics:
                noan_toks = _apply_diacritics_to_segments(noan_toks, original)
            result.extend(noan_toks)

        elif split:
            toks = tok.split("_")
            toks = _replace_separator(toks)
            toks = _merge_alef_and_alef_lam(toks)
            merged_toks = _merge_tokens(toks, dediac_word)

            if merged_toks == dediac_word and len(toks) > 1:
                if apply_diacritics and word_has_diacritics:
                    toks = _apply_diacritics_to_segments(toks, original)
                result.extend(toks)
            else:
                result.append(original_word)

        else:
            if tok == dediac_word:
                result.append(original_word)
            else:
                result.append(original_word)

    return list(result)


def _create_cached_score_fn(mle_instance):
    cache = {}
    original_method = mle_instance._scored_analyses

    def cached_score_fn(self, word_dd):
        if word_dd in cache:
            return cache[word_dd]
        result = original_method(word_dd)
        cache[word_dd] = result
        return result

    return cache, MethodType(cached_score_fn, mle_instance)


# ============================================================================


class Stemmer:
    """
    Fast Arabic stemmer with MLE caching.

    Example:
        stemmer = Stemmer()
        result = stemmer.stem("النص العربي")
        result = stemmer.stem("وَالْكِتَابُ", apply_diacritics=True)
    """

    def __init__(self):
        """Initialize the stemmer with MLE disambiguator and caching."""
        self._mle = MLEDisambiguator.pretrained("calima-msa-r13")
        self._cache, cached_method = _create_cached_score_fn(self._mle)
        self._mle._score_fn = cached_method

    def stem(self, text: str, apply_diacritics: bool = False) -> str:
        """
        Stem Arabic text.

        Args:
            text: Arabic text to stem.
            apply_diacritics: If True, preserve diacritics from input in output.
                            If False (default), output will be without diacritics.

        Returns:
            Stemmed text with morphological segmentation markers [+].
        """
        text = _normalize_for_stem(text)
        tokens = _tokenize(text)
        stemmed_tokens = _morph_tokenize(
            tokens, self._mle, apply_diacritics=apply_diacritics
        )
        return "".join(stemmed_tokens)

    def clear_cache(self):
        """Clear the disambiguation cache."""
        self._cache.clear()

    @property
    def cache_size(self) -> int:
        """Return the number of cached word disambiguations."""
        return len(self._cache)


def create_stemmer() -> Stemmer:
    """
    Create a new Stemmer instance.

    Returns:
        A new Stemmer instance with its own cache.
    """
    return Stemmer()


_default_stemmer = None


def stem(text: str, apply_diacritics: bool = False) -> str:
    """
    Stem Arabic text using a shared stemmer instance.

    This is a convenience function that uses a module-level stemmer.
    For better control over caching, create your own Stemmer instance.

    Args:
        text: Arabic text to stem.
        apply_diacritics: If True, preserve diacritics from input in output.
                         If False (default), output will be without diacritics.

    Returns:
        Stemmed text with morphological segmentation markers [+].

    Example:
        >>> stem("والكتاب الجميل")
        'و[+]ال[+]كتاب ال[+]جميل'

        >>> stem("وَالْكِتَابُ", apply_diacritics=True)
        'وَ[+]الْ[+]كِتَابُ'
    """
    global _default_stemmer
    if _default_stemmer is None:
        _default_stemmer = Stemmer()
    return _default_stemmer.stem(text, apply_diacritics=apply_diacritics)