






import re
from typing import Callable, List, Optional, Union

import tensorflow as tf


try:
		from tensorflow.keras.optimizers.legacy import Adam
except ImportError:
		from tensorflow.keras.optimizers import Adam





class __lowerCamelCase (      tf.keras.optimizers.schedules.LearningRateSchedule):





	"""simple docstring"""






	def __init__(							self       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase = 1.0       ,    UpperCAmelCase = None       ,    ):
							"""simple docstring"""







							super().__init__()
							_UpperCAmelCase									=			initial_learning_rate
							_UpperCAmelCase									=			warmup_steps
							_UpperCAmelCase									=			power
							_UpperCAmelCase									=			decay_schedule_fn
							_UpperCAmelCase									=			name

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







							with tf.name_scope(self.name or 'WarmUp'    ) as name:
													# Implements polynomial warmup. i.e., if global_step < warmup_steps, the
													# learning rate will be `global_step/num_warmup_steps * init_lr`.
													_UpperCAmelCase									=			tf.cast(UpperCAmelCase       ,    tf.floataa    )
													_UpperCAmelCase									=			tf.cast(self.warmup_steps       ,    tf.floataa    )
													_UpperCAmelCase									=			global_step_float / warmup_steps_float
													_UpperCAmelCase									=			self.initial_learning_rate * tf.math.pow(UpperCAmelCase       ,    self.power    )
													return tf.cond(
													    global_step_float < warmup_steps_float       ,    lambda: warmup_learning_rate       ,    lambda: self.decay_schedule_fn(step - self.warmup_steps    )       ,    name=UpperCAmelCase       ,    )

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







							return {
							    "initial_learning_rate": self.initial_learning_rate,
							    "decay_schedule_fn": self.decay_schedule_fn,
							    "warmup_steps": self.warmup_steps,
							    "power": self.power,
							    "name": self.name,
							}



def __A    (      __lowerCAmelCase ,		__lowerCAmelCase ,		__lowerCAmelCase ,		__lowerCAmelCase = 0.0 ,		__lowerCAmelCase = 0.9 ,		__lowerCAmelCase = 0.9_99 ,		__lowerCAmelCase = 1E-8 ,		__lowerCAmelCase = None ,		__lowerCAmelCase = None ,		__lowerCAmelCase = 0.0 ,		__lowerCAmelCase = 1.0 ,		__lowerCAmelCase = None ,		)->  List[Any]:





						"""simple docstring"""



						_UpperCAmelCase									=			tf.keras.optimizers.schedules.PolynomialDecay(
						    initial_learning_rate=__lowerCAmelCase ,		decay_steps=num_train_steps - num_warmup_steps ,		end_learning_rate=init_lr * min_lr_ratio ,		power=__lowerCAmelCase ,		)
						if num_warmup_steps:
												_UpperCAmelCase									=			WarmUp(
												    initial_learning_rate=__lowerCAmelCase ,		decay_schedule_fn=__lowerCAmelCase ,		warmup_steps=__lowerCAmelCase ,		)
						if weight_decay_rate > 0.0:
												_UpperCAmelCase									=			AdamWeightDecay(
												    learning_rate=__lowerCAmelCase ,		weight_decay_rate=__lowerCAmelCase ,		beta_a=__lowerCAmelCase ,		beta_a=__lowerCAmelCase ,		epsilon=__lowerCAmelCase ,		clipnorm=__lowerCAmelCase ,		global_clipnorm=__lowerCAmelCase ,		exclude_from_weight_decay=['LayerNorm', 'layer_norm', 'bias'] ,		include_in_weight_decay=__lowerCAmelCase ,		)
						else:
												_UpperCAmelCase									=			tf.keras.optimizers.Adam(
												    learning_rate=__lowerCAmelCase ,		beta_a=__lowerCAmelCase ,		beta_a=__lowerCAmelCase ,		epsilon=__lowerCAmelCase ,		clipnorm=__lowerCAmelCase ,		global_clipnorm=__lowerCAmelCase ,		)
						# We return the optimizer and the LR scheduler in order to better track the
						# evolution of the LR independently of the optimizer.
						return optimizer, lr_schedule





