




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







   def __init__(      self		:						str					,						_snake_case		:						Tuple					,						_snake_case		:						List[str]=13					,						_snake_case		:						Tuple=7					,						_snake_case		:						Optional[Any]=True					,						_snake_case		:						List[str]=True					,						_snake_case		:						Dict=False					,						_snake_case		:						List[str]=True					,						_snake_case		:						Optional[int]=99					,						_snake_case		:						Optional[int]=32					,						_snake_case		:						Dict=5					,						_snake_case		:						List[Any]=4					,						_snake_case		:						Optional[int]=37					,						_snake_case		:						Optional[Any]="gelu"					,						_snake_case		:						Tuple=0.1					,						_snake_case		:						Tuple=0.1					,						_snake_case		:						int=512					,						_snake_case		:						int=16					,						_snake_case		:						List[Any]=2					,						_snake_case		:						Any=0.0_2					,						_snake_case		:						str=3					,						_snake_case		:						Union[str, Any]=4					,						_snake_case		:						str=None					,						):

        """simple docstring"""


        UpperCAmelCase_          =							parent
        UpperCAmelCase_          =							batch_size
        UpperCAmelCase_          =							seq_length
        UpperCAmelCase_          =							is_training
        UpperCAmelCase_          =							use_input_mask
        UpperCAmelCase_          =							use_token_type_ids
        UpperCAmelCase_          =							use_labels
        UpperCAmelCase_          =							vocab_size
        UpperCAmelCase_          =							hidden_size
        UpperCAmelCase_          =							num_hidden_layers
        UpperCAmelCase_          =							num_attention_heads
        UpperCAmelCase_          =							intermediate_size
        UpperCAmelCase_          =							hidden_act
        UpperCAmelCase_          =							hidden_dropout_prob
        UpperCAmelCase_          =							attention_probs_dropout_prob
        UpperCAmelCase_          =							max_position_embeddings
        UpperCAmelCase_          =							type_vocab_size
        UpperCAmelCase_          =							type_sequence_label_size
        UpperCAmelCase_          =							initializer_range
        UpperCAmelCase_          =							num_labels
        UpperCAmelCase_          =							num_choices
        UpperCAmelCase_          =							scope







   def lowerCamelCase						(      self		:						Optional[Any]):

        """simple docstring"""


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

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

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

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

        UpperCAmelCase_          =							self.get_config()

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







   def lowerCamelCase						(      self		:						str):

        """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 lowerCamelCase						(      self		:						Optional[Any]					,						_snake_case		:						int					,						_snake_case		:						str					,						_snake_case		:						Union[str, Any]					,						_snake_case		:						Dict					,						_snake_case		:						Optional[Any]					,						_snake_case		:						Any					,						_snake_case		:						Optional[Any]):

        """simple docstring"""


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







   def lowerCamelCase						(      self		:						Dict					,						_snake_case		:						Optional[int]					,						_snake_case		:						Dict					,						_snake_case		:						List[Any]					,						_snake_case		:						List[Any]					,						_snake_case		:						Any					,						_snake_case		:						List[Any]					,						_snake_case		:						str					,						_snake_case		:						Optional[int]					,						_snake_case		:						str					,						):

        """simple docstring"""


        UpperCAmelCase_          =							True
        UpperCAmelCase_          =							OpenLlamaModel(_snake_case)
        model.to(_snake_case)
        model.eval()
        UpperCAmelCase_          =							model(
            _snake_case					,						attention_mask=_snake_case					,						encoder_hidden_states=_snake_case					,						encoder_attention_mask=_snake_case					,						)
        UpperCAmelCase_          =							model(
            _snake_case					,						attention_mask=_snake_case					,						encoder_hidden_states=_snake_case					,						)
        UpperCAmelCase_          =							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 lowerCamelCase						(      self		:						Union[str, Any]					,						_snake_case		:						List[Any]					,						_snake_case		:						Tuple					,						_snake_case		:						int					,						_snake_case		:						List[Any]					,						_snake_case		:						Any					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						Optional[int]					,						_snake_case		:						Union[str, Any]					,						):

        """simple docstring"""


        UpperCAmelCase_          =							OpenLlamaForCausalLM(config=_snake_case)
        model.to(_snake_case)
        model.eval()
        UpperCAmelCase_          =							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 lowerCamelCase						(      self		:						List[str]					,						_snake_case		:						Union[str, Any]					,						_snake_case		:						Optional[int]					,						_snake_case		:						Any					,						_snake_case		:						Tuple					,						_snake_case		:						List[Any]					,						_snake_case		:						str					,						_snake_case		:						Tuple					,						_snake_case		:						Tuple					,						_snake_case		:						Dict					,						):

        """simple docstring"""


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

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

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

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

        UpperCAmelCase_          =							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]
        UpperCAmelCase_          =							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
        UpperCAmelCase_          =							ids_tensor((1,)					,						output_from_past.shape[-1]).item()
        UpperCAmelCase_          =							output_from_no_past[:, -3:, random_slice_idx].detach()
        UpperCAmelCase_          =							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 lowerCamelCase						(      self		:						Any):

        """simple docstring"""


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







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







   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


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







   def lowerCamelCase						(      self		:						Tuple):

        """simple docstring"""


        self.config_tester.run_common_tests()







   def lowerCamelCase						(      self		:						str):

        """simple docstring"""


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







   def lowerCamelCase						(      self		:						Dict):

        """simple docstring"""


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







   def lowerCamelCase						(      self		:						Union[str, Any]):

        """simple docstring"""


        UpperCAmelCase_     ,       UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs_for_common()
        UpperCAmelCase_          =							3
        UpperCAmelCase_          =							input_dict['''input_ids''']
        UpperCAmelCase_          =							input_ids.ne(1).to(_snake_case)
        UpperCAmelCase_          =							ids_tensor([self.model_tester.batch_size]					,						self.model_tester.type_sequence_label_size)
        UpperCAmelCase_          =							OpenLlamaForSequenceClassification(_snake_case)
        model.to(_snake_case)
        model.eval()
        UpperCAmelCase_          =							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 lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        UpperCAmelCase_     ,       UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs_for_common()
        UpperCAmelCase_          =							3
        UpperCAmelCase_          =							'''single_label_classification'''
        UpperCAmelCase_          =							input_dict['''input_ids''']
        UpperCAmelCase_          =							input_ids.ne(1).to(_snake_case)
        UpperCAmelCase_          =							ids_tensor([self.model_tester.batch_size]					,						self.model_tester.type_sequence_label_size)
        UpperCAmelCase_          =							OpenLlamaForSequenceClassification(_snake_case)
        model.to(_snake_case)
        model.eval()
        UpperCAmelCase_          =							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 lowerCamelCase						(      self		:						int):

        """simple docstring"""


        UpperCAmelCase_     ,       UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs_for_common()
        UpperCAmelCase_          =							3
        UpperCAmelCase_          =							'''multi_label_classification'''
        UpperCAmelCase_          =							input_dict['''input_ids''']
        UpperCAmelCase_          =							input_ids.ne(1).to(_snake_case)
        UpperCAmelCase_          =							ids_tensor(
            [self.model_tester.batch_size, config.num_labels]					,						self.model_tester.type_sequence_label_size).to(torch.float)
        UpperCAmelCase_          =							OpenLlamaForSequenceClassification(_snake_case)
        model.to(_snake_case)
        model.eval()
        UpperCAmelCase_          =							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 lowerCamelCase						(      self		:						List[Any]):

        """simple docstring"""


        pass







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

        """simple docstring"""


        UpperCAmelCase_     ,       UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs_for_common()
        UpperCAmelCase_          =							ids_tensor([1, 10]					,						config.vocab_size)
        UpperCAmelCase_          =							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
        UpperCAmelCase_          =							OpenLlamaModel(_snake_case)
        original_model.to(_snake_case)
        original_model.eval()
        UpperCAmelCase_          =							original_model(_snake_case).last_hidden_state
        UpperCAmelCase_          =							original_model(_snake_case).last_hidden_state

        set_seed(42)  # Fixed seed at init time so the two models get the same random weights
        UpperCAmelCase_          =							{'''type''': scaling_type, '''factor''': 1_0.0}
        UpperCAmelCase_          =							OpenLlamaModel(_snake_case)
        scaled_model.to(_snake_case)
        scaled_model.eval()
        UpperCAmelCase_          =							scaled_model(_snake_case).last_hidden_state
        UpperCAmelCase_          =							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))








import os

import numpy
import onnx







def 						_snake_case(  SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				) ->    Union[str, Any]:
       lowercase							:     int       	=      a.name
       lowercase							:     Any       	=      b.name

       lowercase							:     Optional[Any]       	=      """"""
       lowercase							:     Dict       	=      """"""

       lowercase							:     int       	=      a == b

       lowercase							:     int       	=      name_a
       lowercase							:     List[str]       	=      name_b

       return res







def 						_snake_case(  SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				) ->    str:
       for i, input_name in enumerate(node_proto.input				):
              if input_name == name:
                     node_proto.input.insert(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)
                     node_proto.input.pop(i + 1				)

       if node_proto.op_type == "If":
              _graph_replace_input_with(node_proto.attribute[0].g					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)
              _graph_replace_input_with(node_proto.attribute[1].g					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)
       if node_proto.op_type == "Loop":
              _graph_replace_input_with(node_proto.attribute[0].g					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)







def 						_snake_case(  SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				) ->    Union[str, Any]:
       for n in graph_proto.node:
              _node_replace_input_with(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)







def 						_snake_case(  SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				) ->    List[Any]:
       lowercase							:     Any       	=      list(model.graph.initializer				)
       lowercase							:     Dict       	=      list(model_without_ext.graph.initializer				)
       for i, ref_i in ind_to_replace:
              assert inits_with_data[i].name == inits[i].name
              assert inits_with_data[ref_i].name == inits[ref_i].name
              assert i > ref_i

              lowercase							:     Union[str, Any]       	=      inits[i].name
              lowercase							:     Dict       	=      inits[ref_i].name

              model_without_ext.graph.initializer.remove(inits[i]				)

              # for n in model.graph.node:
              _graph_replace_input_with(model_without_ext.graph					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)







def 						_snake_case(  SCREAMING_SNAKE_CASE__				) ->    List[str]:
       lowercase							:     Union[str, Any]       	=      os.path.dirname(SCREAMING_SNAKE_CASE__				)
       lowercase							:     Dict       	=      os.path.basename(SCREAMING_SNAKE_CASE__				)

       lowercase							:     str       	=      onnx.load(os.path.join(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)				)

       lowercase							:     List[str]       	=      list(model.graph.initializer				)

       lowercase							:     Tuple       	=      set()
       lowercase							:     int       	=      {}
       lowercase							:     Optional[Any]       	=      []

       lowercase							:     Dict       	=      0

       for i in range(len(SCREAMING_SNAKE_CASE__				)				):
              if i in dup_set:
                     continue

              for j in range(i + 1					,	len(SCREAMING_SNAKE_CASE__				)				):
                     if j in dup_set:
                            continue
                     if _is_equal_tensor_proto(inits[i]					,	inits[j]				):
                            dup_set.add(SCREAMING_SNAKE_CASE__				)
                            dup_set.add(SCREAMING_SNAKE_CASE__				)

                            lowercase							:     int       	=      inits[j].data_type
                            lowercase							:     Optional[int]       	=      numpy.prod(inits[j].dims				)
                            if dtype == 1:
                                   mem_size *= 4
                            elif dtype == 6:
                                   mem_size *= 4
                            elif dtype == 7 or dtype == 11:
                                   mem_size *= 8
                            else:
                                   print("""unexpected data type: """					,	SCREAMING_SNAKE_CASE__				)
                            total_reduced_size += mem_size

                            lowercase							:     Tuple       	=      inits[i].name
                            lowercase							:     int       	=      inits[j].name

                            if name_i in dup_map:
                                   dup_map[name_i].append(SCREAMING_SNAKE_CASE__				)
                            else:
                                   lowercase							:     List[str]       	=      [name_j]
                            ind_to_replace.append((j, i)				)

       print("""total reduced size: """					,	total_reduced_size / 1_024 / 1_024 / 1_024					,	"""GB"""				)

       lowercase							:     str       	=      sorted(SCREAMING_SNAKE_CASE__				)
       _remove_dup_initializers_from_model(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)

       lowercase							:     Optional[Any]       	=      """optimized_""" + model_file_name
       lowercase							:     Dict       	=      os.path.join(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)
       onnx.save(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)

       return new_model
