






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







		def __init__(					self       ,A__       ,A__=13       ,A__=7       ,A__=True       ,A__=True       ,A__=False       ,A__=True       ,A__=99       ,A__=32       ,A__=5       ,A__=4       ,A__=37       ,A__="gelu"       ,A__=0.1       ,A__=0.1       ,A__=512       ,A__=16       ,A__=2       ,A__=0.02       ,A__=3       ,A__=4       ,A__=None       ,):
						_A	:				int   =      parent
						_A	:				Union[str, Any]   =      batch_size
						_A	:				List[Any]   =      seq_length
						_A	:				Any   =      is_training
						_A	:				Any   =      use_input_mask
						_A	:				int   =      use_token_type_ids
						_A	:				Any   =      use_labels
						_A	:				Dict   =      vocab_size
						_A	:				Optional[Any]   =      hidden_size
						_A	:				Optional[int]   =      num_hidden_layers
						_A	:				Tuple   =      num_attention_heads
						_A	:				str   =      intermediate_size
						_A	:				Optional[Any]   =      hidden_act
						_A	:				Dict   =      hidden_dropout_prob
						_A	:				List[Any]   =      attention_probs_dropout_prob
						_A	:				Union[str, Any]   =      max_position_embeddings
						_A	:				Optional[int]   =      type_vocab_size
						_A	:				str   =      type_sequence_label_size
						_A	:				int   =      initializer_range
						_A	:				Any   =      num_labels
						_A	:				Union[str, Any]   =      num_choices
						_A	:				Dict   =      scope







		def  A__      (					self		):
						_A	:				List[Any]   =      ids_tensor([self.batch_size, self.seq_length]       ,self.vocab_size		)

						_A	:				Any   =      None
						if self.use_input_mask:
										_A	:				List[Any]   =      random_attention_mask([self.batch_size, self.seq_length]		)

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

						_A	:				Union[str, Any]   =      None
						_A	:				Union[str, Any]   =      None
						_A	:				Any   =      None
						if self.use_labels:
										_A	:				Dict   =      ids_tensor([self.batch_size]       ,self.type_sequence_label_size		)
										_A	:				Optional[int]   =      ids_tensor([self.batch_size, self.seq_length]       ,self.num_labels		)
										_A	:				List[str]   =      ids_tensor([self.batch_size]       ,self.num_choices		)

						_A	:				List[str]   =      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=A__       ,initializer_range=self.initializer_range       ,use_stable_embedding=A__       ,)







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







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







		def  A__      (					self       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,):
						_A	:				List[Any]   =      OpenLlamaForCausalLM(config=A__		)
						model.to(A__		)
						model.eval()
						_A	:				str   =      model(A__       ,attention_mask=A__       ,labels=A__		)
						self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.seq_length, self.vocab_size)		)







		def  A__      (					self       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,):
						_A	:				Dict   =      True
						_A	:				Tuple   =      True
						_A	:				int   =      OpenLlamaForCausalLM(config=A__		)
						model.to(A__		)
						model.eval()

						# first forward pass
						_A	:				Any   =      model(
						    A__       ,attention_mask=A__       ,encoder_hidden_states=A__       ,encoder_attention_mask=A__       ,use_cache=A__       ,)
						_A	:				Tuple   =      outputs.past_key_values

						# create hypothetical multiple next token and extent to next_input_ids
						_A	:				Union[str, Any]   =      ids_tensor((self.batch_size, 3)       ,config.vocab_size		)
						_A	:				List[str]   =      ids_tensor((self.batch_size, 3)       ,vocab_size=2		)

						# append to next input_ids and
						_A	:				List[str]   =      torch.cat([input_ids, next_tokens]       ,dim=-1		)
						_A	:				Union[str, Any]   =      torch.cat([input_mask, next_mask]       ,dim=-1		)

						_A	:				List[str]   =      model(
						    A__       ,attention_mask=A__       ,encoder_hidden_states=A__       ,encoder_attention_mask=A__       ,output_hidden_states=A__       ,)['''hidden_states'''][0]
						_A	:				Optional[int]   =      model(
						    A__       ,attention_mask=A__       ,encoder_hidden_states=A__       ,encoder_attention_mask=A__       ,past_key_values=A__       ,output_hidden_states=A__       ,)['''hidden_states'''][0]

						# select random slice
						_A	:				List[str]   =      ids_tensor((1,)       ,output_from_past.shape[-1]		).item()
						_A	:				Tuple   =      output_from_no_past[:, -3:, random_slice_idx].detach()
						_A	:				Dict   =      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(A__       ,A__       ,atol=1E-3		)		)







		def  A__      (					self		):
						_A	:				List[Any]   =      self.prepare_config_and_inputs()
						(
						    (
						    _A
						)  ,				(
						    _A
						)  ,				(
						    _A
						)  ,				(
						    _A
						)  ,				(
						    _A
						)  ,				(
						    _A
						)  ,				(
						    _A
						)  ,				
						)	:				Optional[Any]   =      config_and_inputs
						_A	:				int   =      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
						return config, inputs_dict






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







		def  A__      (					self		):
						_A	:				int   =      OpenLlamaModelTester(self		)
						_A	:				List[str]   =      ConfigTester(self       ,config_class=A__       ,hidden_size=37		)







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







		def  A__      (					self		):
						_A	:				Dict   =      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*A__		)







		def  A__      (					self		):
						_A	:				Optional[Any]   =      self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
										_A	:				List[str]   =      type
										self.model_tester.create_and_check_model(*A__		)







		def  A__      (					self		):
						_A  ,				_A	:				List[str]   =      self.model_tester.prepare_config_and_inputs_for_common()
						_A	:				List[Any]   =      3
						_A	:				Optional[int]   =      input_dict['''input_ids''']
						_A	:				Any   =      input_ids.ne(1		).to(A__		)
						_A	:				int   =      ids_tensor([self.model_tester.batch_size]       ,self.model_tester.type_sequence_label_size		)
						_A	:				Any   =      OpenLlamaForSequenceClassification(A__		)
						model.to(A__		)
						model.eval()
						_A	:				Dict   =      model(A__       ,attention_mask=A__       ,labels=A__		)
						self.assertEqual(result.logits.shape       ,(self.model_tester.batch_size, self.model_tester.num_labels)		)







		def  A__      (					self		):
						_A  ,				_A	:				List[str]   =      self.model_tester.prepare_config_and_inputs_for_common()
						_A	:				Optional[Any]   =      3
						_A	:				List[Any]   =      '''single_label_classification'''
						_A	:				Any   =      input_dict['''input_ids''']
						_A	:				Optional[Any]   =      input_ids.ne(1		).to(A__		)
						_A	:				Dict   =      ids_tensor([self.model_tester.batch_size]       ,self.model_tester.type_sequence_label_size		)
						_A	:				List[str]   =      OpenLlamaForSequenceClassification(A__		)
						model.to(A__		)
						model.eval()
						_A	:				Tuple   =      model(A__       ,attention_mask=A__       ,labels=A__		)
						self.assertEqual(result.logits.shape       ,(self.model_tester.batch_size, self.model_tester.num_labels)		)







		def  A__      (					self		):
						_A  ,				_A	:				List[str]   =      self.model_tester.prepare_config_and_inputs_for_common()
						_A	:				Optional[Any]   =      3
						_A	:				Dict   =      '''multi_label_classification'''
						_A	:				List[Any]   =      input_dict['''input_ids''']
						_A	:				Any   =      input_ids.ne(1		).to(A__		)
						_A	:				Any   =      ids_tensor(
						    [self.model_tester.batch_size, config.num_labels]       ,self.model_tester.type_sequence_label_size		).to(torch.float		)
						_A	:				List[Any]   =      OpenLlamaForSequenceClassification(A__		)
						model.to(A__		)
						model.eval()
						_A	:				List[Any]   =      model(A__       ,attention_mask=A__       ,labels=A__		)
						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       ,A__		):
						_A  ,				_A	:				Optional[Any]   =      self.model_tester.prepare_config_and_inputs_for_common()
						_A	:				List[str]   =      ids_tensor([1, 10]       ,config.vocab_size		)
						_A	:				Dict   =      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
						_A	:				List[str]   =      OpenLlamaModel(A__		)
						original_model.to(A__		)
						original_model.eval()
						_A	:				Any   =      original_model(A__		).last_hidden_state
						_A	:				List[Any]   =      original_model(A__		).last_hidden_state

						set_seed(42		)  # Fixed seed at init time so the two models get the same random weights
						_A	:				Any   =      {'''type''': scaling_type, '''factor''': 10.0}
						_A	:				Optional[Any]   =      OpenLlamaModel(A__		)
						scaled_model.to(A__		)
						scaled_model.eval()
						_A	:				str   =      scaled_model(A__		).last_hidden_state
						_A	:				Union[str, Any]   =      scaled_model(A__		).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(A__       ,A__       ,atol=1E-5		)		)
						else:
										self.assertFalse(torch.allclose(A__       ,A__       ,atol=1E-5		)		)

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











