
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						:


     '''simple docstring'''





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



     def  lowerCAmelCase__       (      self  )->					Union[str, Any]:
          UpperCAmelCase__							: Dict 							=						ids_tensor([self.batch_size, self.seq_length] ,			self.vocab_size  )

          UpperCAmelCase__							: Optional[int] 							=						None
          if self.use_input_mask:
               UpperCAmelCase__							: Any 							=						random_attention_mask([self.batch_size, self.seq_length]  )

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

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

          UpperCAmelCase__							: str 							=						self.get_config()

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



     def  lowerCAmelCase__       (      self  )->					List[Any]:
          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 ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase  )->					int:
          UpperCAmelCase__							: Any 							=						OpenLlamaModel(config=snake_case__  )
          model.to(snake_case__  )
          model.eval()
          UpperCAmelCase__							: List[Any] 							=						model(snake_case__ ,			attention_mask=snake_case__  )
          UpperCAmelCase__							: List[str] 							=						model(snake_case__  )
          self.parent.assertEqual(result.last_hidden_state.shape ,			(self.batch_size, self.seq_length, self.hidden_size)  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			)->					str:
          UpperCAmelCase__							: Optional[int] 							=						True
          UpperCAmelCase__							: Tuple 							=						OpenLlamaModel(snake_case__  )
          model.to(snake_case__  )
          model.eval()
          UpperCAmelCase__							: int 							=						model(
              snake_case__ ,			attention_mask=snake_case__ ,			encoder_hidden_states=snake_case__ ,			encoder_attention_mask=snake_case__ ,			)
          UpperCAmelCase__							: Tuple 							=						model(
              snake_case__ ,			attention_mask=snake_case__ ,			encoder_hidden_states=snake_case__ ,			)
          UpperCAmelCase__							: Optional[int] 							=						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 ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			)->					int:
          UpperCAmelCase__							: Tuple 							=						OpenLlamaForCausalLM(config=snake_case__  )
          model.to(snake_case__  )
          model.eval()
          UpperCAmelCase__							: str 							=						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 ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			)->					Tuple:
          UpperCAmelCase__							: int 							=						True
          UpperCAmelCase__							: Optional[Any] 							=						True
          UpperCAmelCase__							: Tuple 							=						OpenLlamaForCausalLM(config=snake_case__  )
          model.to(snake_case__  )
          model.eval()

          # first forward pass
          UpperCAmelCase__							: Tuple 							=						model(
              snake_case__ ,			attention_mask=snake_case__ ,			encoder_hidden_states=snake_case__ ,			encoder_attention_mask=snake_case__ ,			use_cache=snake_case__ ,			)
          UpperCAmelCase__							: Optional[Any] 							=						outputs.past_key_values

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

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

          UpperCAmelCase__							: Tuple 							=						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__							: str 							=						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__							: List[str] 							=						ids_tensor((1,) ,			output_from_past.shape[-1]  ).item()
          UpperCAmelCase__							: int 							=						output_from_no_past[:, -3:, random_slice_idx].detach()
          UpperCAmelCase__							: List[str] 							=						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:
          UpperCAmelCase__							: str 							=						self.prepare_config_and_inputs()
          (
              (
              UpperCAmelCase__						
          )							,	(
              UpperCAmelCase__						
          )							,	(
              UpperCAmelCase__						
          )							,	(
              UpperCAmelCase__						
          )							,	(
              UpperCAmelCase__						
          )							,	(
              UpperCAmelCase__						
          )							,	(
              UpperCAmelCase__						
          )							,							
          )							: int 							=						config_and_inputs
          UpperCAmelCase__							: Optional[int] 							=						{"input_ids": input_ids, "attention_mask": input_mask}
          return config, inputs_dict







