

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     __SCREAMING_SNAKE_CASE   (SCREAMING_SNAKE_CASE__			):
						snake_case_  =      filter(lambda SCREAMING_SNAKE_CASE__			: p.requires_grad    ,						model.parameters()			)
						snake_case_  =      sum([np.prod(p.size()			) for p in model_parameters]			)
						return params


lowerCAmelCase_			      =      logging.getLogger(__name__)


def     __SCREAMING_SNAKE_CASE   (SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__			):
						if metric == "rouge2":
												snake_case_  =      '''{val_avg_rouge2:.4f}-{step_count}'''
						elif metric == "bleu":
												snake_case_  =      '''{val_avg_bleu:.4f}-{step_count}'''
						elif metric == "em":
												snake_case_  =      '''{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.'''			)

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


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





class       snake_case_  (						pl.Callback   ):







			'''simple docstring'''

			def 							snake_case__(	self					:							Dict				,  _UpperCamelCase					:							Union[str, Any]				,  _UpperCamelCase					:							List[str]				) ->List[Any]:
									snake_case_  =      {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					:							Union[str, Any]				,  _UpperCamelCase					:							pl.Trainer				,  _UpperCamelCase					:							pl.LightningModule				,  _UpperCamelCase					:							str				,  _UpperCamelCase					:							List[str]=True				) ->None:
									logger.info(f'''***** {type_path} results at step {trainer.global_step:05d} *****'''				)
									snake_case_  =      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
									snake_case_  =      Path(pl_module.hparams.output_dir				)
									if type_path == "test":
															snake_case_  =      od / '''test_results.txt'''
															snake_case_  =      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.
															snake_case_  =      od / f'''{type_path}_results/{trainer.global_step:05d}.txt'''
															snake_case_  =      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
																					snake_case_  =      metrics[key]
																					if isinstance(_UpperCamelCase				,  torch.Tensor				):
																											snake_case_  =      val.item()
																					snake_case_  =      f'''{key}: {val:.6f}\n'''
																					writer.write(_UpperCamelCase				)

									if not save_generations:
															return

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







			@rank_zero_only
			def 							snake_case__(	self					:							List[str]				,  _UpperCamelCase					:							int				,  _UpperCamelCase					:							str				) ->Optional[Any]:
									try:
															snake_case_  =      pl_module.model.model.num_parameters()
									except AttributeError:
															snake_case_  =      pl_module.model.num_parameters()

									snake_case_  =      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					:							Tuple				,  _UpperCamelCase					:							pl.Trainer				,  _UpperCamelCase					:							pl.LightningModule				) ->str:
									save_json(pl_module.metrics				,  pl_module.metrics_save_path				)
									return self._write_logs(_UpperCamelCase				,  _UpperCamelCase				,  '''test'''				)







			@rank_zero_only
			def 							snake_case__(	self					:							Tuple				,  _UpperCamelCase					:							pl.Trainer				,  _UpperCamelCase					:							Tuple				) ->Dict:
									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 parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel







