tokengeex 1.1.0

TokenGeeX is an efficient tokenizer for code based on UnigramLM and TokenMonster.
Documentation 
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

"""
This script trains a baseline tokenizer model using the HuggingFace or
SentencePiece library. It is used to compare the performance of TokenGeeX with
other tokenization libraries.
"""

import argparse
import glob
from typing import Iterable

import sentencepiece as spm
from tokenizers import Tokenizer, models, trainers
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from tokenizers.pre_tokenizers import ByteLevel as ByteLevelEncoder


def load_samples(input: str, proportion: float) -> Iterable[str]:
    files = glob.glob(f"{input}/*.bin")
    print(f"Found {len(files)} .bin files in {input}")
    # Each .bin file is a 0x00 separated list of UTF-8 samples.
    for file in files:
        with open(file, "rb") as f:
            samples = f.read().decode("utf-8").split("\0")
            print(f"Loaded {len(samples)} samples from {file}")
            yield from samples[: int(len(samples) * proportion)]


def train_huggingface(samples: Iterable[str], output: str, vocab_size: int):
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        max_token_length=32,
        show_progress=True,
    )  # type: ignore

    tokenizer = Tokenizer(
        model=models.BPE(
            byte_fallback=True,
            fuse_unk=False,
            unk_token=None,
        )
    )
    tokenizer.pre_tokenizer = ByteLevelEncoder(
        add_prefix_space=False,
    )  # type: ignore
    tokenizer.decoder = ByteLevelDecoder()  # type: ignore

    tokenizer.train_from_iterator(samples, trainer=trainer)

    tokenizer.save(output)


def train_sentencepiece(samples: Iterable[str], output: str, vocab_size: int):
    spm.SentencePieceTrainer.Train(
        sentence_iterator=samples,
        vocab_size=vocab_size,
        model_prefix=output,
        model_type="bpe",
        input_sentence_size=10000000,
        shuffle_input_sentence=True,
        split_digits=True,
        byte_fallback=True,
        character_coverage=0.999,
        max_sentencepiece_length=32,
        add_dummy_prefix=False,
        allow_whitespace_only_pieces=True,
        split_by_whitespace=False,
        remove_extra_whitespaces=False,
        normalization_rule_name="identity",
        train_extremely_large_corpus=True,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train a BPE tokenizer model.")
    parser.add_argument(
        "-l",
        required=True,
        type=str,
        choices=["huggingface", "sentencepiece"],
        help="The library to use for training the tokenizer.",
    )
    parser.add_argument(
        "-i",
        required=True,
        type=str,
        help="The directory containing the training data.",
    )
    parser.add_argument(
        "-o",
        required=True,
        type=str,
        help="The file to save the trained model to.",
    )
    parser.add_argument(
        "-v",
        required=True,
        type=int,
        help="The size of the vocabulary to use.",
    )
    parser.add_argument(
        "-p",
        default=1.0,
        type=float,
        help="Proportion of samples to use for training (between 0 and 1).",
    )

    args = parser.parse_args()

    lib, input, output, vocab_size, proportion = args.l, args.i, args.o, args.v, args.s

    print(
        f"Training {lib} BPE model with {vocab_size} vocabulary size from {input}. Writing to {output}."
    )

    samples = load_samples(input, proportion)

    if lib == "huggingface":
        train_huggingface(samples, output, vocab_size)
    elif lib == "sentencepiece":
        train_sentencepiece(samples, output, vocab_size)
    else:
        raise ValueError(f"Unknown library: {lib}")