


'''simple docstring'''



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 			_a			(		_lowercase      :			Tuple						):





     '''simple docstring'''






     __UpperCAmelCase			:				Dict													=    filter(lambda _lowercase						: p.requires_grad    ,      model.parameters()						)
     __UpperCAmelCase			:				int													=    sum([np.prod(p.size()						) for p in model_parameters]						)
     return params


__UpperCAmelCase    :List[Any]				       =     logging.getLogger(__name__)




def 			_a			(		_lowercase      :			Tuple    ,      _lowercase      :			Tuple						):





     '''simple docstring'''






     if metric == "rouge2":
          __UpperCAmelCase			:				Any													=    '''{val_avg_rouge2:.4f}-{step_count}'''
     elif metric == "bleu":
          __UpperCAmelCase			:				Tuple													=    '''{val_avg_bleu:.4f}-{step_count}'''
     elif metric == "em":
          __UpperCAmelCase			:				str													=    '''{val_avg_em:.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.'''						)

     __UpperCAmelCase			:				Any													=    ModelCheckpoint(
         dirpath=_lowercase    ,      filename=_lowercase    ,      monitor=F'val_{metric}'    ,      mode='''max'''    ,      save_top_k=3    ,      every_n_epochs=1    ,      )
     return checkpoint_callback




def 			_a			(		_lowercase      :			Union[str, Any]    ,      _lowercase      :			List[str]						):





     '''simple docstring'''






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



