




import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only
from utils_rag import save_json





def        __UpperCAmelCase						(							lowerCamelCase_ :     Dict  )       ->  Optional[int]:







			"""simple docstring"""





			SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      filter(lambda lowerCamelCase_  : p.requires_grad  ,							model.parameters()  )
			SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      sum([np.prod(p.size()  ) for p in model_parameters]  )
			return params


UpperCamelCase__ :   Tuple						=							logging.getLogger(__name__)





def        __UpperCAmelCase						(							lowerCamelCase_ :     int  ,							lowerCamelCase_ :     Tuple  )       ->  List[str]:







			"""simple docstring"""





			if metric == "rouge2":
						SCREAMING_SNAKE_CASE_   : Any					  =      '{val_avg_rouge2:.4f}-{step_count}'
			elif metric == "bleu":
						SCREAMING_SNAKE_CASE_   : List[str]					  =      '{val_avg_bleu:.4f}-{step_count}'
			elif metric == "em":
						SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      '{val_avg_em:.4f}-{step_count}'
			elif metric == "loss":
						SCREAMING_SNAKE_CASE_   : List[Any]					  =      '{val_avg_loss:.4f}-{step_count}'
			else:
						raise NotImplementedError(
						    F'seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this'
						    ' function.'  )

			SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      ModelCheckpoint(
			    dirpath=__UpperCAmelCase  ,							filename=__UpperCAmelCase  ,							monitor=F'val_{metric}'  ,							mode='max'  ,							save_top_k=1  ,							every_n_epochs=1  ,							)

			return checkpoint_callback





def        __UpperCAmelCase						(							lowerCamelCase_ :     Dict  ,							lowerCamelCase_ :     Optional[int]  )       ->  str:







			"""simple docstring"""





			return EarlyStopping(
			    monitor=F'val_{metric}'  ,							mode='min' if 'loss' in metric else 'max'  ,							patience=__UpperCAmelCase  ,							verbose=__UpperCAmelCase  ,							)







class 				lowerCAmelCase_       (					pl.Callback							):






				def 		snake_case	(			self      ,snake_case__      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : Dict					  =      {F'lr_group_{i}': param['lr'] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups )}
							pl_module.logger.log_metrics(lowercase_ )






				@rank_zero_only
				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__=True ):
							logger.info(F'***** {type_path} results at step {trainer.global_step:05d} *****' )
							SCREAMING_SNAKE_CASE_   : Tuple					  =      trainer.callback_metrics
							trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ['log', 'progress_bar', 'preds']} )
							# Log results
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      Path(pl_module.hparams.output_dir )
							if type_path == "test":
										SCREAMING_SNAKE_CASE_   : str					  =      od / 'test_results.txt'
										SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      od / 'test_generations.txt'
							else:
										# this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
										# If people want this it will be easy enough to add back.
										SCREAMING_SNAKE_CASE_   : Optional[int]					  =      od / F'{type_path}_results/{trainer.global_step:05d}.txt'
										SCREAMING_SNAKE_CASE_   : List[str]					  =      od / F'{type_path}_generations/{trainer.global_step:05d}.txt'
										results_file.parent.mkdir(exist_ok=lowercase_ )
										generations_file.parent.mkdir(exist_ok=lowercase_ )
							with open(lowercase_      ,'a+' ) as writer:
										for key in sorted(lowercase_ ):
													if key in ["log", "progress_bar", "preds"]:
																continue
													SCREAMING_SNAKE_CASE_   : Any					  =      metrics[key]
													if isinstance(lowercase_      ,torch.Tensor ):
																SCREAMING_SNAKE_CASE_   : Tuple					  =      val.item()
													SCREAMING_SNAKE_CASE_   : List[str]					  =      F'{key}: {val:.6f}\n'
													writer.write(lowercase_ )

							if not save_generations:
										return

							if "preds" in metrics:
										SCREAMING_SNAKE_CASE_   : str					  =      '\n'.join(metrics['preds'] )
										generations_file.open('w+' ).write(lowercase_ )






				@rank_zero_only
				def 		snake_case	(			self      ,snake_case__      ,snake_case__ ):
							try:
										SCREAMING_SNAKE_CASE_   : Any					  =      pl_module.model.model.num_parameters()
							except AttributeError:
										SCREAMING_SNAKE_CASE_   : Tuple					  =      pl_module.model.num_parameters()

							SCREAMING_SNAKE_CASE_   : Dict					  =      count_trainable_parameters(lowercase_ )
							# mp stands for million parameters
							trainer.logger.log_metrics({'n_params': npars, 'mp': npars / 1E6, 'grad_mp': n_trainable_pars / 1E6} )






				@rank_zero_only
				def 		snake_case	(			self      ,snake_case__      ,snake_case__ ):
							save_json(pl_module.metrics      ,pl_module.metrics_save_path )
							return self._write_logs(lowercase_      ,lowercase_      ,'test' )






				@rank_zero_only
				def 		snake_case	(			self      ,snake_case__      ,snake_case__ ):
							save_json(pl_module.metrics      ,pl_module.metrics_save_path )
							# Uncommenting this will save val generations
							# return self._write_logs(trainer, pl_module, "valid")









