save_wmt_to_disk.py

# by stas00 and sshleifer
import nlp
from tqdm import tqdm

dataset = 'wmt19'
s = 'ru'
t = 'en'

pair = f'{s}-{t}'

ds = nlp.load_dataset(dataset, pair)
save_dir = Path(f'{dataset}-{pair}')
save_dir.mkdir(exist_ok=True)
for split in tqdm(ds.keys()):
    tr_list = list(ds[split])
    data = [x['translation'] for x in tr_list]
    src, tgt = [], []
    for example in data:
        src.append(example[s])
        tgt.append(example[t])
    if split == 'validation':
        split = 'val' # to save to val.source, val.target like summary datasets
    src_path = save_dir.joinpath(f'{split}.source')
    src_path.open('w+').write('\n'.join(src))
    tgt_path = save_dir.joinpath(f'{split}.target')
    tgt_path.open('w+').write('\n'.join(tgt))

To the future reader: You will find the latest version of this script here: https://github.com/huggingface/transformers/blob/master/examples/seq2seq/download_wmt.py