class 			a       (	pl.Callback ):


     """simple docstring"""






     def 				lowerCamelCase__				(       self  :	str							,				snake_case  :	Tuple							,				snake_case  :	Any    )		->							List[str]:
          __UpperCAmelCase			:				Any													=    {f'lr_group_{i}': param['''lr'''] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups    )}
          pl_module.logger.log_metrics(snake_case    )

     @rank_zero_only
     def 				lowerCamelCase__				(       self  :	List[str]							,				snake_case  :	pl.Trainer							,				snake_case  :	pl.LightningModule							,				snake_case  :	str							,				snake_case  :	Any=True    )		->							None:
          logger.info(f'***** {type_path} results at step {trainer.global_step:05d} *****'    )
          __UpperCAmelCase			:				Optional[Any]													=    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
          __UpperCAmelCase			:				Optional[Any]													=    Path(pl_module.hparams.output_dir    )
          if type_path == "test":
               __UpperCAmelCase			:				Optional[int]													=    od / '''test_results.txt'''
               __UpperCAmelCase			:				List[str]													=    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.
               __UpperCAmelCase			:				List[str]													=    od / f'{type_path}_results/{trainer.global_step:05d}.txt'
               __UpperCAmelCase			:				Optional[int]													=    od / f'{type_path}_generations/{trainer.global_step:05d}.txt'
               results_file.parent.mkdir(exist_ok=snake_case    )
               generations_file.parent.mkdir(exist_ok=snake_case    )
          with open(snake_case							,				'''a+'''    ) as writer:
               for key in sorted(snake_case    ):
                    if key in ["log", "progress_bar", "preds"]:
                         continue
                    __UpperCAmelCase			:				Tuple													=    metrics[key]
                    if isinstance(snake_case							,				torch.Tensor    ):
                         __UpperCAmelCase			:				List[str]													=    val.item()
                    __UpperCAmelCase			:				List[Any]													=    f'{key}: {val:.6f}\n'
                    writer.write(snake_case    )

          if not save_generations:
               return

          if "preds" in metrics:
               __UpperCAmelCase			:				Union[str, Any]													=    '''\n'''.join(metrics['''preds''']    )
               generations_file.open('''w+'''    ).write(snake_case    )

     @rank_zero_only
     def 				lowerCamelCase__				(       self  :	Dict							,				snake_case  :	Optional[Any]							,				snake_case  :	Optional[int]    )		->							Tuple:
          try:
               __UpperCAmelCase			:				str													=    pl_module.model.model.num_parameters()
          except AttributeError:
               __UpperCAmelCase			:				str													=    pl_module.model.num_parameters()

          __UpperCAmelCase			:				int													=    count_trainable_parameters(snake_case    )
          # 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 				lowerCamelCase__				(       self  :	List[Any]							,				snake_case  :	pl.Trainer							,				snake_case  :	pl.LightningModule    )		->							Any:
          save_json(pl_module.metrics							,				pl_module.metrics_save_path    )
          return self._write_logs(snake_case							,				snake_case							,				'''test'''    )






     @rank_zero_only
     def 				lowerCamelCase__				(       self  :	List[str]							,				snake_case  :	pl.Trainer							,				snake_case  :	Dict    )		->							Any:
          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'''



import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from transformers import AutoConfig, TFAutoModel, is_tensorflow_text_available, is_tf_available
from transformers.models.bert.tokenization_bert import BertTokenizer
from transformers.testing_utils import require_tensorflow_text, require_tf, slow


if is_tf_available():
      import tensorflow as tf

if is_tensorflow_text_available():
      from transformers.models.bert import TFBertTokenizer


__UpperCAmelCase    :Optional[int]				       =     ["bert-base-uncased", "bert-base-cased"]
__UpperCAmelCase    :str				       =     "hf-internal-testing/tiny-bert-tf-only"

if is_tf_available():



      class 			a       (	tf.keras.Model ):


           """simple docstring"""






           def __init__(       self  :	List[str]							,				snake_case  :	List[str]    )		->							str:
                super().__init__()
                __UpperCAmelCase			:				List[str]													=    tokenizer
                __UpperCAmelCase			:				List[Any]													=    AutoConfig.from_pretrained(snake_case    )
                __UpperCAmelCase			:				int													=    TFAutoModel.from_config(snake_case    )






           def 				lowerCamelCase__				(       self  :	List[Any]							,				snake_case  :	Optional[int]    )		->							Optional[Any]:
                __UpperCAmelCase			:				List[Any]													=    self.tokenizer(snake_case    )
                __UpperCAmelCase			:				Optional[Any]													=    self.bert(**snake_case    )
                return out["pooler_output"]



@require_tf
@require_tensorflow_text
class 			a       (	unittest.TestCase ):


     """simple docstring"""






     def 				lowerCamelCase__				(       self  :	Optional[int]    )		->							List[str]:
          super().setUp()

          __UpperCAmelCase			:				Tuple													=    [
              BertTokenizer.from_pretrained(snake_case    ) for checkpoint in (TOKENIZER_CHECKPOINTS * 2)
          ]  # repeat for when fast_bert_tokenizer=false
          __UpperCAmelCase			:				Any													=    [TFBertTokenizer.from_pretrained(snake_case    ) for checkpoint in TOKENIZER_CHECKPOINTS] + [
              TFBertTokenizer.from_pretrained(snake_case							,				use_fast_bert_tokenizer=snake_case    )
              for checkpoint in TOKENIZER_CHECKPOINTS
          ]
          assert len(self.tokenizers    ) == len(self.tf_tokenizers    )

          __UpperCAmelCase			:				Optional[int]													=    [
              '''This is a straightforward English test sentence.''',
              '''This one has some weird characters\rto\nsee\r\nif  those\u00E9break things.''',
              '''Now we\'re going to add some Chinese: 一 二 三 一二三''',
              '''And some much more rare Chinese: 齉 堃 齉堃''',
              '''Je vais aussi écrire en français pour tester les accents''',
              '''Classical Irish also has some unusual characters, so in they go: Gaelaċ, ꝼ''',
          ]
          __UpperCAmelCase			:				Optional[int]													=    list(zip(self.test_sentences							,				self.test_sentences[::-1]    )    )

     def 				lowerCamelCase__				(       self  :	Optional[int]    )		->							Optional[Any]:
          for tokenizer, tf_tokenizer in zip(self.tokenizers							,				self.tf_tokenizers    ):
               for test_inputs in (self.test_sentences, self.paired_sentences):
                    __UpperCAmelCase			:				Any													=    tokenizer(snake_case							,				return_tensors='''tf'''							,				padding='''longest'''    )
                    __UpperCAmelCase			:				Optional[int]													=    tf_tokenizer(snake_case    )

                    for key in python_outputs.keys():
                         self.assertTrue(tf.reduce_all(python_outputs[key].shape == tf_outputs[key].shape    )    )
                         self.assertTrue(tf.reduce_all(tf.cast(python_outputs[key]							,				tf.intaa    ) == tf_outputs[key]    )    )

     @slow
     def 				lowerCamelCase__				(       self  :	List[Any]    )		->							str:
          for tf_tokenizer in self.tf_tokenizers:
               __UpperCAmelCase			:				Any													=    tf_tokenizer(self.paired_sentences    )
               __UpperCAmelCase			:				Union[str, Any]													=    tf_tokenizer(
                   text=[sentence[0] for sentence in self.paired_sentences]							,				text_pair=[sentence[1] for sentence in self.paired_sentences]							,				)
               for key in merged_outputs.keys():
                    self.assertTrue(tf.reduce_all(tf.cast(merged_outputs[key]							,				tf.intaa    ) == separated_outputs[key]    )    )

     @slow
     def 				lowerCamelCase__				(       self  :	str    )		->							Union[str, Any]:
          for tf_tokenizer in self.tf_tokenizers:
               __UpperCAmelCase			:				Optional[int]													=    tf.function(snake_case    )
               for test_inputs in (self.test_sentences, self.paired_sentences):
                    __UpperCAmelCase			:				int													=    tf.constant(snake_case    )
                    __UpperCAmelCase			:				Tuple													=    compiled_tokenizer(snake_case    )
                    __UpperCAmelCase			:				Optional[int]													=    tf_tokenizer(snake_case    )

                    for key in eager_outputs.keys():
                         self.assertTrue(tf.reduce_all(eager_outputs[key] == compiled_outputs[key]    )    )






     @slow
     def 				lowerCamelCase__				(       self  :	str    )		->							str:
          for tf_tokenizer in self.tf_tokenizers:
               __UpperCAmelCase			:				List[Any]													=    ModelToSave(tokenizer=snake_case    )
               __UpperCAmelCase			:				Union[str, Any]													=    tf.convert_to_tensor(self.test_sentences    )
               __UpperCAmelCase			:				Tuple													=    model(snake_case    )  # Build model with some sample inputs
               with TemporaryDirectory() as tempdir:
                    __UpperCAmelCase			:				Any													=    Path(snake_case    ) / '''saved.model'''
                    model.save(snake_case    )
                    __UpperCAmelCase			:				str													=    tf.keras.models.load_model(snake_case    )
               __UpperCAmelCase			:				Optional[int]													=    loaded_model(snake_case    )
               # We may see small differences because the loaded model is compiled, so we need an epsilon for the test
               self.assertLessEqual(tf.reduce_max(tf.abs(out - loaded_output    )    )							,				1E-5    )