genome_book / tokenizer /train_open_bpe.py
marisming's picture
Upload folder using huggingface_hub
c9c7b01 verified
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.pre_tokenizers import Whitespace
from tokenizers.trainers import BpeTrainer
def train_tokenizer_from_file(input_file_list, vocab_size=50257, output_file="tokenizer.json"):
"""
从文本文件训练分词器。
参数:
- input_file (str): 输入文本文件路径。
- vocab_size (int): 分词器的词汇表大小。
- output_file (str): 训练好的分词器保存路径。
"""
# 初始化分词器
tokenizer = Tokenizer(BPE())
tokenizer.pre_tokenizer = Whitespace()
# 定义 BPE 训练器
trainer = BpeTrainer(
vocab_size=vocab_size,
special_tokens=["<|endoftext|>", "<pad>", "<unk>"]
)
# 开始训练
tokenizer.train(input_file_list, trainer)
tokenizer.save(output_file)
print(f"分词器已保存到 {output_file}")
# 示例:从生成的文本文件训练分词器
train_tokenizer_from_file(["protein_4g.txt","dna_4g.txt","eng_4g.txt"], vocab_size=100000, output_file="gpt2_gene_eng_tokenizer.json")