split if gpus 0 testargs appendfgpusgpus if isapexavailable testargs appendfp16 else testargs appendgpus0 testargs appenddistributedbackendddpcpu testargs appendnumprocesses2 cmd sys executable strpathfinetunerag file resolve testargs executesubprocessasynccmd envself getenv metricssavepath os path joinoutputdir metrics json with openmetricssavepath as f result json loadf return result requiretorchgpu def testfinetunegpuself result self runfinetunegpus1 self assertgreaterequalresulttest0testavgem 0 2 requiretorchmultigpu def testfinetunemultigpuself result self runfinetunegpus2 self assertgreaterequalresulttest0testavgem 0 2 requiretorchgpu requireray def testfinetunegpurayretrievalself result self runfinetunegpus1 distributedretrieverray self assertgreaterequalresulttest0testavgem 0 2 requiretorchmultigpu requireray def testfinetunemultigpurayretrievalself result self runfinetunegpus1 distributedretrieverray self assertgreaterequalresulttest0testavgem 0 2 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
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.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)