mirror of
https://github.com/huggingface/transformers.git
synced 2025-07-15 10:38:23 +06:00
69233cf03b
Move declaration of log streams to before tests, so that results won't get compounded on top of each other
111 lines
3.9 KiB
Python
111 lines
3.9 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import finetune_rag
|
|
from transformers.file_utils import is_apex_available
|
|
from transformers.testing_utils import (
|
|
TestCasePlus,
|
|
execute_subprocess_async,
|
|
require_ray,
|
|
require_torch_gpu,
|
|
require_torch_multi_gpu,
|
|
)
|
|
|
|
|
|
logging.basicConfig(level=logging.DEBUG)
|
|
logger = logging.getLogger()
|
|
|
|
stream_handler = logging.StreamHandler(sys.stdout)
|
|
logger.addHandler(stream_handler)
|
|
|
|
|
|
class RagFinetuneExampleTests(TestCasePlus):
|
|
def _create_dummy_data(self, data_dir):
|
|
os.makedirs(data_dir, exist_ok=True)
|
|
contents = {"source": "What is love ?", "target": "life"}
|
|
n_lines = {"train": 12, "val": 2, "test": 2}
|
|
for split in ["train", "test", "val"]:
|
|
for field in ["source", "target"]:
|
|
content = "\n".join([contents[field]] * n_lines[split])
|
|
with open(os.path.join(data_dir, f"{split}.{field}"), "w") as f:
|
|
f.write(content)
|
|
|
|
def _run_finetune(self, gpus: int, distributed_retriever: str = "pytorch"):
|
|
tmp_dir = self.get_auto_remove_tmp_dir()
|
|
output_dir = os.path.join(tmp_dir, "output")
|
|
data_dir = os.path.join(tmp_dir, "data")
|
|
self._create_dummy_data(data_dir=data_dir)
|
|
|
|
testargs = f"""
|
|
--data_dir {data_dir} \
|
|
--output_dir {output_dir} \
|
|
--model_name_or_path facebook/rag-sequence-base \
|
|
--model_type rag_sequence \
|
|
--do_train \
|
|
--do_predict \
|
|
--n_val -1 \
|
|
--val_check_interval 1.0 \
|
|
--train_batch_size 2 \
|
|
--eval_batch_size 1 \
|
|
--max_source_length 25 \
|
|
--max_target_length 25 \
|
|
--val_max_target_length 25 \
|
|
--test_max_target_length 25 \
|
|
--label_smoothing 0.1 \
|
|
--dropout 0.1 \
|
|
--attention_dropout 0.1 \
|
|
--weight_decay 0.001 \
|
|
--adam_epsilon 1e-08 \
|
|
--max_grad_norm 0.1 \
|
|
--lr_scheduler polynomial \
|
|
--learning_rate 3e-04 \
|
|
--num_train_epochs 1 \
|
|
--warmup_steps 4 \
|
|
--gradient_accumulation_steps 1 \
|
|
--distributed-port 8787 \
|
|
--use_dummy_dataset 1 \
|
|
--distributed_retriever {distributed_retriever} \
|
|
""".split()
|
|
|
|
if gpus > 0:
|
|
testargs.append(f"--gpus={gpus}")
|
|
if is_apex_available():
|
|
testargs.append("--fp16")
|
|
else:
|
|
testargs.append("--gpus=0")
|
|
testargs.append("--distributed_backend=ddp_cpu")
|
|
testargs.append("--num_processes=2")
|
|
|
|
cmd = [sys.executable, str(Path(finetune_rag.__file__).resolve())] + testargs
|
|
execute_subprocess_async(cmd, env=self.get_env())
|
|
|
|
metrics_save_path = os.path.join(output_dir, "metrics.json")
|
|
with open(metrics_save_path) as f:
|
|
result = json.load(f)
|
|
return result
|
|
|
|
@require_torch_gpu
|
|
def test_finetune_gpu(self):
|
|
result = self._run_finetune(gpus=1)
|
|
self.assertGreaterEqual(result["test"][0]["test_avg_em"], 0.2)
|
|
|
|
@require_torch_multi_gpu
|
|
def test_finetune_multigpu(self):
|
|
result = self._run_finetune(gpus=2)
|
|
self.assertGreaterEqual(result["test"][0]["test_avg_em"], 0.2)
|
|
|
|
@require_torch_gpu
|
|
@require_ray
|
|
def test_finetune_gpu_ray_retrieval(self):
|
|
result = self._run_finetune(gpus=1, distributed_retriever="ray")
|
|
self.assertGreaterEqual(result["test"][0]["test_avg_em"], 0.2)
|
|
|
|
@require_torch_multi_gpu
|
|
@require_ray
|
|
def test_finetune_multigpu_ray_retrieval(self):
|
|
result = self._run_finetune(gpus=1, distributed_retriever="ray")
|
|
self.assertGreaterEqual(result["test"][0]["test_avg_em"], 0.2)
|