| import sentencepiece as spm, os | |
| def train_spm(input_path="data/corpus.txt", model_prefix="tokens/oai_spm", vocab_size=32000): | |
| assert os.path.exists(input_path), f"Missing {input_path}" | |
| spm.SentencePieceTrainer.Train( | |
| input=input_path, | |
| model_prefix=model_prefix, | |
| model_type="unigram", | |
| vocab_size=vocab_size, | |
| character_coverage=0.9995, | |
| input_sentence_size=1000000, | |
| shuffle_input_sentence=True, | |
| unk_id=0, bos_id=1, eos_id=2, pad_id=3 | |
| ) | |
| print("Tokenizer trained:", model_prefix+".model") | |
| if __name__ == "__main__": | |
| train_spm() | |