Add fast tokenizer and benchmark classification compatibility

config.json

@@ -8,8 +8,9 @@
     "AutoModelForCausalLM": "modeling_pinyin_code.PinyinCodeForCausalLM",
     "AutoTokenizer": [
       "tokenization_pinyin_code.EncodedMandarinTokenizer",
-      null
-    ]
+      "tokenization_pinyin_code.EncodedMandarinTokenizerFast"
+    ],
+    "AutoModelForSequenceClassification": "modeling_pinyin_code.PinyinCodeForSequenceClassification"
   },
   "block_size": 512,
   "bos_token_id": 2,
@@ -26,9 +27,9 @@
   "n_layer": 8,
   "num_attention_heads": 8,
   "num_hidden_layers": 8,
-  "pad_token_id": 0,
-  "patch_pathlib_utf8_open": true,
-  "transformers_version": "5.10.2",
+  "pad_token_id": 0,
+  "patch_pathlib_utf8_open": true,
+  "transformers_version": "5.10.2",
   "unk_token_id": 1,
   "use_cache": false,
   "vocab_size": 16000

configuration_pinyin_code.py

@@ -1,56 +1,56 @@
-"""Configuration for the Transformers-compatible pinyin-code causal LM."""
-
-from __future__ import annotations
-
-import functools
-import os
-import pathlib
-
-from transformers import PretrainedConfig
-
-
-_UTF8_PATH_OPEN_PATCH_MARKER = "_pinyin_code_utf8_path_open_patch"
-
-
-def install_utf8_path_open_patch() -> None:
-    """Default text-mode ``Path.open`` calls to UTF-8 when encoding is omitted.
-
-    Some external Windows evaluation pipelines call ``Path.open("r")`` on
-    UTF-8 JSONL data before specifying an encoding. The model is loaded before
-    those datasets, so this narrow compatibility shim lets such pipelines read
-    Mandarin evaluation files without repository-side changes. Explicit
-    encodings and binary modes are left untouched.
-    """
-    current_open = pathlib.Path.open
-    if getattr(current_open, _UTF8_PATH_OPEN_PATCH_MARKER, False):
-        return
-
-    @functools.wraps(current_open)
-    def utf8_default_open(
-        self,
-        mode: str = "r",
-        buffering: int = -1,
-        encoding: str | None = None,
-        errors: str | None = None,
-        newline: str | None = None,
-    ):
-        if encoding is None and "b" not in mode:
-            encoding = "utf-8"
-        return current_open(
-            self,
-            mode=mode,
-            buffering=buffering,
-            encoding=encoding,
-            errors=errors,
-            newline=newline,
-        )
-
-    setattr(utf8_default_open, _UTF8_PATH_OPEN_PATCH_MARKER, True)
-    pathlib.Path.open = utf8_default_open
-
-
-class PinyinCodeConfig(PretrainedConfig):
-    """Configuration for the compact GPT-style pinyin-code decoder."""
+"""Configuration for the Transformers-compatible pinyin-code causal LM."""
+
+from __future__ import annotations
+
+import functools
+import os
+import pathlib
+
+from transformers import PretrainedConfig
+
+
+_UTF8_PATH_OPEN_PATCH_MARKER = "_pinyin_code_utf8_path_open_patch"
+
+
+def install_utf8_path_open_patch() -> None:
+    """Default text-mode ``Path.open`` calls to UTF-8 when encoding is omitted.
+
+    Some external Windows evaluation pipelines call ``Path.open("r")`` on
+    UTF-8 JSONL data before specifying an encoding. The model is loaded before
+    those datasets, so this narrow compatibility shim lets such pipelines read
+    Mandarin evaluation files without repository-side changes. Explicit
+    encodings and binary modes are left untouched.
+    """
+    current_open = pathlib.Path.open
+    if getattr(current_open, _UTF8_PATH_OPEN_PATCH_MARKER, False):
+        return
+
+    @functools.wraps(current_open)
+    def utf8_default_open(
+        self,
+        mode: str = "r",
+        buffering: int = -1,
+        encoding: str | None = None,
+        errors: str | None = None,
+        newline: str | None = None,
+    ):
+        if encoding is None and "b" not in mode:
+            encoding = "utf-8"
+        return current_open(
+            self,
+            mode=mode,
+            buffering=buffering,
+            encoding=encoding,
+            errors=errors,
+            newline=newline,
+        )
+
+    setattr(utf8_default_open, _UTF8_PATH_OPEN_PATCH_MARKER, True)
+    pathlib.Path.open = utf8_default_open
+
+
+class PinyinCodeConfig(PretrainedConfig):
+    """Configuration for the compact GPT-style pinyin-code decoder."""
 
     model_type = "pinyin_code"
 
@@ -63,15 +63,15 @@ class PinyinCodeConfig(PretrainedConfig):
         n_embd: int = 256,
         dropout: float = 0.1,
         bos_token_id: int | None = None,
-        eos_token_id: int | None = None,
-        pad_token_id: int | None = None,
-        unk_token_id: int | None = None,
-        patch_pathlib_utf8_open: bool = False,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
+        eos_token_id: int | None = None,
+        pad_token_id: int | None = None,
+        unk_token_id: int | None = None,
+        patch_pathlib_utf8_open: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
             pad_token_id=pad_token_id,
             unk_token_id=unk_token_id,
             **kwargs,
@@ -85,13 +85,13 @@ class PinyinCodeConfig(PretrainedConfig):
         self.num_hidden_layers = n_layer
         self.num_attention_heads = n_head
         self.hidden_size = n_embd
-        self.max_position_embeddings = block_size
-        self.is_decoder = True
-        self.is_encoder_decoder = False
-        self.use_cache = False
-        self.patch_pathlib_utf8_open = patch_pathlib_utf8_open
-        if (
-            patch_pathlib_utf8_open
-            and os.environ.get("PINYIN_CODE_DISABLE_UTF8_OPEN_PATCH") != "1"
-        ):
-            install_utf8_path_open_patch()
+        self.max_position_embeddings = block_size
+        self.is_decoder = True
+        self.is_encoder_decoder = False
+        self.use_cache = False
+        self.patch_pathlib_utf8_open = patch_pathlib_utf8_open
+        if (
+            patch_pathlib_utf8_open
+            and os.environ.get("PINYIN_CODE_DISABLE_UTF8_OPEN_PATCH") != "1"
+        ):
+            install_utf8_path_open_patch()

modeling_pinyin_code.py

@@ -5,9 +5,13 @@ from __future__ import annotations
 import torch
 from torch import nn
 from torch.nn import functional as F
-from transformers import PreTrainedModel
-from transformers.generation import GenerationMixin
-from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    CausalLMOutput,
+    SequenceClassifierOutput,
+)
 
 from .configuration_pinyin_code import PinyinCodeConfig
 
@@ -98,7 +102,7 @@ class TransformerBlock(nn.Module):
         return x
 
 
-class PinyinCodePreTrainedModel(PreTrainedModel):
+class PinyinCodePreTrainedModel(PreTrainedModel):
     """Base class for pinyin-code Transformers models."""
 
     config_class = PinyinCodeConfig
@@ -110,113 +114,113 @@ class PinyinCodePreTrainedModel(PreTrainedModel):
             nn.init.normal_(module.weight, mean=0.0, std=0.02)
             if module.bias is not None:
                 nn.init.zeros_(module.bias)