@require_torch
class   _lowercase						(       __a	,   __a	,   __a	,   unittest.TestCase				):


     '''simple docstring'''





     _A								=							(
         (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
     )
     _A								=							(OpenLlamaForCausalLM,) if is_torch_available() else ()
     _A								=							(
         {
             '''feature-extraction''': OpenLlamaModel,
             '''text-classification''': OpenLlamaForSequenceClassification,
             '''text-generation''': OpenLlamaForCausalLM,
             '''zero-shot''': OpenLlamaForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     _A								=							False
     _A								=							False



     def  lowerCAmelCase__       (      self  )->					Union[str, Any]:
          UpperCAmelCase__							: Dict 							=						OpenLlamaModelTester(self  )
          UpperCAmelCase__							: Tuple 							=						ConfigTester(self ,			config_class=snake_case__ ,			hidden_size=37  )



     def  lowerCAmelCase__       (      self  )->					List[Any]:
          self.config_tester.run_common_tests()



     def  lowerCAmelCase__       (      self  )->					Any:
          UpperCAmelCase__							: str 							=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case__  )



     def  lowerCAmelCase__       (      self  )->					List[Any]:
          UpperCAmelCase__							: Any 							=						self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
               UpperCAmelCase__							: Union[str, Any] 							=						type
               self.model_tester.create_and_check_model(*snake_case__  )



     def  lowerCAmelCase__       (      self  )->					Any:
          UpperCAmelCase__							,	UpperCAmelCase__							: int 							=						self.model_tester.prepare_config_and_inputs_for_common()
          UpperCAmelCase__							: Any 							=						3
          UpperCAmelCase__							: List[Any] 							=						input_dict["input_ids"]
          UpperCAmelCase__							: Optional[int] 							=						input_ids.ne(1  ).to(snake_case__  )
          UpperCAmelCase__							: str 							=						ids_tensor([self.model_tester.batch_size] ,			self.model_tester.type_sequence_label_size  )
          UpperCAmelCase__							: Any 							=						OpenLlamaForSequenceClassification(snake_case__  )
          model.to(snake_case__  )
          model.eval()
          UpperCAmelCase__							: str 							=						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]:
          UpperCAmelCase__							,	UpperCAmelCase__							: Tuple 							=						self.model_tester.prepare_config_and_inputs_for_common()
          UpperCAmelCase__							: Optional[Any] 							=						3
          UpperCAmelCase__							: str 							=						"single_label_classification"
          UpperCAmelCase__							: Union[str, Any] 							=						input_dict["input_ids"]
          UpperCAmelCase__							: List[str] 							=						input_ids.ne(1  ).to(snake_case__  )
          UpperCAmelCase__							: Union[str, Any] 							=						ids_tensor([self.model_tester.batch_size] ,			self.model_tester.type_sequence_label_size  )
          UpperCAmelCase__							: Union[str, Any] 							=						OpenLlamaForSequenceClassification(snake_case__  )
          model.to(snake_case__  )
          model.eval()
          UpperCAmelCase__							: List[Any] 							=						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[Any]:
          UpperCAmelCase__							,	UpperCAmelCase__							: Optional[int] 							=						self.model_tester.prepare_config_and_inputs_for_common()
          UpperCAmelCase__							: Dict 							=						3
          UpperCAmelCase__							: str 							=						"multi_label_classification"
          UpperCAmelCase__							: str 							=						input_dict["input_ids"]
          UpperCAmelCase__							: Optional[int] 							=						input_ids.ne(1  ).to(snake_case__  )
          UpperCAmelCase__							: Tuple 							=						ids_tensor(
              [self.model_tester.batch_size, config.num_labels] ,			self.model_tester.type_sequence_label_size  ).to(torch.float  )
          UpperCAmelCase__							: List[str] 							=						OpenLlamaForSequenceClassification(snake_case__  )
          model.to(snake_case__  )
          model.eval()
          UpperCAmelCase__							: Union[str, Any] 							=						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  )->					str:
          pass



     @parameterized.expand([("linear",), ("dynamic",)]  )
     def  lowerCAmelCase__       (      self ,			__UpperCamelCase  )->					List[str]:
          UpperCAmelCase__							,	UpperCAmelCase__							: List[Any] 							=						self.model_tester.prepare_config_and_inputs_for_common()
          UpperCAmelCase__							: Optional[Any] 							=						ids_tensor([1, 10] ,			config.vocab_size  )
          UpperCAmelCase__							: Optional[int] 							=						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__							: List[str] 							=						OpenLlamaModel(snake_case__  )
          original_model.to(snake_case__  )
          original_model.eval()
          UpperCAmelCase__							: Optional[int] 							=						original_model(snake_case__  ).last_hidden_state
          UpperCAmelCase__							: List[Any] 							=						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__							: int 							=						{"type": scaling_type, "factor": 10.0}
          UpperCAmelCase__							: int 							=						OpenLlamaModel(snake_case__  )
          scaled_model.to(snake_case__  )
          scaled_model.eval()
          UpperCAmelCase__							: Optional[Any] 							=						scaled_model(snake_case__  ).last_hidden_state
          UpperCAmelCase__							: Any 							=						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  )  )





