

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




				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=512      ,   snake_case=16      ,   snake_case=2      ,   snake_case=0.02      ,   snake_case=3      ,   snake_case=4      ,   snake_case=None      ,   ):
											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_							=    hidden_size
											snake_case_							=    num_hidden_layers
											snake_case_							=    num_attention_heads
											snake_case_							=    intermediate_size
											snake_case_							=    hidden_act
											snake_case_							=    hidden_dropout_prob
											snake_case_							=    attention_probs_dropout_prob
											snake_case_							=    max_position_embeddings
											snake_case_							=    type_vocab_size
											snake_case_							=    type_sequence_label_size
											snake_case_							=    initializer_range
											snake_case_							=    num_labels
											snake_case_							=    num_choices
											snake_case_							=    scope




				def     a							(					self       ):
											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_							=    self.get_config()

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




				def     a							(					self       ):
											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=UpperCamelCase__      ,   initializer_range=self.initializer_range      ,   use_stable_embedding=UpperCamelCase__      ,   )




				def     a							(					self      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case       ):
											snake_case_							=    OpenLlamaModel(config=UpperCamelCase__       )
											model.to(UpperCamelCase__       )
											model.eval()
											snake_case_							=    model(UpperCamelCase__      ,   attention_mask=UpperCamelCase__       )
											snake_case_							=    model(UpperCamelCase__       )
											self.parent.assertEqual(result.last_hidden_state.shape      ,   (self.batch_size, self.seq_length, self.hidden_size)       )




				def     a							(					self      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   ):
											snake_case_							=    True
											snake_case_							=    OpenLlamaModel(UpperCamelCase__       )
											model.to(UpperCamelCase__       )
											model.eval()
											snake_case_							=    model(
											    UpperCamelCase__      ,   attention_mask=UpperCamelCase__      ,   encoder_hidden_states=UpperCamelCase__      ,   encoder_attention_mask=UpperCamelCase__      ,   )
											snake_case_							=    model(
											    UpperCamelCase__      ,   attention_mask=UpperCamelCase__      ,   encoder_hidden_states=UpperCamelCase__      ,   )
											snake_case_							=    model(UpperCamelCase__      ,   attention_mask=UpperCamelCase__       )
											self.parent.assertEqual(result.last_hidden_state.shape      ,   (self.batch_size, self.seq_length, self.hidden_size)       )




				def     a							(					self      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   ):
											snake_case_							=    OpenLlamaForCausalLM(config=UpperCamelCase__       )
											model.to(UpperCamelCase__       )
											model.eval()
											snake_case_							=    model(UpperCamelCase__      ,   attention_mask=UpperCamelCase__      ,   labels=UpperCamelCase__       )
											self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.seq_length, self.vocab_size)       )




				def     a							(					self      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   ):
											snake_case_							=    True
											snake_case_							=    True
											snake_case_							=    OpenLlamaForCausalLM(config=UpperCamelCase__       )
											model.to(UpperCamelCase__       )
											model.eval()

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

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

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

											snake_case_							=    model(
											    UpperCamelCase__      ,   attention_mask=UpperCamelCase__      ,   encoder_hidden_states=UpperCamelCase__      ,   encoder_attention_mask=UpperCamelCase__      ,   output_hidden_states=UpperCamelCase__      ,   )["hidden_states"][0]
											snake_case_							=    model(
											    UpperCamelCase__      ,   attention_mask=UpperCamelCase__      ,   encoder_hidden_states=UpperCamelCase__      ,   encoder_attention_mask=UpperCamelCase__      ,   past_key_values=UpperCamelCase__      ,   output_hidden_states=UpperCamelCase__      ,   )["hidden_states"][0]

											# select random slice
											snake_case_							=    ids_tensor((1,)      ,   output_from_past.shape[-1]       ).item()
											snake_case_							=    output_from_no_past[:, -3:, random_slice_idx].detach()
											snake_case_							=    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(UpperCamelCase__      ,   UpperCamelCase__      ,   atol=1e-3       )       )




				def     a							(					self       ):
											snake_case_							=    self.prepare_config_and_inputs()
											(
											    snake_case_ 
											)							=    config_and_inputs
											snake_case_							=    {"input_ids": input_ids, "attention_mask": input_mask}
											return config, inputs_dict