-        elif isinstance(module, nn.Embedding):
-            nn.init.normal_(module.weight, mean=0.0, std=0.02)
-
-
-class PinyinCodeModel(PinyinCodePreTrainedModel):
-    """Base decoder model returned by ``AutoModel``."""
-
-    def __init__(self, config: PinyinCodeConfig, init_weights: bool = True) -> None:
-        super().__init__(config)
-        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
-        self.position_embedding = nn.Embedding(config.block_size, config.n_embd)
-        self.dropout = nn.Dropout(config.dropout)
-        self.blocks = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))
-        self.ln_f = nn.LayerNorm(config.n_embd)
-        if init_weights:
-            self.post_init()
-
-    def get_input_embeddings(self) -> nn.Embedding:
-        return self.token_embedding
-
-    def set_input_embeddings(self, value: nn.Embedding) -> None:
-        self.token_embedding = value
-
-    def forward(
-        self,
-        input_ids: torch.Tensor | None = None,
-        attention_mask: torch.Tensor | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-        position_ids: torch.Tensor | None = None,
-        output_hidden_states: bool | None = None,
-        return_dict: bool | None = None,
-        **kwargs,
-    ) -> BaseModelOutput | tuple:
-        return_dict = True if return_dict is None else return_dict
-        output_hidden_states = (
-            self.config.output_hidden_states
-            if output_hidden_states is None
-            else output_hidden_states
-        )
-
-        if input_ids is None and inputs_embeds is None:
-            raise ValueError("You must provide either input_ids or inputs_embeds")
-        if input_ids is not None and inputs_embeds is not None:
-            raise ValueError("You cannot provide both input_ids and inputs_embeds")
-
-        if inputs_embeds is None:
-            _, seq_len = input_ids.shape
-            if seq_len > self.config.block_size:
-                raise ValueError(
-                    f"Sequence length {seq_len} exceeds block size {self.config.block_size}"
-                )
-            inputs_embeds = self.token_embedding(input_ids)
-        else:
-            seq_len = inputs_embeds.shape[1]
-            if seq_len > self.config.block_size:
-                raise ValueError(
-                    f"Sequence length {seq_len} exceeds block size {self.config.block_size}"
-                )
-
-        if position_ids is None:
-            if attention_mask is not None:
-                position_ids = attention_mask.long().cumsum(dim=-1) - 1
-                position_ids = position_ids.clamp_min(0)
-            else:
-                position_ids = torch.arange(seq_len, device=inputs_embeds.device)
-        position_ids = position_ids[:, -seq_len:] if position_ids.ndim == 2 else position_ids
-
-        x = inputs_embeds + self.position_embedding(position_ids)
-        x = self.dropout(x)
-        all_hidden_states = (x,) if output_hidden_states else None
-        for block in self.blocks:
-            x = block(x, attention_mask=attention_mask)
-            if output_hidden_states:
-                all_hidden_states = all_hidden_states + (x,)
-        hidden_states = self.ln_f(x)
-        if output_hidden_states:
-            all_hidden_states = all_hidden_states + (hidden_states,)
-
-        if not return_dict:
-            output = (hidden_states,)
-            if output_hidden_states:
-                output = output + (all_hidden_states,)
-            return output
-
-        return BaseModelOutput(
-            last_hidden_state=hidden_states,
-            hidden_states=all_hidden_states,
-        )
-
-
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+
+class PinyinCodeModel(PinyinCodePreTrainedModel):
+    """Base decoder model returned by ``AutoModel``."""
+
+    def __init__(self, config: PinyinCodeConfig, init_weights: bool = True) -> None:
+        super().__init__(config)
+        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
+        self.position_embedding = nn.Embedding(config.block_size, config.n_embd)
+        self.dropout = nn.Dropout(config.dropout)
+        self.blocks = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))
+        self.ln_f = nn.LayerNorm(config.n_embd)
+        if init_weights:
+            self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.token_embedding
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.token_embedding = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> BaseModelOutput | tuple:
+        return_dict = True if return_dict is None else return_dict
+        output_hidden_states = (
+            self.config.output_hidden_states
+            if output_hidden_states is None
+            else output_hidden_states
+        )
+
+        if input_ids is None and inputs_embeds is None:
+            raise ValueError("You must provide either input_ids or inputs_embeds")
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot provide both input_ids and inputs_embeds")
+
+        if inputs_embeds is None:
+            _, seq_len = input_ids.shape
+            if seq_len > self.config.block_size:
+                raise ValueError(
+                    f"Sequence length {seq_len} exceeds block size {self.config.block_size}"
+                )
+            inputs_embeds = self.token_embedding(input_ids)
+        else:
+            seq_len = inputs_embeds.shape[1]
+            if seq_len > self.config.block_size:
+                raise ValueError(
+                    f"Sequence length {seq_len} exceeds block size {self.config.block_size}"
+                )
+
+        if position_ids is None:
+            if attention_mask is not None:
+                position_ids = attention_mask.long().cumsum(dim=-1) - 1
+                position_ids = position_ids.clamp_min(0)
+            else:
+                position_ids = torch.arange(seq_len, device=inputs_embeds.device)
+        position_ids = position_ids[:, -seq_len:] if position_ids.ndim == 2 else position_ids
+
+        x = inputs_embeds + self.position_embedding(position_ids)
+        x = self.dropout(x)
+        all_hidden_states = (x,) if output_hidden_states else None
+        for block in self.blocks:
+            x = block(x, attention_mask=attention_mask)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (x,)
+        hidden_states = self.ln_f(x)
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            output = (hidden_states,)
+            if output_hidden_states:
+                output = output + (all_hidden_states,)
+            return output
+
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+        )
+
+
 class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
-    """Compact GPT-style causal language model using the original architecture."""
-
-    _tied_weights_keys = {"lm_head.weight": "token_embedding.weight"}
-    _keys_to_ignore_on_load_missing = [r"lm_head\.weight"]
-
-    def __init__(self, config: PinyinCodeConfig) -> None:
-        super().__init__(config, init_weights=False)
-        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
-        self.post_init()
-        self.tie_weights()
-
-    def get_output_embeddings(self) -> nn.Linear:
-        return self.lm_head
-
-    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
-        self.lm_head = new_embeddings
+    """Compact GPT-style causal language model using the original architecture."""
+
+    _tied_weights_keys = {"lm_head.weight": "token_embedding.weight"}
+    _keys_to_ignore_on_load_missing = [r"lm_head\.weight"]
+
+    def __init__(self, config: PinyinCodeConfig) -> None:
+        super().__init__(config, init_weights=False)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.post_init()
+        self.tie_weights()
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.lm_head = new_embeddings
 
     def tie_weights(self, *args, **kwargs) -> None:
         self.lm_head.weight = self.token_embedding.weight
@@ -242,20 +246,82 @@ class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
             "position_ids": position_ids,
         }
 
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> CausalLMOutput | tuple:
+        return_dict = True if return_dict is None else return_dict
+
+        decoder_outputs = PinyinCodeModel.forward(
+            self,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            position_ids=position_ids,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+        )
+        logits = self.lm_head(decoder_outputs.last_hidden_state)
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits[:, :-1, :].contiguous().view(-1, logits.size(-1)),
+                labels[:, 1:].contiguous().view(-1),
+                ignore_index=-100,
+            )
+
+        if not return_dict:
+            output = (logits,)
+            if decoder_outputs.hidden_states is not None:
+                output = output + (decoder_outputs.hidden_states,)
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=decoder_outputs.hidden_states,
+        )
+
+
+class PinyinCodeForSequenceClassification(PinyinCodeModel):
+    """Sequence classifier using the pinyin-code decoder backbone."""
+
+    _keys_to_ignore_on_load_unexpected = [r"lm_head\.weight"]
+
+    def __init__(self, config: PinyinCodeConfig) -> None:
+        super().__init__(config, init_weights=False)
+        self.num_labels = config.num_labels
+        classifier_dropout = getattr(config, "classifier_dropout", None)
+        if classifier_dropout is None:
+            classifier_dropout = getattr(config, "hidden_dropout_prob", config.dropout)
+        self.score_dropout = nn.Dropout(classifier_dropout)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        self.post_init()
+
     def forward(
         self,
         input_ids: torch.Tensor | None = None,
         attention_mask: torch.Tensor | None = None,
         labels: torch.Tensor | None = None,
-        inputs_embeds: torch.Tensor | None = None,
+        token_type_ids: torch.Tensor | None = None,
         position_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_attentions: bool | None = None,
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
         **kwargs,
-    ) -> CausalLMOutput | tuple:
+    ) -> SequenceClassifierOutput | tuple:
         return_dict = True if return_dict is None else return_dict
 