'''simple docstring'''

from __future__ import annotations

import unittest

from transformers import FunnelConfig, is_tf_available
from transformers.testing_utils import require_tf

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
							import tensorflow as tf

							from transformers import (
							    TFFunnelBaseModel,
							    TFFunnelForMaskedLM,
							    TFFunnelForMultipleChoice,
							    TFFunnelForPreTraining,
							    TFFunnelForQuestionAnswering,
							    TFFunnelForSequenceClassification,
							    TFFunnelForTokenClassification,
							    TFFunnelModel,
							)






class 		a							:


				def __init__(						self       :   Dict			,							lowercase_       :   str			,							lowercase_       :   Union[str, Any]=13			,							lowercase_       :   Union[str, Any]=7			,							lowercase_       :   Tuple=True			,							lowercase_       :   int=True			,							lowercase_       :   Dict=True			,							lowercase_       :   Any=True			,							lowercase_       :   Union[str, Any]=99			,							lowercase_       :   Tuple=[1, 1, 2]			,							lowercase_       :   List[Any]=1			,							lowercase_       :   int=32			,							lowercase_       :   List[Any]=4			,							lowercase_       :   Tuple=8			,							lowercase_       :   Union[str, Any]=37			,							lowercase_       :   Union[str, Any]="gelu_new"			,							lowercase_       :   str=0.1			,							lowercase_       :   Tuple=0.1			,							lowercase_       :   Optional[int]=0.0			,							lowercase_       :   Optional[int]=512			,							lowercase_       :   int=3			,							lowercase_       :   Dict=0.02			,							lowercase_       :   Union[str, Any]=3			,							lowercase_       :   Dict=4			,							lowercase_       :   List[str]=None			,							lowercase_       :   Tuple=False			,							):
									snake_case_		     =							parent
									snake_case_		     =							batch_size
									snake_case_		     =							seq_length
									snake_case_		     =							is_training
									snake_case_		     =							use_input_mask
									snake_case_		     =							use_token_type_ids
									snake_case_		     =							use_labels
									snake_case_		     =							vocab_size
									snake_case_		     =							block_sizes
									snake_case_		     =							num_decoder_layers
									snake_case_		     =							d_model
									snake_case_		     =							n_head
									snake_case_		     =							d_head
									snake_case_		     =							d_inner
									snake_case_		     =							hidden_act
									snake_case_		     =							hidden_dropout
									snake_case_		     =							attention_dropout
									snake_case_		     =							activation_dropout
									snake_case_		     =							max_position_embeddings
									snake_case_		     =							type_vocab_size
									snake_case_		     =							2
									snake_case_		     =							num_labels
									snake_case_		     =							num_choices
									snake_case_		     =							scope
									snake_case_		     =							initializer_std

									# Used in the tests to check the size of the first attention layer
									snake_case_		     =							n_head
									# Used in the tests to check the size of the first hidden state
									snake_case_		     =							self.d_model
									# Used in the tests to check the number of output hidden states/attentions
									snake_case_		     =							sum(self.block_sizes   ) + (0 if base else self.num_decoder_layers)
									# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
									# the last hidden state of the first block (which is the first hidden state of the decoder).
									if not base:
														snake_case_		     =							self.num_hidden_layers + 2


				def     A_					(						self       :   Union[str, Any]   ):
									snake_case_		     =							ids_tensor([self.batch_size, self.seq_length]			,							self.vocab_size   )

									snake_case_		     =							None
									if self.use_input_mask:
														snake_case_		     =							random_attention_mask([self.batch_size, self.seq_length]   )

									snake_case_		     =							None
									if self.use_token_type_ids:
														snake_case_		     =							ids_tensor([self.batch_size, self.seq_length]			,							self.type_vocab_size   )

									snake_case_		     =							None
									snake_case_		     =							None
									snake_case_		     =							None
									if self.use_labels:
														snake_case_		     =							ids_tensor([self.batch_size]			,							self.type_sequence_label_size   )
														snake_case_		     =							ids_tensor([self.batch_size, self.seq_length]			,							self.num_labels   )
														snake_case_		     =							ids_tensor([self.batch_size]			,							self.num_choices   )

									snake_case_		     =							FunnelConfig(
									    vocab_size=self.vocab_size			,							block_sizes=self.block_sizes			,							num_decoder_layers=self.num_decoder_layers			,							d_model=self.d_model			,							n_head=self.n_head			,							d_head=self.d_head			,							d_inner=self.d_inner			,							hidden_act=self.hidden_act			,							hidden_dropout=self.hidden_dropout			,							attention_dropout=self.attention_dropout			,							activation_dropout=self.activation_dropout			,							max_position_embeddings=self.max_position_embeddings			,							type_vocab_size=self.type_vocab_size			,							initializer_std=self.initializer_std			,							)

									return (
									    config,
									    input_ids,
									    token_type_ids,
									    input_mask,
									    sequence_labels,
									    token_labels,
									    choice_labels,
									)


				def     A_					(						self       :   Optional[Any]			,							lowercase_       :   int			,							lowercase_       :   Tuple			,							lowercase_       :   List[str]			,							lowercase_       :   Union[str, Any]			,							lowercase_       :   Optional[int]			,							lowercase_       :   List[Any]			,							lowercase_       :   Optional[int]			,							):
									snake_case_		     =							TFFunnelModel(config=lowercase_   )
									snake_case_		     =							{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									snake_case_		     =							model(lowercase_   )

									snake_case_		     =							[input_ids, input_mask]
									snake_case_		     =							model(lowercase_   )

									snake_case_		     =							model(lowercase_   )
									self.parent.assertEqual(result.last_hidden_state.shape			,							(self.batch_size, self.seq_length, self.d_model)   )

									snake_case_		     =							False
									snake_case_		     =							TFFunnelModel(config=lowercase_   )
									snake_case_		     =							model(lowercase_   )
									self.parent.assertEqual(result.last_hidden_state.shape			,							(self.batch_size, self.seq_length, self.d_model)   )

									snake_case_		     =							False
									snake_case_		     =							TFFunnelModel(config=lowercase_   )
									snake_case_		     =							model(lowercase_   )
									self.parent.assertEqual(result.last_hidden_state.shape			,							(self.batch_size, self.seq_length, self.d_model)   )


				def     A_					(						self       :   List[Any]			,							lowercase_       :   Tuple			,							lowercase_       :   Dict			,							lowercase_       :   List[Any]			,							lowercase_       :   Dict			,							lowercase_       :   int			,							lowercase_       :   int			,							lowercase_       :   List[Any]			,							):
									snake_case_		     =							TFFunnelBaseModel(config=lowercase_   )
									snake_case_		     =							{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									snake_case_		     =							model(lowercase_   )

									snake_case_		     =							[input_ids, input_mask]
									snake_case_		     =							model(lowercase_   )

									snake_case_		     =							model(lowercase_   )
									self.parent.assertEqual(result.last_hidden_state.shape			,							(self.batch_size, 2, self.d_model)   )

									snake_case_		     =							False
									snake_case_		     =							TFFunnelBaseModel(config=lowercase_   )
									snake_case_		     =							model(lowercase_   )
									self.parent.assertEqual(result.last_hidden_state.shape			,							(self.batch_size, 3, self.d_model)   )

									snake_case_		     =							False
									snake_case_		     =							TFFunnelBaseModel(config=lowercase_   )
									snake_case_		     =							model(lowercase_   )
									self.parent.assertEqual(result.last_hidden_state.shape			,							(self.batch_size, 2, self.d_model)   )


				def     A_					(						self       :   Any			,							lowercase_       :   Optional[int]			,							lowercase_       :   str			,							lowercase_       :   int			,							lowercase_       :   Optional[int]			,							lowercase_       :   Optional[int]			,							lowercase_       :   Dict			,							lowercase_       :   Union[str, Any]			,							):
									snake_case_		     =							TFFunnelForPreTraining(config=lowercase_   )
									snake_case_		     =							{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									snake_case_		     =							model(lowercase_   )
									self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.seq_length)   )


				def     A_					(						self       :   Any			,							lowercase_       :   Union[str, Any]			,							lowercase_       :   Tuple			,							lowercase_       :   str			,							lowercase_       :   Tuple			,							lowercase_       :   List[str]			,							lowercase_       :   Union[str, Any]			,							lowercase_       :   Union[str, Any]			,							):
									snake_case_		     =							TFFunnelForMaskedLM(config=lowercase_   )
									snake_case_		     =							{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									snake_case_		     =							model(lowercase_   )
									self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.seq_length, self.vocab_size)   )


				def     A_					(						self       :   Dict			,							lowercase_       :   Dict			,							lowercase_       :   str			,							lowercase_       :   Optional[Any]			,							lowercase_       :   Tuple			,							lowercase_       :   Optional[Any]			,							lowercase_       :   Tuple			,							lowercase_       :   Tuple			,							):
									snake_case_		     =							self.num_labels
									snake_case_		     =							TFFunnelForSequenceClassification(config=lowercase_   )
									snake_case_		     =							{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									snake_case_		     =							model(lowercase_   )
									self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.num_labels)   )


				def     A_					(						self       :   int			,							lowercase_       :   Optional[int]			,							lowercase_       :   Any			,							lowercase_       :   List[str]			,							lowercase_       :   Optional[int]			,							lowercase_       :   str			,							lowercase_       :   str			,							lowercase_       :   int			,							):
									snake_case_		     =							self.num_choices
									snake_case_		     =							TFFunnelForMultipleChoice(config=lowercase_   )
									snake_case_		     =							tf.tile(tf.expand_dims(lowercase_			,							1   )			,							(1, self.num_choices, 1)   )
									snake_case_		     =							tf.tile(tf.expand_dims(lowercase_			,							1   )			,							(1, self.num_choices, 1)   )
									snake_case_		     =							tf.tile(tf.expand_dims(lowercase_			,							1   )			,							(1, self.num_choices, 1)   )
									snake_case_		     =							{
									    '''input_ids''': multiple_choice_inputs_ids,
									    '''attention_mask''': multiple_choice_input_mask,
									    '''token_type_ids''': multiple_choice_token_type_ids,
									}
									snake_case_		     =							model(lowercase_   )
									self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.num_choices)   )


				def     A_					(						self       :   Any			,							lowercase_       :   str			,							lowercase_       :   Dict			,							lowercase_       :   int			,							lowercase_       :   Optional[int]			,							lowercase_       :   List[str]			,							lowercase_       :   Any			,							lowercase_       :   Optional[Any]			,							):
									snake_case_		     =							self.num_labels
									snake_case_		     =							TFFunnelForTokenClassification(config=lowercase_   )
									snake_case_		     =							{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									snake_case_		     =							model(lowercase_   )
									self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.seq_length, self.num_labels)   )


				def     A_					(						self       :   Optional[Any]			,							lowercase_       :   Any			,							lowercase_       :   Dict			,							lowercase_       :   Any			,							lowercase_       :   int			,							lowercase_       :   Tuple			,							lowercase_       :   List[str]			,							lowercase_       :   Optional[Any]			,							):
									snake_case_		     =							TFFunnelForQuestionAnswering(config=lowercase_   )
									snake_case_		     =							{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
									snake_case_		     =							model(lowercase_   )
									self.parent.assertEqual(result.start_logits.shape			,							(self.batch_size, self.seq_length)   )
									self.parent.assertEqual(result.end_logits.shape			,							(self.batch_size, self.seq_length)   )


				def     A_					(						self       :   Any   ):
									snake_case_		     =							self.prepare_config_and_inputs()
									(
									    (
									    snake_case_
									)     ,(
									    snake_case_
									)     ,(
									    snake_case_
									)     ,(
									    snake_case_
									)     ,(
									    snake_case_
									)     ,(
									    snake_case_
									)     ,(
									    snake_case_
									)     ,
									)		     =							config_and_inputs
									snake_case_		     =							{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
									return config, inputs_dict







@require_tf
class 		a							(  _lowerCamelCase  ,	_lowerCamelCase  ,	unittest.TestCase	):
				snake_case_           =						(
				    (
				        TFFunnelModel,
				        TFFunnelForMaskedLM,
				        TFFunnelForPreTraining,
				        TFFunnelForQuestionAnswering,
				        TFFunnelForTokenClassification,
				    )
				    if is_tf_available()
				    else ()
				)
				snake_case_           =						(
				    {
				        "feature-extraction": (TFFunnelBaseModel, TFFunnelModel),
				        "fill-mask": TFFunnelForMaskedLM,
				        "question-answering": TFFunnelForQuestionAnswering,
				        "text-classification": TFFunnelForSequenceClassification,
				        "token-classification": TFFunnelForTokenClassification,
				        "zero-shot": TFFunnelForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				snake_case_           =						False
				snake_case_           =						False


				def     A_					(						self       :   int   ):
									snake_case_		     =							TFFunnelModelTester(self   )
									snake_case_		     =							ConfigTester(self			,							config_class=lowercase_   )


				def     A_					(						self       :   Union[str, Any]   ):
									self.config_tester.run_common_tests()


				def     A_					(						self       :   List[Any]   ):
									snake_case_		     =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowercase_   )


				def     A_					(						self       :   List[str]   ):
									snake_case_		     =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_pretraining(*lowercase_   )


				def     A_					(						self       :   str   ):
									snake_case_		     =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*lowercase_   )


				def     A_					(						self       :   List[str]   ):
									snake_case_		     =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*lowercase_   )


				def     A_					(						self       :   List[Any]   ):
									snake_case_		     =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*lowercase_   )







@require_tf
class 		a							(  _lowerCamelCase  ,	unittest.TestCase	):
				snake_case_           =						(
				    (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
				)
				snake_case_           =						False
				snake_case_           =						False


				def     A_					(						self       :   Union[str, Any]   ):
									snake_case_		     =							TFFunnelModelTester(self			,							base=lowercase_   )
									snake_case_		     =							ConfigTester(self			,							config_class=lowercase_   )


				def     A_					(						self       :   Dict   ):
									self.config_tester.run_common_tests()


				def     A_					(						self       :   List[str]   ):
									snake_case_		     =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_base_model(*lowercase_   )


				def     A_					(						self       :   Union[str, Any]   ):
									snake_case_		     =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*lowercase_   )


				def     A_					(						self       :   str   ):
									snake_case_		     =							self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_multiple_choice(*lowercase_   )

