tokenizer.py

# tokenizer.py
# Enhanced Paradigm-based segmenter with configurable features:
# - Word boundary tokens
# - Null suffixes as tokens
# - Paradigm-specific suffixes and roots

from collections import OrderedDict
from pathlib import Path
from typing import List, Tuple, Optional
import os, json, re
from huggingface_hub import hf_hub_download
from transformers import PreTrainedTokenizerFast

try:
    from .boundary_discovery import anchor_sequences
except ImportError:
    from boundary_discovery import anchor_sequences

def _deserialize_suffixes_from_json(sfx_list):
    out = set()
    for item in sfx_list:
        if isinstance(item, list):
            # JSON nested: [base, nested_list]
            base, nested = item
            out.add((base, frozenset(nested)))
        else:
            out.add(item)  # plain string like "", "ing", "s"
    return out

def _load_paradigms_any(path):
    import json
    with open(path, "r", encoding="utf-8") as f:
        payload = json.load(f)

    # Case A: new schema with top-level dict {"paradigms": [...]}
    if isinstance(payload, dict) and "paradigms" in payload:
        paradigms = []
        for p in payload["paradigms"]:
            stems = set(p["stems"])
            suffixes = _deserialize_suffixes_from_json(p["suffixes"])
            paradigms.append((stems, suffixes))
        meta = payload.get("meta", {})
        return paradigms, meta

    # Case B: older "list of pairs" JSON [[stems, suffixes], ...]
    if isinstance(payload, list) and payload and isinstance(payload[0], list):
        paradigms = []
        for stems, suffixes in payload:
            stems = set(stems)
            # suffixes may be ["", ["er", ["", "s"]], "ing"] or already strings
            norm = _deserialize_suffixes_from_json(suffixes)
            paradigms.append((stems, norm))
        return paradigms, {}

    # Case C: already python-native structure (rare if not using JSON)
    if isinstance(payload, list) and payload and isinstance(payload[0], (list, tuple)) and len(payload[0]) == 2:
        return payload, {}

    raise ValueError("Unrecognized paradigms.json format")

