File size: 5,811 Bytes
b3c2a26 5a98e33 b3c2a26 5a98e33 b3c2a26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | """Train and package the KorByte-128K tokenizer."""
from __future__ import annotations
import hashlib
import importlib.metadata
import json
import platform
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from tokenizers import Regex, Tokenizer, decoders, pre_tokenizers, processors
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from .config import (
CORE_VOCAB_SIZE,
DEFAULT_MODEL_MAX_LENGTH,
TOTAL_VOCAB_SIZE,
Paths,
)
from .special_tokens import added_special_tokens, special_token_strings
# Keep a leading ASCII space with the following unit, but do not bind arbitrary
# punctuation to a word. Six-digit number chunks avoid needless fragmentation
# of dates, standards, and common Korean financial notation.
PRETOKENIZER_PATTERN = (
r" ?[\p{L}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|"
r" ?\p{N}{1,6}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|"
r"\s+(?!\S)|\s+"
)
def _sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def build_tokenizer() -> Tokenizer:
"""Construct an untrained, exactly reversible byte-level BPE pipeline."""
tokenizer = Tokenizer(BPE(unk_token=None, byte_fallback=False, fuse_unk=False))
tokenizer.normalizer = None
tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
[
pre_tokenizers.Split(Regex(PRETOKENIZER_PATTERN), behavior="isolated", invert=False),
pre_tokenizers.ByteLevel(
add_prefix_space=False,
use_regex=False,
trim_offsets=True,
),
]
)
tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
tokenizer.decoder = decoders.ByteLevel()
return tokenizer
def _compatibility_files(tokenizer: Tokenizer, root: Path) -> None:
tokens = special_token_strings()
vocab = tokenizer.get_vocab(with_added_tokens=True)
additional = list(tokens[4:])
tokenizer_config: dict[str, Any] = {
"add_bos_token": False,
"add_eos_token": False,
"additional_special_tokens": additional,
"bos_token": tokens[0],
"clean_up_tokenization_spaces": False,
"eos_token": tokens[1],
"model_max_length": DEFAULT_MODEL_MAX_LENGTH,
"pad_token": tokens[2],
"tokenizer_class": "PreTrainedTokenizerFast",
"unk_token": tokens[3],
}
special_tokens_map = {
"bos_token": tokens[0],
"eos_token": tokens[1],
"pad_token": tokens[2],
"unk_token": tokens[3],
"additional_special_tokens": additional,
}
added_tokens_map = {token: vocab[token] for token in tokens}
(root / "tokenizer_config.json").write_text(
json.dumps(tokenizer_config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
)
(root / "special_tokens_map.json").write_text(
json.dumps(special_tokens_map, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
)
(root / "added_tokens.json").write_text(
json.dumps(added_tokens_map, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
)
def train_tokenizer(root: Path, *, min_frequency: int = 3) -> dict[str, Any]:
"""Train from the prepared corpus and write a complete Hub-compatible snapshot."""
paths = Paths(root)
if not paths.corpus.is_file():
raise FileNotFoundError(f"Prepared corpus not found: {paths.corpus}")
if not paths.corpus_manifest.is_file():
raise FileNotFoundError(f"Corpus manifest not found: {paths.corpus_manifest}")
tokenizer = build_tokenizer()
trainer = BpeTrainer(
vocab_size=CORE_VOCAB_SIZE,
min_frequency=min_frequency,
show_progress=True,
special_tokens=[],
initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
max_token_length=64,
)
tokenizer.train([str(paths.corpus)], trainer)
core_size = tokenizer.get_vocab_size(with_added_tokens=False)
if core_size != CORE_VOCAB_SIZE:
raise RuntimeError(f"Expected {CORE_VOCAB_SIZE} core tokens, trained {core_size}")
added_count = tokenizer.add_special_tokens(added_special_tokens())
if added_count != TOTAL_VOCAB_SIZE - CORE_VOCAB_SIZE:
raise RuntimeError(f"Expected 256 added tokens, got {added_count}")
total_size = tokenizer.get_vocab_size(with_added_tokens=True)
if total_size != TOTAL_VOCAB_SIZE:
raise RuntimeError(f"Expected {TOTAL_VOCAB_SIZE} total tokens, got {total_size}")
tokenizer.save(str(root / "tokenizer.json"), pretty=True)
tokenizer.model.save(str(root))
_compatibility_files(tokenizer, root)
corpus_manifest = json.loads(paths.corpus_manifest.read_text(encoding="utf-8"))
build_manifest = {
"schema_version": 1,
"created_at": datetime.now(UTC).isoformat(),
"algorithm": "byte-level-bpe",
"normalizer": None,
"pretokenizer_pattern": PRETOKENIZER_PATTERN,
"core_vocab_size": core_size,
"special_token_count": added_count,
"total_vocab_size": total_size,
"min_frequency": min_frequency,
"max_token_length": 64,
"corpus_sha256": corpus_manifest["sha256"],
"tokenizer_sha256": _sha256(root / "tokenizer.json"),
"python": platform.python_version(),
"platform": platform.platform(),
"packages": {
package: importlib.metadata.version(package)
for package in ("datasets", "huggingface-hub", "tokenizers", "transformers")
},
}
paths.build_manifest.write_text(
json.dumps(build_manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
)
return build_manifest
|