-        decoder_outputs = PinyinCodeModel.forward(
+        outputs = PinyinCodeModel.forward(
             self,
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -263,25 +329,61 @@ class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
             position_ids=position_ids,
             output_hidden_states=output_hidden_states,
             return_dict=True,
+            **kwargs,
         )
-        logits = self.lm_head(decoder_outputs.last_hidden_state)
-
-        loss = None
-        if labels is not None:
-            loss = F.cross_entropy(
-                logits[:, :-1, :].contiguous().view(-1, logits.size(-1)),
-                labels[:, 1:].contiguous().view(-1),
-                ignore_index=-100,
-            )
-
+
+        hidden_states = outputs.last_hidden_state
+        batch_size, sequence_length, _ = hidden_states.shape
+        if attention_mask is not None:
+            sequence_indices = attention_mask.to(hidden_states.device).long().sum(dim=-1) - 1
+        else:
+            sequence_indices = torch.full(
+                (batch_size,),
+                sequence_length - 1,
+                dtype=torch.long,
+                device=hidden_states.device,
+            )
+        sequence_indices = sequence_indices.clamp(min=0, max=sequence_length - 1)
+        batch_indices = torch.arange(batch_size, device=hidden_states.device)
+        pooled_hidden_states = hidden_states[batch_indices, sequence_indices]
+        logits = self.classifier(self.score_dropout(pooled_hidden_states))
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            if getattr(self.config, "problem_type", None) is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif labels.dtype in (torch.long, torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                if self.num_labels == 1:
+                    loss = F.mse_loss(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = F.mse_loss(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss = F.cross_entropy(
+                    logits.view(-1, self.num_labels),
+                    labels.view(-1),
+                )
+            elif self.config.problem_type == "multi_label_classification":
+                loss = F.binary_cross_entropy_with_logits(logits, labels)
+
         if not return_dict:
             output = (logits,)
-            if decoder_outputs.hidden_states is not None:
-                output = output + (decoder_outputs.hidden_states,)
+            if outputs.hidden_states is not None:
+                output = output + (outputs.hidden_states,)
+            attentions = getattr(outputs, "attentions", None)
+            if attentions is not None:
+                output = output + (attentions,)
             return ((loss,) + output) if loss is not None else output
 
-        return CausalLMOutput(
+        return SequenceClassifierOutput(
             loss=loss,
             logits=logits,
-            hidden_states=decoder_outputs.hidden_states,
+            hidden_states=outputs.hidden_states,
+            attentions=getattr(outputs, "attentions", None),
         )

preprocessing/extract_babylm_zho.py

@@ -0,0 +1,109 @@
+"""Extract BabyLM Chinese data from Hugging Face into JSONL.
+
+The dataset is gated. Accept the terms at:
+https://huggingface.co/datasets/BabyLM-community/babylm-zho
+
+Then log in with `huggingface-cli login` or set `HF_TOKEN`.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+from typing import Any, Iterable
+
+
+DATASET_ID = "BabyLM-community/babylm-zho"
+
+
+def hf_token() -> str | None:
+    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+
+
+def load_rows(split: str, streaming: bool) -> Iterable[dict[str, Any]]:
+    try:
+        from datasets import load_dataset
+    except ImportError as exc:
+        raise SystemExit(
+            "Missing dependency: install it with `py -m pip install datasets`."
+        ) from exc
+
+    kwargs: dict[str, Any] = {
+        "path": DATASET_ID,
+        "split": split,
+        "streaming": streaming,
+    }
+    token = hf_token()
+    if token:
+        kwargs["token"] = token
+
+    return load_dataset(**kwargs)
+
+
+def keep_row(row: dict[str, Any], args: argparse.Namespace) -> bool:
+    if args.category and row.get("category") not in args.category:
+        return False
+    if args.script and row.get("script") not in args.script:
+        return False
+    if args.language and row.get("language") not in args.language:
+        return False
+    return True
+
+
+def clean_row(row: dict[str, Any], text_only: bool) -> dict[str, Any]:
+    if text_only:
+        return {
+            "doc_id": row.get("doc-id"),
+            "text": row.get("text", ""),
+        }
+    return dict(row)
+
+
+def extract(args: argparse.Namespace) -> int:
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+
+    written = 0
+    with args.output.open("w", encoding="utf-8", newline="\n") as out:
+        for row in load_rows(args.split, args.streaming):
+            if args.max_docs is not None and written >= args.max_docs:
+                break
+            if not keep_row(row, args):
+                continue
+
+            out.write(json.dumps(clean_row(row, args.text_only), ensure_ascii=False))
+            out.write("\n")
+            written += 1
+
+    return written
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=f"Extract {DATASET_ID} from Hugging Face into JSONL."
+    )
+    parser.add_argument("--split", default="train")
+    parser.add_argument("--output", type=Path, default=Path("data/babylm_zho.jsonl"))
+    parser.add_argument("--max-docs", type=int, default=None)
+    parser.add_argument("--category", action="append")
+    parser.add_argument("--script", action="append")
+    parser.add_argument("--language", action="append")
+    parser.add_argument("--text-only", action="store_true")
+    parser.add_argument(
+        "--streaming",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Stream by default so the full dataset is not downloaded first.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    count = extract(args)
+    print(f"Wrote {count:,} records to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

preprocessing/preprocess.py

@@ -1,333 +1,393 @@
-"""Preprocess BabyLM Mandarin JSONL into word-preserving pinyin initial codes."""
-
-from __future__ import annotations
-
-import argparse
-import json
-import logging
-import re
-import sys
-import unicodedata
-from pathlib import Path
-from typing import Any, Iterable, Literal
-
-import jieba
-from pypinyin import Style, pinyin
-
-
-LABELS = {
-    "题干": "<QUESTION>",
-    "选项": "<OPTIONS>",
-    "答案": "<ANSWER>",
-    "解析": "<EXPLANATION>",
-}
-
-PUNCTUATION = set("。，、？！：；.,?!:;()[]{}<>《》【】“”\"'‘’「」『』—-~…/\\")
-CHINESE_RE = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]+")
-Transliteration = Literal["pinyin-code", "pinyin-initial", "hanzi"]
-LATIN_LETTER = r"A-Za-zÀ-ÖØ-öø-ÿĀ-ſƀ-ɏɐ-ʯ"
-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"}
-
-# Match protected markers before ordinary words so tokens like <ANSWER> survive
-# the later English/punctuation handling as a single vocabulary item.
-TOKEN_RE = re.compile(
-    r"<[A-Z_]+>|"
-    r"[\u3400-\u4dbf\u4e00-\u9fff]+|"
-    rf"{LATIN_ALNUM_PATTERN}|"
-    r"\S"
-)
-
-
-def require_dependencies() -> None:
-    """Fail early with a concise install hint and quiet jieba startup logging."""
-    if jieba is None or Style is None or pinyin is None:
-        raise SystemExit(
-            "Missing dependency: install with `py -m pip install jieba pypinyin`."
-        )
-    jieba.setLogLevel(logging.WARNING)
-
-
-def normalize_text(text: str) -> str:
-    """Replace task-specific surface forms with stable special tokens.
-
-    This happens before tokenization so multi-character patterns such as
-    ``$$...$$`` and empty brackets cannot be split into punctuation pieces.
-    """
-    text = unicodedata.normalize("NFKC", text)
-    text = URL_RE.sub(" <URL> ", text)
-    text = re.sub(r"\$\$.*?\$\$", " <MATH> ", text, flags=re.DOTALL)
-    text = re.sub(r"[（(]\s*[）)]", " <BLANK> ", text)
-
-    for label, marker in LABELS.items():
-        text = re.sub(rf"{label}\s*[:：]", f" {marker} ", text)
-
-    text = re.sub(
-        rf"(?<![{LATIN_LETTER}])yes(?![{LATIN_LETTER}])",
-        " <YES> ",
-        text,
-        flags=re.I,
-    )
-    text = re.sub(
-        rf"(?<![{LATIN_LETTER}])no(?![{LATIN_LETTER}])",
-        " <NO> ",
-        text,
-        flags=re.I,
-    )
-    text = re.sub(
-        rf"(?<![{LATIN_LETTER}])[ABCD](?=\s*[:：.．、\)])",
-        r" \g<0> ",
-        text,
-    )
-    text = re.sub(
-        rf"(?<![{LATIN_LETTER}0-9])[-+]?\d+(?:[.,]\d+)*(?:%|％)?"
-        rf"(?![{LATIN_LETTER}0-9])",
-        " <NUM> ",
-        text,
-    )
-
-    text = text.replace("（", "(").replace("）", ")")
-    text = re.sub(r"\s+", " ", text)
-    return text.strip()
-
-
-def latin_token_to_model_token(token: str) -> str:
-    """Normalize non-Mandarin alphanumeric tokens without losing option labels."""
-    upper = token.upper()
-    return upper if upper in {"A", "B", "C", "D"} else token.lower()
-
-
-def should_preserve_fallback_token(token: str) -> bool:
-    """Return true for visible non-Hanzi letters, punctuation, and symbols."""
-    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
-
-
-def split_tone3_syllable(syllable: str) -> tuple[str, int]:
-    """Return the plain pinyin syllable and its tone number.
-
-    ``Style.TONE3`` writes tones as final digits, but neutral tone syllables have
-    no digit. Treat those digitless cases as fifth tone.
-    """
-    match = re.fullmatch(r"([a-züv]+)([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:
-    """Map pinyin syllable length to the requested 0-4 digit offset."""
-    return min(max(len(syllable), 1), 5) - 1
-
-
-def syllable_to_initial_code(syllable: str) -> str:
-    """Convert one pinyin syllable with tone into ``initial + digit``.
-
-    Tone controls initial casing and whether the digit starts from 0 or 5:
-    tones 1/3/5 use uppercase initials, while tones 2/4 use lowercase initials.
-    Syllable length then adds the 0-4 offset that makes the final digit.
-    """
-    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:
-    """Convert one pinyin syllable with tone into its lowercase first letter."""
-    plain, _ = split_tone3_syllable(syllable)
-    return plain[:1].lower()
-
-
-def chinese_word_to_initial_codes(word: str) -> str:
-    """Convert one already-segmented Chinese word to one compact code token.
-
-    The important representation choice is preserved: jieba decides the word
-    boundary, and all syllable codes inside that word are concatenated. For
-    example, ``我们`` becomes ``W6M7`` rather than ``W6 M7``.
-    """
-    syllables = pinyin(word, style=Style.TONE3, heteronym=False, errors="ignore")
-    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)
-
-
-def chinese_word_to_initial_letters(word: str) -> str:
-    """Convert one already-segmented Chinese word to lowercase pinyin initials."""
-    syllables = pinyin(word, style=Style.TONE3, heteronym=False, errors="ignore")
-    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 chinese_word_to_transliteration(word: str, transliteration: Transliteration) -> str:
-    """Convert one segmented Chinese word using the requested transliteration."""
-    if transliteration == "pinyin-code":
-        return chinese_word_to_initial_codes(word)
-    if transliteration == "pinyin-initial":
-        return chinese_word_to_initial_letters(word)
-    if transliteration == "hanzi":
-        return word
-    raise ValueError(f"Unsupported transliteration: {transliteration}")
-
-
-def tokenize_chinese_span(
-    text: str,
-    transliteration: Transliteration = "pinyin-code",
-    use_jieba: bool = True,
-) -> Iterable[str]:
-    """Emit one token per jieba word or per Hanzi character."""
-    words = jieba.cut(text, cut_all=False) if use_jieba else text
-    for word in words:
-        word = word.strip()
-        if not word:
-            continue
-        if CHINESE_RE.search(word):
-            token = chinese_word_to_transliteration(word, transliteration)
-            if token:
-                yield token
-
-
-def process_text(
-    text: str,
-    transliteration: Transliteration = "pinyin-code",
-    use_jieba: bool = True,
-) -> str:
-    """Convert one raw document string into the final space-separated token line."""
-    tokens: list[str] = []
-    for part in TOKEN_RE.findall(normalize_text(text)):
-        if part.startswith("<") and part.endswith(">"):
-            tokens.append(part)
-        elif CHINESE_RE.fullmatch(part):
-            tokens.extend(tokenize_chinese_span(part, transliteration, use_jieba))
-        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("<NUM>")
-        elif should_preserve_fallback_token(part):
-            tokens.append(part.lower())
-
-    return " ".join(tokens)
-
-
-def hanzi_to_encoded(text: str, use_jieba: bool = True) -> str:
-    """Convert normal Hanzi/Mandarin text to the compact initial+digit encoding."""
-    require_dependencies()
-    return process_text(text, "pinyin-code", use_jieba)
-
-
-def read_jsonl(path: Path) -> Iterable[dict[str, Any]]:
-    """Yield JSON objects from UTF-8 JSONL, tolerating a leading BOM if present."""
-    with path.open("r", encoding="utf-8-sig") as handle:
-        for line_number, line in enumerate(handle, start=1):
-            line = line.strip()
-            if not line:
-                continue
-            try:
-                yield json.loads(line)
-            except json.JSONDecodeError as exc:
-                raise ValueError(f"Invalid JSON on line {line_number}: {exc}") from exc
-
-
-def preprocess_file(
-    input_path: Path,
-    output_path: Path,
-    preview_count: int,
-    transliteration: Transliteration = "pinyin-code",
-    use_jieba: bool = True,
-) -> int:
-    """Stream input documents to output while retaining a small preview buffer."""
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-    written = 0
-    previews: list[tuple[str, str]] = []
-
-    with output_path.open("w", encoding="utf-8", newline="\n") as out:
-        for obj in read_jsonl(input_path):
-            text = str(obj.get("text", ""))
-            processed = process_text(text, transliteration, use_jieba)
-            out.write(processed)
-            out.write("\n")
-            written += 1
-
-            if len(previews) < preview_count:
-                previews.append((text, processed))
-
-    for original, processed in previews:
-        print("ORIGINAL:")
-        print(original)
-        print("PROCESSED:")
-        print(processed)
-        print()
-
-    return written
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description="Convert BabyLM Mandarin JSONL text fields to model-ready tokens."
-    )
-    parser.add_argument("--input", type=Path, required=True)
-    parser.add_argument("--output", type=Path, required=True)
-    parser.add_argument("--preview", type=int, default=3)
-    parser.add_argument(
-        "--transliteration",
-        choices=("pinyin-code", "pinyin-initial", "hanzi"),
-        default="pinyin-code",
-        help=(
-            "Mandarin transliteration to emit: 'pinyin-code' keeps the original "
-            "tone/length code, while 'pinyin-initial' emits lowercase pinyin "
-            "first letters only, and 'hanzi' keeps segmented Mandarin as Hanzi."
-        ),
-    )
-    parser.add_argument(
-        "--jieba",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help=(
-            "Use jieba word segmentation for Chinese spans. Disable with "
-            "--no-jieba to emit one token per Hanzi character before "
-            "transliteration."
-        ),
-    )
-    return parser.parse_args()
-
-
-def main() -> None:
-    require_dependencies()
-    if hasattr(sys.stdout, "reconfigure"):
-        sys.stdout.reconfigure(encoding="utf-8")
-    args = parse_args()
-    count = preprocess_file(
-        args.input,
-        args.output,
-        args.preview,
-        args.transliteration,
-        args.jieba,
-    )
-    print(f"Wrote {count:,} processed documents to {args.output}")
-
-
-if __name__ == "__main__":
-    main()
+"""Preprocess BabyLM Mandarin JSONL into word-preserving pinyin initial codes."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import re
+import sys
+import unicodedata
+from multiprocessing import Pool
+from pathlib import Path
+from typing import Any, Iterable, Literal
+
+import jieba
+from pypinyin import Style, pinyin
+
+
+LABELS = {
+    "题干": "<QUESTION>",
+    "选项": "<OPTIONS>",
+    "答案": "<ANSWER>",
+    "解析": "<EXPLANATION>",
+}
+
+PUNCTUATION = set("。，、？！：；.,?!:;()[]{}<>《》【】“”\"'‘’「」『』—-~…/\\")
+CHINESE_RE = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]+")
+Transliteration = Literal["pinyin-code", "pinyin-initial", "hanzi"]
+LATIN_LETTER = r"A-Za-zÀ-ÖØ-öø-ÿĀ-ſƀ-ɏɐ-ʯ"
+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"}
+
+# Match protected markers before ordinary words so tokens like <ANSWER> survive
+# the later English/punctuation handling as a single vocabulary item.
+TOKEN_RE = re.compile(
+    r"<[A-Z_]+>|"
+    r"[\u3400-\u4dbf\u4e00-\u9fff]+|"
+    rf"{LATIN_ALNUM_PATTERN}|"
+    r"\S"
+)
+
+
+def require_dependencies() -> None:
+    """Fail early with a concise install hint and quiet jieba startup logging."""
+    if jieba is None or Style is None or pinyin is None:
+        raise SystemExit(
+            "Missing dependency: install with `py -m pip install jieba pypinyin`."
+        )
+    jieba.setLogLevel(logging.WARNING)
+
+
+def normalize_text(text: str) -> str:
+    """Replace task-specific surface forms with stable special tokens.
+
+    This happens before tokenization so multi-character patterns such as
+    ``$$...$$`` and empty brackets cannot be split into punctuation pieces.
+    """
+    text = unicodedata.normalize("NFKC", text)
+    text = URL_RE.sub(" <URL> ", text)
+    text = re.sub(r"\$\$.*?\$\$", " <MATH> ", text, flags=re.DOTALL)
+    text = re.sub(r"[（(]\s*[）)]", " <BLANK> ", text)
+
+    for label, marker in LABELS.items():
+        text = re.sub(rf"{label}\s*[:：]", f" {marker} ", text)
+
+    text = re.sub(
+        rf"(?<![{LATIN_LETTER}])yes(?![{LATIN_LETTER}])",
+        " <YES> ",
+        text,
+        flags=re.I,
+    )
+    text = re.sub(
+        rf"(?<![{LATIN_LETTER}])no(?![{LATIN_LETTER}])",
+        " <NO> ",
+        text,
+        flags=re.I,
+    )
+    text = re.sub(
+        rf"(?<![{LATIN_LETTER}])[ABCD](?=\s*[:：.．、\)])",
+        r" \g<0> ",
+        text,
+    )
+    text = re.sub(
+        rf"(?<![{LATIN_LETTER}0-9])[-+]?\d+(?:[.,]\d+)*(?:%|％)?"
+        rf"(?![{LATIN_LETTER}0-9])",
+        " <NUM> ",
+        text,
+    )
+
+    text = text.replace("（", "(").replace("）", ")")
+    text = re.sub(r"\s+", " ", text)
+    return text.strip()
+
+
+def latin_token_to_model_token(token: str) -> str:
+    """Normalize non-Mandarin alphanumeric tokens without losing option labels."""
+    upper = token.upper()
+    return upper if upper in {"A", "B", "C", "D"} else token.lower()
+
+
+def should_preserve_fallback_token(token: str) -> bool:
+    """Return true for visible non-Hanzi letters, punctuation, and symbols."""
+    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
+
+
+def split_tone3_syllable(syllable: str) -> tuple[str, int]:
+    """Return the plain pinyin syllable and its tone number.
+
+    ``Style.TONE3`` writes tones as final digits, but neutral tone syllables have
+    no digit. Treat those digitless cases as fifth tone.
+    """
+    match = re.fullmatch(r"([a-züv]+)([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:
+    """Map pinyin syllable length to the requested 0-4 digit offset."""
+    return min(max(len(syllable), 1), 5) - 1
+
+
+def syllable_to_initial_code(syllable: str) -> str:
+    """Convert one pinyin syllable with tone into ``initial + digit``.
+
+    Tone controls initial casing and whether the digit starts from 0 or 5:
+    tones 1/3/5 use uppercase initials, while tones 2/4 use lowercase initials.
+    Syllable length then adds the 0-4 offset that makes the final digit.
+    """
+    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:
+    """Convert one pinyin syllable with tone into its lowercase first letter."""
+    plain, _ = split_tone3_syllable(syllable)
+    return plain[:1].lower()
+
+
+def chinese_word_to_initial_codes(word: str) -> str:
+    """Convert one already-segmented Chinese word to one compact code token.
+
+    The important representation choice is preserved: jieba decides the word
+    boundary, and all syllable codes inside that word are concatenated. For
+    example, ``我们`` becomes ``W6M7`` rather than ``W6 M7``.
+    """
+    syllables = pinyin(word, style=Style.TONE3, heteronym=False, errors="ignore")
+    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)
+
+
+def chinese_word_to_initial_letters(word: str) -> str:
+    """Convert one already-segmented Chinese word to lowercase pinyin initials."""
+    syllables = pinyin(word, style=Style.TONE3, heteronym=False, errors="ignore")
+    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 chinese_word_to_transliteration(word: str, transliteration: Transliteration) -> str:
+    """Convert one segmented Chinese word using the requested transliteration."""
+    if transliteration == "pinyin-code":
+        return chinese_word_to_initial_codes(word)
+    if transliteration == "pinyin-initial":
+        return chinese_word_to_initial_letters(word)
+    if transliteration == "hanzi":
+        return word
+    raise ValueError(f"Unsupported transliteration: {transliteration}")
+
+
+def tokenize_chinese_span(
+    text: str,
+    transliteration: Transliteration = "pinyin-code",
+    use_jieba: bool = True,
+) -> Iterable[str]:
+    """Emit one token per jieba word or per Hanzi character."""
+    words = jieba.cut(text, cut_all=False) if use_jieba else text
+    for word in words:
+        word = word.strip()
+        if not word:
+            continue
+        if CHINESE_RE.search(word):
+            token = chinese_word_to_transliteration(word, transliteration)
+            if token:
+                yield token
+
+
+def process_text(
+    text: str,
+    transliteration: Transliteration = "pinyin-code",
+    use_jieba: bool = True,
+) -> str:
+    """Convert one raw document string into the final space-separated token line."""
+    tokens: list[str] = []
+    for part in TOKEN_RE.findall(normalize_text(text)):
+        if part.startswith("<") and part.endswith(">"):
+            tokens.append(part)
+        elif CHINESE_RE.fullmatch(part):
+            tokens.extend(tokenize_chinese_span(part, transliteration, use_jieba))
+        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("<NUM>")
+        elif should_preserve_fallback_token(part):
+            tokens.append(part.lower())
+
+    return " ".join(tokens)
+
+
+def hanzi_to_encoded(text: str, use_jieba: bool = True) -> str:
+    """Convert normal Hanzi/Mandarin text to the compact initial+digit encoding."""
+    require_dependencies()
+    return process_text(text, "pinyin-code", use_jieba)
+
+
+def read_jsonl(path: Path) -> Iterable[dict[str, Any]]:
+    """Yield JSON objects from UTF-8 JSONL, tolerating a leading BOM if present."""
+    with path.open("r", encoding="utf-8-sig") as handle:
+        for line_number, line in enumerate(handle, start=1):
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                yield json.loads(line)
+            except json.JSONDecodeError as exc:
+                raise ValueError(f"Invalid JSON on line {line_number}: {exc}") from exc
+
+
+def iter_processing_tasks(
+    input_path: Path,
+    transliteration: Transliteration,
+    use_jieba: bool,
+) -> Iterable[tuple[str, Transliteration, bool]]:
+    """Yield independent document tasks without loading the full corpus."""
+    for obj in read_jsonl(input_path):
+        yield str(obj.get("text", "")), transliteration, use_jieba
+
+
+def process_task(
+    task: tuple[str, Transliteration, bool],
+) -> tuple[str, str]:
+    """Process one document and return its original and encoded text."""
+    text, transliteration, use_jieba = task
+    return text, process_text(text, transliteration, use_jieba)
+
+
+def preprocess_file(
+    input_path: Path,
+    output_path: Path,
+    preview_count: int,
+    transliteration: Transliteration = "pinyin-code",
+    use_jieba: bool = True,
+    workers: int = 1,
+    chunksize: int = 32,
+) -> int:
+    """Stream documents to output, optionally using ordered multiprocessing."""
+    if workers < 1:
+        raise ValueError("workers must be at least 1")
+    if chunksize < 1:
+        raise ValueError("chunksize must be at least 1")
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    written = 0
+    previews: list[tuple[str, str]] = []
+    tasks = iter_processing_tasks(input_path, transliteration, use_jieba)
+
+    with output_path.open("w", encoding="utf-8", newline="\n") as out:
+        if workers == 1:
+            for text, processed in map(process_task, tasks):
+                out.write(processed)
+                out.write("\n")
+                written += 1
+
+                if len(previews) < preview_count:
+                    previews.append((text, processed))
+        else:
+            with Pool(
+                processes=workers,
+                initializer=require_dependencies,
+            ) as pool:
+                results = pool.imap(
+                    process_task,
+                    tasks,
+                    chunksize=chunksize,
+                )
+
+                for text, processed in results:
+                    out.write(processed)
+                    out.write("\n")
+                    written += 1
+
+                    if len(previews) < preview_count:
+                        previews.append((text, processed))
+
+    for original, processed in previews:
+        print("ORIGINAL:")
+        print(original)
+        print("PROCESSED:")
+        print(processed)
+        print()
+
+    return written
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Convert BabyLM Mandarin JSONL text fields to model-ready tokens."
+    )
+    parser.add_argument("--input", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--preview", type=int, default=3)
+    parser.add_argument(
+        "--transliteration",
+        choices=("pinyin-code", "pinyin-initial", "hanzi"),
+        default="pinyin-code",
+        help=(
+            "Mandarin transliteration to emit: 'pinyin-code' keeps the original "
+            "tone/length code, while 'pinyin-initial' emits lowercase pinyin "
+            "first letters only, and 'hanzi' keeps segmented Mandarin as Hanzi."
+        ),
+    )
+    parser.add_argument(
+        "--jieba",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help=(
+            "Use jieba word segmentation for Chinese spans. Disable with "
+            "--no-jieba to emit one token per Hanzi character before "
+            "transliteration."
+        ),
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=1,
+        help=(
+            "Number of document-processing worker processes. "
+            "The default of 1 preserves sequential behavior."
+        ),
+    )
+    parser.add_argument(
+        "--chunksize",
+        type=int,
+        default=32,
+        help="Documents assigned to each worker per multiprocessing batch.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    require_dependencies()
+    if hasattr(sys.stdout, "reconfigure"):
+        sys.stdout.reconfigure(encoding="utf-8")
+    args = parse_args()
+    count = preprocess_file(
+        args.input,
+        args.output,
+        args.preview,
+        args.transliteration,
+        args.jieba,
+        args.workers,
+        args.chunksize,
+    )
+    print(f"Wrote {count:,} processed documents to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

preprocessing/probability_matrix.py

@@ -0,0 +1,173 @@
+"""Build a bucketed tone/length probability matrix from tone statistics JSON."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from collections import Counter
+from pathlib import Path
+from typing import Any
+
+
+TONE_BUCKETS = {
+    "1": "tone_1",
+    "2": "tone_2",
+    "4": "tone_4",
+    "3": "tone_3_or_5",
+    "5": "tone_3_or_5",
+    "blank": "tone_3_or_5",
+}
+TONE_ORDER = ("tone_1", "tone_2", "tone_4", "tone_3_or_5")
+
+LENGTH_ORDER = ("length_1_or_2", "length_3", "length_4_to_6")
+
+
+def load_stats(path: Path) -> dict[str, Any]:
+    return json.loads(path.read_text(encoding="utf-8-sig"))
+
+
+def length_bucket(length: int) -> str:
+    if length in {1, 2}:
+        return "length_1_or_2"
+    if length == 3:
+        return "length_3"
+    if length in {4, 5, 6}:
+        return "length_4_to_6"
+    raise ValueError(f"Unsupported pinyin length: {length}")
+
+
+def bucketed_marginals(stats: dict[str, Any]) -> tuple[Counter[str], Counter[str], int]:
+    tone_counts: Counter[str] = Counter()
+    for tone, values in stats.get("tones", {}).items():
+        bucket = TONE_BUCKETS.get(str(tone))
+        if bucket is None:
+            raise ValueError(f"Unsupported tone bucket: {tone}")
+        tone_counts[bucket] += int(values["count"])
+
+    length_counts: Counter[str] = Counter()
+    for values in stats.get("lengths", []):
+        bucket = length_bucket(int(values["length"]))
+        length_counts[bucket] += int(values["count"])
+
+    total = int(stats.get("total_unique_cases") or sum(tone_counts.values()))
+    return tone_counts, length_counts, total
+
+
+def exact_matrix_from_cases(stats: dict[str, Any]) -> tuple[dict[str, dict[str, int]], int]:
+    matrix = {tone: {length: 0 for length in LENGTH_ORDER} for tone in TONE_ORDER}
+    for case in stats["unique_cases"]:
+        tone_bucket = TONE_BUCKETS.get(str(case["tone"]))
+        if tone_bucket is None:
+            raise ValueError(f"Unsupported tone bucket: {case['tone']}")
+        len_bucket = length_bucket(int(case["length"]))
+        matrix[tone_bucket][len_bucket] += 1
+    total = sum(sum(row.values()) for row in matrix.values())
+    return matrix, total
+
+
+def estimated_matrix_from_marginals(
+    tone_counts: Counter[str],
+    length_counts: Counter[str],
+    total: int,
+) -> dict[str, dict[str, float]]:
+    if total == 0:
+        return {tone: {length: 0.0 for length in LENGTH_ORDER} for tone in TONE_ORDER}
+
+    return {
+        tone: {
+            length: tone_counts[tone] * length_counts[length] / total
+            for length in LENGTH_ORDER
+        }
+        for tone in TONE_ORDER
+    }
+
+
+def matrix_payload(stats: dict[str, Any]) -> dict[str, Any]:
+    if "unique_cases" in stats:
+        counts, total = exact_matrix_from_cases(stats)
+        method = "exact_from_unique_cases"
+    else:
+        tone_counts, length_counts, total = bucketed_marginals(stats)
+        counts = estimated_matrix_from_marginals(tone_counts, length_counts, total)
+        method = "estimated_from_marginals_assuming_independence"
+
+    probabilities = {
+        tone: {
+            length: round((counts[tone][length] / total) if total else 0.0, 6)
+            for length in LENGTH_ORDER
+        }
+        for tone in TONE_ORDER
+    }
+
+    percentages = {
+        tone: {
+            length: round(probabilities[tone][length] * 100, 4)
+            for length in LENGTH_ORDER
+        }
+        for tone in TONE_ORDER
+    }
+
+    return {
+        "method": method,
+        "total_unique_cases": total,
+        "tone_buckets": list(TONE_ORDER),
+        "length_buckets": list(LENGTH_ORDER),
+        "counts": counts,
+        "probabilities": probabilities,
+        "percentages": percentages,
+    }
+
+
+def print_cli_matrix(payload: dict[str, Any]) -> None:
+    print(f"method: {payload['method']}")
+    print(f"total_unique_cases: {payload['total_unique_cases']}")
+    print()
+
+    headers = ["tone \\ length", *LENGTH_ORDER]
+    rows = []
+    for tone in TONE_ORDER:
+        row = [tone]
+        for length in LENGTH_ORDER:
+            row.append(f"{payload['percentages'][tone][length]:.4f}%")
+        rows.append(row)
+
+    widths = [
+        max(len(str(row[index])) for row in [headers, *rows])
+        for index in range(len(headers))
+    ]
+    print(" | ".join(value.ljust(widths[index]) for index, value in enumerate(headers)))
+    print("-+-".join("-" * width for width in widths))
+    for row in rows:
+        print(" | ".join(value.ljust(widths[index]) for index, value in enumerate(row)))
+
+
+def write_output(path: Path, payload: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(
+        json.dumps(payload, ensure_ascii=False, indent=2),
+        encoding="utf-8",
+        newline="\n",
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Create a bucketed tone/length probability matrix."
+    )
+    parser.add_argument("--input", type=Path, default=Path("data/10k_statistics.jsonl"))
+    parser.add_argument("--output", type=Path)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    payload = matrix_payload(load_stats(args.input))
+    print_cli_matrix(payload)
+
+    if args.output:
+        write_output(args.output, payload)
+        print(f"\nWrote matrix JSON to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

preprocessing/split_long_sentencepiece_lines.py

@@ -0,0 +1,147 @@
+"""Split long processed-text lines for SentencePiece training.
+
+SentencePiece BPE can fail on very long lines when whitespace splitting is
+disabled. This utility writes a tokenizer-training copy of a processed corpus
+where long documents are split at whitespace boundaries while preserving all
+tokens.
+"""
+
+from __future__ import annotations
+
+import argparse
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable
+
+
+@dataclass
+class SplitStats:
+    input_lines: int = 0
+    output_lines: int = 0
+    split_lines: int = 0
+    max_input_chars: int = 0
+    max_output_chars: int = 0
+
+
+def split_line_at_whitespace(line: str, max_chars: int) -> list[str]:
+    """Return chunks no longer than max_chars, splitting only between tokens."""
+    line = line.strip()
+    if not line:
+        return []
+    if len(line) <= max_chars:
+        return [line]
+
+    chunks: list[str] = []
+    current: list[str] = []
+    current_len = 0
+
+    for token in line.split():
+        token_len = len(token)
+        if token_len > max_chars:
+            raise ValueError(
+                f"Token longer than --max-chars ({token_len} > {max_chars}): {token[:80]!r}"
+            )
+
+        next_len = token_len if not current else current_len + 1 + token_len
+        if current and next_len > max_chars:
+            chunks.append(" ".join(current))
+            current = [token]
+            current_len = token_len
+        else:
+            current.append(token)
+            current_len = next_len
+
+    if current:
+        chunks.append(" ".join(current))
+
+    return chunks
+
+
+def iter_split_lines(input_path: Path, max_chars: int, stats: SplitStats) -> Iterable[str]:
+    """Yield split lines and update stats."""
+    with input_path.open("r", encoding="utf-8-sig") as handle:
+        for raw_line in handle:
+            line = raw_line.rstrip("\n\r")
+            stats.input_lines += 1
+            stats.max_input_chars = max(stats.max_input_chars, len(line))
+
+            chunks = split_line_at_whitespace(line, max_chars)
+            if len(chunks) > 1:
+                stats.split_lines += 1
+
+            for chunk in chunks:
+                stats.output_lines += 1
+                stats.max_output_chars = max(stats.max_output_chars, len(chunk))
+                yield chunk
+
+
+def split_file(input_path: Path, output_path: Path, max_chars: int, dry_run: bool) -> SplitStats:
+    """Split input_path into output_path and return summary statistics."""
+    if max_chars <= 0:
+        raise ValueError("--max-chars must be greater than zero")
+
+    stats = SplitStats()
+    split_lines = iter_split_lines(input_path, max_chars, stats)
+
+    if dry_run:
+        for _ in split_lines:
+            pass
+        return stats
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with output_path.open("w", encoding="utf-8", newline="\n") as out:
+        for line in split_lines:
+            out.write(line)
+            out.write("\n")
+
+    return stats
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Create a SentencePiece-training copy of processed text by splitting "
+            "long lines at whitespace boundaries."
+        )
+    )
+    parser.add_argument(
+        "--input",
+        type=Path,
+        required=True,
+        help="Processed .txt file with one document per line.",
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        required=True,
+        help="Output .txt file to use for SentencePiece training.",
+    )
+    parser.add_argument(
+        "--max-chars",
+        type=int,
+        default=60000,
+        help="Maximum characters per output line. Keep below 65535 for SentencePiece BPE.",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Report what would be written without creating the output file.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    stats = split_file(args.input, args.output, args.max_chars, args.dry_run)
+
+    action = "Would write" if args.dry_run else "Wrote"
+    print(f"{action} {stats.output_lines:,} lines from {stats.input_lines:,} input lines")
+    print(f"Split input lines: {stats.split_lines:,}")
+    print(f"Max input chars: {stats.max_input_chars:,}")
+    print(f"Max output chars: {stats.max_output_chars:,}")
+    if not args.dry_run:
+        print(f"Output: {args.output}")
+
+
+if __name__ == "__main__":
+    main()

preprocessing/statistics_to_latex.py

@@ -0,0 +1,190 @@
+"""Render pinyin tone/length summary statistics as a publication-ready LaTeX table."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any
+
+
+TONE_LABELS = {
+    "1": "Tone 1",
+    "2": "Tone 2",
+    "3": "Tone 3",
+    "4": "Tone 4",
+    "blank": "Neutral/unmarked",
+}
+TONE_ORDER = ("1", "2", "3", "4", "blank")
+
+
+def load_stats(path: Path) -> dict[str, Any]:
+    """Load the JSON summary payload written by ``tone_statistics.py``."""
+    try:
+        return json.loads(path.read_text(encoding="utf-8-sig"))
+    except json.JSONDecodeError as exc:
+        raise ValueError(f"Invalid statistics JSON in {path}: {exc}") from exc
+
+
+def latex_escape(text: str) -> str:
+    """Escape free text inserted into LaTeX prose fields."""
+    replacements = {
+        "\\": r"\textbackslash{}",
+        "&": r"\&",
+        "%": r"\%",
+        "$": r"\$",
+        "#": r"\#",
+        "_": r"\_",
+        "{": r"\{",
+        "}": r"\}",
+        "~": r"\textasciitilde{}",
+        "^": r"\textasciicircum{}",
+    }
+    return "".join(replacements.get(char, char) for char in text)
+
+
+def validate_stats(stats: dict[str, Any]) -> int:
+    """Validate required summary fields and return the number of unique cases."""
+    try:
+        total = int(stats["total_unique_cases"])
+        tones = stats["tones"]
+        lengths = stats["lengths"]
+    except (KeyError, TypeError, ValueError) as exc:
+        raise ValueError("Statistics payload is missing required summary fields.") from exc
+
+    missing_tones = [tone for tone in TONE_ORDER if tone not in tones]
+    if missing_tones:
+        raise ValueError(f"Missing tone categories: {', '.join(missing_tones)}")
+
+    tone_total = sum(int(tones[tone]["count"]) for tone in TONE_ORDER)
+    length_total = sum(int(item["count"]) for item in lengths)
+    if tone_total != total or length_total != total:
+        raise ValueError(
+            "Statistics totals do not agree: "
+            f"reported={total}, tones={tone_total}, lengths={length_total}."
+        )
+    return total
+
+
+def format_row(
+    tone_label: str = "",
+    tone_count: str = "",
+    tone_percent: str = "",
+    length_label: str = "",
+    length_count: str = "",
+    length_percent: str = "",
+) -> str:
+    return (
+        f"{tone_label} & {tone_count} & {tone_percent} & "
+        f"{length_label} & {length_count} & {length_percent} \\\\"
+    )
+
+
+def render_table(
+    stats: dict[str, Any],
+    caption: str,
+    label: str,
+) -> str:
+    """Return a standalone LaTeX table environment for summary statistics."""
+    total = validate_stats(stats)
+    count_format = f"{len(str(total))}.0"
+    tone_rows = [
+        (
+            TONE_LABELS[tone],
+            str(int(stats["tones"][tone]["count"])),
+            f'{float(stats["tones"][tone]["percentage"]):.2f}',
+        )
+        for tone in TONE_ORDER
+    ]
+    length_rows = [
+        (
+            str(int(item["length"])),
+            str(int(item["count"])),
+            f'{float(item["percentage"]):.2f}',
+        )
+        for item in sorted(stats["lengths"], key=lambda item: int(item["length"]))
+    ]
+
+    rows: list[str] = []
+    data_row_count = max(len(tone_rows), len(length_rows))
+    for index in range(data_row_count):
+        tone = tone_rows[index] if index < len(tone_rows) else ("", "", "")
+        length = length_rows[index] if index < len(length_rows) else ("", "", "")
+        rows.append(format_row(*tone, *length))
+    rows.append(r"\addlinespace")
+    rows.append(format_row("Total", str(total), "100.00", "Total", str(total), "100.00"))
+
+    body = "\n".join(rows)
+    return f"""% Requires \\usepackage{{booktabs}}
+% Requires \\usepackage{{siunitx}}
+\\begin{{table}}[t]
+\\centering
+\\caption{{{latex_escape(caption)}}}
+\\label{{{label}}}
+\\begin{{tabular}}{{
+    l
+    S[table-format={count_format}]
+    S[table-format=3.2]
+    @{{\\hspace{{1.75em}}}}
+    l
+    S[table-format={count_format}]
+    S[table-format=3.2]
+}}
+\\toprule
+\\multicolumn{{3}}{{c}}{{Tone category}} &
+\\multicolumn{{3}}{{c}}{{Syllable length (letters)}} \\\\
+\\cmidrule(lr){{1-3}} \\cmidrule(lr){{4-6}}
+{{Category}} & {{Count}} & {{Percent (\\%)}} &
+{{Length}} & {{Count}} & {{Percent (\\%)}} \\\\
+\\midrule
+{body}
+\\bottomrule
+\\end{{tabular}}
+\\vspace{{2pt}}
+
+\\begin{{minipage}}{{0.95\\linewidth}}
+\\footnotesize\\emph{{Note.}} Percentages are calculated over {total:,} unique
+character--contextual-pinyin cases; syllable length excludes the final tone
+digit. Neutral/unmarked denotes readings without an explicit tone digit.
+\\end{{minipage}}
+\\end{{table}}
+"""
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Convert pinyin summary statistics JSON to a LaTeX table."
+    )
+    parser.add_argument(
+        "--input",
+        type=Path,
+        default=Path("data/10k_statistics.jsonl"),
+        help="Summary statistics JSON file.",
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=Path("tables/10k_statistics_table.tex"),
+        help="Path for the generated LaTeX table.",
+    )
+    parser.add_argument(
+        "--caption",
+        default=(
+            "Distribution of tone categories and syllable lengths among unique "
+            "contextual pinyin cases in the 10k BabyLM-Zho sample."
+        ),
+    )
+    parser.add_argument("--label", default="tab:10k-pinyin-statistics")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    table = render_table(load_stats(args.input), args.caption, args.label)
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(table, encoding="utf-8", newline="\n")
+    print(f"Wrote LaTeX table to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

preprocessing/tone_statistics.py

@@ -0,0 +1,285 @@
+"""Count unique Mandarin character-to-pinyin tone and length cases in JSONL."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import re
+import sys
+from collections import Counter
+from pathlib import Path
+from typing import Any, Iterable
+
+import jieba
+from pypinyin import Style, pinyin
+
+
+CHINESE_RE = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]+")
+TONE_RE = re.compile(r"[1-4]$")
+TONE_ORDER = ("1", "2", "3", "4", "blank")
+
+
+def require_dependencies() -> None:
+    """Fail early with a concise install hint and quiet jieba startup logging."""
+    if jieba is None or Style is None or pinyin is None:
+        raise SystemExit(
+            "Missing dependency: install with `py -m pip install jieba pypinyin`."
+        )
+    jieba.setLogLevel(logging.WARNING)
+
+
+def read_jsonl(path: Path) -> Iterable[dict[str, Any]]:
+    """Yield JSON objects from UTF-8 JSONL, tolerating a leading BOM if present."""
+    with path.open("r", encoding="utf-8-sig") as handle:
+        for line_number, line in enumerate(handle, start=1):
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                yield json.loads(line)
+            except json.JSONDecodeError as exc:
+                raise ValueError(f"Invalid JSON on line {line_number}: {exc}") from exc
+
+
+def chinese_spans(text: str) -> Iterable[str]:
+    """Yield contiguous Chinese spans from mixed text."""
+    yield from CHINESE_RE.findall(text)
+
+
+def pinyin_tone_case(character: str, pinyin_with_tone: str) -> tuple[str, str, str]:
+    """Return the unique mapping key plus the tone bucket for one character."""
+    tone_match = TONE_RE.search(pinyin_with_tone)
+    tone = tone_match.group(0) if tone_match else "blank"
+    return character, pinyin_with_tone.lower(), tone
+
+
+def pinyin_length(pinyin_with_tone: str) -> int:
+    """Return pinyin syllable length after removing a final tone digit."""
+    return len(TONE_RE.sub("", pinyin_with_tone))
+
+
+def iter_contextual_cases(text: str) -> Iterable[tuple[str, str, str]]:
+    """Segment Chinese text with jieba and yield contextual pinyin cases.
+
+    pypinyin receives whole jieba words instead of isolated characters, which lets
+    it choose readings for common contextual forms such as polyphonic characters.
+    """
+    for span in chinese_spans(text):
+        for word in jieba.cut(span, cut_all=False):
+            word = word.strip()
+            if not word:
+                continue
+
+            pronunciations = pinyin(
+                word,
+                style=Style.TONE3,
+                heteronym=False,
+                neutral_tone_with_five=False,
+                errors="ignore",
+            )
+            chinese_chars = [char for char in word if CHINESE_RE.fullmatch(char)]
+
+            for character, syllable in zip(chinese_chars, pronunciations):
+                if not syllable:
+                    continue
+                yield pinyin_tone_case(character, syllable[0])
+
+
+def collect_unique_cases(input_path: Path, text_field: str) -> set[tuple[str, str, str]]:
+    """Collect unique character/pinyin/tone cases from a JSONL file."""
+    unique_cases: set[tuple[str, str, str]] = set()
+    for obj in read_jsonl(input_path):
+        text = str(obj.get(text_field, ""))
+        unique_cases.update(iter_contextual_cases(text))
+    return unique_cases
+
+
+def summarize(unique_cases: set[tuple[str, str, str]]) -> dict[str, Any]:
+    """Build counts and percentages for unique cases grouped by tone and length."""
+    tone_counts = Counter(tone for _, _, tone in unique_cases)
+    length_counts = Counter(pinyin_length(syllable) for _, syllable, _ in unique_cases)
+    total = len(unique_cases)
+
+    tones = {}
+    for tone in TONE_ORDER:
+        count = tone_counts[tone]
+        percentage = (count / total * 100) if total else 0.0
+        tones[tone] = {
+            "count": count,
+            "percentage": round(percentage, 4),
+        }
+
+    lengths = [
+        {
+            "length": length,
+            "count": count,
+            "percentage": round((count / total * 100) if total else 0.0, 4),
+        }
+        for length, count in sorted(length_counts.items())
+    ]
+
+    return {
+        "total_unique_cases": total,
+        "tones": tones,
+        "lengths": lengths,
+    }
+
+
+def write_case_list(
+    output_path: Path,
+    unique_cases: set[tuple[str, str, str]],
+    summary: dict[str, Any],
+) -> None:
+    """Write summary plus all unique mappings for later inspection."""
+    payload = {
+        **summary,
+        "unique_cases": [
+            {
+                "character": char,
+                "pinyin": syllable,
+                "tone": tone,
+                "length": pinyin_length(syllable),
+            }
+            for char, syllable, tone in sorted(unique_cases)
+        ],
+    }
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text(
+        json.dumps(payload, ensure_ascii=False, indent=2),
+        encoding="utf-8",
+        newline="\n",
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Segment JSONL text with jieba, convert to contextual pinyin, and "
+            "count unique Chinese character/pinyin cases by tone and length."
+        )
+    )
+    parser.add_argument("--input", type=Path, required=True)
+    parser.add_argument("--output", type=Path)
+    parser.add_argument("--text-field", default="text")
+    parser.add_argument(
+        "--show-cases",
+        action="store_true",
+        help="Print every unique character/pinyin/tone case after the summary.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    require_dependencies()
+    if hasattr(sys.stdout, "reconfigure"):
+        sys.stdout.reconfigure(encoding="utf-8")
+
+    args = parse_args()
+    unique_cases = collect_unique_cases(args.input, args.text_field)
+    summary = summarize(unique_cases)
+
+    print(json.dumps(summary, ensure_ascii=False, indent=2))
+
+    if args.show_cases:
+        for char, syllable, tone in sorted(unique_cases):
+            print(f"{char}\t{syllable}\t{tone}")
+
+    if args.output:
+        write_case_list(args.output, unique_cases, summary)
+        print(f"Wrote tone statistics to {args.output}")
+
+
+if __name__ == "__main__":
+    main()
+
+
+"""
+10k_babylm_zho.jsonl
+{
+  "total_unique_cases": 5426,
+  "tones": {
+    "1": {
+      "count": 1370,
+      "percentage": 25.2488
+    },
+    "2": {
+      "count": 1352,
+      "percentage": 24.9171
+    },
+    "3": {
+      "count": 903,
+      "percentage": 16.6421
+    },
+    "4": {
+      "count": 1754,
+      "percentage": 32.3258
+    },
+    "blank": {
+      "count": 47,
+      "percentage": 0.8662
+    }
+  },
+  "lengths": [
+    {
+      "length": 1,
+      "count": 29,
+      "percentage": 0.5345
+    },
+    {
+      "length": 2,
+      "count": 1523,
+      "percentage": 28.0686
+    },
+    {
+      "length": 3,
+      "count": 2108,
+      "percentage": 38.85
+    },
+    {
+      "length": 4,
+      "count": 1449,
+      "percentage": 26.7048
+    },
+    {
+      "length": 5,
+      "count": 298,
+      "percentage": 5.4921
+    },
+    {
+      "length": 6,
+      "count": 19,
+      "percentage": 0.3502
+    }
+  ]
+}
+"""
+
+"""
+babylm_zho.jsonl
+{
+  "total_unique_cases": 8231,
+  "tones": {
+    "1": {
+      "count": 2092,
+      "percentage": 25.4161
+    },
+    "2": {
+      "count": 2095,
+      "percentage": 25.4526
+    },
+    "3": {
+      "count": 1350,
+      "percentage": 16.4014
+    },
+    "4": {
+      "count": 2629,
+      "percentage": 31.9402
+    },
+    "blank": {
+      "count": 65,
+      "percentage": 0.7897
+    }
+  }
+}
+"""

tokenization_pinyin_code.py

@@ -2,16 +2,19 @@
 
 from __future__ import annotations
 
-import logging
-import math
-import re
-import shutil
-import unicodedata
-from pathlib import Path
-from typing import Any
+import logging
+import math
+import re
+import shutil
+import unicodedata
+from pathlib import Path
+from typing import Any
 
-import sentencepiece as spm
-from transformers import PreTrainedTokenizer
+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]")
@@ -78,7 +81,7 @@ def should_preserve_fallback_token(token: str) -> bool:
     return True
 
 
-class PinyinCodeTokenizer(PreTrainedTokenizer):
+class PinyinCodeTokenizer(PreTrainedTokenizer):
     """Slow tokenizer that preserves the existing SentencePiece model."""
 
     vocab_files_names = {"vocab_file": "tokenizer.model"}
@@ -277,150 +280,150 @@ class PinyinCodeTokenizer(PreTrainedTokenizer):
 
         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)
+    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"]
+            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,
             )
-
-        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,
-        )
+
+        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
@@ -430,43 +433,43 @@ class PinyinCodeTokenizer(PreTrainedTokenizer):
             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,
-        )
+    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:
@@ -544,5 +547,160 @@ class PinyinCodeTokenizer(PreTrainedTokenizer):
         return (str(output_path),)
 
 
-class EncodedMandarinTokenizer(PinyinCodeTokenizer):
-    """Tokenizer wrapper that hides Hanzi-to-encoded-Mandarin preprocessing."""
+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),)

tokenizer_config.json

@@ -1,53 +1,20 @@
 {
-  "added_tokens_decoder": {
-    "0": {
-      "content": "<pad>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "1": {
-      "content": "<unk>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "2": {
-      "content": "<s>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "3": {
-      "content": "</s>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    }
-  },
-  "auto_map": {
-    "AutoTokenizer": [
-      "tokenization_pinyin_code.EncodedMandarinTokenizer",
-      null
-    ]
-  },
-  "backend": "custom",
+  "backend": "tokenizers",
   "bos_token": "<s>",
   "eos_token": "</s>",
   "jieba": true,
   "model_max_length": 512,
   "pad_token": "<pad>",
   "pinyin_format": "pinyin-code",
-  "tokenizer_class": "EncodedMandarinTokenizer",
+  "tokenizer_class": "EncodedMandarinTokenizerFast",
   "transliteration": "pinyin-code",
   "unk_token": "<unk>",
-  "use_jieba": true
+  "use_jieba": true,
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_pinyin_code.EncodedMandarinTokenizer",
+      "tokenization_pinyin_code.EncodedMandarinTokenizerFast"
+    ]
+  },
+  "tokenizer_kind": "sentencepiece"
 }