# ----------------------------
# Enhanced Paradigm-based segmenter
# ----------------------------
class EnhancedParadigmFinderSegmenter:
    def __init__(self, paradigms, config):
        self.paradigms = paradigms
        self.config = config
        self.lowercase = config.get("lowercase", True)
        self.space_punct = config.get("space_punct", True)
        self.use_word_boundaries = config.get("use_word_boundaries", False)
        self.null_suffixes_as_tokens = config.get("null_suffixes_as_tokens", False)
        self.paradigm_specific_suffixes = config.get("paradigm_specific_suffixes", False)
        self.paradigm_specific_roots = config.get("paradigm_specific_roots", False)
        self.word_boundary_token = config.get("word_boundary_token", "▁")
        self.null_suffix_token = config.get("null_suffix_token", "ε")
        self.paradigm_token_format = config.get("paradigm_token_format", "{token}_p{paradigm_idx}")
        self.fallback_mode = config.get("fallback_mode", "none")
        self.boundaries_discovery = config.get("boundaries_discovery", False)
        self.boundary_discovery_mode = config.get("boundary_discovery_mode", "space_free_only")
        self.boundary_space_marker = config.get("boundary_space_marker", "_")
        self.boundary_min_sequence_length = config.get("boundary_min_sequence_length", 2)
        self.segment_cache_size = max(0, int(config.get("segment_cache_size", 200000)))
        self._segment_cache = OrderedDict() if self.segment_cache_size > 0 else None
        self.space_free_lexicon_meta = config.get("space_free_lexicon", {})
        self.language_zero_morphemes = set(config.get("language_zero_morphemes", {}).values())
        self._space_free_candidates_by_initial = self._build_space_free_candidates(
            self.space_free_lexicon_meta
        )
        self._candidates_by_initial = {}
        for p_idx, (stems, suffixes) in enumerate(self.paradigms):
            for stem in stems:
                if not stem:
                    continue
                initial = stem[0]
                if initial not in self._candidates_by_initial:
                    self._candidates_by_initial[initial] = {}
                self._candidates_by_initial[initial].setdefault(p_idx, []).append((stem, suffixes))
        for initial, paradigms_by_rank in self._candidates_by_initial.items():
            ordered = []
            for p_idx in sorted(paradigms_by_rank):
                stems_for_rank = sorted(
                    paradigms_by_rank[p_idx],
                    key=lambda item: (-len(item[0]), item[0]),
                )
                ordered.append((p_idx, stems_for_rank))
            self._candidates_by_initial[initial] = ordered
        if self.fallback_mode not in {"none", "suffix"}:
            raise ValueError(f"Unsupported fallback_mode: {self.fallback_mode}")
        if self.boundary_discovery_mode not in {"space_free_only", "all"}:
            raise ValueError(f"Unsupported boundary_discovery_mode: {self.boundary_discovery_mode}")

    @staticmethod
    def _is_han_char(ch: str) -> bool:
        return (
            "\u3400" <= ch <= "\u4dbf"
            or "\u4e00" <= ch <= "\u9fff"
            or "\uf900" <= ch <= "\ufaff"
        )

    def _contains_han(self, text: str) -> bool:
        return any(self._is_han_char(ch) for ch in text)

    def _is_zero_suffix_marker(self, suffix) -> bool:
        return isinstance(suffix, str) and suffix in self.language_zero_morphemes

    def _build_space_free_candidates(self, lexicon_meta):
        candidates_by_initial = {}
        if not isinstance(lexicon_meta, dict):
            return candidates_by_initial

        languages = lexicon_meta.get("languages", {})
        for lang_meta in languages.values():
            for token in lang_meta.get("tokens", []):
                if not token or any(ch.isspace() for ch in token):
                    continue
                initial = token[0]
                candidates_by_initial.setdefault(initial, set()).add(token)

        for initial, tokens in list(candidates_by_initial.items()):
            candidates_by_initial[initial] = sorted(tokens, key=lambda tok: (-len(tok), tok))
        return candidates_by_initial

    def _segment_cache_get(self, word: str, fallback: bool, top_k: int) -> Optional[List[str]]:
        if self._segment_cache is None:
            return None
        key = (word, fallback, top_k)
        cached = self._segment_cache.get(key)
        if cached is None:
            return None
        self._segment_cache.move_to_end(key)
        return list(cached)

    def _segment_cache_put(self, word: str, fallback: bool, top_k: int, pieces: List[str]) -> List[str]:
        if self._segment_cache is None:
            return pieces
        key = (word, fallback, top_k)
        self._segment_cache[key] = tuple(pieces)
        self._segment_cache.move_to_end(key)
        if len(self._segment_cache) > self.segment_cache_size:
            self._segment_cache.popitem(last=False)
        return pieces

    def _format_token(self, token: str, paradigm_idx: Optional[int], apply_label: bool) -> str:
        if not apply_label or paradigm_idx is None or not token:
            return token
        return self.paradigm_token_format.format(token=token, paradigm_idx=paradigm_idx)

    def _match_suffixes(self, suffixes, remainder: str, paradigm_idx: int) -> List[List[str]]:
        matches = []

        for suffix in sorted(
            suffixes,
            key=lambda s: (
                0 if isinstance(s, str) else 1,
                -(0 if self._is_zero_suffix_marker(s) else (len(s) if isinstance(s, str) else len(s[0]))),
                "" if self._is_zero_suffix_marker(s) else (s if isinstance(s, str) else s[0]),
            ),
        ):
            if isinstance(suffix, (tuple, list)):
                base, nested = suffix
                if remainder.startswith(base):
                    sub = remainder[len(base):]
                    for nested_match in self._match_suffixes(nested, sub, paradigm_idx):
                        piece = self._format_token(
                            base,
                            paradigm_idx,
                            apply_label=self.paradigm_specific_suffixes,
                        )
                        if piece:
                            matches.append([piece] + nested_match)
                        else:
                            matches.append(nested_match)
            elif self._is_zero_suffix_marker(suffix) and remainder == "":
                if self.null_suffixes_as_tokens:
                    matches.append([
                        self._format_token(
                            self.null_suffix_token,
                            paradigm_idx,
                            apply_label=self.paradigm_specific_suffixes,
                        )
                    ])
                else:
                    matches.append([])
            elif remainder == suffix:
                if suffix:
                    matches.append([
                        self._format_token(
                            suffix,
                            paradigm_idx,
                            apply_label=self.paradigm_specific_suffixes,
                        )
                    ])
                elif self.null_suffixes_as_tokens:
                    matches.append([
                        self._format_token(
                            self.null_suffix_token,
                            paradigm_idx,
                            apply_label=self.paradigm_specific_suffixes,
                        )
                    ])
                else:
                    matches.append([])

        matches.sort(key=lambda parts: (-len(parts), tuple(parts)))
        return matches

    def _best_full_match(self, word: str) -> Optional[List[str]]:
        initial_candidates = self._candidates_by_initial.get(word[0], []) if word else []

        for p_idx, stems_for_rank in initial_candidates:
            best_match = None
            best_score = None

            for stem, suffixes in stems_for_rank:
                if not word.startswith(stem):
                    continue

                remainder = word[len(stem):]
                suffix_matches = self._match_suffixes(suffixes, remainder, p_idx)
                if not suffix_matches:
                    continue

                root = self._format_token(
                    stem,
                    p_idx,
                    apply_label=self.paradigm_specific_roots,
                )
                for suffix_parts in suffix_matches:
                    score = (len(stem), len(suffix_parts))
                    if best_score is None or score > best_score:
                        best_score = score
                        best_match = [root] + suffix_parts

            if best_match is not None:
                return best_match

        return None

    def _preprocess(self, text: str) -> str:
        s = text
        if self.lowercase:
            s = s.lower()
        if self.space_punct:
            s = re.sub(r"([^\w\s'])", r" \1 ", s)
        s = re.sub(r"\s+", " ", s).strip()
        return s

    def _prepare_units(self, raw_text: str) -> List[str]:
        if self._space_free_candidates_by_initial and self._contains_han(raw_text):
            return self._preprocess(raw_text).split()
        if not self.boundaries_discovery:
            return self._preprocess(raw_text).split()
        if self.boundary_discovery_mode == "space_free_only" and any(ch.isspace() for ch in raw_text.strip()):
            return self._preprocess(raw_text).split()
        return anchor_sequences(
            raw_text.lower() if self.lowercase else raw_text,
            space_marker=self.boundary_space_marker,
            min_sequence_length=self.boundary_min_sequence_length,
        )

    def _segment_space_free_word(self, word: str) -> List[str]:
        pieces = []
        idx = 0

        while idx < len(word):
            initial = word[idx]
            best = None

            for candidate in self._space_free_candidates_by_initial.get(initial, []):
                if word.startswith(candidate, idx):
                    best = candidate
                    break

            if best is None:
                best = initial

            pieces.append(best)
            idx += len(best)

        return pieces

    def _segment_word(self, word: str, fallback=True, top_k=20) -> List[str]:
        """Enhanced segmentation with deterministic paradigm selection."""
        cached = self._segment_cache_get(word, fallback, top_k)
        if cached is not None:
            return cached

        if self._space_free_candidates_by_initial and self._contains_han(word):
            return self._segment_cache_put(word, fallback, top_k, self._segment_space_free_word(word))

        full_match = self._best_full_match(word)
        if full_match is not None:
            return self._segment_cache_put(word, fallback, top_k, full_match)

        if fallback and self.fallback_mode == "suffix":
            candidates = self.paradigms[:top_k]
            longest = ""

            def collect_flat(sfx):
                for s in sfx:
                    if isinstance(s, (tuple, list)):
                        yield s[0]
                        yield from collect_flat(s[1])
                    else:
                        yield s
            for _, suffixes in candidates:
                for suffix in collect_flat(suffixes):
                    if word.endswith(suffix) and len(suffix) > len(longest):
                        longest = suffix
            if longest:
                stem = word[:-len(longest)]
                return self._segment_cache_put(word, fallback, top_k, [stem, longest])

        return self._segment_cache_put(word, fallback, top_k, [word])

    def segment_to_tokens(self, raw_text: str, fallback=True, top_k=20) -> List[str]:
        words = self._prepare_units(raw_text)
        segmented = []

        for word_idx, word in enumerate(words):
            if self.use_word_boundaries and word_idx > 0:
                segmented.append(self.word_boundary_token)
            segmented.extend(self._segment_word(word, fallback=fallback, top_k=top_k))

        return segmented

    def segment_with_alignment(self, raw_text: str) -> Tuple[str, List[Optional[int]]]:
        """
        Return the labeled segmentation string used by both training and inference.
        The alignment list is kept as a placeholder because the wrapper currently
        delegates offsets to the fast tokenizer directly.
        """
        segmented_tokens = self.segment_to_tokens(raw_text, fallback=True)
        segmented_text = " ".join(segmented_tokens)
        return segmented_text, [None] * len(segmented_text)