"""simple docstring"""




from pathlib import PurePosixPath
from typing import Optional

import fsspec
from fsspec import AbstractFileSystem
from huggingface_hub.hf_api import DatasetInfo

from ..utils.file_utils import get_authentication_headers_for_url
from ..utils.hub import hf_hub_url







class   _lowercase						(       lowerCAmelCase_				):


     '''simple docstring'''





     _A								=							''
     _A								=							'hf-legacy'  # "hf://"" is reserved for hffs



     def __init__(      self ,			__UpperCamelCase = None ,			__UpperCamelCase = None ,			**__UpperCamelCase ,			)->					str:
          super().__init__(self ,			**__UpperCamelCase  )
          UpperCAmelCase__							: int 							=						repo_info
          UpperCAmelCase__							: Optional[int] 							=						token
          UpperCAmelCase__							: Optional[Any] 							=						None



     def  lowerCAmelCase__       (      self  )->					Optional[Any]:
          if self.dir_cache is None:
               UpperCAmelCase__							: str 							=						{}
               for hf_file in self.repo_info.siblings:
                    # TODO(QL): add sizes
                    UpperCAmelCase__							: str 							=						{
                        "name": hf_file.rfilename,
                        "size": None,
                        "type": "file",
                    }
                    self.dir_cache.update(
                        {
                            str(__UpperCamelCase  ): {"name": str(__UpperCamelCase  ), "size": None, "type": "directory"}
                            for d in list(PurePosixPath(hf_file.rfilename  ).parents  )[:-1]
                        }  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase = "rb" ,			**__UpperCamelCase ,			)->					List[Any]:
          if not isinstance(self.repo_info ,			__UpperCamelCase  ):
               raise NotImplementedError(F"Open is only implemented for dataset repositories, but got {self.repo_info}"  )
          UpperCAmelCase__							: Union[str, Any] 							=						hf_hub_url(self.repo_info.id ,			__UpperCamelCase ,			revision=self.repo_info.sha  )
          return fsspec.open(
              __UpperCamelCase ,			mode=__UpperCamelCase ,			headers=get_authentication_headers_for_url(__UpperCamelCase ,			use_auth_token=self.token  ) ,			client_kwargs={"trust_env": True} ,			).open()



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			**__UpperCamelCase  )->					List[str]:
          self._get_dirs()
          UpperCAmelCase__							: Union[str, Any] 							=						self._strip_protocol(__UpperCamelCase  )
          if path in self.dir_cache:
               return self.dir_cache[path]
          else:
               raise FileNotFoundError(__UpperCamelCase  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase=False ,			**__UpperCamelCase  )->					str:
          self._get_dirs()
          UpperCAmelCase__							: str 							=						PurePosixPath(path.strip("/"  )  )
          UpperCAmelCase__							: Optional[Any] 							=						{}
          for p, f in self.dir_cache.items():
               UpperCAmelCase__							: Optional[int] 							=						PurePosixPath(p.strip("/"  )  )
               UpperCAmelCase__							: Dict 							=						p.parent
               if root == path:
                    UpperCAmelCase__							: Tuple 							=						f
          UpperCAmelCase__							: List[Any] 							=						list(paths.values()  )
          if detail:
               return out
          else:
               return sorted(f["name"] for f in out  )



