






import shutil
import tempfile
import unittest
from unittest.mock import patch

from transformers import (
    DefaultFlowCallback,
    IntervalStrategy,
    PrinterCallback,
    ProgressCallback,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    is_torch_available,
)
from transformers.testing_utils import require_torch


if is_torch_available():
		from transformers.trainer import DEFAULT_CALLBACKS

		from .test_trainer import RegressionDataset, RegressionModelConfig, RegressionPreTrainedModel





class __lowerCamelCase (      snake_case__):





	"""simple docstring"""






	def __init__(							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			[]

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							self.events.append('on_init_end'    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							self.events.append('on_train_begin'    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							self.events.append('on_train_end'    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							self.events.append('on_epoch_begin'    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							self.events.append('on_epoch_end'    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							self.events.append('on_step_begin'    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							self.events.append('on_step_end'    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							self.events.append('on_evaluate'    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							self.events.append('on_predict'    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							self.events.append('on_save'    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							self.events.append('on_log'    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							self.events.append('on_prediction_step'    )





@require_torch
class __lowerCamelCase (      unittest.TestCase):





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			tempfile.mkdtemp()

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							shutil.rmtree(self.output_dir    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase=0       ,    UpperCAmelCase=0       ,    UpperCAmelCase=64       ,    UpperCAmelCase=64       ,    UpperCAmelCase=None       ,    UpperCAmelCase=False       ,    **UpperCAmelCase    ):
							"""simple docstring"""







							_UpperCAmelCase									=			RegressionDataset(length=UpperCAmelCase    )
							_UpperCAmelCase									=			RegressionDataset(length=UpperCAmelCase    )
							_UpperCAmelCase									=			RegressionModelConfig(a=UpperCAmelCase       ,    b=UpperCAmelCase    )
							_UpperCAmelCase									=			RegressionPreTrainedModel(UpperCAmelCase    )

							_UpperCAmelCase									=			TrainingArguments(self.output_dir       ,    disable_tqdm=UpperCAmelCase       ,    report_to=[]       ,    **UpperCAmelCase    )
							return Trainer(
							    UpperCAmelCase       ,    UpperCAmelCase       ,    train_dataset=UpperCAmelCase       ,    eval_dataset=UpperCAmelCase       ,    callbacks=UpperCAmelCase       ,    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase    ):
							"""simple docstring"""







							self.assertEqual(len(UpperCAmelCase    )       ,    len(UpperCAmelCase    )    )

							# Order doesn't matter
							_UpperCAmelCase									=			sorted(UpperCAmelCase       ,    key=lambda UpperCAmelCase    : cb.__name__ if isinstance(UpperCAmelCase       ,    UpperCAmelCase    ) else cb.__class__.__name__    )
							_UpperCAmelCase									=			sorted(UpperCAmelCase       ,    key=lambda UpperCAmelCase    : cb.__name__ if isinstance(UpperCAmelCase       ,    UpperCAmelCase    ) else cb.__class__.__name__    )

							for cba, cba in zip(UpperCAmelCase       ,    UpperCAmelCase    ):
													if isinstance(UpperCAmelCase       ,    UpperCAmelCase    ) and isinstance(UpperCAmelCase       ,    UpperCAmelCase    ):
																			self.assertEqual(UpperCAmelCase       ,    UpperCAmelCase    )
													elif isinstance(UpperCAmelCase       ,    UpperCAmelCase    ) and not isinstance(UpperCAmelCase       ,    UpperCAmelCase    ):
																			self.assertEqual(UpperCAmelCase       ,    cba.__class__    )
													elif not isinstance(UpperCAmelCase       ,    UpperCAmelCase    ) and isinstance(UpperCAmelCase       ,    UpperCAmelCase    ):
																			self.assertEqual(cba.__class__       ,    UpperCAmelCase    )
													else:
																			self.assertEqual(UpperCAmelCase       ,    UpperCAmelCase    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase    ):
							"""simple docstring"""







							_UpperCAmelCase									=			['on_init_end', 'on_train_begin']
							_UpperCAmelCase									=			0
							_UpperCAmelCase									=			len(trainer.get_eval_dataloader()    )
							_UpperCAmelCase									=			['on_prediction_step'] * len(trainer.get_eval_dataloader()    ) + ['on_log', 'on_evaluate']
							for _ in range(trainer.state.num_train_epochs    ):
													expected_events.append('on_epoch_begin'    )
													for _ in range(UpperCAmelCase    ):
																			step += 1
																			expected_events += ["on_step_begin", "on_step_end"]
																			if step % trainer.args.logging_steps == 0:
																									expected_events.append('on_log'    )
																			if trainer.args.evaluation_strategy == IntervalStrategy.STEPS and step % trainer.args.eval_steps == 0:
																									expected_events += evaluation_events.copy()
																			if step % trainer.args.save_steps == 0:
																									expected_events.append('on_save'    )
													expected_events.append('on_epoch_end'    )
													if trainer.args.evaluation_strategy == IntervalStrategy.EPOCH:
																			expected_events += evaluation_events.copy()
							expected_events += ["on_log", "on_train_end"]
							return expected_events

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			self.get_trainer()
							_UpperCAmelCase									=			DEFAULT_CALLBACKS.copy() + [ProgressCallback]
							self.check_callbacks_equality(trainer.callback_handler.callbacks       ,    UpperCAmelCase    )

							# Callbacks passed at init are added to the default callbacks
							_UpperCAmelCase									=			self.get_trainer(callbacks=[MyTestTrainerCallback]    )
							expected_callbacks.append(UpperCAmelCase    )
							self.check_callbacks_equality(trainer.callback_handler.callbacks       ,    UpperCAmelCase    )

							# TrainingArguments.disable_tqdm controls if use ProgressCallback or PrinterCallback
							_UpperCAmelCase									=			self.get_trainer(disable_tqdm=UpperCAmelCase    )
							_UpperCAmelCase									=			DEFAULT_CALLBACKS.copy() + [PrinterCallback]
							self.check_callbacks_equality(trainer.callback_handler.callbacks       ,    UpperCAmelCase    )

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			DEFAULT_CALLBACKS.copy() + [ProgressCallback]
							_UpperCAmelCase									=			self.get_trainer()

							# We can add, pop, or remove by class name
							trainer.remove_callback(UpperCAmelCase    )
							expected_callbacks.remove(UpperCAmelCase    )
							self.check_callbacks_equality(trainer.callback_handler.callbacks       ,    UpperCAmelCase    )

							_UpperCAmelCase									=			self.get_trainer()
							_UpperCAmelCase									=			trainer.pop_callback(UpperCAmelCase    )
							self.assertEqual(cb.__class__       ,    UpperCAmelCase    )
							self.check_callbacks_equality(trainer.callback_handler.callbacks       ,    UpperCAmelCase    )

							trainer.add_callback(UpperCAmelCase    )
							expected_callbacks.insert(0       ,    UpperCAmelCase    )
							self.check_callbacks_equality(trainer.callback_handler.callbacks       ,    UpperCAmelCase    )

							# We can also add, pop, or remove by instance
							_UpperCAmelCase									=			self.get_trainer()
							_UpperCAmelCase									=			trainer.callback_handler.callbacks[0]
							trainer.remove_callback(UpperCAmelCase    )
							expected_callbacks.remove(UpperCAmelCase    )
							self.check_callbacks_equality(trainer.callback_handler.callbacks       ,    UpperCAmelCase    )

							_UpperCAmelCase									=			self.get_trainer()
							_UpperCAmelCase									=			trainer.callback_handler.callbacks[0]
							_UpperCAmelCase									=			trainer.pop_callback(UpperCAmelCase    )
							self.assertEqual(UpperCAmelCase       ,    UpperCAmelCase    )
							self.check_callbacks_equality(trainer.callback_handler.callbacks       ,    UpperCAmelCase    )

							trainer.add_callback(UpperCAmelCase    )
							expected_callbacks.insert(0       ,    UpperCAmelCase    )
							self.check_callbacks_equality(trainer.callback_handler.callbacks       ,    UpperCAmelCase    )

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							import warnings

							# XXX: for now ignore scatter_gather warnings in this test since it's not relevant to what's being tested
							warnings.simplefilter(action='ignore'       ,    category=UpperCAmelCase    )

							_UpperCAmelCase									=			self.get_trainer(callbacks=[MyTestTrainerCallback]    )
							trainer.train()
							_UpperCAmelCase									=			trainer.callback_handler.callbacks[-2].events
							self.assertEqual(UpperCAmelCase       ,    self.get_expected_events(UpperCAmelCase    )    )

							# Independent log/save/eval
							_UpperCAmelCase									=			self.get_trainer(callbacks=[MyTestTrainerCallback]       ,    logging_steps=5    )
							trainer.train()
							_UpperCAmelCase									=			trainer.callback_handler.callbacks[-2].events
							self.assertEqual(UpperCAmelCase       ,    self.get_expected_events(UpperCAmelCase    )    )

							_UpperCAmelCase									=			self.get_trainer(callbacks=[MyTestTrainerCallback]       ,    save_steps=5    )
							trainer.train()
							_UpperCAmelCase									=			trainer.callback_handler.callbacks[-2].events
							self.assertEqual(UpperCAmelCase       ,    self.get_expected_events(UpperCAmelCase    )    )

							_UpperCAmelCase									=			self.get_trainer(callbacks=[MyTestTrainerCallback]       ,    eval_steps=5       ,    evaluation_strategy='steps'    )
							trainer.train()
							_UpperCAmelCase									=			trainer.callback_handler.callbacks[-2].events
							self.assertEqual(UpperCAmelCase       ,    self.get_expected_events(UpperCAmelCase    )    )

							_UpperCAmelCase									=			self.get_trainer(callbacks=[MyTestTrainerCallback]       ,    evaluation_strategy='epoch'    )
							trainer.train()
							_UpperCAmelCase									=			trainer.callback_handler.callbacks[-2].events
							self.assertEqual(UpperCAmelCase       ,    self.get_expected_events(UpperCAmelCase    )    )

							# A bit of everything
							_UpperCAmelCase									=			self.get_trainer(
							    callbacks=[MyTestTrainerCallback]       ,    logging_steps=3       ,    save_steps=10       ,    eval_steps=5       ,    evaluation_strategy='steps'       ,    )
							trainer.train()
							_UpperCAmelCase									=			trainer.callback_handler.callbacks[-2].events
							self.assertEqual(UpperCAmelCase       ,    self.get_expected_events(UpperCAmelCase    )    )

							# warning should be emitted for duplicated callbacks
							with patch('transformers.trainer_callback.logger.warning'    ) as warn_mock:
													_UpperCAmelCase									=			self.get_trainer(
													    callbacks=[MyTestTrainerCallback, MyTestTrainerCallback]       ,    )
													assert str(UpperCAmelCase    ) in warn_mock.call_args[0][0]








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)
_a      	=  logging.getLogger()

_a      	=  logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)





class __lowerCamelCase (      snake_case__):





	"""simple docstring"""






	def        UpperCamelCase    (							self       ,    UpperCAmelCase    ):
							"""simple docstring"""







							os.makedirs(UpperCAmelCase       ,    exist_ok=UpperCAmelCase    )
							_UpperCAmelCase									=			{'source': 'What is love ?', 'target': 'life'}
							_UpperCAmelCase									=			{'train': 12, 'val': 2, 'test': 2}
							for split in ["train", "test", "val"]:
													for field in ["source", "target"]:
																			_UpperCAmelCase									=			'\n'.join([contents[field]] * n_lines[split]    )
																			with open(os.path.join(UpperCAmelCase       ,    F"""{split}.{field}"""    )       ,    'w'    ) as f:
																									f.write(UpperCAmelCase    )

	def        UpperCamelCase    (							self       ,    UpperCAmelCase       ,    UpperCAmelCase = "pytorch"    ):
							"""simple docstring"""







							_UpperCAmelCase									=			self.get_auto_remove_tmp_dir()
							_UpperCAmelCase									=			os.path.join(UpperCAmelCase       ,    'output'    )
							_UpperCAmelCase									=			os.path.join(UpperCAmelCase       ,    'data'    )
							self._create_dummy_data(data_dir=UpperCAmelCase    )

							_UpperCAmelCase									=			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'    )

							_UpperCAmelCase									=			[sys.executable, str(Path(finetune_rag.__file__    ).resolve()    )] + testargs
							execute_subprocess_async(UpperCAmelCase       ,    env=self.get_env()    )

							_UpperCAmelCase									=			os.path.join(UpperCAmelCase       ,    'metrics.json'    )
							with open(UpperCAmelCase    ) as f:
													_UpperCAmelCase									=			json.load(UpperCAmelCase    )
							return result

	@require_torch_gpu
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			self._run_finetune(gpus=1    )
							self.assertGreaterEqual(result['test'][0]['test_avg_em']       ,    0.2    )

	@require_torch_multi_gpu
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			self._run_finetune(gpus=2    )
							self.assertGreaterEqual(result['test'][0]['test_avg_em']       ,    0.2    )

	@require_torch_gpu
	@require_ray
	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			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        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			self._run_finetune(gpus=1       ,    distributed_retriever='ray'    )
							self.assertGreaterEqual(result['test'][0]['test_avg_em']       ,    0.2    )