# ----------------------------
# Offset remapping helper
# ----------------------------
def remap_offsets_to_raw(offsets: List[Tuple[int,int]], pre2raw: List[Optional[int]]) -> List[Tuple[int,int]]:
    mapped = []
    L = len(pre2raw)
    for s,e in offsets:
        s = max(0, min(s, L)); e = max(0, min(e, L))
        rs = re_ = None
        t = s
        while t < e and rs is None:
            if pre2raw[t] is not None: rs = pre2raw[t]
            t += 1
        t = e - 1
        while t >= s and re_ is None:
            if pre2raw[t] is not None: re_ = pre2raw[t] + 1
            t -= 1
        mapped.append((rs if rs is not None else 0, re_ if re_ is not None else 0))
    return mapped

# ----------------------------
# Public wrapper
# ----------------------------
class EnhancedParadigmTokenizerWrapper(PreTrainedTokenizerFast):
    slow_tokenizer_class = None

    @staticmethod
    def _resolve_assets_dir(name_or_path, tok_file) -> Optional[str]:
        candidates = []
        if tok_file:
            candidates.append(Path(tok_file).parent)

        if name_or_path and os.path.isdir(name_or_path):
            candidates.append(Path(name_or_path).resolve())

        module_dir = Path(__file__).resolve().parent
        commit_hash = module_dir.name
        repo_id = name_or_path if isinstance(name_or_path, str) else None
        if repo_id and "/" in repo_id:
            repo_cache_name = f"models--{repo_id.replace('/', '--')}"
            for parent in module_dir.parents:
                candidates.append(parent / "hub" / repo_cache_name / "snapshots" / commit_hash)

        for parent in module_dir.parents:
            hub_root = parent / "hub"
            if not hub_root.is_dir():
                continue
            candidates.extend(hub_root.glob(f"models--*--*/snapshots/{commit_hash}"))

        seen = set()
        for candidate in candidates:
            candidate = Path(candidate)
            key = str(candidate)
            if key in seen:
                continue
            seen.add(key)
            if (candidate / "tokenizer.json").is_file() and (candidate / "paradigms.json").is_file():
                return str(candidate)
        return None

    @staticmethod
    def _download_required_assets(name_or_path, revision, cache_dir=None, local_files_only=False) -> Optional[str]:
        if not isinstance(name_or_path, str) or "/" not in name_or_path:
            return None

        required_files = [
            "tokenizer.json",
            "paradigms.json",
            "preprocess_config.json",
            "tokenizer_config.json",
        ]
        downloaded = []
        for filename in required_files:
            try:
                downloaded.append(
                    hf_hub_download(
                        repo_id=name_or_path,
                        filename=filename,
                        revision=revision,
                        cache_dir=cache_dir,
                        local_files_only=local_files_only,
                    )
                )
            except Exception:
                if filename in {"paradigms.json", "tokenizer.json"}:
                    return None

        snapshot_candidates = []
        for path in downloaded:
            candidate = Path(path).parent
            snapshot_candidates.append(candidate)
            for parent in Path(path).parents:
                snapshot_candidates.append(parent / "snapshots" / str(revision))

        seen = set()
        for candidate in snapshot_candidates:
            candidate = Path(candidate)
            key = str(candidate)
            if key in seen:
                continue
            seen.add(key)
            if (candidate / "tokenizer.json").is_file() and (candidate / "paradigms.json").is_file():
                return str(candidate)
        return None

    def __init__(self, *args, **kwargs):
        # Ensure fast tokenizer is loaded directly (no slow->fast conversion)
        name_or_path = kwargs.get("name_or_path", None)
        if name_or_path is None and len(args) > 0 and isinstance(args[0], str):
            name_or_path = args[0]
        cache_dir = kwargs.get("cache_dir")
        local_files_only = bool(kwargs.get("local_files_only", False))

        if "tokenizer_file" not in kwargs and "tokenizer_object" not in kwargs and name_or_path is not None:
            tf = os.path.join(name_or_path, "tokenizer.json")
            if os.path.isfile(tf):
                kwargs["tokenizer_file"] = tf
            else:
                commit_hash = Path(__file__).resolve().parent.name
                downloaded_assets_dir = self._download_required_assets(
                    name_or_path=name_or_path,
                    revision=commit_hash,
                    cache_dir=cache_dir,
                    local_files_only=local_files_only,
                )
                if downloaded_assets_dir is None and local_files_only:
                    downloaded_assets_dir = self._download_required_assets(
                        name_or_path=name_or_path,
                        revision=commit_hash,
                        cache_dir=cache_dir,
                        local_files_only=False,
                    )
                if downloaded_assets_dir is not None:
                    kwargs["tokenizer_file"] = str(Path(downloaded_assets_dir) / "tokenizer.json")

        super().__init__(*args, **kwargs)

        tok_file = kwargs.get("tokenizer_file", getattr(self, "tokenizer_file", None))
        hf_dir = self._resolve_assets_dir(name_or_path, tok_file)
        if hf_dir is None:
            commit_hash = Path(__file__).resolve().parent.name
            hf_dir = self._download_required_assets(
                name_or_path=name_or_path,
                revision=commit_hash,
                cache_dir=cache_dir,
                local_files_only=local_files_only,
            )
            if hf_dir is None and local_files_only:
                hf_dir = self._download_required_assets(
                    name_or_path=name_or_path,
                    revision=commit_hash,
                    cache_dir=cache_dir,
                    local_files_only=False,
                )
        if hf_dir is None:
            raise FileNotFoundError(
                "Could not resolve local tokenizer assets directory from tokenizer_file "
                f"or name_or_path={name_or_path!r}"
            )

        # Load paradigms
        ppath = os.path.join(hf_dir, "paradigms.json")
        if not os.path.exists(ppath):
            raise FileNotFoundError(f"Missing paradigms.json in {hf_dir}")
        self.paradigms, self.paradigms_meta = _load_paradigms_any(ppath)

        # Load configuration
        self.config = {}
        cpath = os.path.join(hf_dir, "tokenizer_config.json")
        if os.path.exists(cpath):
            with open(cpath, "r", encoding="utf-8") as f:
                try:
                    self.config.update(json.load(f))
                except json.JSONDecodeError:
                    pass

        # Load preprocessing flags
        pre_cfg = {"lowercase": True, "space_punct": True}
        pre_cpath = os.path.join(hf_dir, "preprocess_config.json")
        if os.path.exists(pre_cpath):
            with open(pre_cpath, "r", encoding="utf-8") as f:
                pre_cfg.update(json.load(f))

        # Merge configs
        full_config = {**pre_cfg, **self.config}
        if self.paradigms_meta.get("space_free_lexicon"):
            full_config["space_free_lexicon"] = self.paradigms_meta["space_free_lexicon"]
        if self.paradigms_meta.get("language_zero_morphemes"):
            full_config["language_zero_morphemes"] = self.paradigms_meta["language_zero_morphemes"]

        self.segmenter = EnhancedParadigmFinderSegmenter(
            paradigms=self.paradigms,
            config=full_config,
        )

    def _segment_input(self, value):
        if isinstance(value, str):
            seg, _ = self.segmenter.segment_with_alignment(value)
            return seg
        if isinstance(value, (list, tuple)):
            segs = []
            for item in value:
                if not isinstance(item, str):
                    raise TypeError("batched inputs must contain only strings")
                seg, _ = self.segmenter.segment_with_alignment(item)
                segs.append(seg)
            return segs
        raise TypeError("text inputs must be str or List[str]/Tuple[str]")

    # ---- main entry point ----
    def __call__(self, text, text_pair=None, **kwargs):
        seg_text = self._segment_input(text)
        if text_pair is None:
            return super().__call__(seg_text, **kwargs)

        seg_text_pair = self._segment_input(text_pair)
        return super().__call__(seg_text, text_pair=seg_text_pair, **kwargs)

    def tokenize(self, text, **kwargs):
        # Intercept manual .tokenize() calls to ensure segmentation happens first
        if isinstance(text, str):
            return super().tokenize(self._segment_input(text), **kwargs)
        elif isinstance(text, list):
            # Tokenize each string separately, then flatten (matches HF behavior)
            out = []
            for t in text:
                out.extend(super().tokenize(self._segment_input(t), **kwargs))
            return out
        else:
            raise TypeError("tokenize() expects str or List[str]")

    def encode(self, text, text_pair=None, **kwargs):
        seg_text = self._segment_input(text)
        if text_pair is None:
            return super().encode(seg_text, **kwargs)
        return super().encode(seg_text, text_pair=self._segment_input(text_pair), **kwargs)

    def encode_plus(self, text, text_pair=None, **kwargs):
        seg_text = self._segment_input(text)
        if text_pair is None:
            return super().encode_plus(seg_text, **kwargs)
        return super().encode_plus(
            seg_text,
            text_pair=self._segment_input(text_pair),
            **kwargs,
        )