timorobrecht commited on
Commit
bfd7047
·
verified ·
1 Parent(s): 5457b54

Upload folder using huggingface_hub

Browse files
README.md ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ library_name: transformers
3
+ pipeline_tag: text-generation
4
+ tags:
5
+ - causal-lm
6
+ - trust-remote-code
7
+ - sentencepiece
8
+ ---
9
+
10
+ # Pinyin-Code Causal LM
11
+
12
+ This repository contains a custom Transformers causal language model.
13
+ External evaluation repositories should load it with
14
+ `trust_remote_code=True` and use the `causal` backend.
15
+
16
+ ## Dependencies
17
+
18
+ Install the runtime dependencies before loading the model:
19
+
20
+ ```bash
21
+ pip install torch transformers safetensors sentencepiece pypinyin jieba
22
+ ```
23
+
24
+ `sentencepiece` is required for `AutoTokenizer`. `pypinyin` is required
25
+ for raw Mandarin-to-pinyin tokenization. `jieba` is required when
26
+ `use_jieba` is true; this export was created with `use_jieba=true`.
27
+
28
+ ## Loading
29
+
30
+ ```python
31
+ from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoTokenizer
32
+
33
+ model_path = "PATH_OR_REPO_ID"
34
+
35
+ config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
36
+ tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
37
+ base_model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
38
+ model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
39
+ ```
40
+
41
+ ## Evaluation
42
+
43
+ Configure external evaluators with:
44
+
45
+ - model path: this local folder or Hugging Face repo ID
46
+ - backend: `causal`
47
+ - trust remote code: enabled
48
+
49
+ The tokenizer accepts raw text through standard calls such as
50
+ `tokenizer(text)`, `tokenizer(text, add_special_tokens=False)`, and
51
+ `tokenizer(texts, padding=True, truncation=True, return_tensors="pt")`.
52
+
53
+ Export metadata:
54
+
55
+ - transliteration: `pinyin-code`
56
+ - use_jieba: `true`
config.json ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "PinyinCodeForCausalLM"
4
+ ],
5
+ "auto_map": {
6
+ "AutoConfig": "configuration_pinyin_code.PinyinCodeConfig",
7
+ "AutoModel": "modeling_pinyin_code.PinyinCodeModel",
8
+ "AutoModelForCausalLM": "modeling_pinyin_code.PinyinCodeForCausalLM",
9
+ "AutoTokenizer": [
10
+ "tokenization_pinyin_code.EncodedMandarinTokenizer",
11
+ null
12
+ ]
13
+ },
14
+ "block_size": 512,
15
+ "bos_token_id": 2,
16
+ "dropout": 0.1,
17
+ "dtype": "float32",
18
+ "eos_token_id": 3,
19
+ "evaluation_backend": "causal",
20
+ "hidden_size": 512,
21
+ "is_decoder": true,
22
+ "max_position_embeddings": 512,
23
+ "model_type": "pinyin_code",
24
+ "n_embd": 512,
25
+ "n_head": 8,
26
+ "n_layer": 8,
27
+ "num_attention_heads": 8,
28
+ "num_hidden_layers": 8,
29
+ "pad_token_id": 0,
30
+ "transformers_version": "5.10.2",
31
+ "unk_token_id": 1,
32
+ "use_cache": false,
33
+ "vocab_size": 16000
34
+ }
configuration_pinyin_code.py ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Configuration for the Transformers-compatible pinyin-code causal LM."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from transformers import PretrainedConfig
6
+
7
+
8
+ class PinyinCodeConfig(PretrainedConfig):
9
+ """Configuration for the compact GPT-style pinyin-code decoder."""
10
+
11
+ model_type = "pinyin_code"
12
+
13
+ def __init__(
14
+ self,
15
+ vocab_size: int = 8000,
16
+ block_size: int = 128,
17
+ n_layer: int = 6,
18
+ n_head: int = 8,
19
+ n_embd: int = 256,
20
+ dropout: float = 0.1,
21
+ bos_token_id: int | None = None,
22
+ eos_token_id: int | None = None,
23
+ pad_token_id: int | None = None,
24
+ unk_token_id: int | None = None,
25
+ **kwargs,
26
+ ) -> None:
27
+ super().__init__(
28
+ bos_token_id=bos_token_id,
29
+ eos_token_id=eos_token_id,
30
+ pad_token_id=pad_token_id,
31
+ unk_token_id=unk_token_id,
32
+ **kwargs,
33
+ )
34
+ self.vocab_size = vocab_size
35
+ self.block_size = block_size
36
+ self.n_layer = n_layer
37
+ self.n_head = n_head
38
+ self.n_embd = n_embd
39
+ self.dropout = dropout
40
+ self.num_hidden_layers = n_layer
41
+ self.num_attention_heads = n_head
42
+ self.hidden_size = n_embd
43
+ self.max_position_embeddings = block_size
44
+ self.is_decoder = True
45
+ self.is_encoder_decoder = False
46
+ self.use_cache = False
full_chinese_spm.vocab ADDED
The diff for this file is too large to render. See raw diff
 
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 2,
3
+ "eos_token_id": 3,
4
+ "pad_token_id": 0,
5
+ "transformers_version": "5.10.2"
6
+ }
hf/__init__.py ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Hugging Face Transformers integration for the pinyin-code language model."""
2
+
3
+ from .configuration_pinyin_code import PinyinCodeConfig
4
+ from .modeling_pinyin_code import PinyinCodeForCausalLM, PinyinCodeModel
5
+ from .tokenization_pinyin_code import EncodedMandarinTokenizer, PinyinCodeTokenizer
6
+
7
+ __all__ = [
8
+ "EncodedMandarinTokenizer",
9
+ "PinyinCodeConfig",
10
+ "PinyinCodeForCausalLM",
11
+ "PinyinCodeModel",
12
+ "PinyinCodeTokenizer",
13
+ ]
hf/configuration_pinyin_code.py ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Configuration for the Transformers-compatible pinyin-code causal LM."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from transformers import PretrainedConfig
6
+
7
+
8
+ class PinyinCodeConfig(PretrainedConfig):
9
+ """Configuration for the compact GPT-style pinyin-code decoder."""
10
+
11
+ model_type = "pinyin_code"
12
+
13
+ def __init__(
14
+ self,
15
+ vocab_size: int = 8000,
16
+ block_size: int = 128,
17
+ n_layer: int = 6,
18
+ n_head: int = 8,
19
+ n_embd: int = 256,
20
+ dropout: float = 0.1,
21
+ bos_token_id: int | None = None,
22
+ eos_token_id: int | None = None,
23
+ pad_token_id: int | None = None,
24
+ unk_token_id: int | None = None,
25
+ **kwargs,
26
+ ) -> None:
27
+ super().__init__(
28
+ bos_token_id=bos_token_id,
29
+ eos_token_id=eos_token_id,
30
+ pad_token_id=pad_token_id,
31
+ unk_token_id=unk_token_id,
32
+ **kwargs,
33
+ )
34
+ self.vocab_size = vocab_size
35
+ self.block_size = block_size
36
+ self.n_layer = n_layer
37
+ self.n_head = n_head
38
+ self.n_embd = n_embd
39
+ self.dropout = dropout
40
+ self.num_hidden_layers = n_layer
41
+ self.num_attention_heads = n_head
42
+ self.hidden_size = n_embd
43
+ self.max_position_embeddings = block_size
44
+ self.is_decoder = True
45
+ self.is_encoder_decoder = False
46
+ self.use_cache = False
hf/modeling_pinyin_code.py ADDED
@@ -0,0 +1,262 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Transformers-compatible implementation of the pinyin-code causal LM."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import torch
6
+ from torch import nn
7
+ from torch.nn import functional as F
8
+ from transformers import PreTrainedModel
9
+ from transformers.generation import GenerationMixin
10
+ from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
11
+
12
+ from .configuration_pinyin_code import PinyinCodeConfig
13
+
14
+
15
+ class CausalSelfAttention(nn.Module):
16
+ """Multi-head masked self-attention matching the original training module."""
17
+
18
+ def __init__(self, config: PinyinCodeConfig) -> None:
19
+ super().__init__()
20
+ if config.n_embd % config.n_head != 0:
21
+ raise ValueError("n_embd must be divisible by n_head")
22
+
23
+ self.n_head = config.n_head
24
+ self.head_dim = config.n_embd // config.n_head
25
+ self.dropout_p = config.dropout
26
+ self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd)
27
+ self.proj = nn.Linear(config.n_embd, config.n_embd)
28
+ self.resid_dropout = nn.Dropout(config.dropout)
29
+
30
+ def forward(self, x: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
31
+ batch_size, seq_len, embd = x.shape
32
+ q, k, v = self.qkv(x).split(embd, dim=2)
33
+
34
+ q = q.view(batch_size, seq_len, self.n_head, self.head_dim).transpose(1, 2)
35
+ k = k.view(batch_size, seq_len, self.n_head, self.head_dim).transpose(1, 2)
36
+ v = v.view(batch_size, seq_len, self.n_head, self.head_dim).transpose(1, 2)
37
+
38
+ dropout_p = self.dropout_p if self.training else 0.0
39
+ if attention_mask is not None:
40
+ causal_mask = torch.ones(
41
+ seq_len,
42
+ seq_len,
43
+ device=x.device,
44
+ dtype=torch.bool,
45
+ ).tril()
46
+ key_mask = attention_mask[:, None, None, :seq_len].to(dtype=torch.bool)
47
+ attn_mask = causal_mask.view(1, 1, seq_len, seq_len) & key_mask
48
+ y = F.scaled_dot_product_attention(
49
+ q,
50
+ k,
51
+ v,
52
+ attn_mask=attn_mask,
53
+ dropout_p=dropout_p,
54
+ is_causal=False,
55
+ )
56
+ else:
57
+ y = F.scaled_dot_product_attention(
58
+ q,
59
+ k,
60
+ v,
61
+ dropout_p=dropout_p,
62
+ is_causal=True,
63
+ )
64
+
65
+ y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, embd)
66
+ return self.resid_dropout(self.proj(y))
67
+
68
+
69
+ class FeedForward(nn.Module):
70
+ """Transformer MLP block."""
71
+
72
+ def __init__(self, config: PinyinCodeConfig) -> None:
73
+ super().__init__()
74
+ self.net = nn.Sequential(
75
+ nn.Linear(config.n_embd, 4 * config.n_embd),
76
+ nn.GELU(),
77
+ nn.Linear(4 * config.n_embd, config.n_embd),
78
+ nn.Dropout(config.dropout),
79
+ )
80
+
81
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
82
+ return self.net(x)
83
+
84
+
85
+ class TransformerBlock(nn.Module):
86
+ """Pre-norm Transformer block."""
87
+
88
+ def __init__(self, config: PinyinCodeConfig) -> None:
89
+ super().__init__()
90
+ self.ln_1 = nn.LayerNorm(config.n_embd)
91
+ self.attn = CausalSelfAttention(config)
92
+ self.ln_2 = nn.LayerNorm(config.n_embd)
93
+ self.mlp = FeedForward(config)
94
+
95
+ def forward(self, x: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
96
+ x = x + self.attn(self.ln_1(x), attention_mask=attention_mask)
97
+ x = x + self.mlp(self.ln_2(x))
98
+ return x
99
+
100
+
101
+ class PinyinCodePreTrainedModel(PreTrainedModel):
102
+ """Base class for pinyin-code Transformers models."""
103
+
104
+ config_class = PinyinCodeConfig
105
+ base_model_prefix = "pinyin_code"
106
+ supports_gradient_checkpointing = False
107
+
108
+ def _init_weights(self, module: nn.Module) -> None:
109
+ if isinstance(module, nn.Linear):
110
+ nn.init.normal_(module.weight, mean=0.0, std=0.02)
111
+ if module.bias is not None:
112
+ nn.init.zeros_(module.bias)
113
+ elif isinstance(module, nn.Embedding):
114
+ nn.init.normal_(module.weight, mean=0.0, std=0.02)
115
+
116
+
117
+ class PinyinCodeModel(PinyinCodePreTrainedModel):
118
+ """Base decoder model returned by ``AutoModel``."""
119
+
120
+ def __init__(self, config: PinyinCodeConfig, init_weights: bool = True) -> None:
121
+ super().__init__(config)
122
+ self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
123
+ self.position_embedding = nn.Embedding(config.block_size, config.n_embd)
124
+ self.dropout = nn.Dropout(config.dropout)
125
+ self.blocks = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))
126
+ self.ln_f = nn.LayerNorm(config.n_embd)
127
+ if init_weights:
128
+ self.post_init()
129
+
130
+ def get_input_embeddings(self) -> nn.Embedding:
131
+ return self.token_embedding
132
+
133
+ def set_input_embeddings(self, value: nn.Embedding) -> None:
134
+ self.token_embedding = value
135
+
136
+ def forward(
137
+ self,
138
+ input_ids: torch.Tensor | None = None,
139
+ attention_mask: torch.Tensor | None = None,
140
+ inputs_embeds: torch.Tensor | None = None,
141
+ position_ids: torch.Tensor | None = None,
142
+ return_dict: bool | None = None,
143
+ **kwargs,
144
+ ) -> BaseModelOutput | tuple:
145
+ return_dict = True if return_dict is None else return_dict
146
+
147
+ if input_ids is None and inputs_embeds is None:
148
+ raise ValueError("You must provide either input_ids or inputs_embeds")
149
+ if input_ids is not None and inputs_embeds is not None:
150
+ raise ValueError("You cannot provide both input_ids and inputs_embeds")
151
+
152
+ if inputs_embeds is None:
153
+ _, seq_len = input_ids.shape
154
+ if seq_len > self.config.block_size:
155
+ raise ValueError(
156
+ f"Sequence length {seq_len} exceeds block size {self.config.block_size}"
157
+ )
158
+ inputs_embeds = self.token_embedding(input_ids)
159
+ else:
160
+ seq_len = inputs_embeds.shape[1]
161
+ if seq_len > self.config.block_size:
162
+ raise ValueError(
163
+ f"Sequence length {seq_len} exceeds block size {self.config.block_size}"
164
+ )
165
+
166
+ if position_ids is None:
167
+ if attention_mask is not None:
168
+ position_ids = attention_mask.long().cumsum(dim=-1) - 1
169
+ position_ids = position_ids.clamp_min(0)
170
+ else:
171
+ position_ids = torch.arange(seq_len, device=inputs_embeds.device)
172
+ position_ids = position_ids[:, -seq_len:] if position_ids.ndim == 2 else position_ids
173
+
174
+ x = inputs_embeds + self.position_embedding(position_ids)
175
+ x = self.dropout(x)
176
+ for block in self.blocks:
177
+ x = block(x, attention_mask=attention_mask)
178
+ hidden_states = self.ln_f(x)
179
+
180
+ if not return_dict:
181
+ return (hidden_states,)
182
+
183
+ return BaseModelOutput(last_hidden_state=hidden_states)
184
+
185
+
186
+ class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
187
+ """Compact GPT-style causal language model using the original architecture."""
188
+
189
+ _tied_weights_keys = {"lm_head.weight": "token_embedding.weight"}
190
+ _keys_to_ignore_on_load_missing = [r"lm_head\.weight"]
191
+
192
+ def __init__(self, config: PinyinCodeConfig) -> None:
193
+ super().__init__(config, init_weights=False)
194
+ self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
195
+ self.post_init()
196
+ self.tie_weights()
197
+
198
+ def get_output_embeddings(self) -> nn.Linear:
199
+ return self.lm_head
200
+
201
+ def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
202
+ self.lm_head = new_embeddings
203
+
204
+ def tie_weights(self, *args, **kwargs) -> None:
205
+ self.lm_head.weight = self.token_embedding.weight
206
+
207
+ def prepare_inputs_for_generation(
208
+ self,
209
+ input_ids: torch.Tensor,
210
+ past_key_values=None,
211
+ attention_mask: torch.Tensor | None = None,
212
+ **kwargs,
213
+ ) -> dict:
214
+ if input_ids.shape[1] > self.config.block_size:
215
+ input_ids = input_ids[:, -self.config.block_size :]
216
+ if attention_mask is not None:
217
+ attention_mask = attention_mask[:, -self.config.block_size :]
218
+ position_ids = None
219
+ if attention_mask is not None:
220
+ position_ids = attention_mask.long().cumsum(dim=-1) - 1
221
+ position_ids = position_ids.clamp_min(0)
222
+ return {
223
+ "input_ids": input_ids,
224
+ "attention_mask": attention_mask,
225
+ "position_ids": position_ids,
226
+ }
227
+
228
+ def forward(
229
+ self,
230
+ input_ids: torch.Tensor | None = None,
231
+ attention_mask: torch.Tensor | None = None,
232
+ labels: torch.Tensor | None = None,
233
+ inputs_embeds: torch.Tensor | None = None,
234
+ position_ids: torch.Tensor | None = None,
235
+ return_dict: bool | None = None,
236
+ **kwargs,
237
+ ) -> CausalLMOutput | tuple:
238
+ return_dict = True if return_dict is None else return_dict
239
+
240
+ decoder_outputs = PinyinCodeModel.forward(
241
+ self,
242
+ input_ids=input_ids,
243
+ attention_mask=attention_mask,
244
+ inputs_embeds=inputs_embeds,
245
+ position_ids=position_ids,
246
+ return_dict=True,
247
+ )
248
+ logits = self.lm_head(decoder_outputs.last_hidden_state)
249
+
250
+ loss = None
251
+ if labels is not None:
252
+ loss = F.cross_entropy(
253
+ logits[:, :-1, :].contiguous().view(-1, logits.size(-1)),
254
+ labels[:, 1:].contiguous().view(-1),
255
+ ignore_index=-100,
256
+ )
257
+
258
+ if not return_dict:
259
+ output = (logits,)
260
+ return ((loss,) + output) if loss is not None else output
261
+
262
+ return CausalLMOutput(loss=loss, logits=logits)
hf/tokenization_pinyin_code.py ADDED
@@ -0,0 +1,402 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SentencePiece tokenizer wrapper for pinyin-code Transformers models."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import logging
6
+ import re
7
+ import shutil
8
+ import unicodedata
9
+ from pathlib import Path
10
+ from typing import Any
11
+
12
+ import sentencepiece as spm
13
+ from transformers import PreTrainedTokenizer
14
+
15
+
16
+ CHINESE_RE = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]")
17
+ CHINESE_SPAN_RE = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]+")
18
+ PINYIN_CODE_TOKEN_RE = re.compile(
19
+ r"(?<![A-Za-z0-9])[A-Za-z]\d(?:[A-Za-z]\d)*(?![A-Za-z0-9])"
20
+ )
21
+ SPECIAL_MARKER_RE = re.compile(r"<[A-Z_]+>")
22
+ PUNCTUATION = set(
23
+ "\u3002\uff0c\u3001\uff1f\uff01\uff1a\uff1b.,?!:;()[]{}<>\u300a\u300b"
24
+ "\u3010\u3011\u201c\u201d\"'\u2018\u2019\u300c\u300d\u300e\u300f"
25
+ "\u2014-~\u2026/\\"
26
+ )
27
+ LATIN_LETTER = (
28
+ r"A-Za-z\u00c0-\u00d6\u00d8-\u00f6\u00f8-\u00ff"
29
+ r"\u0100-\u017f\u0180-\u024f\u0250-\u02af"
30
+ )
31
+ LATIN_ALNUM_PATTERN = (
32
+ rf"(?:[{LATIN_LETTER}][{LATIN_LETTER}0-9]*"
33
+ rf"(?:[-_][{LATIN_LETTER}0-9]+)*|"
34
+ rf"[0-9]+[{LATIN_LETTER}][{LATIN_LETTER}0-9]*"
35
+ rf"(?:[-_][{LATIN_LETTER}0-9]+)*)"
36
+ )
37
+ LATIN_ALNUM_RE = re.compile(LATIN_ALNUM_PATTERN)
38
+ URL_RE = re.compile(r"\b(?:https?://\S*|www\.\S+)", flags=re.I)
39
+ DISCARDED_UNICODE_CATEGORIES = {"Cc", "Cf", "Co", "Cs", "Cn"}
40
+ TOKEN_RE = re.compile(
41
+ r"<[A-Z_]+>|"
42
+ r"[\u3400-\u4dbf\u4e00-\u9fff]+|"
43
+ rf"{LATIN_ALNUM_PATTERN}|"
44
+ r"\S"
45
+ )
46
+ LABELS = {
47
+ "\u9898\u5e72": "<QUESTION>",
48
+ "\u9009\u9879": "<OPTIONS>",
49
+ "\u7b54\u6848": "<ANSWER>",
50
+ "\u89e3\u6790": "<EXPLANATION>",
51
+ }
52
+ PINYIN_FORMAT_ALIASES = {
53
+ "code": "pinyin-code",
54
+ "codes": "pinyin-code",
55
+ "pinyin-code": "pinyin-code",
56
+ "initial": "pinyin-initial",
57
+ "initials": "pinyin-initial",
58
+ "pinyin-initial": "pinyin-initial",
59
+ "hanzi": "hanzi",
60
+ }
61
+
62
+
63
+ def latin_token_to_model_token(token: str) -> str:
64
+ upper = token.upper()
65
+ return upper if upper in {"A", "B", "C", "D"} else token.lower()
66
+
67
+
68
+ def should_preserve_fallback_token(token: str) -> bool:
69
+ if token == "\ufffd":
70
+ return False
71
+ for char in token:
72
+ category = unicodedata.category(char)
73
+ if category in DISCARDED_UNICODE_CATEGORIES:
74
+ return False
75
+ if category[0] not in {"L", "P", "S"}:
76
+ return False
77
+ return True
78
+
79
+
80
+ class PinyinCodeTokenizer(PreTrainedTokenizer):
81
+ """Slow tokenizer that preserves the existing SentencePiece model."""
82
+
83
+ vocab_files_names = {"vocab_file": "tokenizer.model"}
84
+ model_input_names = ["input_ids", "attention_mask"]
85
+
86
+ def __init__(
87
+ self,
88
+ vocab_file: str,
89
+ add_bos_token: bool = False,
90
+ add_eos_token: bool = False,
91
+ transliteration: str = "pinyin-code",
92
+ pinyin_format: str | None = None,
93
+ use_jieba: bool = True,
94
+ jieba: bool | None = None,
95
+ **kwargs,
96
+ ) -> None:
97
+ self.vocab_file = vocab_file
98
+ self.sp_model = spm.SentencePieceProcessor(model_file=vocab_file)
99
+ self.add_bos_token = add_bos_token
100
+ self.add_eos_token = add_eos_token
101
+ self.transliteration = self._normalize_transliteration(
102
+ pinyin_format or transliteration
103
+ )
104
+ self.use_jieba = use_jieba if jieba is None else jieba
105
+
106
+ kwargs.setdefault("unk_token", self._piece_or_none(self.sp_model.unk_id()))
107
+ kwargs.setdefault("bos_token", self._piece_or_none(self.sp_model.bos_id()))
108
+ kwargs.setdefault("eos_token", self._piece_or_none(self.sp_model.eos_id()))
109
+ kwargs.setdefault("pad_token", self._piece_or_none(self.sp_model.pad_id()))
110
+ kwargs.setdefault("transliteration", self.transliteration)
111
+ kwargs.setdefault("pinyin_format", self.transliteration)
112
+ kwargs.setdefault("use_jieba", self.use_jieba)
113
+ kwargs.setdefault("jieba", self.use_jieba)
114
+ super().__init__(**kwargs)
115
+
116
+ def _normalize_transliteration(self, value: str) -> str:
117
+ normalized = PINYIN_FORMAT_ALIASES.get(value.lower())
118
+ if normalized is None:
119
+ allowed = ", ".join(sorted(set(PINYIN_FORMAT_ALIASES.values())))
120
+ raise ValueError(f"Unsupported transliteration {value!r}; choose from {allowed}")
121
+ return normalized
122
+
123
+ def _piece_or_none(self, token_id: int) -> str | None:
124
+ if token_id is None or token_id < 0:
125
+ return None
126
+ return self.sp_model.id_to_piece(token_id)
127
+
128
+ def _looks_preprocessed(self, text: str) -> bool:
129
+ if SPECIAL_MARKER_RE.search(text):
130
+ return True
131
+ if self.transliteration == "pinyin-code" and PINYIN_CODE_TOKEN_RE.search(text):
132
+ return True
133
+ return False
134
+
135
+ def _preprocess_raw_text(self, text: str) -> str:
136
+ if not CHINESE_RE.search(text) and self._looks_preprocessed(text):
137
+ return text
138
+ try:
139
+ from preprocessing.preprocess import (
140
+ hanzi_to_encoded,
141
+ process_text,
142
+ require_dependencies,
143
+ )
144
+ except ImportError:
145
+ return self._fallback_process_text(text)
146
+
147
+ require_dependencies()
148
+ if self.transliteration == "pinyin-code":
149
+ return hanzi_to_encoded(text, self.use_jieba)
150
+ return process_text(text, self.transliteration, self.use_jieba)
151
+
152
+ def _fallback_process_text(self, text: str) -> str:
153
+ if self.use_jieba:
154
+ try:
155
+ import jieba
156
+ except ImportError as exc:
157
+ raise ImportError(
158
+ "Tokenizing raw Mandarin benchmark text with jieba segmentation "
159
+ "requires jieba. Install the model dependencies before running "
160
+ "lm_eval."
161
+ ) from exc
162
+
163
+ jieba.setLogLevel(logging.WARNING)
164
+ else:
165
+ jieba = None
166
+
167
+ if self.transliteration != "hanzi":
168
+ try:
169
+ from pypinyin import Style, pinyin
170
+ except ImportError as exc:
171
+ raise ImportError(
172
+ "Tokenizing raw Mandarin benchmark text as pinyin requires pypinyin. "
173
+ "Install the model dependencies before running lm_eval."
174
+ ) from exc
175
+
176
+
177
+ def normalize_text(value: str) -> str:
178
+ value = unicodedata.normalize("NFKC", value)
179
+ value = URL_RE.sub(" <URL> ", value)
180
+ value = re.sub(r"\$\$.*?\$\$", " <MATH> ", value, flags=re.DOTALL)
181
+ value = re.sub(r"[\uff08(]\s*[\uff09)]", " <BLANK> ", value)
182
+ for label, marker in LABELS.items():
183
+ value = re.sub(rf"{label}\s*[:\uff1a]", f" {marker} ", value)
184
+ value = re.sub(
185
+ rf"(?<![{LATIN_LETTER}])yes(?![{LATIN_LETTER}])",
186
+ " <YES> ",
187
+ value,
188
+ flags=re.I,
189
+ )
190
+ value = re.sub(
191
+ rf"(?<![{LATIN_LETTER}])no(?![{LATIN_LETTER}])",
192
+ " <NO> ",
193
+ value,
194
+ flags=re.I,
195
+ )
196
+ value = re.sub(
197
+ rf"(?<![{LATIN_LETTER}])[ABCD](?=\s*[:\uff1a.\uff0e\u3001\)])",
198
+ r" \g<0> ",
199
+ value,
200
+ )
201
+ value = re.sub(
202
+ rf"(?<![{LATIN_LETTER}0-9])[-+]?\d+(?:[.,]\d+)*(?:%|\uff05)?"
203
+ rf"(?![{LATIN_LETTER}0-9])",
204
+ " <NUM> ",
205
+ value,
206
+ )
207
+ value = value.replace("\uff08", "(").replace("\uff09", ")")
208
+ return re.sub(r"\s+", " ", value).strip()
209
+
210
+ def split_tone3_syllable(syllable: str) -> tuple[str, int]:
211
+ match = re.fullmatch(r"([a-z\u00fcv]+)([1-5]?)", syllable.lower())
212
+ if not match:
213
+ return syllable, 5
214
+ plain, tone = match.groups()
215
+ return plain, int(tone or "5")
216
+
217
+ def length_digit_offset(syllable: str) -> int:
218
+ return min(max(len(syllable), 1), 5) - 1
219
+
220
+ def syllable_to_initial_code(syllable: str) -> str:
221
+ plain, tone = split_tone3_syllable(syllable)
222
+ if not plain:
223
+ return ""
224
+ tone_offset = 5 if tone in {3, 4, 5} else 0
225
+ digit = tone_offset + length_digit_offset(plain)
226
+ initial = plain[0].upper() if tone in {1, 3, 5} else plain[0].lower()
227
+ return f"{initial}{digit}"
228
+
229
+ def syllable_to_initial_letter(syllable: str) -> str:
230
+ plain, _ = split_tone3_syllable(syllable)
231
+ return plain[:1].lower()
232
+
233
+ def convert_word(word: str) -> str:
234
+ if self.transliteration == "hanzi":
235
+ return word
236
+ syllables = pinyin(word, style=Style.TONE3, heteronym=False, errors="ignore")
237
+ if self.transliteration == "pinyin-code":
238
+ codes = [
239
+ syllable_to_initial_code(item[0])
240
+ for item in syllables
241
+ if item and item[0]
242
+ ]
243
+ return "".join(code for code in codes if code)
244
+ initials = [
245
+ syllable_to_initial_letter(item[0])
246
+ for item in syllables
247
+ if item and item[0]
248
+ ]
249
+ return "".join(initial for initial in initials if initial)
250
+
251
+ def tokenize_chinese_span(value: str) -> list[str]:
252
+ tokens = []
253
+ words = jieba.cut(value, cut_all=False) if self.use_jieba else value
254
+ for word in words:
255
+ word = word.strip()
256
+ if word and CHINESE_SPAN_RE.search(word):
257
+ token = convert_word(word)
258
+ if token:
259
+ tokens.append(token)
260
+ return tokens
261
+
262
+ tokens = []
263
+ for part in TOKEN_RE.findall(normalize_text(text)):
264
+ if part.startswith("<") and part.endswith(">"):
265
+ tokens.append(part)
266
+ elif CHINESE_SPAN_RE.fullmatch(part):
267
+ tokens.extend(tokenize_chinese_span(part))
268
+ elif part in PUNCTUATION:
269
+ tokens.append(part)
270
+ elif LATIN_ALNUM_RE.fullmatch(part):
271
+ tokens.append(latin_token_to_model_token(part))
272
+ elif part.isdigit():
273
+ tokens.append("<NUM>")
274
+ elif should_preserve_fallback_token(part):
275
+ tokens.append(part.lower())
276
+
277
+ return " ".join(tokens)
278
+
279
+ def _preprocess_tokenizer_input(self, value: Any) -> Any:
280
+ if value is None:
281
+ return None
282
+ if isinstance(value, str):
283
+ return self._preprocess_raw_text(value)
284
+ if isinstance(value, tuple):
285
+ return tuple(self._preprocess_tokenizer_input(item) for item in value)
286
+ if isinstance(value, list):
287
+ return [self._preprocess_tokenizer_input(item) for item in value]
288
+ return value
289
+
290
+ def __call__(self, text=None, text_pair=None, *args, **kwargs):
291
+ if "text_target" in kwargs:
292
+ kwargs["text_target"] = self._preprocess_tokenizer_input(kwargs["text_target"])
293
+ if "text_pair_target" in kwargs:
294
+ kwargs["text_pair_target"] = self._preprocess_tokenizer_input(
295
+ kwargs["text_pair_target"]
296
+ )
297
+
298
+ text = self._preprocess_tokenizer_input(text)
299
+ text_pair = self._preprocess_tokenizer_input(text_pair)
300
+ if text_pair is None:
301
+ return super().__call__(text, *args, **kwargs)
302
+ return super().__call__(text, text_pair, *args, **kwargs)
303
+
304
+ def encode(self, text, text_pair=None, add_special_tokens=True, *args, **kwargs):
305
+ kwargs["add_special_tokens"] = add_special_tokens
306
+ text = self._preprocess_tokenizer_input(text)
307
+ text_pair = self._preprocess_tokenizer_input(text_pair)
308
+ if text_pair is None:
309
+ return super().encode(text, *args, **kwargs)
310
+ return super().encode(text, text_pair, *args, **kwargs)
311
+
312
+ def encode_plus(self, text, text_pair=None, *args, **kwargs):
313
+ text = self._preprocess_tokenizer_input(text)
314
+ text_pair = self._preprocess_tokenizer_input(text_pair)
315
+ if text_pair is None:
316
+ return super().encode_plus(text, *args, **kwargs)
317
+ return super().encode_plus(text, text_pair, *args, **kwargs)
318
+
319
+ def batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs):
320
+ batch_text_or_text_pairs = self._preprocess_tokenizer_input(
321
+ batch_text_or_text_pairs
322
+ )
323
+ return super().batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs)
324
+
325
+ @property
326
+ def vocab_size(self) -> int:
327
+ return self.sp_model.get_piece_size()
328
+
329
+ def get_vocab(self) -> dict[str, int]:
330
+ vocab = {self.sp_model.id_to_piece(i): i for i in range(self.vocab_size)}
331
+ vocab.update(self.added_tokens_encoder)
332
+ return vocab
333
+
334
+ def _tokenize(self, text: str) -> list[str]:
335
+ text = self._preprocess_raw_text(text)
336
+ return self.sp_model.encode(text, out_type=str)
337
+
338
+ def _convert_token_to_id(self, token: str) -> int:
339
+ return self.sp_model.piece_to_id(token)
340
+
341
+ def _convert_id_to_token(self, index: int) -> str:
342
+ return self.sp_model.id_to_piece(index)
343
+
344
+ def convert_tokens_to_string(self, tokens: list[str]) -> str:
345
+ return self.sp_model.decode(tokens)
346
+
347
+ def build_inputs_with_special_tokens(
348
+ self,
349
+ token_ids_0: list[int],
350
+ token_ids_1: list[int] | None = None,
351
+ ) -> list[int]:
352
+ output = list(token_ids_0)
353
+ if self.add_bos_token and self.bos_token_id is not None:
354
+ output = [self.bos_token_id] + output
355
+ if self.add_eos_token and self.eos_token_id is not None:
356
+ output = output + [self.eos_token_id]
357
+ if token_ids_1 is not None:
358
+ output += list(token_ids_1)
359
+ if self.add_eos_token and self.eos_token_id is not None:
360
+ output.append(self.eos_token_id)
361
+ return output
362
+
363
+ def get_special_tokens_mask(
364
+ self,
365
+ token_ids_0: list[int],
366
+ token_ids_1: list[int] | None = None,
367
+ already_has_special_tokens: bool = False,
368
+ ) -> list[int]:
369
+ if already_has_special_tokens:
370
+ special_ids = set(self.all_special_ids)
371
+ return [1 if token_id in special_ids else 0 for token_id in token_ids_0]
372
+
373
+ mask = [0] * len(token_ids_0)
374
+ if self.add_bos_token and self.bos_token_id is not None:
375
+ mask = [1] + mask
376
+ if self.add_eos_token and self.eos_token_id is not None:
377
+ mask = mask + [1]
378
+ if token_ids_1 is not None:
379
+ mask += [0] * len(token_ids_1)
380
+ if self.add_eos_token and self.eos_token_id is not None:
381
+ mask.append(1)
382
+ return mask
383
+
384
+ def create_token_type_ids_from_sequences(
385
+ self,
386
+ token_ids_0: list[int],
387
+ token_ids_1: list[int] | None = None,
388
+ ) -> list[int]:
389
+ return [0] * len(self.build_inputs_with_special_tokens(token_ids_0, token_ids_1))
390
+
391
+ def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
392
+ output_name = "tokenizer.model"
393
+ if filename_prefix:
394
+ output_name = f"{filename_prefix}-{output_name}"
395
+ output_path = Path(save_directory) / output_name
396
+ if Path(self.vocab_file).resolve() != output_path.resolve():
397
+ shutil.copyfile(self.vocab_file, output_path)
398
+ return (str(output_path),)
399
+
400
+
401
+ class EncodedMandarinTokenizer(PinyinCodeTokenizer):
402
+ """Tokenizer wrapper that hides Hanzi-to-encoded-Mandarin preprocessing."""
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:382172cacb9a6b44b81cf9d26be5a43ada59ae52b8c5e10acc89723d57f37aa2
3
+ size 134706152
modeling_pinyin_code.py ADDED
@@ -0,0 +1,262 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Transformers-compatible implementation of the pinyin-code causal LM."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import torch
6
+ from torch import nn
7
+ from torch.nn import functional as F
8
+ from transformers import PreTrainedModel
9
+ from transformers.generation import GenerationMixin
10
+ from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
11
+
12
+ from .configuration_pinyin_code import PinyinCodeConfig
13
+
14
+
15
+ class CausalSelfAttention(nn.Module):
16
+ """Multi-head masked self-attention matching the original training module."""
17
+
18
+ def __init__(self, config: PinyinCodeConfig) -> None:
19
+ super().__init__()
20
+ if config.n_embd % config.n_head != 0:
21
+ raise ValueError("n_embd must be divisible by n_head")
22
+
23
+ self.n_head = config.n_head
24
+ self.head_dim = config.n_embd // config.n_head
25
+ self.dropout_p = config.dropout
26
+ self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd)
27
+ self.proj = nn.Linear(config.n_embd, config.n_embd)
28
+ self.resid_dropout = nn.Dropout(config.dropout)
29
+
30
+ def forward(self, x: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
31
+ batch_size, seq_len, embd = x.shape
32
+ q, k, v = self.qkv(x).split(embd, dim=2)
33
+
34
+ q = q.view(batch_size, seq_len, self.n_head, self.head_dim).transpose(1, 2)
35
+ k = k.view(batch_size, seq_len, self.n_head, self.head_dim).transpose(1, 2)
36
+ v = v.view(batch_size, seq_len, self.n_head, self.head_dim).transpose(1, 2)
37
+
38
+ dropout_p = self.dropout_p if self.training else 0.0
39
+ if attention_mask is not None:
40
+ causal_mask = torch.ones(
41
+ seq_len,
42
+ seq_len,
43
+ device=x.device,
44
+ dtype=torch.bool,
45
+ ).tril()
46
+ key_mask = attention_mask[:, None, None, :seq_len].to(dtype=torch.bool)
47
+ attn_mask = causal_mask.view(1, 1, seq_len, seq_len) & key_mask
48
+ y = F.scaled_dot_product_attention(
49
+ q,
50
+ k,
51
+ v,
52
+ attn_mask=attn_mask,
53
+ dropout_p=dropout_p,
54
+ is_causal=False,
55
+ )
56
+ else:
57
+ y = F.scaled_dot_product_attention(
58
+ q,
59
+ k,
60
+ v,
61
+ dropout_p=dropout_p,
62
+ is_causal=True,
63
+ )
64
+
65
+ y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, embd)
66
+ return self.resid_dropout(self.proj(y))
67
+
68
+
69
+ class FeedForward(nn.Module):
70
+ """Transformer MLP block."""
71
+
72
+ def __init__(self, config: PinyinCodeConfig) -> None:
73
+ super().__init__()
74
+ self.net = nn.Sequential(
75
+ nn.Linear(config.n_embd, 4 * config.n_embd),
76
+ nn.GELU(),
77
+ nn.Linear(4 * config.n_embd, config.n_embd),
78
+ nn.Dropout(config.dropout),
79
+ )
80
+
81
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
82
+ return self.net(x)
83
+
84
+
85
+ class TransformerBlock(nn.Module):
86
+ """Pre-norm Transformer block."""
87
+
88
+ def __init__(self, config: PinyinCodeConfig) -> None:
89
+ super().__init__()
90
+ self.ln_1 = nn.LayerNorm(config.n_embd)
91
+ self.attn = CausalSelfAttention(config)
92
+ self.ln_2 = nn.LayerNorm(config.n_embd)
93
+ self.mlp = FeedForward(config)
94
+
95
+ def forward(self, x: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
96
+ x = x + self.attn(self.ln_1(x), attention_mask=attention_mask)
97
+ x = x + self.mlp(self.ln_2(x))
98
+ return x
99
+
100
+
101
+ class PinyinCodePreTrainedModel(PreTrainedModel):
102
+ """Base class for pinyin-code Transformers models."""
103
+
104
+ config_class = PinyinCodeConfig
105
+ base_model_prefix = "pinyin_code"
106
+ supports_gradient_checkpointing = False
107
+
108
+ def _init_weights(self, module: nn.Module) -> None:
109
+ if isinstance(module, nn.Linear):
110
+ nn.init.normal_(module.weight, mean=0.0, std=0.02)
111
+ if module.bias is not None:
112
+ nn.init.zeros_(module.bias)
113
+ elif isinstance(module, nn.Embedding):
114
+ nn.init.normal_(module.weight, mean=0.0, std=0.02)
115
+
116
+
117
+ class PinyinCodeModel(PinyinCodePreTrainedModel):
118
+ """Base decoder model returned by ``AutoModel``."""
119
+
120
+ def __init__(self, config: PinyinCodeConfig, init_weights: bool = True) -> None:
121
+ super().__init__(config)
122
+ self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
123
+ self.position_embedding = nn.Embedding(config.block_size, config.n_embd)
124
+ self.dropout = nn.Dropout(config.dropout)
125
+ self.blocks = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))
126
+ self.ln_f = nn.LayerNorm(config.n_embd)
127
+ if init_weights:
128
+ self.post_init()
129
+
130
+ def get_input_embeddings(self) -> nn.Embedding:
131
+ return self.token_embedding
132
+
133
+ def set_input_embeddings(self, value: nn.Embedding) -> None:
134
+ self.token_embedding = value
135
+
136
+ def forward(
137
+ self,
138
+ input_ids: torch.Tensor | None = None,
139
+ attention_mask: torch.Tensor | None = None,
140
+ inputs_embeds: torch.Tensor | None = None,
141
+ position_ids: torch.Tensor | None = None,
142
+ return_dict: bool | None = None,
143
+ **kwargs,
144
+ ) -> BaseModelOutput | tuple:
145
+ return_dict = True if return_dict is None else return_dict
146
+
147
+ if input_ids is None and inputs_embeds is None:
148
+ raise ValueError("You must provide either input_ids or inputs_embeds")
149
+ if input_ids is not None and inputs_embeds is not None:
150
+ raise ValueError("You cannot provide both input_ids and inputs_embeds")
151
+
152
+ if inputs_embeds is None:
153
+ _, seq_len = input_ids.shape
154
+ if seq_len > self.config.block_size:
155
+ raise ValueError(
156
+ f"Sequence length {seq_len} exceeds block size {self.config.block_size}"
157
+ )
158
+ inputs_embeds = self.token_embedding(input_ids)
159
+ else:
160
+ seq_len = inputs_embeds.shape[1]
161
+ if seq_len > self.config.block_size:
162
+ raise ValueError(
163
+ f"Sequence length {seq_len} exceeds block size {self.config.block_size}"
164
+ )
165
+
166
+ if position_ids is None:
167
+ if attention_mask is not None:
168
+ position_ids = attention_mask.long().cumsum(dim=-1) - 1
169
+ position_ids = position_ids.clamp_min(0)
170
+ else:
171
+ position_ids = torch.arange(seq_len, device=inputs_embeds.device)
172
+ position_ids = position_ids[:, -seq_len:] if position_ids.ndim == 2 else position_ids
173
+
174
+ x = inputs_embeds + self.position_embedding(position_ids)
175
+ x = self.dropout(x)
176
+ for block in self.blocks:
177
+ x = block(x, attention_mask=attention_mask)
178
+ hidden_states = self.ln_f(x)
179
+
180
+ if not return_dict:
181
+ return (hidden_states,)
182
+
183
+ return BaseModelOutput(last_hidden_state=hidden_states)
184
+
185
+
186
+ class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
187
+ """Compact GPT-style causal language model using the original architecture."""
188
+
189
+ _tied_weights_keys = {"lm_head.weight": "token_embedding.weight"}
190
+ _keys_to_ignore_on_load_missing = [r"lm_head\.weight"]
191
+
192
+ def __init__(self, config: PinyinCodeConfig) -> None:
193
+ super().__init__(config, init_weights=False)
194
+ self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
195
+ self.post_init()
196
+ self.tie_weights()
197
+
198
+ def get_output_embeddings(self) -> nn.Linear:
199
+ return self.lm_head
200
+
201
+ def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
202
+ self.lm_head = new_embeddings
203
+
204
+ def tie_weights(self, *args, **kwargs) -> None:
205
+ self.lm_head.weight = self.token_embedding.weight
206
+
207
+ def prepare_inputs_for_generation(
208
+ self,
209
+ input_ids: torch.Tensor,
210
+ past_key_values=None,
211
+ attention_mask: torch.Tensor | None = None,
212
+ **kwargs,
213
+ ) -> dict:
214
+ if input_ids.shape[1] > self.config.block_size:
215
+ input_ids = input_ids[:, -self.config.block_size :]
216
+ if attention_mask is not None:
217
+ attention_mask = attention_mask[:, -self.config.block_size :]
218
+ position_ids = None
219
+ if attention_mask is not None:
220
+ position_ids = attention_mask.long().cumsum(dim=-1) - 1
221
+ position_ids = position_ids.clamp_min(0)
222
+ return {
223
+ "input_ids": input_ids,
224
+ "attention_mask": attention_mask,
225
+ "position_ids": position_ids,
226
+ }
227
+
228
+ def forward(
229
+ self,
230
+ input_ids: torch.Tensor | None = None,
231
+ attention_mask: torch.Tensor | None = None,
232
+ labels: torch.Tensor | None = None,
233
+ inputs_embeds: torch.Tensor | None = None,
234
+ position_ids: torch.Tensor | None = None,
235
+ return_dict: bool | None = None,
236
+ **kwargs,
237
+ ) -> CausalLMOutput | tuple:
238
+ return_dict = True if return_dict is None else return_dict
239
+
240
+ decoder_outputs = PinyinCodeModel.forward(
241
+ self,
242
+ input_ids=input_ids,
243
+ attention_mask=attention_mask,
244
+ inputs_embeds=inputs_embeds,
245
+ position_ids=position_ids,
246
+ return_dict=True,
247
+ )
248
+ logits = self.lm_head(decoder_outputs.last_hidden_state)
249
+
250
+ loss = None
251
+ if labels is not None:
252
+ loss = F.cross_entropy(
253
+ logits[:, :-1, :].contiguous().view(-1, logits.size(-1)),
254
+ labels[:, 1:].contiguous().view(-1),
255
+ ignore_index=-100,
256
+ )
257
+
258
+ if not return_dict:
259
+ output = (logits,)
260
+ return ((loss,) + output) if loss is not None else output
261
+
262
+ return CausalLMOutput(loss=loss, logits=logits)
preprocessing/__init__.py ADDED
File without changes
preprocessing/preprocess.py ADDED
@@ -0,0 +1,333 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Preprocess BabyLM Mandarin JSONL into word-preserving pinyin initial codes."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import argparse
6
+ import json
7
+ import logging
8
+ import re
9
+ import sys
10
+ import unicodedata
11
+ from pathlib import Path
12
+ from typing import Any, Iterable, Literal
13
+
14
+ import jieba
15
+ from pypinyin import Style, pinyin
16
+
17
+
18
+ LABELS = {
19
+ "题干": "<QUESTION>",
20
+ "选项": "<OPTIONS>",
21
+ "答案": "<ANSWER>",
22
+ "解析": "<EXPLANATION>",
23
+ }
24
+
25
+ PUNCTUATION = set("。,、?!:;.,?!:;()[]{}<>《》【】“”\"'‘’「」『』—-~…/\\")
26
+ CHINESE_RE = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]+")
27
+ Transliteration = Literal["pinyin-code", "pinyin-initial", "hanzi"]
28
+ LATIN_LETTER = r"A-Za-zÀ-ÖØ-öø-ÿĀ-ſƀ-ɏɐ-ʯ"
29
+ LATIN_ALNUM_PATTERN = (
30
+ rf"(?:[{LATIN_LETTER}][{LATIN_LETTER}0-9]*"
31
+ rf"(?:[-_][{LATIN_LETTER}0-9]+)*|"
32
+ rf"[0-9]+[{LATIN_LETTER}][{LATIN_LETTER}0-9]*"
33
+ rf"(?:[-_][{LATIN_LETTER}0-9]+)*)"
34
+ )
35
+ LATIN_ALNUM_RE = re.compile(LATIN_ALNUM_PATTERN)
36
+ URL_RE = re.compile(r"\b(?:https?://\S*|www\.\S+)", flags=re.I)
37
+ DISCARDED_UNICODE_CATEGORIES = {"Cc", "Cf", "Co", "Cs", "Cn"}
38
+
39
+ # Match protected markers before ordinary words so tokens like <ANSWER> survive
40
+ # the later English/punctuation handling as a single vocabulary item.
41
+ TOKEN_RE = re.compile(
42
+ r"<[A-Z_]+>|"
43
+ r"[\u3400-\u4dbf\u4e00-\u9fff]+|"
44
+ rf"{LATIN_ALNUM_PATTERN}|"
45
+ r"\S"
46
+ )
47
+
48
+
49
+ def require_dependencies() -> None:
50
+ """Fail early with a concise install hint and quiet jieba startup logging."""
51
+ if jieba is None or Style is None or pinyin is None:
52
+ raise SystemExit(
53
+ "Missing dependency: install with `py -m pip install jieba pypinyin`."
54
+ )
55
+ jieba.setLogLevel(logging.WARNING)
56
+
57
+
58
+ def normalize_text(text: str) -> str:
59
+ """Replace task-specific surface forms with stable special tokens.
60
+
61
+ This happens before tokenization so multi-character patterns such as
62
+ ``$$...$$`` and empty brackets cannot be split into punctuation pieces.
63
+ """
64
+ text = unicodedata.normalize("NFKC", text)
65
+ text = URL_RE.sub(" <URL> ", text)
66
+ text = re.sub(r"\$\$.*?\$\$", " <MATH> ", text, flags=re.DOTALL)
67
+ text = re.sub(r"[((]\s*[))]", " <BLANK> ", text)
68
+
69
+ for label, marker in LABELS.items():
70
+ text = re.sub(rf"{label}\s*[::]", f" {marker} ", text)
71
+
72
+ text = re.sub(
73
+ rf"(?<![{LATIN_LETTER}])yes(?![{LATIN_LETTER}])",
74
+ " <YES> ",
75
+ text,
76
+ flags=re.I,
77
+ )
78
+ text = re.sub(
79
+ rf"(?<![{LATIN_LETTER}])no(?![{LATIN_LETTER}])",
80
+ " <NO> ",
81
+ text,
82
+ flags=re.I,
83
+ )
84
+ text = re.sub(
85
+ rf"(?<![{LATIN_LETTER}])[ABCD](?=\s*[::..、\)])",
86
+ r" \g<0> ",
87
+ text,
88
+ )
89
+ text = re.sub(
90
+ rf"(?<![{LATIN_LETTER}0-9])[-+]?\d+(?:[.,]\d+)*(?:%|%)?"
91
+ rf"(?![{LATIN_LETTER}0-9])",
92
+ " <NUM> ",
93
+ text,
94
+ )
95
+
96
+ text = text.replace("(", "(").replace(")", ")")
97
+ text = re.sub(r"\s+", " ", text)
98
+ return text.strip()
99
+
100
+
101
+ def latin_token_to_model_token(token: str) -> str:
102
+ """Normalize non-Mandarin alphanumeric tokens without losing option labels."""
103
+ upper = token.upper()
104
+ return upper if upper in {"A", "B", "C", "D"} else token.lower()
105
+
106
+
107
+ def should_preserve_fallback_token(token: str) -> bool:
108
+ """Return true for visible non-Hanzi letters, punctuation, and symbols."""
109
+ if token == "\ufffd":
110
+ return False
111
+ for char in token:
112
+ category = unicodedata.category(char)
113
+ if category in DISCARDED_UNICODE_CATEGORIES:
114
+ return False
115
+ if category[0] not in {"L", "P", "S"}:
116
+ return False
117
+ return True
118
+
119
+
120
+ def split_tone3_syllable(syllable: str) -> tuple[str, int]:
121
+ """Return the plain pinyin syllable and its tone number.
122
+
123
+ ``Style.TONE3`` writes tones as final digits, but neutral tone syllables have
124
+ no digit. Treat those digitless cases as fifth tone.
125
+ """
126
+ match = re.fullmatch(r"([a-züv]+)([1-5]?)", syllable.lower())
127
+ if not match:
128
+ return syllable, 5
129
+
130
+ plain, tone = match.groups()
131
+ return plain, int(tone or "5")
132
+
133
+
134
+ def length_digit_offset(syllable: str) -> int:
135
+ """Map pinyin syllable length to the requested 0-4 digit offset."""
136
+ return min(max(len(syllable), 1), 5) - 1
137
+
138
+
139
+ def syllable_to_initial_code(syllable: str) -> str:
140
+ """Convert one pinyin syllable with tone into ``initial + digit``.
141
+
142
+ Tone controls initial casing and whether the digit starts from 0 or 5:
143
+ tones 1/3/5 use uppercase initials, while tones 2/4 use lowercase initials.
144
+ Syllable length then adds the 0-4 offset that makes the final digit.
145
+ """
146
+ plain, tone = split_tone3_syllable(syllable)
147
+ if not plain:
148
+ return ""
149
+
150
+ tone_offset = 5 if tone in {3, 4, 5} else 0
151
+ digit = tone_offset + length_digit_offset(plain)
152
+ initial = plain[0].upper() if tone in {1, 3, 5} else plain[0].lower()
153
+ return f"{initial}{digit}"
154
+
155
+
156
+ def syllable_to_initial_letter(syllable: str) -> str:
157
+ """Convert one pinyin syllable with tone into its lowercase first letter."""
158
+ plain, _ = split_tone3_syllable(syllable)
159
+ return plain[:1].lower()
160
+
161
+
162
+ def chinese_word_to_initial_codes(word: str) -> str:
163
+ """Convert one already-segmented Chinese word to one compact code token.
164
+
165
+ The important representation choice is preserved: jieba decides the word
166
+ boundary, and all syllable codes inside that word are concatenated. For
167
+ example, ``我们`` becomes ``W6M7`` rather than ``W6 M7``.
168
+ """
169
+ syllables = pinyin(word, style=Style.TONE3, heteronym=False, errors="ignore")
170
+ codes = [syllable_to_initial_code(item[0]) for item in syllables if item and item[0]]
171
+ return "".join(code for code in codes if code)
172
+
173
+
174
+ def chinese_word_to_initial_letters(word: str) -> str:
175
+ """Convert one already-segmented Chinese word to lowercase pinyin initials."""
176
+ syllables = pinyin(word, style=Style.TONE3, heteronym=False, errors="ignore")
177
+ initials = [
178
+ syllable_to_initial_letter(item[0]) for item in syllables if item and item[0]
179
+ ]
180
+ return "".join(initial for initial in initials if initial)
181
+
182
+
183
+ def chinese_word_to_transliteration(word: str, transliteration: Transliteration) -> str:
184
+ """Convert one segmented Chinese word using the requested transliteration."""
185
+ if transliteration == "pinyin-code":
186
+ return chinese_word_to_initial_codes(word)
187
+ if transliteration == "pinyin-initial":
188
+ return chinese_word_to_initial_letters(word)
189
+ if transliteration == "hanzi":
190
+ return word
191
+ raise ValueError(f"Unsupported transliteration: {transliteration}")
192
+
193
+
194
+ def tokenize_chinese_span(
195
+ text: str,
196
+ transliteration: Transliteration = "pinyin-code",
197
+ use_jieba: bool = True,
198
+ ) -> Iterable[str]:
199
+ """Emit one token per jieba word or per Hanzi character."""
200
+ words = jieba.cut(text, cut_all=False) if use_jieba else text
201
+ for word in words:
202
+ word = word.strip()
203
+ if not word:
204
+ continue
205
+ if CHINESE_RE.search(word):
206
+ token = chinese_word_to_transliteration(word, transliteration)
207
+ if token:
208
+ yield token
209
+
210
+
211
+ def process_text(
212
+ text: str,
213
+ transliteration: Transliteration = "pinyin-code",
214
+ use_jieba: bool = True,
215
+ ) -> str:
216
+ """Convert one raw document string into the final space-separated token line."""
217
+ tokens: list[str] = []
218
+ for part in TOKEN_RE.findall(normalize_text(text)):
219
+ if part.startswith("<") and part.endswith(">"):
220
+ tokens.append(part)
221
+ elif CHINESE_RE.fullmatch(part):
222
+ tokens.extend(tokenize_chinese_span(part, transliteration, use_jieba))
223
+ elif part in PUNCTUATION:
224
+ tokens.append(part)
225
+ elif LATIN_ALNUM_RE.fullmatch(part):
226
+ tokens.append(latin_token_to_model_token(part))
227
+ elif part.isdigit():
228
+ tokens.append("<NUM>")
229
+ elif should_preserve_fallback_token(part):
230
+ tokens.append(part.lower())
231
+
232
+ return " ".join(tokens)
233
+
234
+
235
+ def hanzi_to_encoded(text: str, use_jieba: bool = True) -> str:
236
+ """Convert normal Hanzi/Mandarin text to the compact initial+digit encoding."""
237
+ require_dependencies()
238
+ return process_text(text, "pinyin-code", use_jieba)
239
+
240
+
241
+ def read_jsonl(path: Path) -> Iterable[dict[str, Any]]:
242
+ """Yield JSON objects from UTF-8 JSONL, tolerating a leading BOM if present."""
243
+ with path.open("r", encoding="utf-8-sig") as handle:
244
+ for line_number, line in enumerate(handle, start=1):
245
+ line = line.strip()
246
+ if not line:
247
+ continue
248
+ try:
249
+ yield json.loads(line)
250
+ except json.JSONDecodeError as exc:
251
+ raise ValueError(f"Invalid JSON on line {line_number}: {exc}") from exc
252
+
253
+
254
+ def preprocess_file(
255
+ input_path: Path,
256
+ output_path: Path,
257
+ preview_count: int,
258
+ transliteration: Transliteration = "pinyin-code",
259
+ use_jieba: bool = True,
260
+ ) -> int:
261
+ """Stream input documents to output while retaining a small preview buffer."""
262
+ output_path.parent.mkdir(parents=True, exist_ok=True)
263
+ written = 0
264
+ previews: list[tuple[str, str]] = []
265
+
266
+ with output_path.open("w", encoding="utf-8", newline="\n") as out:
267
+ for obj in read_jsonl(input_path):
268
+ text = str(obj.get("text", ""))
269
+ processed = process_text(text, transliteration, use_jieba)
270
+ out.write(processed)
271
+ out.write("\n")
272
+ written += 1
273
+
274
+ if len(previews) < preview_count:
275
+ previews.append((text, processed))
276
+
277
+ for original, processed in previews:
278
+ print("ORIGINAL:")
279
+ print(original)
280
+ print("PROCESSED:")
281
+ print(processed)
282
+ print()
283
+
284
+ return written
285
+
286
+
287
+ def parse_args() -> argparse.Namespace:
288
+ parser = argparse.ArgumentParser(
289
+ description="Convert BabyLM Mandarin JSONL text fields to model-ready tokens."
290
+ )
291
+ parser.add_argument("--input", type=Path, required=True)
292
+ parser.add_argument("--output", type=Path, required=True)
293
+ parser.add_argument("--preview", type=int, default=3)
294
+ parser.add_argument(
295
+ "--transliteration",
296
+ choices=("pinyin-code", "pinyin-initial", "hanzi"),
297
+ default="pinyin-code",
298
+ help=(
299
+ "Mandarin transliteration to emit: 'pinyin-code' keeps the original "
300
+ "tone/length code, while 'pinyin-initial' emits lowercase pinyin "
301
+ "first letters only, and 'hanzi' keeps segmented Mandarin as Hanzi."
302
+ ),
303
+ )
304
+ parser.add_argument(
305
+ "--jieba",
306
+ action=argparse.BooleanOptionalAction,
307
+ default=True,
308
+ help=(
309
+ "Use jieba word segmentation for Chinese spans. Disable with "
310
+ "--no-jieba to emit one token per Hanzi character before "
311
+ "transliteration."
312
+ ),
313
+ )
314
+ return parser.parse_args()
315
+
316
+
317
+ def main() -> None:
318
+ require_dependencies()
319
+ if hasattr(sys.stdout, "reconfigure"):
320
+ sys.stdout.reconfigure(encoding="utf-8")
321
+ args = parse_args()
322
+ count = preprocess_file(
323
+ args.input,
324
+ args.output,
325
+ args.preview,
326
+ args.transliteration,
327
+ args.jieba,
328
+ )
329
+ print(f"Wrote {count:,} processed documents to {args.output}")
330
+
331
+
332
+ if __name__ == "__main__":
333
+ main()
special_tokens_map.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": "<s>",
3
+ "eos_token": "</s>",
4
+ "pad_token": "<pad>",
5
+ "unk_token": "<unk>"
6
+ }
tokenization_pinyin_code.py ADDED
@@ -0,0 +1,402 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SentencePiece tokenizer wrapper for pinyin-code Transformers models."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import logging
6
+ import re
7
+ import shutil
8
+ import unicodedata
9
+ from pathlib import Path
10
+ from typing import Any
11
+
12
+ import sentencepiece as spm
13
+ from transformers import PreTrainedTokenizer
14
+
15
+
16
+ CHINESE_RE = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]")
17
+ CHINESE_SPAN_RE = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]+")
18
+ PINYIN_CODE_TOKEN_RE = re.compile(
19
+ r"(?<![A-Za-z0-9])[A-Za-z]\d(?:[A-Za-z]\d)*(?![A-Za-z0-9])"
20
+ )
21
+ SPECIAL_MARKER_RE = re.compile(r"<[A-Z_]+>")
22
+ PUNCTUATION = set(
23
+ "\u3002\uff0c\u3001\uff1f\uff01\uff1a\uff1b.,?!:;()[]{}<>\u300a\u300b"
24
+ "\u3010\u3011\u201c\u201d\"'\u2018\u2019\u300c\u300d\u300e\u300f"
25
+ "\u2014-~\u2026/\\"
26
+ )
27
+ LATIN_LETTER = (
28
+ r"A-Za-z\u00c0-\u00d6\u00d8-\u00f6\u00f8-\u00ff"
29
+ r"\u0100-\u017f\u0180-\u024f\u0250-\u02af"
30
+ )
31
+ LATIN_ALNUM_PATTERN = (
32
+ rf"(?:[{LATIN_LETTER}][{LATIN_LETTER}0-9]*"
33
+ rf"(?:[-_][{LATIN_LETTER}0-9]+)*|"
34
+ rf"[0-9]+[{LATIN_LETTER}][{LATIN_LETTER}0-9]*"
35
+ rf"(?:[-_][{LATIN_LETTER}0-9]+)*)"
36
+ )
37
+ LATIN_ALNUM_RE = re.compile(LATIN_ALNUM_PATTERN)
38
+ URL_RE = re.compile(r"\b(?:https?://\S*|www\.\S+)", flags=re.I)
39
+ DISCARDED_UNICODE_CATEGORIES = {"Cc", "Cf", "Co", "Cs", "Cn"}
40
+ TOKEN_RE = re.compile(
41
+ r"<[A-Z_]+>|"
42
+ r"[\u3400-\u4dbf\u4e00-\u9fff]+|"
43
+ rf"{LATIN_ALNUM_PATTERN}|"
44
+ r"\S"
45
+ )
46
+ LABELS = {
47
+ "\u9898\u5e72": "<QUESTION>",
48
+ "\u9009\u9879": "<OPTIONS>",
49
+ "\u7b54\u6848": "<ANSWER>",
50
+ "\u89e3\u6790": "<EXPLANATION>",
51
+ }
52
+ PINYIN_FORMAT_ALIASES = {
53
+ "code": "pinyin-code",
54
+ "codes": "pinyin-code",
55
+ "pinyin-code": "pinyin-code",
56
+ "initial": "pinyin-initial",
57
+ "initials": "pinyin-initial",
58
+ "pinyin-initial": "pinyin-initial",
59
+ "hanzi": "hanzi",
60
+ }
61
+
62
+
63
+ def latin_token_to_model_token(token: str) -> str:
64
+ upper = token.upper()
65
+ return upper if upper in {"A", "B", "C", "D"} else token.lower()
66
+
67
+
68
+ def should_preserve_fallback_token(token: str) -> bool:
69
+ if token == "\ufffd":
70
+ return False
71
+ for char in token:
72
+ category = unicodedata.category(char)
73
+ if category in DISCARDED_UNICODE_CATEGORIES:
74
+ return False
75
+ if category[0] not in {"L", "P", "S"}:
76
+ return False
77
+ return True
78
+
79
+
80
+ class PinyinCodeTokenizer(PreTrainedTokenizer):
81
+ """Slow tokenizer that preserves the existing SentencePiece model."""
82
+
83
+ vocab_files_names = {"vocab_file": "tokenizer.model"}
84
+ model_input_names = ["input_ids", "attention_mask"]
85
+
86
+ def __init__(
87
+ self,
88
+ vocab_file: str,
89
+ add_bos_token: bool = False,
90
+ add_eos_token: bool = False,
91
+ transliteration: str = "pinyin-code",
92
+ pinyin_format: str | None = None,
93
+ use_jieba: bool = True,
94
+ jieba: bool | None = None,
95
+ **kwargs,
96
+ ) -> None:
97
+ self.vocab_file = vocab_file
98
+ self.sp_model = spm.SentencePieceProcessor(model_file=vocab_file)
99
+ self.add_bos_token = add_bos_token
100
+ self.add_eos_token = add_eos_token
101
+ self.transliteration = self._normalize_transliteration(
102
+ pinyin_format or transliteration
103
+ )
104
+ self.use_jieba = use_jieba if jieba is None else jieba
105
+
106
+ kwargs.setdefault("unk_token", self._piece_or_none(self.sp_model.unk_id()))
107
+ kwargs.setdefault("bos_token", self._piece_or_none(self.sp_model.bos_id()))
108
+ kwargs.setdefault("eos_token", self._piece_or_none(self.sp_model.eos_id()))
109
+ kwargs.setdefault("pad_token", self._piece_or_none(self.sp_model.pad_id()))
110
+ kwargs.setdefault("transliteration", self.transliteration)
111
+ kwargs.setdefault("pinyin_format", self.transliteration)
112
+ kwargs.setdefault("use_jieba", self.use_jieba)
113
+ kwargs.setdefault("jieba", self.use_jieba)
114
+ super().__init__(**kwargs)
115
+
116
+ def _normalize_transliteration(self, value: str) -> str:
117
+ normalized = PINYIN_FORMAT_ALIASES.get(value.lower())
118
+ if normalized is None:
119
+ allowed = ", ".join(sorted(set(PINYIN_FORMAT_ALIASES.values())))
120
+ raise ValueError(f"Unsupported transliteration {value!r}; choose from {allowed}")
121
+ return normalized
122
+
123
+ def _piece_or_none(self, token_id: int) -> str | None:
124
+ if token_id is None or token_id < 0:
125
+ return None
126
+ return self.sp_model.id_to_piece(token_id)
127
+
128
+ def _looks_preprocessed(self, text: str) -> bool:
129
+ if SPECIAL_MARKER_RE.search(text):
130
+ return True
131
+ if self.transliteration == "pinyin-code" and PINYIN_CODE_TOKEN_RE.search(text):
132
+ return True
133
+ return False
134
+
135
+ def _preprocess_raw_text(self, text: str) -> str:
136
+ if not CHINESE_RE.search(text) and self._looks_preprocessed(text):
137
+ return text
138
+ try:
139
+ from preprocessing.preprocess import (
140
+ hanzi_to_encoded,
141
+ process_text,
142
+ require_dependencies,
143
+ )
144
+ except ImportError:
145
+ return self._fallback_process_text(text)
146
+
147
+ require_dependencies()
148
+ if self.transliteration == "pinyin-code":
149
+ return hanzi_to_encoded(text, self.use_jieba)
150
+ return process_text(text, self.transliteration, self.use_jieba)
151
+
152
+ def _fallback_process_text(self, text: str) -> str:
153
+ if self.use_jieba:
154
+ try:
155
+ import jieba
156
+ except ImportError as exc:
157
+ raise ImportError(
158
+ "Tokenizing raw Mandarin benchmark text with jieba segmentation "
159
+ "requires jieba. Install the model dependencies before running "
160
+ "lm_eval."
161
+ ) from exc
162
+
163
+ jieba.setLogLevel(logging.WARNING)
164
+ else:
165
+ jieba = None
166
+
167
+ if self.transliteration != "hanzi":
168
+ try:
169
+ from pypinyin import Style, pinyin
170
+ except ImportError as exc:
171
+ raise ImportError(
172
+ "Tokenizing raw Mandarin benchmark text as pinyin requires pypinyin. "
173
+ "Install the model dependencies before running lm_eval."
174
+ ) from exc
175
+
176
+
177
+ def normalize_text(value: str) -> str:
178
+ value = unicodedata.normalize("NFKC", value)
179
+ value = URL_RE.sub(" <URL> ", value)
180
+ value = re.sub(r"\$\$.*?\$\$", " <MATH> ", value, flags=re.DOTALL)
181
+ value = re.sub(r"[\uff08(]\s*[\uff09)]", " <BLANK> ", value)
182
+ for label, marker in LABELS.items():
183
+ value = re.sub(rf"{label}\s*[:\uff1a]", f" {marker} ", value)
184
+ value = re.sub(
185
+ rf"(?<![{LATIN_LETTER}])yes(?![{LATIN_LETTER}])",
186
+ " <YES> ",
187
+ value,
188
+ flags=re.I,
189
+ )
190
+ value = re.sub(
191
+ rf"(?<![{LATIN_LETTER}])no(?![{LATIN_LETTER}])",
192
+ " <NO> ",
193
+ value,
194
+ flags=re.I,
195
+ )
196
+ value = re.sub(
197
+ rf"(?<![{LATIN_LETTER}])[ABCD](?=\s*[:\uff1a.\uff0e\u3001\)])",
198
+ r" \g<0> ",
199
+ value,
200
+ )
201
+ value = re.sub(
202
+ rf"(?<![{LATIN_LETTER}0-9])[-+]?\d+(?:[.,]\d+)*(?:%|\uff05)?"
203
+ rf"(?![{LATIN_LETTER}0-9])",
204
+ " <NUM> ",
205
+ value,
206
+ )
207
+ value = value.replace("\uff08", "(").replace("\uff09", ")")
208
+ return re.sub(r"\s+", " ", value).strip()
209
+
210
+ def split_tone3_syllable(syllable: str) -> tuple[str, int]:
211
+ match = re.fullmatch(r"([a-z\u00fcv]+)([1-5]?)", syllable.lower())
212
+ if not match:
213
+ return syllable, 5
214
+ plain, tone = match.groups()
215
+ return plain, int(tone or "5")
216
+
217
+ def length_digit_offset(syllable: str) -> int:
218
+ return min(max(len(syllable), 1), 5) - 1
219
+
220
+ def syllable_to_initial_code(syllable: str) -> str:
221
+ plain, tone = split_tone3_syllable(syllable)
222
+ if not plain:
223
+ return ""
224
+ tone_offset = 5 if tone in {3, 4, 5} else 0
225
+ digit = tone_offset + length_digit_offset(plain)
226
+ initial = plain[0].upper() if tone in {1, 3, 5} else plain[0].lower()
227
+ return f"{initial}{digit}"
228
+
229
+ def syllable_to_initial_letter(syllable: str) -> str:
230
+ plain, _ = split_tone3_syllable(syllable)
231
+ return plain[:1].lower()
232
+
233
+ def convert_word(word: str) -> str:
234
+ if self.transliteration == "hanzi":
235
+ return word
236
+ syllables = pinyin(word, style=Style.TONE3, heteronym=False, errors="ignore")
237
+ if self.transliteration == "pinyin-code":
238
+ codes = [
239
+ syllable_to_initial_code(item[0])
240
+ for item in syllables
241
+ if item and item[0]
242
+ ]
243
+ return "".join(code for code in codes if code)
244
+ initials = [
245
+ syllable_to_initial_letter(item[0])
246
+ for item in syllables
247
+ if item and item[0]
248
+ ]
249
+ return "".join(initial for initial in initials if initial)
250
+
251
+ def tokenize_chinese_span(value: str) -> list[str]:
252
+ tokens = []
253
+ words = jieba.cut(value, cut_all=False) if self.use_jieba else value
254
+ for word in words:
255
+ word = word.strip()
256
+ if word and CHINESE_SPAN_RE.search(word):
257
+ token = convert_word(word)
258
+ if token:
259
+ tokens.append(token)
260
+ return tokens
261
+
262
+ tokens = []
263
+ for part in TOKEN_RE.findall(normalize_text(text)):
264
+ if part.startswith("<") and part.endswith(">"):
265
+ tokens.append(part)
266
+ elif CHINESE_SPAN_RE.fullmatch(part):
267
+ tokens.extend(tokenize_chinese_span(part))
268
+ elif part in PUNCTUATION:
269
+ tokens.append(part)
270
+ elif LATIN_ALNUM_RE.fullmatch(part):
271
+ tokens.append(latin_token_to_model_token(part))
272
+ elif part.isdigit():
273
+ tokens.append("<NUM>")
274
+ elif should_preserve_fallback_token(part):
275
+ tokens.append(part.lower())
276
+
277
+ return " ".join(tokens)
278
+
279
+ def _preprocess_tokenizer_input(self, value: Any) -> Any:
280
+ if value is None:
281
+ return None
282
+ if isinstance(value, str):
283
+ return self._preprocess_raw_text(value)
284
+ if isinstance(value, tuple):
285
+ return tuple(self._preprocess_tokenizer_input(item) for item in value)
286
+ if isinstance(value, list):
287
+ return [self._preprocess_tokenizer_input(item) for item in value]
288
+ return value
289
+
290
+ def __call__(self, text=None, text_pair=None, *args, **kwargs):
291
+ if "text_target" in kwargs:
292
+ kwargs["text_target"] = self._preprocess_tokenizer_input(kwargs["text_target"])
293
+ if "text_pair_target" in kwargs:
294
+ kwargs["text_pair_target"] = self._preprocess_tokenizer_input(
295
+ kwargs["text_pair_target"]
296
+ )
297
+
298
+ text = self._preprocess_tokenizer_input(text)
299
+ text_pair = self._preprocess_tokenizer_input(text_pair)
300
+ if text_pair is None:
301
+ return super().__call__(text, *args, **kwargs)
302
+ return super().__call__(text, text_pair, *args, **kwargs)
303
+
304
+ def encode(self, text, text_pair=None, add_special_tokens=True, *args, **kwargs):
305
+ kwargs["add_special_tokens"] = add_special_tokens
306
+ text = self._preprocess_tokenizer_input(text)
307
+ text_pair = self._preprocess_tokenizer_input(text_pair)
308
+ if text_pair is None:
309
+ return super().encode(text, *args, **kwargs)
310
+ return super().encode(text, text_pair, *args, **kwargs)
311
+
312
+ def encode_plus(self, text, text_pair=None, *args, **kwargs):
313
+ text = self._preprocess_tokenizer_input(text)
314
+ text_pair = self._preprocess_tokenizer_input(text_pair)
315
+ if text_pair is None:
316
+ return super().encode_plus(text, *args, **kwargs)
317
+ return super().encode_plus(text, text_pair, *args, **kwargs)
318
+
319
+ def batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs):
320
+ batch_text_or_text_pairs = self._preprocess_tokenizer_input(
321
+ batch_text_or_text_pairs
322
+ )
323
+ return super().batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs)
324
+
325
+ @property
326
+ def vocab_size(self) -> int:
327
+ return self.sp_model.get_piece_size()
328
+
329
+ def get_vocab(self) -> dict[str, int]:
330
+ vocab = {self.sp_model.id_to_piece(i): i for i in range(self.vocab_size)}
331
+ vocab.update(self.added_tokens_encoder)
332
+ return vocab
333
+
334
+ def _tokenize(self, text: str) -> list[str]:
335
+ text = self._preprocess_raw_text(text)
336
+ return self.sp_model.encode(text, out_type=str)
337
+
338
+ def _convert_token_to_id(self, token: str) -> int:
339
+ return self.sp_model.piece_to_id(token)
340
+
341
+ def _convert_id_to_token(self, index: int) -> str:
342
+ return self.sp_model.id_to_piece(index)
343
+
344
+ def convert_tokens_to_string(self, tokens: list[str]) -> str:
345
+ return self.sp_model.decode(tokens)
346
+
347
+ def build_inputs_with_special_tokens(
348
+ self,
349
+ token_ids_0: list[int],
350
+ token_ids_1: list[int] | None = None,
351
+ ) -> list[int]:
352
+ output = list(token_ids_0)
353
+ if self.add_bos_token and self.bos_token_id is not None:
354
+ output = [self.bos_token_id] + output
355
+ if self.add_eos_token and self.eos_token_id is not None:
356
+ output = output + [self.eos_token_id]
357
+ if token_ids_1 is not None:
358
+ output += list(token_ids_1)
359
+ if self.add_eos_token and self.eos_token_id is not None:
360
+ output.append(self.eos_token_id)
361
+ return output
362
+
363
+ def get_special_tokens_mask(
364
+ self,
365
+ token_ids_0: list[int],
366
+ token_ids_1: list[int] | None = None,
367
+ already_has_special_tokens: bool = False,
368
+ ) -> list[int]:
369
+ if already_has_special_tokens:
370
+ special_ids = set(self.all_special_ids)
371
+ return [1 if token_id in special_ids else 0 for token_id in token_ids_0]
372
+
373
+ mask = [0] * len(token_ids_0)
374
+ if self.add_bos_token and self.bos_token_id is not None:
375
+ mask = [1] + mask
376
+ if self.add_eos_token and self.eos_token_id is not None:
377
+ mask = mask + [1]
378
+ if token_ids_1 is not None:
379
+ mask += [0] * len(token_ids_1)
380
+ if self.add_eos_token and self.eos_token_id is not None:
381
+ mask.append(1)
382
+ return mask
383
+
384
+ def create_token_type_ids_from_sequences(
385
+ self,
386
+ token_ids_0: list[int],
387
+ token_ids_1: list[int] | None = None,
388
+ ) -> list[int]:
389
+ return [0] * len(self.build_inputs_with_special_tokens(token_ids_0, token_ids_1))
390
+
391
+ def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
392
+ output_name = "tokenizer.model"
393
+ if filename_prefix:
394
+ output_name = f"{filename_prefix}-{output_name}"
395
+ output_path = Path(save_directory) / output_name
396
+ if Path(self.vocab_file).resolve() != output_path.resolve():
397
+ shutil.copyfile(self.vocab_file, output_path)
398
+ return (str(output_path),)
399
+
400
+
401
+ class EncodedMandarinTokenizer(PinyinCodeTokenizer):
402
+ """Tokenizer wrapper that hides Hanzi-to-encoded-Mandarin preprocessing."""
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:60cf469adc184103e3c73309b971a667d4ec04211cf33527e0341716ccd3f5ee
3
+ size 285315
tokenizer_config.json ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "added_tokens_decoder": {
3
+ "0": {
4
+ "content": "<pad>",
5
+ "lstrip": false,
6
+ "normalized": false,
7
+ "rstrip": false,
8
+ "single_word": false,
9
+ "special": true
10
+ },
11
+ "1": {
12
+ "content": "<unk>",
13
+ "lstrip": false,
14
+ "normalized": false,
15
+ "rstrip": false,
16
+ "single_word": false,
17
+ "special": true
18
+ },
19
+ "2": {
20
+ "content": "<s>",
21
+ "lstrip": false,
22
+ "normalized": false,
23
+ "rstrip": false,
24
+ "single_word": false,
25
+ "special": true
26
+ },
27
+ "3": {
28
+ "content": "</s>",
29
+ "lstrip": false,
30
+ "normalized": false,
31
+ "rstrip": false,
32
+ "single_word": false,
33
+ "special": true
34
+ }
35
+ },
36
+ "auto_map": {
37
+ "AutoTokenizer": [
38
+ "tokenization_pinyin_code.EncodedMandarinTokenizer",
39
+ null
40
+ ]
41
+ },
42
+ "backend": "custom",
43
+ "bos_token": "<s>",
44
+ "eos_token": "</s>",
45
+ "jieba": true,
46
+ "model_max_length": 512,
47
+ "pad_token": "<pad>",
48
+ "pinyin_format": "pinyin-code",
49
+ "tokenizer_class": "EncodedMandarinTokenizer",
50
+ "transliteration": "pinyin-code",
51
+ "unk_token": "<unk>",
52
+ "use_jieba": true
53
+ }
training_metadata.json ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "epoch": 5,
3
+ "evaluation_backend": "causal",
4
+ "global_step": 11159,
5
+ "jieba": true,
6
+ "source_checkpoint": "models\\full_chinese_gpu3.1\\best.pt",
7
+ "transliteration": "pinyin-code",
8
+ "use_jieba": true,
9
+ "validation_loss": 5.517338045504915
10
+ }