ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: 修复特殊token 问题

Browse Source
This commit is contained in:
ViperEkura
2026-04-02 16:16:02 +08:00
parent e01ec081b3
commit f44ad6912e
28 changed files with 334 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pipeline/tokenizer.py
+87 -44
View File
@@ -7,100 +7,143 @@ from typing import List, Union, Optional, Tuple, Iterator
class BpeTokenizer:
def __init__(self, path: Optional[str] = None):
self._control_tokens = ["<bos>", "<eos>", "<pad>"]
self._special_tokens = ["<|im_start|>", "<|im_end|>"]
self._control_tokens = [
"<｜begin▁of▁sentence｜>",
"<｜end▁of▁sentence｜>",
"<｜▁pad▁｜>",
]
self._special_tokens = ["<｜im▁start｜>", "<｜im▁end｜>"]
model = BPE()
self._tokenizer = Tokenizer(model)
self._tokenizer.normalizer = normalizers.Sequence([
normalizers.NFC(),
normalizers.Strip()
])
self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
pre_tokenizers.UnicodeScripts(),
pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True)
])
self._tokenizer.normalizer = normalizers.Sequence(
[normalizers.NFC(), normalizers.Strip()]
)
self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
[
pre_tokenizers.UnicodeScripts(),
pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True),
]
)
self._tokenizer.decoder = decoders.ByteLevel()
self._tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
if path is not None:
self._tokenizer = Tokenizer.from_file(path)
def _prepare_trainer(self, vocab_size: int, min_freq: int, reserved_token_size: int, max_token_length: int = 18) -> Tuple[BpeTrainer, int, List[str]]:
def _prepare_trainer(
self,
vocab_size: int,
min_freq: int,
reserved_token_size: int,
max_token_length: int = 18,
) -> Tuple[BpeTrainer, int, List[str]]:
assert reserved_token_size > len(self._special_tokens)
reserved_tokens = [f"<|reserve{i:02d}|>" for i in range(reserved_token_size - len(self._special_tokens))]
detail_vocab_size = vocab_size - (len(reserved_tokens) + len(self._special_tokens))
reserved_tokens = [
f"<｜reserve{i:02d}｜>"
for i in range(reserved_token_size - len(self._special_tokens))
]
detail_vocab_size = vocab_size - (
len(reserved_tokens) + len(self._special_tokens)
)
alphabet = pre_tokenizers.ByteLevel.alphabet()
min_size = len(alphabet) + len(self._control_tokens)
assert detail_vocab_size > min_size
trainer = BpeTrainer(
vocab_size=detail_vocab_size,
min_frequency=min_freq,
limit_alphabet=detail_vocab_size // 6,
max_token_length=max_token_length,
special_tokens=self._control_tokens,
special_tokens=self._control_tokens + self._special_tokens,
initial_alphabet=alphabet,
show_progress=True,
)
return trainer, detail_vocab_size, reserved_tokens
def train(self, files: List[str], vocab_size: int, min_freq: int, reserved_token_size: int = 100) -> None:
def train(
self,
files: List[str],
vocab_size: int,
min_freq: int,
reserved_token_size: int = 100,
) -> None:
trainer, _, reserved_tokens = self._prepare_trainer(
vocab_size=vocab_size,
min_freq=min_freq,
reserved_token_size=reserved_token_size
reserved_token_size=reserved_token_size,
)
self._tokenizer.train(files=files, trainer=trainer)
self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens)
def train_from_iterator(self, iterator: Iterator[str], vocab_size: int, min_freq: int, reserved_token_size: int = 100) -> None:
self._tokenizer.add_special_tokens(
self._control_tokens + self._special_tokens + reserved_tokens
)
def train_from_iterator(
self,
iterator: Iterator[str],
vocab_size: int,
min_freq: int,
reserved_token_size: int = 100,
) -> None:
trainer, _, reserved_tokens = self._prepare_trainer(
vocab_size=vocab_size,
min_freq=min_freq,
reserved_token_size=reserved_token_size
reserved_token_size=reserved_token_size,
)
self._tokenizer.train_from_iterator(iterator=iterator, trainer=trainer)
self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens)
self._tokenizer.add_special_tokens(
self._control_tokens + self._special_tokens + reserved_tokens
)
def save(self, path: str) -> None:
self._tokenizer.save(path)
def load(self, path: str) -> None:
self._tokenizer = Tokenizer.from_file(path)
def encode(self, tokens: Union[str, List[str]], out_ids: bool = True, add_special_tokens: bool = False) -> Union[List[int], List[str], List[List[int]], List[List[str]]]:
def encode(
self,
tokens: Union[str, List[str]],
out_ids: bool = True,
add_special_tokens: bool = False,
) -> Union[List[int], List[str], List[List[int]], List[List[str]]]:
if isinstance(tokens, str):
encoded: Encoding = self._tokenizer.encode(tokens, add_special_tokens=add_special_tokens)
encoded: Encoding = self._tokenizer.encode(
tokens, add_special_tokens=add_special_tokens
)
return encoded.ids if out_ids else encoded.tokens
elif isinstance(tokens, list):
encoded_list: List[Encoding] = self._tokenizer.encode_batch(tokens, add_special_tokens=add_special_tokens)
return [encoded.ids if out_ids else encoded.tokens for encoded in encoded_list]
encoded_list: List[Encoding] = self._tokenizer.encode_batch(
tokens, add_special_tokens=add_special_tokens
)
return [
encoded.ids if out_ids else encoded.tokens for encoded in encoded_list
]
def decode(self, tokens: List[int], skip_special_tokens: bool=True) -> str:
def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
def __len__(self) -> int:
return self._tokenizer.get_vocab_size()
@property
def stop_ids(self) -> List[int]:
stop_token = self._control_tokens + self._special_tokens
stop_ids = [self._tokenizer.token_to_id(token) for token in stop_token]
return stop_ids
@property
def bos_id(self) -> int:
return self._tokenizer.token_to_id("<bos>")
return self._tokenizer.token_to_id("<｜begin▁of▁sentence｜>")
@property
def eos_id(self) -> int:
return self._tokenizer.token_to_id("<eos>")
return self._tokenizer.token_to_id("<｜end▁of▁sentence｜>")
@property
def pad_id(self) -> int:
return self._tokenizer.token_to_id("<pad>")
return self._tokenizer.token_to_id("<｜▁pad▁｜>")