File size: 1,506 Bytes
295d147 | 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 | import argparse
from pathlib import Path
from tokenizers import Tokenizer
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from tokenizers.models import BPE
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.trainers import BpeTrainer
SPECIAL_TOKENS = [
"<|pad|>",
"<|unk|>",
"<|bos|>",
"<|eos|>",
"<|user|>",
"<|assistant|>",
"<|system|>",
"<|godot|>",
"<|gdscript|>",
"<|code|>",
"<|error|>",
"<|patch|>",
]
def collect_files():
files = [str(p) for p in Path("data/clean").glob("*.txt")]
files += [str(p) for p in Path("data/instructions").glob("*.jsonl")]
return files
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--vocab-size", type=int, default=16000)
parser.add_argument("--out", default="tokenizer/tokenizer.json")
args = parser.parse_args()
files = collect_files()
if not files:
raise SystemExit("No training files found in data/clean or data/instructions")
tokenizer = Tokenizer(BPE(unk_token="<|unk|>"))
tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
tokenizer.decoder = ByteLevelDecoder()
trainer = BpeTrainer(
vocab_size=args.vocab_size,
min_frequency=2,
special_tokens=SPECIAL_TOKENS,
)
tokenizer.train(files, trainer)
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
tokenizer.save(args.out)
print(f"Saved {args.out}")
if __name__ == "__main__":
main()
|