


'''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 __snake_case      (   UpperCAmelCase_		:		List[Any]	):
 lowerCamelCase_  				=						filter(lambda UpperCAmelCase_	: p.requires_grad      ,       model.parameters()	)
 lowerCamelCase_  				=						sum([np.prod(p.size()	) for p in model_parameters]	)
 return params


a_						:       List[str]      			=				logging.getLogger(__name__)





def __snake_case      (   UpperCAmelCase_		:		Tuple      ,       UpperCAmelCase_		:		int	):
 if metric == "rouge2":
  lowerCamelCase_  				=						"{val_avg_rouge2:.4f}-{step_count}"
 elif metric == "bleu":
  lowerCamelCase_  				=						"{val_avg_bleu:.4f}-{step_count}"
 elif metric == "em":
  lowerCamelCase_  				=						"{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."	)

 lowerCamelCase_  				=						ModelCheckpoint(
     dirpath=UpperCAmelCase_      ,       filename=UpperCAmelCase_      ,       monitor=F'''val_{metric}'''      ,       mode="max"      ,       save_top_k=3      ,       every_n_epochs=1      ,       )
 return checkpoint_callback





def __snake_case      (   UpperCAmelCase_		:		int      ,       UpperCAmelCase_		:		List[str]	):
 return EarlyStopping(
     monitor=F'''val_{metric}'''      ,       mode="min" if "loss" in metric else "max"      ,       patience=UpperCAmelCase_      ,       verbose=UpperCAmelCase_      ,       )







class   snake_case					(					pl.Callback    ):





  """simple docstring"""




  def 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						{f'''lr_group_{i}''': param["lr"] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups  )}
   pl_module.logger.log_metrics(UpperCamelCase  )






  @rank_zero_only
  def 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase=True  ):







   """simple docstring"""
   logger.info(f'''***** {type_path} results at step {trainer.global_step:05d} *****'''  )
   lowerCamelCase_  				=						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
   lowerCamelCase_  				=						Path(pl_module.hparams.output_dir  )
   if type_path == "test":
    lowerCamelCase_  				=						od / "test_results.txt"
    lowerCamelCase_  				=						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.
    lowerCamelCase_  				=						od / f'''{type_path}_results/{trainer.global_step:05d}.txt'''
    lowerCamelCase_  				=						od / f'''{type_path}_generations/{trainer.global_step:05d}.txt'''
    results_file.parent.mkdir(exist_ok=UpperCamelCase  )
    generations_file.parent.mkdir(exist_ok=UpperCamelCase  )
   with open(UpperCamelCase    ,						"a+"  ) as writer:
    for key in sorted(UpperCamelCase  ):
     if key in ["log", "progress_bar", "preds"]:
      continue
     lowerCamelCase_  				=						metrics[key]
     if isinstance(UpperCamelCase    ,						torch.Tensor  ):
      lowerCamelCase_  				=						val.item()
     lowerCamelCase_  				=						f'''{key}: {val:.6f}\n'''
     writer.write(UpperCamelCase  )

   if not save_generations:
    return

   if "preds" in metrics:
    lowerCamelCase_  				=						"\n".join(metrics["preds"]  )
    generations_file.open("w+"  ).write(UpperCamelCase  )






  @rank_zero_only
  def 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   try:
    lowerCamelCase_  				=						pl_module.model.model.num_parameters()
   except AttributeError:
    lowerCamelCase_  				=						pl_module.model.num_parameters()

   lowerCamelCase_  				=						count_trainable_parameters(UpperCamelCase  )
   # 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    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   save_json(pl_module.metrics    ,						pl_module.metrics_save_path  )
   return self._write_logs(UpperCamelCase    ,						UpperCamelCase    ,						"test"  )







  @rank_zero_only
  def 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   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

import numpy as np

from transformers import DistilBertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
     import jax.numpy as jnp

     from transformers.models.distilbert.modeling_flax_distilbert import (
         FlaxDistilBertForMaskedLM,
         FlaxDistilBertForMultipleChoice,
         FlaxDistilBertForQuestionAnswering,
         FlaxDistilBertForSequenceClassification,
         FlaxDistilBertForTokenClassification,
         FlaxDistilBertModel,
     )





class   snake_case					(					unittest.TestCase    ):





  """simple docstring"""




  def __init__(					self    ,						UpperCamelCase    ,						UpperCamelCase=13    ,						UpperCamelCase=7    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						UpperCamelCase=99    ,						UpperCamelCase=32    ,						UpperCamelCase=5    ,						UpperCamelCase=4    ,						UpperCamelCase=37    ,						UpperCamelCase="gelu"    ,						UpperCamelCase=0.1    ,						UpperCamelCase=0.1    ,						UpperCamelCase=512    ,						UpperCamelCase=16    ,						UpperCamelCase=2    ,						UpperCamelCase=0.02    ,						UpperCamelCase=4    ,						):







   """simple docstring"""
   lowerCamelCase_  				=						parent
   lowerCamelCase_  				=						batch_size
   lowerCamelCase_  				=						seq_length
   lowerCamelCase_  				=						is_training
   lowerCamelCase_  				=						use_attention_mask
   lowerCamelCase_  				=						use_token_type_ids
   lowerCamelCase_  				=						use_labels
   lowerCamelCase_  				=						vocab_size
   lowerCamelCase_  				=						hidden_size
   lowerCamelCase_  				=						num_hidden_layers
   lowerCamelCase_  				=						num_attention_heads
   lowerCamelCase_  				=						intermediate_size
   lowerCamelCase_  				=						hidden_act
   lowerCamelCase_  				=						hidden_dropout_prob
   lowerCamelCase_  				=						attention_probs_dropout_prob
   lowerCamelCase_  				=						max_position_embeddings
   lowerCamelCase_  				=						type_vocab_size
   lowerCamelCase_  				=						type_sequence_label_size
   lowerCamelCase_  				=						initializer_range
   lowerCamelCase_  				=						num_choices






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size  )

   lowerCamelCase_  				=						None
   if self.use_attention_mask:
    lowerCamelCase_  				=						random_attention_mask([self.batch_size, self.seq_length]  )

   lowerCamelCase_  				=						DistilBertConfig(
       vocab_size=self.vocab_size    ,						dim=self.hidden_size    ,						n_layers=self.num_hidden_layers    ,						n_heads=self.num_attention_heads    ,						hidden_dim=self.intermediate_size    ,						hidden_act=self.hidden_act    ,						dropout=self.hidden_dropout_prob    ,						attention_dropout=self.attention_probs_dropout_prob    ,						max_position_embeddings=self.max_position_embeddings    ,						initializer_range=self.initializer_range    ,						tie_weights_=UpperCamelCase    ,						)

   return config, input_ids, attention_mask







  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.prepare_config_and_inputs()
   lowerCamelCase_				,lowerCamelCase_				,lowerCamelCase_  				=						config_and_inputs
   lowerCamelCase_  				=						{"input_ids": input_ids, "attention_mask": attention_mask}
   return config, inputs_dict





@require_flax
class   snake_case					(					lowercase    ,		unittest.TestCase    ):





  """simple docstring"""




  _lowerCamelCase						=						(
      (
          FlaxDistilBertModel,
          FlaxDistilBertForMaskedLM,
          FlaxDistilBertForMultipleChoice,
          FlaxDistilBertForQuestionAnswering,
          FlaxDistilBertForSequenceClassification,
          FlaxDistilBertForTokenClassification,
          FlaxDistilBertForQuestionAnswering,
      )
      if is_flax_available()
      else ()
  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						FlaxDistilBertModelTester(self  )







  @slow
  def 						snake_case	(					self  ):







   """simple docstring"""
   for model_class_name in self.all_model_classes:
    lowerCamelCase_  				=						model_class_name.from_pretrained("distilbert-base-uncased"  )
    lowerCamelCase_  				=						model(np.ones((1, 1)  )  )
    self.assertIsNotNone(UpperCamelCase  )





@require_flax
class   snake_case					(					unittest.TestCase    ):





  """simple docstring"""




  @slow
  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						FlaxDistilBertModel.from_pretrained("distilbert-base-uncased"  )
   lowerCamelCase_  				=						np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]  )
   lowerCamelCase_  				=						np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]  )
   lowerCamelCase_  				=						model(UpperCamelCase    ,						attention_mask=UpperCamelCase  )[0]
   lowerCamelCase_  				=						(1, 11, 768)
   self.assertEqual(output.shape    ,						UpperCamelCase  )
   lowerCamelCase_  				=						np.array([[[-0.1_639, 0.3_299, 0.1_648], [-0.1_746, 0.3_289, 0.1_710], [-0.1_884, 0.3_357, 0.1_810]]]  )

   self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]    ,						UpperCamelCase    ,						atol=1e-4  )  )



