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        lowerCamelCase__				:





			def __init__(       self	:	Any  ,  lowercase__	:	List[Any]  ,  lowercase__	:	int=13  ,  lowercase__	:	str=7  ,  lowercase__	:	Optional[Any]=True  ,  lowercase__	:	Optional[int]=True  ,  lowercase__	:	List[Any]=False  ,  lowercase__	:	int=True  ,  lowercase__	:	int=99  ,  lowercase__	:	Any=32  ,  lowercase__	:	Optional[int]=5  ,  lowercase__	:	Optional[int]=4  ,  lowercase__	:	List[str]=37  ,  lowercase__	:	int="gelu"  ,  lowercase__	:	List[Any]=0.1  ,  lowercase__	:	str=0.1  ,  lowercase__	:	str=5_12  ,  lowercase__	:	Union[str, Any]=16  ,  lowercase__	:	List[Any]=2  ,  lowercase__	:	List[Any]=0.0_2  ,  lowercase__	:	int=3  ,  lowercase__	:	Optional[Any]=4  ,  lowercase__	:	List[Any]=None  ,  ):
										_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   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]						):
										_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   SCREAMING_SNAKE_CASE__     (       self	:	str						):
										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=lowercase__  ,  initializer_range=self.initializer_range  ,  use_stable_embedding=lowercase__  ,  )





			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict  ,  lowercase__	:	Dict  ,  lowercase__	:	Tuple  ,  lowercase__	:	Any  ,  lowercase__	:	Any  ,  lowercase__	:	List[str]  ,  lowercase__	:	str  ,  lowercase__	:	List[Any]						):
										_lowerCAmelCase							=							OpenLlamaModel(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__						)
										_lowerCAmelCase							=							model(lowercase__						)
										self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[int]  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	List[str]  ,  lowercase__	:	Any  ,  lowercase__	:	List[str]  ,  lowercase__	:	int  ,  lowercase__	:	Dict  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	Optional[Any]  ,  ):
										_lowerCAmelCase							=							True
										_lowerCAmelCase							=							OpenLlamaModel(lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(
										    lowercase__  ,  attention_mask=lowercase__  ,  encoder_hidden_states=lowercase__  ,  encoder_attention_mask=lowercase__  ,  )
										_lowerCAmelCase							=							model(
										    lowercase__  ,  attention_mask=lowercase__  ,  encoder_hidden_states=lowercase__  ,  )
										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__						)
										self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[int]  ,  lowercase__	:	int  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	Dict  ,  lowercase__	:	List[str]  ,  lowercase__	:	List[str]  ,  lowercase__	:	Dict  ,  lowercase__	:	Optional[int]  ,  ):
										_lowerCAmelCase							=							OpenLlamaForCausalLM(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__  ,  labels=lowercase__						)
										self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.vocab_size)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict  ,  lowercase__	:	List[Any]  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	str  ,  lowercase__	:	Any  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	int  ,  ):
										_lowerCAmelCase							=							True
										_lowerCAmelCase							=							True
										_lowerCAmelCase							=							OpenLlamaForCausalLM(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()

										# first forward pass
										_lowerCAmelCase							=							model(
										    lowercase__  ,  attention_mask=lowercase__  ,  encoder_hidden_states=lowercase__  ,  encoder_attention_mask=lowercase__  ,  use_cache=lowercase__  ,  )
										_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(
										    lowercase__  ,  attention_mask=lowercase__  ,  encoder_hidden_states=lowercase__  ,  encoder_attention_mask=lowercase__  ,  output_hidden_states=lowercase__  ,  )["""hidden_states"""][0]
										_lowerCAmelCase							=							model(
										    lowercase__  ,  attention_mask=lowercase__  ,  encoder_hidden_states=lowercase__  ,  encoder_attention_mask=lowercase__  ,  past_key_values=lowercase__  ,  output_hidden_states=lowercase__  ,  )["""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(lowercase__  ,  lowercase__  ,  atol=1e-3						)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[int]						):
										_lowerCAmelCase							=							self.prepare_config_and_inputs()
										(
										    _lowerCAmelCase			
										)							=							config_and_inputs
										_lowerCAmelCase							=							{"""input_ids""": input_ids, """attention_mask""": input_mask}
										return config, inputs_dict







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





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]						):
										_lowerCAmelCase							=							OpenLlamaModelTester(self						)
										_lowerCAmelCase							=							ConfigTester(self  ,  config_class=lowercase__  ,  hidden_size=37						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										self.config_tester.run_common_tests()





			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	str						):
										_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(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs_for_common()
										_lowerCAmelCase							=							3
										_lowerCAmelCase							=							input_dict["""input_ids"""]
										_lowerCAmelCase							=							input_ids.ne(1						).to(lowercase__						)
										_lowerCAmelCase							=							ids_tensor([self.model_tester.batch_size]  ,  self.model_tester.type_sequence_label_size						)
										_lowerCAmelCase							=							OpenLlamaForSequenceClassification(lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__  ,  labels=lowercase__						)
										self.assertEqual(result.logits.shape  ,  (self.model_tester.batch_size, self.model_tester.num_labels)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Any						):
										_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(lowercase__						)
										_lowerCAmelCase							=							ids_tensor([self.model_tester.batch_size]  ,  self.model_tester.type_sequence_label_size						)
										_lowerCAmelCase							=							OpenLlamaForSequenceClassification(lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__  ,  labels=lowercase__						)
										self.assertEqual(result.logits.shape  ,  (self.model_tester.batch_size, self.model_tester.num_labels)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[int]						):
										_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(lowercase__						)
										_lowerCAmelCase							=							ids_tensor(
										    [self.model_tester.batch_size, config.num_labels]  ,  self.model_tester.type_sequence_label_size						).to(torch.float						)
										_lowerCAmelCase							=							OpenLlamaForSequenceClassification(lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__  ,  labels=lowercase__						)
										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   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										pass





			@parameterized.expand([('linear',), ('dynamic',)]						)
			def   SCREAMING_SNAKE_CASE__     (       self	:	int  ,  lowercase__	:	Union[str, Any]						):
										_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(lowercase__						)
										original_model.to(lowercase__						)
										original_model.eval()
										_lowerCAmelCase							=							original_model(lowercase__						).last_hidden_state
										_lowerCAmelCase							=							original_model(lowercase__						).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""": 1_0.0}
										_lowerCAmelCase							=							OpenLlamaModel(lowercase__						)
										scaled_model.to(lowercase__						)
										scaled_model.eval()
										_lowerCAmelCase							=							scaled_model(lowercase__						).last_hidden_state
										_lowerCAmelCase							=							scaled_model(lowercase__						).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(lowercase__  ,  lowercase__  ,  atol=1e-5						)						)
										else:
																	self.assertFalse(torch.allclose(lowercase__  ,  lowercase__  ,  atol=1e-5						)						)

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












