数据集引入

安装依赖

pip install transformers datasets sentencepiece

数据预处理

from datasets import load_dataset
from utils import prepare_training_data
import sentencepiece as spm
import os

# 需要改成项目路径
os.chdir('your disk position') 
data_dir = 'data'
if not os.path.exists(data_dir):
    os.makedirs(data_dir)
# 加载WMT 2014 英语-德语数据集
dataset = load_dataset('wmt14', 'de-en')
prepare_training_data(dataset, 'WMT_de_en')
# 生成 BPM token
spm.SentencePieceTrainer.Train(
    '--input=data/WMT_de_en_train.txt '
    '--model_prefix=bpe '
    '--vocab_size=37000 '
    '--model_type=bpe '
    '--bos_id=1 --eos_id=2 --unk_id=3 --pad_id=0 '
    '--character_coverage=0.9995 '
    '--shuffle_input_sentence=true'
)
# 最后将生成的BPE文件放到data文件夹下