@require_torch
class  lowercase				(       __snake_case	,  __snake_case	,  __snake_case	,  unittest.TestCase			):
				__SCREAMING_SNAKE_CASE					: Tuple       			=			(
				    (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
				)
				__SCREAMING_SNAKE_CASE					: Any       			=			(OpenLlamaForCausalLM,) if is_torch_available() else ()
				__SCREAMING_SNAKE_CASE					: Optional[int]       			=			(
				    {
				        """feature-extraction""": OpenLlamaModel,
				        """text-classification""": OpenLlamaForSequenceClassification,
				        """text-generation""": OpenLlamaForCausalLM,
				        """zero-shot""": OpenLlamaForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				__SCREAMING_SNAKE_CASE					: List[str]       			=			False
				__SCREAMING_SNAKE_CASE					: List[Any]       			=			False




				def     a							(					self       ):
											snake_case_							=    OpenLlamaModelTester(self       )
											snake_case_							=    ConfigTester(self      ,   config_class=UpperCamelCase__      ,   hidden_size=37       )




				def     a							(					self       ):
											self.config_tester.run_common_tests()




				def     a							(					self       ):
											snake_case_							=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*UpperCamelCase__       )




				def     a							(					self       ):
											snake_case_							=    self.model_tester.prepare_config_and_inputs()
											for type in ["absolute", "relative_key", "relative_key_query"]:
																		snake_case_							=    type
																		self.model_tester.create_and_check_model(*UpperCamelCase__       )




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




				def     a							(					self       ):
											snake_case_							=    self.model_tester.prepare_config_and_inputs_for_common()
											snake_case_							=    3
											snake_case_							=    "single_label_classification"
											snake_case_							=    input_dict["input_ids"]
											snake_case_							=    input_ids.ne(1       ).to(UpperCamelCase__       )
											snake_case_							=    ids_tensor([self.model_tester.batch_size]      ,   self.model_tester.type_sequence_label_size       )
											snake_case_							=    OpenLlamaForSequenceClassification(UpperCamelCase__       )
											model.to(UpperCamelCase__       )
											model.eval()
											snake_case_							=    model(UpperCamelCase__      ,   attention_mask=UpperCamelCase__      ,   labels=UpperCamelCase__       )
											self.assertEqual(result.logits.shape      ,   (self.model_tester.batch_size, self.model_tester.num_labels)       )




				def     a							(					self       ):
											snake_case_							=    self.model_tester.prepare_config_and_inputs_for_common()
											snake_case_							=    3
											snake_case_							=    "multi_label_classification"
											snake_case_							=    input_dict["input_ids"]
											snake_case_							=    input_ids.ne(1       ).to(UpperCamelCase__       )
											snake_case_							=    ids_tensor(
											    [self.model_tester.batch_size, config.num_labels]      ,   self.model_tester.type_sequence_label_size       ).to(torch.float       )
											snake_case_							=    OpenLlamaForSequenceClassification(UpperCamelCase__       )
											model.to(UpperCamelCase__       )
											model.eval()
											snake_case_							=    model(UpperCamelCase__      ,   attention_mask=UpperCamelCase__      ,   labels=UpperCamelCase__       )
											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     a							(					self       ):
											pass




				@parameterized.expand([('linear',), ('dynamic',)]       )
				def     a							(					self      ,   snake_case       ):
											snake_case_							=    self.model_tester.prepare_config_and_inputs_for_common()
											snake_case_							=    ids_tensor([1, 10]      ,   config.vocab_size       )
											snake_case_							=    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
											snake_case_							=    OpenLlamaModel(UpperCamelCase__       )
											original_model.to(UpperCamelCase__       )
											original_model.eval()
											snake_case_							=    original_model(UpperCamelCase__       ).last_hidden_state
											snake_case_							=    original_model(UpperCamelCase__       ).last_hidden_state

											set_seed(42       )  # Fixed seed at init time so the two models get the same random weights
											snake_case_							=    {"type": scaling_type, "factor": 10.0}
											snake_case_							=    OpenLlamaModel(UpperCamelCase__       )
											scaled_model.to(UpperCamelCase__       )
											scaled_model.eval()
											snake_case_							=    scaled_model(UpperCamelCase__       ).last_hidden_state
											snake_case_							=    scaled_model(UpperCamelCase__       ).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(UpperCamelCase__      ,   UpperCamelCase__      ,   atol=1e-5       )       )
											else:
																		self.assertFalse(torch.allclose(UpperCamelCase__      ,   UpperCamelCase__      ,   atol=1e-5       )       )

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







"""simple docstring"""


from numpy import exp, pi, sqrt






def     a__    (    __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE = 0.0					, __SCREAMING_SNAKE_CASE = 1.0 )	->	int:
		return 1 / sqrt(2 * pi * sigma**2 ) * exp(-((x - mu) ** 2) / (2 * sigma**2) )


if __name__ == "__main__":
			import doctest

			doctest.testmod()


