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()