93 lines
3.0 KiB
Python
93 lines
3.0 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
import tarfile
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
import pytorch_lightning as pl
|
|
import wget
|
|
|
|
from transformers.testing_utils import slow
|
|
|
|
from .finetune import SummarizationModule, main
|
|
from .test_seq2seq_examples import CUDA_AVAILABLE, MBART_TINY
|
|
from .utils import load_json
|
|
|
|
|
|
def fetch_and_save_wmt_100():
|
|
# TODO(SS): DELETEME
|
|
DATA_URL = "https://s3.amazonaws.com/datasets.huggingface.co/translation/wmt_en_ro_test_data.tgz"
|
|
|
|
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
|
|
|
|
|
|
@slow
|
|
@pytest.mark.skipif(not CUDA_AVAILABLE, reason="too slow to run on CPU")
|
|
def test_train_mbart_cc25_enro_script():
|
|
data_dir = "examples/seq2seq/test_data/wmt_en_ro"
|
|
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 = 1 # 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}", "--learning_rate=3e-1"]
|
|
)
|
|
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
|
|
model = main(args)
|
|
contents = os.listdir(output_dir)
|
|
# ckpt_name = "val_avg_rouge2=0.0000-step_count=2.ckpt" # "val_avg_rouge2=0.0000-epoch=1.ckpt" #
|
|
# "epoch=1-val_avg_rouge2=0.0000.ckpt"
|
|
contents = {os.path.basename(p) for p in contents}
|
|
|
|
# self.assertIn(ckpt_name, contents)
|
|
|
|
assert "test_generations.txt" in contents
|
|
assert "test_results.txt" in contents
|
|
|
|
metrics = load_json(model.metrics_save_path)
|
|
first_step_stats = metrics["val"][0]
|
|
last_step_stats = metrics["val"][-1]
|
|
assert last_step_stats["val_avg_gen_time"] >= 0.01
|
|
assert 1.0 >= last_step_stats["val_avg_gen_time"]
|
|
assert first_step_stats["val_avg_bleu"] < last_step_stats["val_avg_bleu"]
|
|
|
|
assert isinstance(last_step_stats[f"val_avg_{model.val_metric}"], float)
|
|
# desired_n_evals = int(args_d["max_epochs"] * (1 / args_d["val_check_interval"]) + 1)
|
|
# assert len(metrics["val"]) == desired_n_evals
|
|
assert len(metrics["test"]) == 1
|