class __lowerCamelCase (      snake_case__):





	"""simple docstring"""






	def __init__(							self       ,    UpperCAmelCase = 0.0_01       ,    UpperCAmelCase = 0.9       ,    UpperCAmelCase = 0.9_99       ,    UpperCAmelCase = 1e-7       ,    UpperCAmelCase = False       ,    UpperCAmelCase = 0.0       ,    UpperCAmelCase = None       ,    UpperCAmelCase = None       ,    UpperCAmelCase = "AdamWeightDecay"       ,    **UpperCAmelCase       ,    ):
							"""simple docstring"""







							super().__init__(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    )
							_UpperCAmelCase									=			weight_decay_rate
							_UpperCAmelCase									=			include_in_weight_decay
							_UpperCAmelCase									=			exclude_from_weight_decay

	@classmethod
	def        UpperCamelCase    (							cls       ,    UpperCAmelCase    ):
							"""simple docstring"""







							_UpperCAmelCase									=			{'WarmUp': WarmUp}
							return super(UpperCAmelCase       ,    cls    ).from_config(UpperCAmelCase       ,    custom_objects=UpperCAmelCase    )

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







							super(UpperCAmelCase       ,    self    )._prepare_local(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			tf.constant(
							    self.weight_decay_rate       ,    name='adam_weight_decay_rate'    )

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







							_UpperCAmelCase									=			self._do_use_weight_decay(var.name    )
							if do_decay:
													return var.assign_sub(
													    learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]['weight_decay_rate']       ,    use_locking=self._use_locking       ,    )
							return tf.no_op()

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







							_UpperCAmelCase					, _UpperCAmelCase									=			list(zip(*UpperCAmelCase    )    )
							return super(UpperCAmelCase       ,    self    ).apply_gradients(zip(UpperCAmelCase       ,    UpperCAmelCase    )       ,    name=UpperCAmelCase       ,    **UpperCAmelCase    )

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







							if apply_state is None:
													return self._decayed_lr_t[var_dtype], {}

							_UpperCAmelCase									=			apply_state or {}
							_UpperCAmelCase									=			apply_state.get((var_device, var_dtype)    )
							if coefficients is None:
													_UpperCAmelCase									=			self._fallback_apply_state(UpperCAmelCase       ,    UpperCAmelCase    )
													_UpperCAmelCase									=			coefficients

							return coefficients["lr_t"], {"apply_state": apply_state}

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







							_UpperCAmelCase					, _UpperCAmelCase									=			self._get_lr(var.device       ,    var.dtype.base_dtype       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			self._decay_weights_op(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase    )
							with tf.control_dependencies([decay]    ):
													return super(UpperCAmelCase       ,    self    )._resource_apply_dense(UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    )

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







							_UpperCAmelCase					, _UpperCAmelCase									=			self._get_lr(var.device       ,    var.dtype.base_dtype       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			self._decay_weights_op(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase    )
							with tf.control_dependencies([decay]    ):
													return super(UpperCAmelCase       ,    self    )._resource_apply_sparse(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    **UpperCAmelCase    )

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







							_UpperCAmelCase									=			super().get_config()
							config.update({'weight_decay_rate': self.weight_decay_rate}    )
							return config

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







							if self.weight_decay_rate == 0:
													return False

							if self._include_in_weight_decay:
													for r in self._include_in_weight_decay:
																			if re.search(UpperCAmelCase       ,    UpperCAmelCase    ) is not None:
																									return True

							if self._exclude_from_weight_decay:
													for r in self._exclude_from_weight_decay:
																			if re.search(UpperCAmelCase       ,    UpperCAmelCase    ) is not None:
																									return False
							return True





class __lowerCamelCase (      snake_case__):





	"""simple docstring"""






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







							_UpperCAmelCase									=			[]
							_UpperCAmelCase									=			None

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







							if self._accum_steps is None:
													_UpperCAmelCase									=			tf.Variable(
													    tf.constant(0       ,    dtype=tf.intaa    )       ,    trainable=UpperCAmelCase       ,    synchronization=tf.VariableSynchronization.ON_READ       ,    aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA       ,    )

							return self._accum_steps.value()

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







							if not self._gradients:
													raise ValueError('The accumulator should be called first to initialize the gradients'    )
							return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]

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







							if not self._gradients:
													_UpperCAmelCase									=			self.step  # Create the step variable.
													self._gradients.extend(
													    [
													        tf.Variable(
													            tf.zeros_like(UpperCAmelCase    )       ,    trainable=UpperCAmelCase       ,    synchronization=tf.VariableSynchronization.ON_READ       ,    aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA       ,    )
													        if gradient is not None
													        else gradient
													        for gradient in gradients
													    ]    )
							if len(UpperCAmelCase    ) != len(self._gradients    ):
													raise ValueError(F"""Expected {len(self._gradients    )} gradients, but got {len(UpperCAmelCase    )}"""    )

							for accum_gradient, gradient in zip(self._gradients       ,    UpperCAmelCase    ):
													if accum_gradient is not None and gradient is not None:
																			accum_gradient.assign_add(UpperCAmelCase    )

							self._accum_steps.assign_add(1    )

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







							if not self._gradients:
													return
							self._accum_steps.assign(0    )
							for gradient in self._gradients:
													if gradient is not None:
																			gradient.assign(tf.zeros_like(UpperCAmelCase    )    )








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    )
