"""SentencePiece tokenizer wrapper for pinyin-code Transformers models.""" from __future__ import annotations import logging import math import re import shutil import unicodedata from pathlib import Path from typing import Any import sentencepiece as spm from tokenizers import Tokenizer, decoders, normalizers from tokenizers.models import BPE from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast from transformers.tokenization_utils_base import generate_merges CHINESE_RE = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]") CHINESE_SPAN_RE = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]+") PINYIN_CODE_TOKEN_RE = re.compile( r"(?") PUNCTUATION = set( "\u3002\uff0c\u3001\uff1f\uff01\uff1a\uff1b.,?!:;()[]{}<>\u300a\u300b" "\u3010\u3011\u201c\u201d\"'\u2018\u2019\u300c\u300d\u300e\u300f" "\u2014-~\u2026/\\" ) LATIN_LETTER = ( r"A-Za-z\u00c0-\u00d6\u00d8-\u00f6\u00f8-\u00ff" r"\u0100-\u017f\u0180-\u024f\u0250-\u02af" ) LATIN_ALNUM_PATTERN = ( rf"(?:[{LATIN_LETTER}][{LATIN_LETTER}0-9]*" rf"(?:[-_][{LATIN_LETTER}0-9]+)*|" rf"[0-9]+[{LATIN_LETTER}][{LATIN_LETTER}0-9]*" rf"(?:[-_][{LATIN_LETTER}0-9]+)*)" ) LATIN_ALNUM_RE = re.compile(LATIN_ALNUM_PATTERN) URL_RE = re.compile(r"\b(?:https?://\S*|www\.\S+)", flags=re.I) DISCARDED_UNICODE_CATEGORIES = {"Cc", "Cf", "Co", "Cs", "Cn"} TOKEN_RE = re.compile( r"<[A-Z_]+>|" r"[\u3400-\u4dbf\u4e00-\u9fff]+|" rf"{LATIN_ALNUM_PATTERN}|" r"\S" ) LABELS = { "\u9898\u5e72": "", "\u9009\u9879": "", "\u7b54\u6848": "", "\u89e3\u6790": "", } PINYIN_FORMAT_ALIASES = { "code": "pinyin-code", "codes": "pinyin-code", "pinyin-code": "pinyin-code", "initial": "pinyin-initial", "initials": "pinyin-initial", "pinyin-initial": "pinyin-initial", "hanzi": "hanzi", } def latin_token_to_model_token(token: str) -> str: upper = token.upper() return upper if upper in {"A", "B", "C", "D"} else token.lower() def should_preserve_fallback_token(token: str) -> bool: if token == "\ufffd": return False for char in token: category = unicodedata.category(char) if category in DISCARDED_UNICODE_CATEGORIES: return False if category[0] not in {"L", "P", "S"}: return False return True class PinyinCodeTokenizer(PreTrainedTokenizer): """Slow tokenizer that preserves the existing SentencePiece model.""" vocab_files_names = {"vocab_file": "tokenizer.model"} model_input_names = ["input_ids", "attention_mask"] def __init__( self, vocab_file: str, add_bos_token: bool = False, add_eos_token: bool = False, transliteration: str = "pinyin-code", pinyin_format: str | None = None, use_jieba: bool = True, jieba: bool | None = None, **kwargs, ) -> None: self.vocab_file = vocab_file self.sp_model = spm.SentencePieceProcessor(model_file=vocab_file) self.add_bos_token = add_bos_token self.add_eos_token = add_eos_token self.transliteration = self._normalize_transliteration( pinyin_format or transliteration ) self.use_jieba = use_jieba if jieba is None else jieba kwargs.setdefault("unk_token", self._piece_or_none(self.sp_model.unk_id())) kwargs.setdefault("bos_token", self._piece_or_none(self.sp_model.bos_id())) kwargs.setdefault("eos_token", self._piece_or_none(self.sp_model.eos_id())) kwargs.setdefault("pad_token", self._piece_or_none(self.sp_model.pad_id())) kwargs.setdefault("transliteration", self.transliteration) kwargs.setdefault("pinyin_format", self.transliteration) kwargs.setdefault("use_jieba", self.use_jieba) kwargs.setdefault("jieba", self.use_jieba) super().__init__(**kwargs) def _normalize_transliteration(self, value: str) -> str: normalized = PINYIN_FORMAT_ALIASES.get(value.lower()) if normalized is None: allowed = ", ".join(sorted(set(PINYIN_FORMAT_ALIASES.values()))) raise ValueError(f"Unsupported transliteration {value!r}; choose from {allowed}") return normalized def _piece_or_none(self, token_id: int) -> str | None: if token_id is None or token_id < 0: return None return self.sp_model.id_to_piece(token_id) def _looks_preprocessed(self, text: str) -> bool: if SPECIAL_MARKER_RE.search(text): return True if self.transliteration == "pinyin-code" and PINYIN_CODE_TOKEN_RE.search(text): return True return False def _preprocess_raw_text(self, text: str) -> str: if not CHINESE_RE.search(text) and self._looks_preprocessed(text): return text try: from preprocessing.preprocess import ( hanzi_to_encoded, process_text, require_dependencies, ) except ImportError: return self._fallback_process_text(text) require_dependencies() if self.transliteration == "pinyin-code": return hanzi_to_encoded(text, self.use_jieba) return process_text(text, self.transliteration, self.use_jieba) def _fallback_process_text(self, text: str) -> str: if self.use_jieba: try: import jieba except ImportError as exc: raise ImportError( "Tokenizing raw Mandarin benchmark text with jieba segmentation " "requires jieba. Install the model dependencies before running " "lm_eval." ) from exc jieba.setLogLevel(logging.WARNING) else: jieba = None if self.transliteration != "hanzi": try: from pypinyin import Style, pinyin except ImportError as exc: raise ImportError( "Tokenizing raw Mandarin benchmark text as pinyin requires pypinyin. " "Install the model dependencies before running lm_eval." ) from exc def normalize_text(value: str) -> str: value = unicodedata.normalize("NFKC", value) value = URL_RE.sub(" ", value) value = re.sub(r"\$\$.*?\$\$", " ", value, flags=re.DOTALL) value = re.sub(r"[\uff08(]\s*[\uff09)]", " ", value) for label, marker in LABELS.items(): value = re.sub(rf"{label}\s*[:\uff1a]", f" {marker} ", value) value = re.sub( rf"(? ", value, flags=re.I, ) value = re.sub( rf"(? ", value, flags=re.I, ) value = re.sub( rf"(? ", value, ) value = re.sub( rf"(? ", value, ) value = value.replace("\uff08", "(").replace("\uff09", ")") return re.sub(r"\s+", " ", value).strip() def split_tone3_syllable(syllable: str) -> tuple[str, int]: match = re.fullmatch(r"([a-z\u00fcv]+)([1-5]?)", syllable.lower()) if not match: return syllable, 5 plain, tone = match.groups() return plain, int(tone or "5") def length_digit_offset(syllable: str) -> int: return min(max(len(syllable), 1), 5) - 1 def syllable_to_initial_code(syllable: str) -> str: plain, tone = split_tone3_syllable(syllable) if not plain: return "" tone_offset = 5 if tone in {3, 4, 5} else 0 digit = tone_offset + length_digit_offset(plain) initial = plain[0].upper() if tone in {1, 3, 5} else plain[0].lower() return f"{initial}{digit}" def syllable_to_initial_letter(syllable: str) -> str: plain, _ = split_tone3_syllable(syllable) return plain[:1].lower() def convert_word(word: str) -> str: if self.transliteration == "hanzi": return word syllables = pinyin(word, style=Style.TONE3, heteronym=False, errors="ignore") if self.transliteration == "pinyin-code": codes = [ syllable_to_initial_code(item[0]) for item in syllables if item and item[0] ] return "".join(code for code in codes if code) initials = [ syllable_to_initial_letter(item[0]) for item in syllables if item and item[0] ] return "".join(initial for initial in initials if initial) def tokenize_chinese_span(value: str) -> list[str]: tokens = [] words = jieba.cut(value, cut_all=False) if self.use_jieba else value for word in words: word = word.strip() if word and CHINESE_SPAN_RE.search(word): token = convert_word(word) if token: tokens.append(token) return tokens tokens = [] for part in TOKEN_RE.findall(normalize_text(text)): if part.startswith("<") and part.endswith(">"): tokens.append(part) elif CHINESE_SPAN_RE.fullmatch(part): tokens.extend(tokenize_chinese_span(part)) elif part in PUNCTUATION: tokens.append(part) elif LATIN_ALNUM_RE.fullmatch(part): tokens.append(latin_token_to_model_token(part)) elif part.isdigit(): tokens.append("") elif should_preserve_fallback_token(part): tokens.append(part.lower()) return " ".join(tokens) def _preprocess_tokenizer_input(self, value: Any) -> Any: if value is None: return None if isinstance(value, str): return self._preprocess_raw_text(value) if isinstance(value, tuple): return tuple(self._preprocess_tokenizer_input(item) for item in value) if isinstance(value, list): return [self._preprocess_tokenizer_input(item) for item in value] return value def _non_content_token_ids(self) -> set[int]: return { token_id for token_id in ( self.pad_token_id, self.bos_token_id, self.eos_token_id, self.cls_token_id, self.sep_token_id, self.mask_token_id, ) if token_id is not None } def _offset_source_text(self, value: Any, is_split_into_words: bool = False) -> str: if value is None: return "" if isinstance(value, str): return value if isinstance(value, tuple): return " ".join(self._offset_source_text(item) for item in value) if isinstance(value, list): separator = " " if is_split_into_words else "" return separator.join(self._offset_source_text(item) for item in value) return str(value) def _synthetic_offset_mapping(self, text: Any, input_ids: Any, is_split_into_words: bool = False) -> list[tuple[int, int]]: """Return slow-tokenizer-compatible offsets for evaluators that require them. SentencePiece offsets are not available for this Python tokenizer because raw Mandarin text is preprocessed into pinyin-code before encoding. These spans conservatively distribute non-special tokens across the original text so suffix/completion masking code can run without requiring a fast tokenizer. """ ids = input_ids.tolist() if hasattr(input_ids, "tolist") else list(input_ids) source = self._offset_source_text(text, is_split_into_words=is_split_into_words) source_length = len(source) if not ids: return [] if source_length == 0: return [(0, 0) for _ in ids] non_content_ids = self._non_content_token_ids() content_positions = [ index for index, token_id in enumerate(ids) if int(token_id) not in non_content_ids ] if not content_positions: return [(0, 0) for _ in ids] offsets = [(0, 0) for _ in ids] count = len(content_positions) for ordinal, position in enumerate(content_positions): start = math.floor(ordinal * source_length / count) end = math.ceil((ordinal + 1) * source_length / count) if end <= start: end = min(source_length, start + 1) offsets[position] = (start, end) return offsets def _with_optional_offsets( self, encoding, original_text: Any, return_offsets_mapping: bool, is_split_into_words: bool = False, return_tensors: str | None = None, ): if not return_offsets_mapping: return encoding input_ids = encoding["input_ids"] tensor_input = hasattr(input_ids, "ndim") input_ids_list = input_ids.tolist() if tensor_input else input_ids is_batched = False if tensor_input: is_batched = input_ids.ndim > 1 elif input_ids_list and isinstance(input_ids_list[0], list): is_batched = True if is_batched: if isinstance(original_text, list) and not is_split_into_words: texts = original_text else: texts = [original_text] * len(input_ids_list) offsets = [ self._synthetic_offset_mapping(text, ids, is_split_into_words=is_split_into_words) for text, ids in zip(texts, input_ids_list) ] else: offsets = self._synthetic_offset_mapping( original_text, input_ids_list, is_split_into_words=is_split_into_words, ) if return_tensors == "pt" or tensor_input: try: import torch offsets = torch.tensor(offsets, dtype=torch.long) except ImportError: pass encoding["offset_mapping"] = offsets return encoding def __call__(self, text=None, text_pair=None, *args, **kwargs): original_text = text return_offsets_mapping = bool(kwargs.pop("return_offsets_mapping", False)) is_split_into_words = bool(kwargs.get("is_split_into_words", False)) return_tensors = kwargs.get("return_tensors") if "text_target" in kwargs: kwargs["text_target"] = self._preprocess_tokenizer_input(kwargs["text_target"]) if "text_pair_target" in kwargs: kwargs["text_pair_target"] = self._preprocess_tokenizer_input( kwargs["text_pair_target"] ) text = self._preprocess_tokenizer_input(text) text_pair = self._preprocess_tokenizer_input(text_pair) if text_pair is None: encoding = super().__call__(text, *args, **kwargs) else: encoding = super().__call__(text, text_pair, *args, **kwargs) return self._with_optional_offsets( encoding, original_text, return_offsets_mapping, is_split_into_words=is_split_into_words, return_tensors=return_tensors, ) def encode(self, text, text_pair=None, add_special_tokens=True, *args, **kwargs): kwargs["add_special_tokens"] = add_special_tokens text = self._preprocess_tokenizer_input(text) text_pair = self._preprocess_tokenizer_input(text_pair) if text_pair is None: return super().encode(text, *args, **kwargs) return super().encode(text, text_pair, *args, **kwargs) def encode_plus(self, text, text_pair=None, *args, **kwargs): original_text = text return_offsets_mapping = bool(kwargs.pop("return_offsets_mapping", False)) is_split_into_words = bool(kwargs.get("is_split_into_words", False)) return_tensors = kwargs.get("return_tensors") text = self._preprocess_tokenizer_input(text) text_pair = self._preprocess_tokenizer_input(text_pair) if text_pair is None: encoding = super().encode_plus(text, *args, **kwargs) else: encoding = super().encode_plus(text, text_pair, *args, **kwargs) return self._with_optional_offsets( encoding, original_text, return_offsets_mapping, is_split_into_words=is_split_into_words, return_tensors=return_tensors, ) def batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs): original_batch = batch_text_or_text_pairs return_offsets_mapping = bool(kwargs.pop("return_offsets_mapping", False)) is_split_into_words = bool(kwargs.get("is_split_into_words", False)) return_tensors = kwargs.get("return_tensors") batch_text_or_text_pairs = self._preprocess_tokenizer_input( batch_text_or_text_pairs ) encoding = super().batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs) return self._with_optional_offsets( encoding, original_batch, return_offsets_mapping, is_split_into_words=is_split_into_words, return_tensors=return_tensors, ) @property def vocab_size(self) -> int: return self.sp_model.get_piece_size() def get_vocab(self) -> dict[str, int]: vocab = {self.sp_model.id_to_piece(i): i for i in range(self.vocab_size)} vocab.update(self.added_tokens_encoder) return vocab def _tokenize(self, text: str) -> list[str]: text = self._preprocess_raw_text(text) return self.sp_model.encode(text, out_type=str) def _convert_token_to_id(self, token: str) -> int: return self.sp_model.piece_to_id(token) def _convert_id_to_token(self, index: int) -> str: return self.sp_model.id_to_piece(index) def convert_tokens_to_string(self, tokens: list[str]) -> str: return self.sp_model.decode(tokens) def build_inputs_with_special_tokens( self, token_ids_0: list[int], token_ids_1: list[int] | None = None, ) -> list[int]: output = list(token_ids_0) if self.add_bos_token and self.bos_token_id is not None: output = [self.bos_token_id] + output if self.add_eos_token and self.eos_token_id is not None: output = output + [self.eos_token_id] if token_ids_1 is not None: output += list(token_ids_1) if self.add_eos_token and self.eos_token_id is not None: output.append(self.eos_token_id) return output def get_special_tokens_mask( self, token_ids_0: list[int], token_ids_1: list[int] | None = None, already_has_special_tokens: bool = False, ) -> list[int]: if already_has_special_tokens: special_ids = set(self.all_special_ids) return [1 if token_id in special_ids else 0 for token_id in token_ids_0] mask = [0] * len(token_ids_0) if self.add_bos_token and self.bos_token_id is not None: mask = [1] + mask if self.add_eos_token and self.eos_token_id is not None: mask = mask + [1] if token_ids_1 is not None: mask += [0] * len(token_ids_1) if self.add_eos_token and self.eos_token_id is not None: mask.append(1) return mask def create_token_type_ids_from_sequences( self, token_ids_0: list[int], token_ids_1: list[int] | None = None, ) -> list[int]: return [0] * len(self.build_inputs_with_special_tokens(token_ids_0, token_ids_1)) def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]: output_name = "tokenizer.model" if filename_prefix: output_name = f"{filename_prefix}-{output_name}" output_path = Path(save_directory) / output_name if Path(self.vocab_file).resolve() != output_path.resolve(): shutil.copyfile(self.vocab_file, output_path) return (str(output_path),) class EncodedMandarinTokenizer(PinyinCodeTokenizer): """Tokenizer wrapper that hides Hanzi-to-encoded-Mandarin preprocessing.""" def build_sentencepiece_bpe_backend(vocab_file: str) -> Tokenizer: """Build a tokenizers backend equivalent to the trained SentencePiece BPE.""" processor = spm.SentencePieceProcessor(model_file=vocab_file) vocab = { processor.id_to_piece(index): index for index in range(processor.get_piece_size()) } tokenizer = Tokenizer( BPE( vocab=vocab, merges=generate_merges(vocab), unk_token=processor.id_to_piece(processor.unk_id()), fuse_unk=False, ) ) tokenizer.normalizer = normalizers.Sequence( [normalizers.Prepend("▁"), normalizers.Replace(" ", "▁")] ) tokenizer.decoder = decoders.Sequence([decoders.Replace("▁", " ")]) return tokenizer class EncodedMandarinTokenizerFast(PreTrainedTokenizerFast): """Fast tokenizer preserving the raw-Hanzi pinyin-code preprocessing path.""" vocab_files_names = { "vocab_file": "tokenizer.model", "tokenizer_file": "tokenizer.json", } slow_tokenizer_class = EncodedMandarinTokenizer model_input_names = ["input_ids", "attention_mask"] def __init__( self, vocab_file: str | None = None, tokenizer_file: str | None = None, add_bos_token: bool = False, add_eos_token: bool = False, transliteration: str = "pinyin-code", pinyin_format: str | None = None, use_jieba: bool = True, jieba: bool | None = None, **kwargs, ) -> None: if vocab_file is None: raise ValueError("EncodedMandarinTokenizerFast requires tokenizer.model") self.vocab_file = vocab_file self.sp_model = spm.SentencePieceProcessor(model_file=vocab_file) self.transliteration = PinyinCodeTokenizer._normalize_transliteration( self, pinyin_format or transliteration, ) self.use_jieba = use_jieba if jieba is None else jieba kwargs.setdefault("unk_token", self._piece_or_none(self.sp_model.unk_id())) kwargs.setdefault("bos_token", self._piece_or_none(self.sp_model.bos_id())) kwargs.setdefault("eos_token", self._piece_or_none(self.sp_model.eos_id())) kwargs.setdefault("pad_token", self._piece_or_none(self.sp_model.pad_id())) kwargs.setdefault("transliteration", self.transliteration) kwargs.setdefault("pinyin_format", self.transliteration) kwargs.setdefault("use_jieba", self.use_jieba) kwargs.setdefault("jieba", self.use_jieba) if tokenizer_file is None: kwargs["tokenizer_object"] = build_sentencepiece_bpe_backend(vocab_file) super().__init__( vocab_file=vocab_file, tokenizer_file=tokenizer_file, add_bos_token=add_bos_token, add_eos_token=add_eos_token, **kwargs, ) _normalize_transliteration = PinyinCodeTokenizer._normalize_transliteration _piece_or_none = PinyinCodeTokenizer._piece_or_none _looks_preprocessed = PinyinCodeTokenizer._looks_preprocessed _preprocess_raw_text = PinyinCodeTokenizer._preprocess_raw_text _fallback_process_text = PinyinCodeTokenizer._fallback_process_text def _preprocess_tokenizer_input(self, value: Any) -> Any: if value is None: return None if isinstance(value, str): return self._preprocess_raw_text(value) if isinstance(value, tuple): return tuple(self._preprocess_tokenizer_input(item) for item in value) if isinstance(value, list): return [self._preprocess_tokenizer_input(item) for item in value] return value def __call__(self, text=None, text_pair=None, *args, **kwargs): if "text_target" in kwargs: kwargs["text_target"] = self._preprocess_tokenizer_input(kwargs["text_target"]) if "text_pair_target" in kwargs: kwargs["text_pair_target"] = self._preprocess_tokenizer_input( kwargs["text_pair_target"] ) text = self._preprocess_tokenizer_input(text) text_pair = self._preprocess_tokenizer_input(text_pair) if text_pair is None: return super().__call__(text, *args, **kwargs) return super().__call__(text, text_pair, *args, **kwargs) def encode(self, text, text_pair=None, add_special_tokens=True, *args, **kwargs): kwargs["add_special_tokens"] = add_special_tokens text = self._preprocess_tokenizer_input(text) text_pair = self._preprocess_tokenizer_input(text_pair) if text_pair is None: return super().encode(text, *args, **kwargs) return super().encode(text, text_pair, *args, **kwargs) def encode_plus(self, text, text_pair=None, *args, **kwargs): text = self._preprocess_tokenizer_input(text) text_pair = self._preprocess_tokenizer_input(text_pair) if text_pair is None: return super().encode_plus(text, *args, **kwargs) return super().encode_plus(text, text_pair, *args, **kwargs) def batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs): batch_text_or_text_pairs = self._preprocess_tokenizer_input( batch_text_or_text_pairs ) return super().batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs) def build_inputs_with_special_tokens( self, token_ids_0: list[int], token_ids_1: list[int] | None = None, ) -> list[int]: output = list(token_ids_0) if self.add_bos_token and self.bos_token_id is not None: output = [self.bos_token_id] + output if self.add_eos_token and self.eos_token_id is not None: output = output + [self.eos_token_id] if token_ids_1 is not None: output += list(token_ids_1) if self.add_eos_token and self.eos_token_id is not None: output.append(self.eos_token_id) return output def save_vocabulary( self, save_directory: str, filename_prefix: str | None = None, ) -> tuple[str]: output_name = "tokenizer.model" if filename_prefix: output_name = f"{filename_prefix}-{output_name}" output_path = Path(save_directory) / output_name if Path(self.vocab_file).resolve() != output_path.resolve(): shutil.copyfile(self.vocab_file, output_path) return (str(output_path),)