


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)
lowerCAmelCase__			:      Optional[int]								=logging.getLogger()

lowerCAmelCase__			:      str								=logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







class 	__lowercase    (A__  ):





							"""simple docstring"""







							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__       ):





										"""simple docstring"""




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






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__ = "pytorch"       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Any					=					self.get_auto_remove_tmp_dir()
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					os.path.join(_lowerCamelCase			,						'output'       )
										SCREAMING_SNAKE_CASE_		:     List[str]					=					os.path.join(_lowerCamelCase			,						'data'       )
										self._create_dummy_data(data_dir=_lowerCamelCase       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					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'       )

										SCREAMING_SNAKE_CASE_		:     List[str]					=					[sys.executable, str(Path(finetune_rag.__file__       ).resolve()       )] + testargs
										execute_subprocess_async(_lowerCamelCase			,						env=self.get_env()       )

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					os.path.join(_lowerCamelCase			,						'metrics.json'       )
										with open(_lowerCamelCase       ) as f:
													SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					json.load(_lowerCamelCase       )
										return result






							@require_torch_gpu
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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






							@require_torch_multi_gpu
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[Any]					=					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"""




										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					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"""




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







"""simple docstring"""






# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings('ignore', category=UserWarning, module='torch.optim.lr_scheduler')





class 				a__					:



     def __init__(							self :List[str]      ,  _lowerCamelCase :Any      ,  _lowerCamelCase :Optional[Any]      ,  _lowerCamelCase :bool = True      ,  _lowerCamelCase :bool = False   ):







      '''simple docstring'''



      UpperCamelCase_     :					List[Any]   	=scheduler
      UpperCamelCase_     :					Any   	=optimizers if isinstance(_lowerCamelCase      ,  (list, tuple)   ) else [optimizers]
      UpperCamelCase_     :					Dict   	=split_batches
      UpperCamelCase_     :					str   	=step_with_optimizer
      UpperCamelCase_     :					Dict   	=GradientState()



     def  lowerCamelCase_						(							self :Optional[int]      ,  *_lowerCamelCase :str      ,  **_lowerCamelCase :Union[str, Any]   ):







      '''simple docstring'''



      if not self.step_with_optimizer:
       # No link between scheduler and optimizer -> just step
       self.scheduler.step(*_lowerCamelCase      ,  **_lowerCamelCase   )
       return

      # Otherwise, first make sure the optimizer was stepped.
      if not self.gradient_state.sync_gradients:
       if self.gradient_state.adjust_scheduler:
        self.scheduler._step_count += 1
       return

      for opt in self.optimizers:
       if opt.step_was_skipped:
        return
      if self.split_batches:
       # Split batches -> the training dataloader batch size is not changed so one step per training step
       self.scheduler.step(*_lowerCamelCase      ,  **_lowerCamelCase   )
      else:
       # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
       # num_processes steps per training step
       UpperCamelCase_     :					Optional[Any]   	=AcceleratorState().num_processes
       for _ in range(_lowerCamelCase   ):
        # Special case when using OneCycle and `drop_last` was not used
        if hasattr(self.scheduler      ,  'total_steps'   ):
         if self.scheduler._step_count <= self.scheduler.total_steps:
          self.scheduler.step(*_lowerCamelCase      ,  **_lowerCamelCase   )
        else:
         self.scheduler.step(*_lowerCamelCase      ,  **_lowerCamelCase   )



     def  lowerCamelCase_						(							self :Dict   ):







      '''simple docstring'''



      return self.scheduler.get_last_lr()



     def  lowerCamelCase_						(							self :Tuple   ):







      '''simple docstring'''



      return self.scheduler.state_dict()



     def  lowerCamelCase_						(							self :Any      ,  _lowerCamelCase :Any   ):







      '''simple docstring'''



      self.scheduler.load_state_dict(_lowerCamelCase   )



     def  lowerCamelCase_						(							self :Any   ):







      '''simple docstring'''



      return self.scheduler.get_lr()




     def  lowerCamelCase_						(							self :Optional[int]      ,  *_lowerCamelCase :Union[str, Any]      ,  **_lowerCamelCase :Optional[int]   ):







      '''simple docstring'''



      return self.scheduler.print_lr(*_lowerCamelCase      ,  **_lowerCamelCase   )