import logging
import os
from logging import (
    CRITICAL,  # NOQA
    DEBUG,  # NOQA
    ERROR,  # NOQA
    FATAL,  # NOQA
    INFO,  # NOQA
    NOTSET,  # NOQA
    WARN,  # NOQA
    WARNING,  # NOQA
)
from typing import Optional

from tqdm import auto as tqdm_lib


a__		     =     {
    '''debug''': logging.DEBUG,
    '''info''': logging.INFO,
    '''warning''': logging.WARNING,
    '''error''': logging.ERROR,
    '''critical''': logging.CRITICAL,
}

a__		     =     logging.WARNING





def 				A__ ()     ->						Any:
					__UpperCamelCase					:	Union[str, Any]             =      os.getenv("""DATASETS_VERBOSITY"""		,	snake_case				)
					if env_level_str:
										if env_level_str in log_levels:
															return log_levels[env_level_str]
										else:
															logging.getLogger().warning(
															    F'''Unknown option DATASETS_VERBOSITY={env_level_str}, '''
															    F'''has to be one of: { ", ".join(log_levels.keys()				) }'''				)
					return _default_log_level





def 				A__ ()     ->						str:
					return __name__.split("""."""				)[0]





def 				A__ ()     ->						logging.Logger:
					return logging.getLogger(_get_library_name()				)