class       a							:



 def __init__(     self				,       _snake_case				,       _snake_case=13				,       _snake_case=7				,       _snake_case=True				,       _snake_case=True				,       _snake_case=False				,       _snake_case=True				,       _snake_case=99				,       _snake_case=32				,       _snake_case=5				,       _snake_case=4				,       _snake_case=37				,       _snake_case="gelu"				,       _snake_case=0.1				,       _snake_case=0.1				,       _snake_case=5_12				,       _snake_case=16				,       _snake_case=2				,       _snake_case=0.02				,       _snake_case=3				,       _snake_case=4				,       _snake_case=None				,       ):







  """simple docstring"""






  lowerCAmelCase    				= parent
  lowerCAmelCase    				= batch_size
  lowerCAmelCase    				= seq_length
  lowerCAmelCase    				= is_training
  lowerCAmelCase    				= use_input_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_labels
  lowerCAmelCase    				= num_choices
  lowerCAmelCase    				= scope



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= ids_tensor([self.batch_size, self.seq_length]				,       self.vocab_size     )

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

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

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

  lowerCAmelCase    				= self.get_config()

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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  return OpenLlamaConfig(
      vocab_size=self.vocab_size				,       hidden_size=self.hidden_size				,       num_hidden_layers=self.num_hidden_layers				,       num_attention_heads=self.num_attention_heads				,       intermediate_size=self.intermediate_size				,       hidden_act=self.hidden_act				,       hidden_dropout_prob=self.hidden_dropout_prob				,       attention_probs_dropout_prob=self.attention_probs_dropout_prob				,       max_position_embeddings=self.max_position_embeddings				,       type_vocab_size=self.type_vocab_size				,       is_decoder=_snake_case				,       initializer_range=self.initializer_range				,       use_stable_embedding=_snake_case				,       )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenLlamaModel(config=_snake_case     )
  model.to(_snake_case     )
  model.eval()
  lowerCAmelCase    				= model(_snake_case				,       attention_mask=_snake_case     )
  lowerCAmelCase    				= model(_snake_case     )
  self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       ):







  """simple docstring"""






  lowerCAmelCase    				= True
  lowerCAmelCase    				= OpenLlamaModel(_snake_case     )
  model.to(_snake_case     )
  model.eval()
  lowerCAmelCase    				= model(
      _snake_case				,       attention_mask=_snake_case				,       encoder_hidden_states=_snake_case				,       encoder_attention_mask=_snake_case				,       )
  lowerCAmelCase    				= model(
      _snake_case				,       attention_mask=_snake_case				,       encoder_hidden_states=_snake_case				,       )
  lowerCAmelCase    				= model(_snake_case				,       attention_mask=_snake_case     )
  self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       ):







  """simple docstring"""






  lowerCAmelCase    				= OpenLlamaForCausalLM(config=_snake_case     )
  model.to(_snake_case     )
  model.eval()
  lowerCAmelCase    				= model(_snake_case				,       attention_mask=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       ):







  """simple docstring"""






  lowerCAmelCase    				= True
  lowerCAmelCase    				= True
  lowerCAmelCase    				= OpenLlamaForCausalLM(config=_snake_case     )
  model.to(_snake_case     )
  model.eval()

  # first forward pass
  lowerCAmelCase    				= model(
      _snake_case				,       attention_mask=_snake_case				,       encoder_hidden_states=_snake_case				,       encoder_attention_mask=_snake_case				,       use_cache=_snake_case				,       )
  lowerCAmelCase    				= outputs.past_key_values

  # create hypothetical multiple next token and extent to next_input_ids
  lowerCAmelCase    				= ids_tensor((self.batch_size, 3)				,       config.vocab_size     )
  lowerCAmelCase    				= ids_tensor((self.batch_size, 3)				,       vocab_size=2     )

  # append to next input_ids and
  lowerCAmelCase    				= torch.cat([input_ids, next_tokens]				,       dim=-1     )
  lowerCAmelCase    				= torch.cat([input_mask, next_mask]				,       dim=-1     )

  lowerCAmelCase    				= model(
      _snake_case				,       attention_mask=_snake_case				,       encoder_hidden_states=_snake_case				,       encoder_attention_mask=_snake_case				,       output_hidden_states=_snake_case				,       )['hidden_states'][0]
  lowerCAmelCase    				= model(
      _snake_case				,       attention_mask=_snake_case				,       encoder_hidden_states=_snake_case				,       encoder_attention_mask=_snake_case				,       past_key_values=_snake_case				,       output_hidden_states=_snake_case				,       )['hidden_states'][0]

  # select random slice
  lowerCAmelCase    				= ids_tensor((1,)				,       output_from_past.shape[-1]     ).item()
  lowerCAmelCase    				= output_from_no_past[:, -3:, random_slice_idx].detach()
  lowerCAmelCase    				= output_from_past[:, :, random_slice_idx].detach()

  self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]     )

  # test that outputs are equal for slice
  self.parent.assertTrue(torch.allclose(_snake_case				,       _snake_case				,       atol=1E-3     )     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.prepare_config_and_inputs()
  (
      (
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,
  )    				= config_and_inputs
  lowerCAmelCase    				= {'input_ids': input_ids, 'attention_mask': input_mask}
  return config, inputs_dict







@require_torch
class       a							(   a__   , a__   , a__   , unittest.TestCase  ):
 snake_case__             =       (
     (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
 )
 snake_case__             =       (OpenLlamaForCausalLM,) if is_torch_available() else ()
 snake_case__             =       (
     {
         '''feature-extraction''': OpenLlamaModel,
         '''text-classification''': OpenLlamaForSequenceClassification,
         '''text-generation''': OpenLlamaForCausalLM,
         '''zero-shot''': OpenLlamaForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 snake_case__             =       False
 snake_case__             =       False



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenLlamaModelTester(self     )
  lowerCAmelCase    				= ConfigTester(self				,       config_class=_snake_case				,       hidden_size=37     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  self.config_tester.run_common_tests()



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_model(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  for type in ["absolute", "relative_key", "relative_key_query"]:
   lowerCAmelCase    				= type
   self.model_tester.create_and_check_model(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase				,lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs_for_common()
  lowerCAmelCase    				= 3
  lowerCAmelCase    				= input_dict['input_ids']
  lowerCAmelCase    				= input_ids.ne(1     ).to(_snake_case     )
  lowerCAmelCase    				= ids_tensor([self.model_tester.batch_size]				,       self.model_tester.type_sequence_label_size     )
  lowerCAmelCase    				= OpenLlamaForSequenceClassification(_snake_case     )
  model.to(_snake_case     )
  model.eval()
  lowerCAmelCase    				= model(_snake_case				,       attention_mask=_snake_case				,       labels=_snake_case     )
  self.assertEqual(result.logits.shape				,       (self.model_tester.batch_size, self.model_tester.num_labels)     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase				,lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs_for_common()
  lowerCAmelCase    				= 3
  lowerCAmelCase    				= 'single_label_classification'
  lowerCAmelCase    				= input_dict['input_ids']
  lowerCAmelCase    				= input_ids.ne(1     ).to(_snake_case     )
  lowerCAmelCase    				= ids_tensor([self.model_tester.batch_size]				,       self.model_tester.type_sequence_label_size     )
  lowerCAmelCase    				= OpenLlamaForSequenceClassification(_snake_case     )
  model.to(_snake_case     )
  model.eval()
  lowerCAmelCase    				= model(_snake_case				,       attention_mask=_snake_case				,       labels=_snake_case     )
  self.assertEqual(result.logits.shape				,       (self.model_tester.batch_size, self.model_tester.num_labels)     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase				,lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs_for_common()
  lowerCAmelCase    				= 3
  lowerCAmelCase    				= 'multi_label_classification'
  lowerCAmelCase    				= input_dict['input_ids']
  lowerCAmelCase    				= input_ids.ne(1     ).to(_snake_case     )
  lowerCAmelCase    				= ids_tensor(
      [self.model_tester.batch_size, config.num_labels]				,       self.model_tester.type_sequence_label_size     ).to(torch.float     )
  lowerCAmelCase    				= OpenLlamaForSequenceClassification(_snake_case     )
  model.to(_snake_case     )
  model.eval()
  lowerCAmelCase    				= model(_snake_case				,       attention_mask=_snake_case				,       labels=_snake_case     )
  self.assertEqual(result.logits.shape				,       (self.model_tester.batch_size, self.model_tester.num_labels)     )



 @unittest.skip('Open-Llama buffers include complex numbers, which breaks this test'     )
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  pass



 @parameterized.expand([('linear',), ('dynamic',)]     )
 def    UpperCamelCase__  (     self				,       _snake_case     ):







  """simple docstring"""






  lowerCAmelCase				,lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs_for_common()
  lowerCAmelCase    				= ids_tensor([1, 10]				,       config.vocab_size     )
  lowerCAmelCase    				= ids_tensor([1, int(config.max_position_embeddings * 1.5     )]				,       config.vocab_size     )

  set_seed(42     )  # Fixed seed at init time so the two models get the same random weights
  lowerCAmelCase    				= OpenLlamaModel(_snake_case     )
  original_model.to(_snake_case     )
  original_model.eval()
  lowerCAmelCase    				= original_model(_snake_case     ).last_hidden_state
  lowerCAmelCase    				= original_model(_snake_case     ).last_hidden_state

  set_seed(42     )  # Fixed seed at init time so the two models get the same random weights
  lowerCAmelCase    				= {'type': scaling_type, 'factor': 10.0}
  lowerCAmelCase    				= OpenLlamaModel(_snake_case     )
  scaled_model.to(_snake_case     )
  scaled_model.eval()
  lowerCAmelCase    				= scaled_model(_snake_case     ).last_hidden_state
  lowerCAmelCase    				= scaled_model(_snake_case     ).last_hidden_state

  # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
  # maximum sequence length, so the outputs for the short input should match.
  if scaling_type == "dynamic":
   self.assertTrue(torch.allclose(_snake_case				,       _snake_case				,       atol=1E-5     )     )
  else:
   self.assertFalse(torch.allclose(_snake_case				,       _snake_case				,       atol=1E-5     )     )

  # The output should be different for long inputs
  self.assertFalse(torch.allclose(_snake_case				,       _snake_case				,       atol=1E-5     )     )





