usrbinenv python 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license download a dataset using the datasets package and save it to the format expected by finetune py format of savedir train source train target val source val target test source test target args srclang str source language tgtlang str target language dataset str wmt16 wmt17 etc wmt16 is a good start as it s small to get the full list run import datasets printd id for d in datasets listdatasets if wmt in d id savedir str where to save the datasets defaults to f datasetsrclangtgtlang usage downloadwmtdataset ro en dataset wmt16 saves to wmt16roen to save to val source val target like summary datasets reader is the bottleneck so writing one record at a time doesn t slow things down usr bin env python 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license download a dataset using the datasets package and save it to the format expected by finetune py format of save_dir train source train target val source val target test source test target args src_lang str source language tgt_lang str target language dataset str wmt16 wmt17 etc wmt16 is a good start as it s small to get the full list run import datasets print d id for d in datasets list_datasets if wmt in d id save_dir str where to save the datasets defaults to f dataset src_lang tgt_lang usage download_wmt_dataset ro en dataset wmt16 saves to wmt16 ro en to save to val source val target like summary datasets reader is the bottleneck so writing one record at a time doesn t slow things down
from pathlib import Path

import fire
from tqdm import tqdm


def download_wmt_dataset(src_lang="ro", tgt_lang="en", dataset="wmt16", save_dir=None) -> None:
    
    try:
        import datasets
    except (ModuleNotFoundError, ImportError):
        raise ImportError("run pip install datasets")
    pair = f"{src_lang}-{tgt_lang}"
    print(f"Converting {dataset}-{pair}")
    ds = datasets.load_dataset(dataset, pair)
    if save_dir is None:
        save_dir = f"{dataset}-{pair}"
    save_dir = Path(save_dir)
    save_dir.mkdir(exist_ok=True)

    for split in ds.keys():
        print(f"Splitting {split} with {ds[split].num_rows} records")

        
        fn = "val" if split == "validation" else split
        src_path = save_dir.joinpath(f"{fn}.source")
        tgt_path = save_dir.joinpath(f"{fn}.target")
        src_fp = src_path.open("w+")
        tgt_fp = tgt_path.open("w+")

        
        for x in tqdm(ds[split]):
            ex = x["translation"]
            src_fp.write(ex[src_lang] + "\n")
            tgt_fp.write(ex[tgt_lang] + "\n")

    print(f"Saved {dataset} dataset to {save_dir}")


if __name__ == "__main__":
    fire.Fire(download_wmt_dataset)