def 				A__ ()     ->						None:
					# Apply our default configuration to the library root logger.
					__UpperCamelCase					:	Any             =      _get_library_root_logger()
					library_root_logger.setLevel(_get_default_logging_level()				)





def 				A__ ()     ->						None:
					__UpperCamelCase					:	List[Any]             =      _get_library_root_logger()
					library_root_logger.setLevel(logging.NOTSET				)





def 				A__ (snake_case   :		Optional[str] = None				)     ->						logging.Logger:
					if name is None:
										__UpperCamelCase					:	Optional[int]             =      _get_library_name()
					return logging.getLogger(snake_case				)





def 				A__ ()     ->						int:
					return _get_library_root_logger().getEffectiveLevel()





def 				A__ (snake_case   :		int				)     ->						None:
					_get_library_root_logger().setLevel(snake_case				)





def 				A__ ()     ->						Dict:
					return set_verbosity(snake_case				)





def 				A__ ()     ->						Any:
					return set_verbosity(snake_case				)





def 				A__ ()     ->						Dict:
					return set_verbosity(snake_case				)





def 				A__ ()     ->						str:
					return set_verbosity(snake_case				)





def 				A__ ()     ->						None:
					__UpperCamelCase					:	str             =      False





def 				A__ ()     ->						None:
					__UpperCamelCase					:	List[str]             =      True


# Configure the library root logger at the module level (singleton-like)
_configure_library_root_logger()


class    SCREAMING_SNAKE_CASE_				:



	"""simple docstring"""


	def __init__(					self		:						Optional[int]					,     *lowerCAmelCase		:						str					,     **lowerCAmelCase		:						List[str]	)	->       Union[str, Any]:  # pylint: disable=unused-argument







						"""simple docstring"""







						__UpperCamelCase					:	List[str]             =      args[0] if args else None






	def __iter__(					self		:						List[Any]	)	->       List[Any]:







						"""simple docstring"""







						return iter(self._iterator	)






	def __getattr__(					self		:						int					,     lowerCAmelCase		:						List[str]	)	->       Dict:







						"""simple docstring"""







						def empty_fn(*lowerCAmelCase		:						List[Any]					,     **lowerCAmelCase		:						Dict	):  # pylint: disable=unused-argument
											return

						return empty_fn






	def __enter__(					self		:						Tuple	)	->       Any:







						"""simple docstring"""







						return self






	def __exit__(					self		:						str					,     lowerCAmelCase		:						Dict					,     lowerCAmelCase		:						Dict					,     lowerCAmelCase		:						List[str]	)	->       Tuple:







						"""simple docstring"""







						return






a__		     =     True


class    SCREAMING_SNAKE_CASE_				:



	"""simple docstring"""


	def __call__(					self		:						str					,     *lowerCAmelCase		:						List[Any]					,     lowerCAmelCase		:						Dict=False					,     **lowerCAmelCase		:						Tuple	)	->       Optional[Any]:







						"""simple docstring"""







						if _tqdm_active and not disable:
											return tqdm_lib.tqdm(*lowerCAmelCase					,     **lowerCAmelCase	)
						else:
											return EmptyTqdm(*lowerCAmelCase					,     **lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Any					,     *lowerCAmelCase		:						List[str]					,     **lowerCAmelCase		:						str	)	->       Dict:







						"""simple docstring"""







						__UpperCamelCase					:	Any             =      None
						if _tqdm_active:
											return tqdm_lib.tqdm.set_lock(*lowerCAmelCase					,     **lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Dict	)	->       Union[str, Any]:







						"""simple docstring"""







						if _tqdm_active:
											return tqdm_lib.tqdm.get_lock()






a__		     =     _tqdm_cls()





def 				A__ ()     ->						bool:
					global _tqdm_active
					return bool(_tqdm_active				)





def 				A__ ()     ->						List[str]:
					global _tqdm_active
					__UpperCamelCase					:	Tuple             =      True







def 				A__ ()     ->						int:
					global _tqdm_active
					__UpperCamelCase					:	Any             =      False