import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				Optional[Any]   =      []
						_A	:				str   =      []

						for i in range(self.num_layers		):
										_A	:				Union[str, Any]   =      self.in_channels if i == 0 else self.out_channels

										_A	:				str   =      FlaxResnetBlockaD(
										    in_channels=A__       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

										_A	:				List[str]   =      FlaxTransformeraDModel(
										    in_channels=self.out_channels       ,n_heads=self.num_attention_heads       ,d_head=self.out_channels // self.num_attention_heads       ,depth=1       ,use_linear_projection=self.use_linear_projection       ,only_cross_attention=self.only_cross_attention       ,use_memory_efficient_attention=self.use_memory_efficient_attention       ,dtype=self.dtype       ,)
										attentions.append(A__		)

						_A	:				List[str]   =      resnets
						_A	:				Any   =      attentions

						if self.add_downsample:
										_A	:				Any   =      FlaxDownsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						_A	:				List[str]   =      ()

						for resnet, attn in zip(self.resnets       ,self.attentions		):
										_A	:				Optional[Any]   =      resnet(A__       ,A__       ,deterministic=A__		)
										_A	:				Tuple   =      attn(A__       ,A__       ,deterministic=A__		)
										output_states += (hidden_states,)

						if self.add_downsample:
										_A	:				Optional[int]   =      self.downsamplers_a(A__		)
										output_states += (hidden_states,)

						return hidden_states, output_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				List[Any]   =      []

						for i in range(self.num_layers		):
										_A	:				int   =      self.in_channels if i == 0 else self.out_channels

										_A	:				List[Any]   =      FlaxResnetBlockaD(
										    in_channels=A__       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)
						_A	:				Tuple   =      resnets

						if self.add_downsample:
										_A	:				Optional[Any]   =      FlaxDownsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__=True		):
						_A	:				List[Any]   =      ()

						for resnet in self.resnets:
										_A	:				Optional[int]   =      resnet(A__       ,A__       ,deterministic=A__		)
										output_states += (hidden_states,)

						if self.add_downsample:
										_A	:				List[Any]   =      self.downsamplers_a(A__		)
										output_states += (hidden_states,)

						return hidden_states, output_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				List[str]   =      []
						_A	:				Optional[Any]   =      []

						for i in range(self.num_layers		):
										_A	:				Dict   =      self.in_channels if (i == self.num_layers - 1) else self.out_channels
										_A	:				List[str]   =      self.prev_output_channel if i == 0 else self.out_channels

										_A	:				Optional[int]   =      FlaxResnetBlockaD(
										    in_channels=resnet_in_channels + res_skip_channels       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

										_A	:				Union[str, Any]   =      FlaxTransformeraDModel(
										    in_channels=self.out_channels       ,n_heads=self.num_attention_heads       ,d_head=self.out_channels // self.num_attention_heads       ,depth=1       ,use_linear_projection=self.use_linear_projection       ,only_cross_attention=self.only_cross_attention       ,use_memory_efficient_attention=self.use_memory_efficient_attention       ,dtype=self.dtype       ,)
										attentions.append(A__		)

						_A	:				Dict   =      resnets
						_A	:				int   =      attentions

						if self.add_upsample:
										_A	:				str   =      FlaxUpsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__       ,A__       ,A__=True		):
						for resnet, attn in zip(self.resnets       ,self.attentions		):
										# pop res hidden states
										_A	:				List[Any]   =      res_hidden_states_tuple[-1]
										_A	:				int   =      res_hidden_states_tuple[:-1]
										_A	:				List[Any]   =      jnp.concatenate((hidden_states, res_hidden_states)       ,axis=-1		)

										_A	:				str   =      resnet(A__       ,A__       ,deterministic=A__		)
										_A	:				Optional[int]   =      attn(A__       ,A__       ,deterministic=A__		)

						if self.add_upsample:
										_A	:				Union[str, Any]   =      self.upsamplers_a(A__		)

						return hidden_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				Optional[Any]   =      []

						for i in range(self.num_layers		):
										_A	:				Any   =      self.in_channels if (i == self.num_layers - 1) else self.out_channels
										_A	:				Optional[Any]   =      self.prev_output_channel if i == 0 else self.out_channels

										_A	:				List[Any]   =      FlaxResnetBlockaD(
										    in_channels=resnet_in_channels + res_skip_channels       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

						_A	:				int   =      resnets

						if self.add_upsample:
										_A	:				List[str]   =      FlaxUpsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						for resnet in self.resnets:
										# pop res hidden states
										_A	:				Tuple   =      res_hidden_states_tuple[-1]
										_A	:				Optional[int]   =      res_hidden_states_tuple[:-1]
										_A	:				Any   =      jnp.concatenate((hidden_states, res_hidden_states)       ,axis=-1		)

										_A	:				Optional[Any]   =      resnet(A__       ,A__       ,deterministic=A__		)

						if self.add_upsample:
										_A	:				Tuple   =      self.upsamplers_a(A__		)

						return hidden_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						# there is always at least one resnet
						_A	:				List[str]   =      [
						    FlaxResnetBlockaD(
						        in_channels=self.in_channels       ,out_channels=self.in_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
						]

						_A	:				Union[str, Any]   =      []

						for _ in range(self.num_layers		):
										_A	:				Any   =      FlaxTransformeraDModel(
										    in_channels=self.in_channels       ,n_heads=self.num_attention_heads       ,d_head=self.in_channels // self.num_attention_heads       ,depth=1       ,use_linear_projection=self.use_linear_projection       ,use_memory_efficient_attention=self.use_memory_efficient_attention       ,dtype=self.dtype       ,)
										attentions.append(A__		)

										_A	:				Any   =      FlaxResnetBlockaD(
										    in_channels=self.in_channels       ,out_channels=self.in_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

						_A	:				List[str]   =      resnets
						_A	:				str   =      attentions







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						_A	:				Optional[int]   =      self.resnets[0](A__       ,A__		)
						for attn, resnet in zip(self.attentions       ,self.resnets[1:]		):
										_A	:				Dict   =      attn(A__       ,A__       ,deterministic=A__		)
										_A	:				Any   =      resnet(A__       ,A__       ,deterministic=A__		)

						return hidden_states



