khmer_segmentation.py

"""
Khmer word-segmentation pipeline that must run before SentencePiece training and
before every future encode/decode call. See khmer_segmentation_sentencepiece_approach.md
for the root-cause analysis this implements.

The tokenizer is the PAIR (this module, khmer_sp.model) — not the .model file alone.
"""
import re
import json
from typing import List, Optional, Tuple

from khmernltk import word_tokenize
import wordninja

# Khmer Unicode block: U+1780-U+17FF
NON_KHMER_BLOCK = re.compile(r"[^ក-៿\s]+")

# Distinct sentinel pairs per masker, built with chr() rather than embedding raw private-use
# characters in source (which some tool layers silently strip to nothing). If both maskers
# shared one sentinel pair, restore() for one type would wrongly try to resolve the other
# type's tokens whenever both masks are active on the same text (IndexError).
GAZ_OPEN, GAZ_CLOSE = chr(0xE000), chr(0xE001)
NK_OPEN, NK_CLOSE = chr(0xE002), chr(0xE003)
LAT_OPEN, LAT_CLOSE = chr(0xE004), chr(0xE005)
_GAZ_PATTERN = re.compile(GAZ_OPEN + r"(\d+)" + GAZ_CLOSE)
_NK_PATTERN = re.compile(NK_OPEN + r"(\d+)" + NK_CLOSE)
_LAT_SENTINEL_PATTERN = re.compile(f"{LAT_OPEN}(\\d+){LAT_CLOSE}")
_ALNUM_ONLY = re.compile(r"[A-Za-z0-9]+")


class GazetteerMasker:
    """Mask loanwords/acronyms before segmentation so khmer-nltk doesn't shatter them."""

    MIN_ENTRY_LENGTH = 4

    def __init__(self, entries: List[str]):
        bad = [e for e in entries if len(e) < self.MIN_ENTRY_LENGTH]
        if bad:
            raise ValueError(f"Gazetteer entries below min length {self.MIN_ENTRY_LENGTH}: {bad}")
        # Longest-first: a short entry must never win a match inside a longer one.
        self.entries = sorted(set(entries), key=len, reverse=True)
        self.pattern = re.compile("|".join(re.escape(e) for e in self.entries)) if self.entries else None

    @classmethod
    def from_file(cls, gazetteer_path: str) -> "GazetteerMasker":
        with open(gazetteer_path, encoding="utf-8") as f:
            return cls(json.load(f))

    def mask(self, text: str) -> Tuple[str, List[str]]:
        if not self.pattern:
            return text, []
        spans: List[str] = []

        def repl(m):
            spans.append(m.group(0))
            return f"{GAZ_OPEN}{len(spans) - 1}{GAZ_CLOSE}"

        return self.pattern.sub(repl, text), spans

    def restore(self, tokens: List[str], spans: List[str]) -> List[str]:
        # .sub(), not .match(): khmer-nltk can fuse a masked span directly onto adjacent
        # Khmer text with no space (e.g. "ការ" + mask -> one token "ការ3"), and
        # .match() only fires when the sentinel is the ENTIRE token — otherwise it silently
        # leaves the raw sentinel + index number in the output instead of the real span.
        return [_GAZ_PATTERN.sub(lambda m: spans[int(m.group(1))], t) for t in tokens]


class NonKhmerMasker:
    """Mask non-Khmer script (English, digits, punctuation) before segmentation."""

    def mask(self, text: str) -> Tuple[str, List[str]]:
        spans: List[str] = []

        def repl(m):
            spans.append(m.group(0))
            return f"{NK_OPEN}{len(spans) - 1}{NK_CLOSE}"

        return NON_KHMER_BLOCK.sub(repl, text), spans

    def restore(self, tokens: List[str], spans: List[str]) -> List[str]:
        # See GazetteerMasker.restore() — same fused-token bug, same .sub() fix.
        return [_NK_PATTERN.sub(lambda m: spans[int(m.group(1))], t) for t in tokens]


