""" 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(" ", "")