# 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, )