class LatinSplitter:
    """
    Decompose glued English/Latin runs (e.g. `tryAImodel` -> `try AI model`) using
    frequency-based `wordninja` segmentation. This is a precision improvement for
    code-switched text, NOT a fix for the Khmer fusion bug — ordinary English already
    has spaces almost everywhere, so it was never causing the systemic ln(V) plateau.

    Three safeguards were found necessary by testing (not assumed):
    - `wordninja` alone mangles domain terms outside its generic English corpus, e.g.
      ACLEDA -> AC LED A, COVID-19 -> C OVID 19, Bakong -> Ba kong, 5G -> 5 G. Entries in
      `exceptions` are protected (masked out) before wordninja ever sees that substring.
    - Matching exceptions case-INsensitively creates its own collisions: a 3-char entry
      like ABA (bank name) would match the substring "aba" inside the ordinary word
      "database", corrupting it to "dat aba se". Matching case-sensitively avoids this
      since real acronyms/brand names are written in a distinctive case.
    - `wordninja` silently DROPS characters outside [A-Za-z0-9] instead of erroring —
      confirmed failure: "$5.99USD" -> "5 99 USD" (the "$" and the decimal point both
      vanish). Only spans that are purely alphanumeric are passed to wordninja; anything
      containing punctuation/currency symbols/hyphens is left untouched to avoid data loss.
    """

    def __init__(self, exceptions: List[str]):
        # Longest-first, same rule as GazetteerMasker: a short entry must never win a
        # match inside a longer one.
        self.exceptions = sorted(set(exceptions), key=len, reverse=True)
        self.pattern = re.compile("|".join(re.escape(e) for e in self.exceptions)) if self.exceptions else None

    @classmethod
    def from_file(cls, exceptions_path: str) -> "LatinSplitter":
        with open(exceptions_path, encoding="utf-8") as f:
            return cls(json.load(f))

    def _maybe_split(self, part: str) -> List[str]:
        if not _ALNUM_ONLY.fullmatch(part):
            return [part]
        words = wordninja.split(part)
        return words if words else [part]

    def split(self, span: str) -> str:
        """One non-Khmer span -> space-joined decomposition (protected terms kept whole)."""
        if not self.pattern:
            return " ".join(self._maybe_split(span))

        protected: List[str] = []

        def repl(m):
            protected.append(m.group(0))
            return f"{LAT_OPEN}{len(protected) - 1}{LAT_CLOSE}"

        masked = self.pattern.sub(repl, span)

        out: List[str] = []
        for part in re.split(f"({LAT_OPEN}\\d+{LAT_CLOSE})", masked):
            if not part:
                continue
            m = _LAT_SENTINEL_PATTERN.fullmatch(part)
            if m:
                out.append(protected[int(m.group(1))])
            else:
                out.extend(self._maybe_split(part))
        return " ".join(out)


def segment_line(
    line: str,
    gazetteer: GazetteerMasker,
    non_khmer: NonKhmerMasker,
    latin_splitter: Optional[LatinSplitter] = None,
) -> str:
    """Raw text -> space-joined, word-segmented text (ready for SentencePiece)."""
    masked, nk_spans = non_khmer.mask(line)
    if latin_splitter is not None:
        nk_spans = [latin_splitter.split(s) for s in nk_spans]
    masked, gz_spans = gazetteer.mask(masked)
    tokens = word_tokenize(masked)
    # khmer-nltk emits literal whitespace as its own tokens; drop them here since we
    # re-insert exactly one space between every remaining token below. Keeping them would
    # double/triple spacing at every point the original text already had a space.
    tokens = [t for t in tokens if t.strip() != ""]
    tokens = gazetteer.restore(tokens, gz_spans)
    tokens = non_khmer.restore(tokens, nk_spans)
    return " ".join(tokens)


class KhmerTokenizer:
    """
    Segmentation pipeline + SentencePiece model, wrapped together.
    Use this (not a bare SentencePieceProcessor) for every future encode/decode —
    the segmentation step must never drift out of sync between training and inference.
    """

    def __init__(self, sp_model_path: str, gazetteer_path: str, latin_exceptions_path: Optional[str] = None):
        import sentencepiece as spm

        self.sp = spm.SentencePieceProcessor(model_file=sp_model_path)
        self.gazetteer = GazetteerMasker.from_file(gazetteer_path)
        self.non_khmer = NonKhmerMasker()
        self.latin_splitter = LatinSplitter.from_file(latin_exceptions_path) if latin_exceptions_path else None

    def preprocess(self, text: str) -> str:
        return segment_line(text, self.gazetteer, self.non_khmer, self.latin_splitter)

    def encode(self, text: str, out_type=int):
        return self.sp.encode(self.preprocess(text), out_type=out_type)

    def decode(self, ids) -> str:
        # Strip the artificial word-boundary spaces introduced for training;
        # natural Khmer orthography does not space every word.
        return self.sp.decode(ids).replace(" ", "")