61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import pytorch_lightning as pl
|
|
import torch
|
|
|
|
from .finetune import SummarizationModule, main
|
|
from .test_seq2seq_examples import make_test_data_dir, MBART_TINY, CUDA_AVAILABLE
|
|
import wget
|
|
import tarfile
|
|
WMT_100_URL = 'asd'
|
|
DATA_URL = 'https://s3.amazonaws.com/datasets.huggingface.co/translation/wmt_en_ro_test_data.tgz'
|
|
def fetch_and_save_wmt_100():
|
|
dest_dir = 'wmt_100_dir'
|
|
if os.path.isdir(dest_dir):
|
|
return dest_dir
|
|
filename = wget.download(DATA_URL)
|
|
tarball = tarfile.TarFile(filename)
|
|
tarball.extractall(path=dest_dir)
|
|
os.remove(filename)
|
|
return dest_dir
|
|
|
|
|
|
|
|
def test_train_mbart_cc25_enro_script():
|
|
data_dir = fetch_and_save_wmt_100()
|
|
env_vars_to_replace = {
|
|
"$MAX_LEN": 200,
|
|
"$BS": 4,
|
|
"$GAS": 1,
|
|
"$ENRO_DIR": data_dir,
|
|
"facebook/mbart-large-cc25": MBART_TINY,
|
|
}
|
|
|
|
# Clean up bash script
|
|
bash_script = Path("examples/seq2seq/train_mbart_cc25_enro.sh").open().read().split("finetune.py")[1].strip()
|
|
bash_script = bash_script.replace("\\\n", "").strip().replace("$@", "")
|
|
for k, v in env_vars_to_replace.items():
|
|
bash_script = bash_script.replace(k, str(v))
|
|
output_dir = tempfile.mkdtemp(prefix="output")
|
|
|
|
if CUDA_AVAILABLE:
|
|
gpus = torch.cuda.device_count()
|
|
else:
|
|
bash_script = bash_script.replace("--fp16", "")
|
|
gpus = 0
|
|
|
|
testargs = ["finetune.py"] + bash_script.split() + [f"--output_dir={output_dir}", f"--gpus={gpus}"]
|
|
with patch.object(sys, "argv", testargs):
|
|
parser = argparse.ArgumentParser()
|
|
parser = pl.Trainer.add_argparse_args(parser)
|
|
parser = SummarizationModule.add_model_specific_args(parser, os.getcwd())
|
|
args = parser.parse_args()
|
|
# assert args.gpus == gpus THIS BREAKS
|
|
# args.gpus = gpus
|
|
main(args)
|