File size: 2,165 Bytes
0a83496 | 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 | from pathlib import Path
from datasets import interleave_datasets, load_dataset
from tokenizers import ByteLevelBPETokenizer
from tqdm import tqdm
NUM_DOCUMENTS = 500_000
VOCAB_SIZE = 32_000
SEED = 42
OUTPUT_FILE = Path(__file__).with_name("hanse_tokenizer.json")
SOURCES = (
("HuggingFaceFW/fineweb-edu", "sample-10BT", 0.40),
("HuggingFaceFW/fineweb-2", "deu_Latn", 0.35),
("HuggingFaceFW/finewiki", "de", 0.20),
("HuggingFaceFW/finewiki", "en", 0.05),
)
SPECIAL_TOKENS = (
"<|pad|>",
"<|bos|>",
"<|eos|>",
"<|unk|>",
"<|system|>",
"<|user|>",
"<|assistant|>",
"<|tool|>",
"<|tool_result|>",
"<|end_of_turn|>",
)
def training_corpus(dataset):
accepted = 0
with tqdm(total=NUM_DOCUMENTS, desc="Feeding documents", unit="docs", dynamic_ncols=True) as progress:
for example in dataset:
text = example.get("text")
if not isinstance(text, str) or not (text := text.strip()):
continue
yield text
accepted += 1
progress.update()
if accepted == NUM_DOCUMENTS:
return
raise RuntimeError(f"Dataset exhausted after {accepted:,} usable documents")
def main() -> None:
streams = [
load_dataset(name, config, split="train", streaming=True)
for name, config, _ in SOURCES
]
dataset = interleave_datasets(
streams,
probabilities=[probability for _, _, probability in SOURCES],
seed=SEED,
stopping_strategy="all_exhausted",
)
tokenizer = ByteLevelBPETokenizer()
tokenizer.train_from_iterator(
training_corpus(dataset),
vocab_size=VOCAB_SIZE,
min_frequency=2,
special_tokens=list(SPECIAL_TOKENS),
length=NUM_DOCUMENTS,
show_progress=True,
)
assert tokenizer.get_vocab_size() == VOCAB_SIZE
assert [tokenizer.token_to_id(token) for token in SPECIAL_TOKENS] == list(range(len(SPECIAL_TOKENS)))
tokenizer.save(str(OUTPUT_FILE))
print(f"[+] Saved {tokenizer.get_vocab_size():,}-token tokenizer to {OUTPUT_FILE}")
if __name__ == "__main__":
main()
|