

"""simple docstring"""


import copy

from ...configuration_utils import PretrainedConfig
from ...utils import add_start_docstrings


__A						   =			r"\n    [`RagConfig`] stores the configuration of a *RagModel*. Configuration objects inherit from [`PretrainedConfig`] and\n    can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information.\n\n    Args:\n        title_sep (`str`, *optional*, defaults to  `\" / \"`):\n            Separator inserted between the title and the text of the retrieved document when calling [`RagRetriever`].\n        doc_sep (`str`, *optional*, defaults to  `\" // \"`):\n            Separator inserted between the text of the retrieved document and the original input when calling\n            [`RagRetriever`].\n        n_docs (`int`, *optional*, defaults to 5):\n            Number of documents to retrieve.\n        max_combined_length (`int`, *optional*, defaults to 300):\n            Max length of contextualized input returned by [`~RagRetriever.__call__`].\n        retrieval_vector_size (`int`, *optional*, defaults to 768):\n            Dimensionality of the document embeddings indexed by [`RagRetriever`].\n        retrieval_batch_size (`int`, *optional*, defaults to 8):\n            Retrieval batch size, defined as the number of queries issues concurrently to the faiss index encapsulated\n            [`RagRetriever`].\n        dataset (`str`, *optional*, defaults to `\"wiki_dpr\"`):\n            A dataset identifier of the indexed dataset in HuggingFace Datasets (list all available datasets and ids\n            using `datasets.list_datasets()`).\n        dataset_split (`str`, *optional*, defaults to `\"train\"`)\n            Which split of the `dataset` to load.\n        index_name (`str`, *optional*, defaults to `\"compressed\"`)\n            The index name of the index associated with the `dataset`. One can choose between `\"legacy\"`, `\"exact\"` and\n            `\"compressed\"`.\n        index_path (`str`, *optional*)\n            The path to the serialized faiss index on disk.\n        passages_path (`str`, *optional*):\n            A path to text passages compatible with the faiss index. Required if using\n            [`~models.rag.retrieval_rag.LegacyIndex`]\n        use_dummy_dataset (`bool`, *optional*, defaults to `False`)\n            Whether to load a \"dummy\" variant of the dataset specified by `dataset`.\n        label_smoothing (`float`, *optional*, defaults to 0.0):\n            Only relevant if `return_loss` is set to `True`. Controls the `epsilon` parameter value for label smoothing\n            in the loss calculation. If set to 0, no label smoothing is performed.\n        do_marginalize (`bool`, *optional*, defaults to `False`):\n            If `True`, the logits are marginalized over all documents by making use of\n            `torch.nn.functional.log_softmax`.\n        reduce_loss (`bool`, *optional*, defaults to `False`):\n            Whether or not to reduce the NLL loss using the `torch.Tensor.sum` operation.\n        do_deduplication (`bool`, *optional*, defaults to `True`):\n            Whether or not to deduplicate the generations from different context documents for a given input. Has to be\n            set to `False` if used while training with distributed backend.\n        exclude_bos_score (`bool`, *optional*, defaults to `False`):\n            Whether or not to disregard the BOS token when computing the loss.\n        output_retrieved(`bool`, *optional*, defaults to `False`):\n            If set to `True`, `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and\n            `context_attention_mask` are returned. See returned tensors for more detail.\n        use_cache (`bool`, *optional*, defaults to `True`):\n            Whether or not the model should return the last key/values attentions (not used by all models).\n        forced_eos_token_id (`int`, *optional*):\n            The id of the token to force as the last generated token when `max_length` is reached. Usually set to\n            `eos_token_id`.\n"





@add_start_docstrings(__snake_case  )
class     snake_case	(				__snake_case  ):
				SCREAMING_SNAKE_CASE_   :      Any         =  """rag"""
				SCREAMING_SNAKE_CASE_   :      List[Any]         =  True

				def __init__(	self							:				Union[str, Any]     ,  UpperCamelCase__							:				Optional[Any]=None     ,  UpperCamelCase__							:				str=True     ,  UpperCamelCase__							:				List[str]=None     ,  UpperCamelCase__							:				Optional[Any]=None     ,  UpperCamelCase__							:				Dict=None     ,  UpperCamelCase__							:				int=None     ,  UpperCamelCase__							:				List[str]=None     ,  UpperCamelCase__							:				Optional[int]=" / "     ,  UpperCamelCase__							:				Any=" // "     ,  UpperCamelCase__							:				int=5     ,  UpperCamelCase__							:				Union[str, Any]=3_0_0     ,  UpperCamelCase__							:				str=7_6_8     ,  UpperCamelCase__							:				Union[str, Any]=8     ,  UpperCamelCase__							:				Tuple="wiki_dpr"     ,  UpperCamelCase__							:				Optional[Any]="train"     ,  UpperCamelCase__							:				str="compressed"     ,  UpperCamelCase__							:				Any=None     ,  UpperCamelCase__							:				int=None     ,  UpperCamelCase__							:				Any=False     ,  UpperCamelCase__							:				List[str]=False     ,  UpperCamelCase__							:				Optional[int]=0.0     ,  UpperCamelCase__							:				Any=True     ,  UpperCamelCase__							:				List[Any]=False     ,  UpperCamelCase__							:				str=False     ,  UpperCamelCase__							:				Tuple=False     ,  UpperCamelCase__							:				int=True     ,  UpperCamelCase__							:				List[Any]=None     ,  **UpperCamelCase__							:				str     ,  )->   List[str]:



						'''simple docstring'''


						super().__init__(
						    bos_token_id=UpperCamelCase__     ,  pad_token_id=UpperCamelCase__     ,  eos_token_id=UpperCamelCase__     ,  decoder_start_token_id=UpperCamelCase__     ,  forced_eos_token_id=UpperCamelCase__     ,  is_encoder_decoder=UpperCamelCase__     ,  prefix=UpperCamelCase__     ,  vocab_size=UpperCamelCase__     ,  **UpperCamelCase__     ,  )
						assert (
						    "question_encoder" in kwargs and "generator" in kwargs
						), "Config has to be initialized with question_encoder and generator config"
						__lowerCAmelCase:       Dict          =			kwargs.pop("question_encoder")
						__lowerCAmelCase:       Optional[Any]          =			question_encoder_config.pop("model_type")
						__lowerCAmelCase:       Optional[int]          =			kwargs.pop("generator")
						__lowerCAmelCase:       Optional[Any]          =			decoder_config.pop("model_type")

						from ..auto.configuration_auto import AutoConfig

						__lowerCAmelCase:       Optional[Any]          =			AutoConfig.for_model(UpperCamelCase__     ,  **UpperCamelCase__)
						__lowerCAmelCase:       Any          =			AutoConfig.for_model(UpperCamelCase__     ,  **UpperCamelCase__)

						__lowerCAmelCase:       Optional[int]          =			reduce_loss
						__lowerCAmelCase:       List[str]          =			label_smoothing
						__lowerCAmelCase:       str          =			exclude_bos_score
						__lowerCAmelCase:       Any          =			do_marginalize

						__lowerCAmelCase:       List[Any]          =			title_sep
						__lowerCAmelCase:       Tuple          =			doc_sep
						__lowerCAmelCase:       List[Any]          =			n_docs
						__lowerCAmelCase:       Union[str, Any]          =			max_combined_length

						__lowerCAmelCase:       Optional[int]          =			dataset
						__lowerCAmelCase:       List[str]          =			dataset_split
						__lowerCAmelCase:       Dict          =			index_name

						__lowerCAmelCase:       str          =			retrieval_vector_size
						__lowerCAmelCase:       str          =			retrieval_batch_size
						__lowerCAmelCase:       List[Any]          =			passages_path
						__lowerCAmelCase:       int          =			index_path
						__lowerCAmelCase:       Optional[int]          =			use_dummy_dataset

						__lowerCAmelCase:       Union[str, Any]          =			output_retrieved

						__lowerCAmelCase:       str          =			do_deduplication

						__lowerCAmelCase:       List[str]          =			use_cache

						if self.forced_eos_token_id is None:
								__lowerCAmelCase:       Any          =			getattr(self.generator     ,  "forced_eos_token_id"     ,  UpperCamelCase__)

				@classmethod
				def 							lowercase_     (	cls							:				Optional[Any]     ,  UpperCamelCase__							:				PretrainedConfig     ,  UpperCamelCase__							:				PretrainedConfig     ,  **UpperCamelCase__							:				Optional[Any])->   PretrainedConfig:



						'''simple docstring'''


						return cls(question_encoder=question_encoder_config.to_dict()     ,  generator=generator_config.to_dict()     ,  **UpperCamelCase__)




				def 							lowercase_     (	self							:				Union[str, Any])->   Optional[int]:



						'''simple docstring'''


						__lowerCAmelCase:       Optional[Any]          =			copy.deepcopy(self.__dict__)
						__lowerCAmelCase:       str          =			self.question_encoder.to_dict()
						__lowerCAmelCase:       Tuple          =			self.generator.to_dict()
						__lowerCAmelCase:       int          =			self.__class__.model_type
						return output





"""simple docstring"""


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							:				Any     ,  UpperCamelCase__							:				int     ,  UpperCamelCase__							:				str=1_3     ,  UpperCamelCase__							:				Tuple=7     ,  UpperCamelCase__							:				Optional[Any]=True     ,  UpperCamelCase__							:				Optional[Any]=True     ,  UpperCamelCase__							:				Any=False     ,  UpperCamelCase__							:				str=True     ,  UpperCamelCase__							:				List[Any]=9_9     ,  UpperCamelCase__							:				Optional[int]=3_2     ,  UpperCamelCase__							:				List[Any]=5     ,  UpperCamelCase__							:				Union[str, Any]=4     ,  UpperCamelCase__							:				List[str]=3_7     ,  UpperCamelCase__							:				Optional[Any]="gelu"     ,  UpperCamelCase__							:				int=0.1     ,  UpperCamelCase__							:				Dict=0.1     ,  UpperCamelCase__							:				Dict=5_1_2     ,  UpperCamelCase__							:				List[Any]=1_6     ,  UpperCamelCase__							:				Optional[int]=2     ,  UpperCamelCase__							:				Dict=0.02     ,  UpperCamelCase__							:				int=3     ,  UpperCamelCase__							:				Any=4     ,  UpperCamelCase__							:				Optional[int]=None     ,  )->   Optional[Any]:



						'''simple docstring'''


						__lowerCAmelCase:       Tuple          =			parent
						__lowerCAmelCase:       Optional[int]          =			batch_size
						__lowerCAmelCase:       int          =			seq_length
						__lowerCAmelCase:       Any          =			is_training
						__lowerCAmelCase:       List[Any]          =			use_input_mask
						__lowerCAmelCase:       Any          =			use_token_type_ids
						__lowerCAmelCase:       Dict          =			use_labels
						__lowerCAmelCase:       Union[str, Any]          =			vocab_size
						__lowerCAmelCase:       Union[str, Any]          =			hidden_size
						__lowerCAmelCase:       int          =			num_hidden_layers
						__lowerCAmelCase:       List[Any]          =			num_attention_heads
						__lowerCAmelCase:       int          =			intermediate_size
						__lowerCAmelCase:       Optional[Any]          =			hidden_act
						__lowerCAmelCase:       Optional[Any]          =			hidden_dropout_prob
						__lowerCAmelCase:       Optional[int]          =			attention_probs_dropout_prob
						__lowerCAmelCase:       Any          =			max_position_embeddings
						__lowerCAmelCase:       Optional[int]          =			type_vocab_size
						__lowerCAmelCase:       str          =			type_sequence_label_size
						__lowerCAmelCase:       int          =			initializer_range
						__lowerCAmelCase:       Dict          =			num_labels
						__lowerCAmelCase:       Dict          =			num_choices
						__lowerCAmelCase:       str          =			scope

				def 							lowercase_     (	self							:				Optional[Any])->   Optional[Any]:



						'''simple docstring'''


						__lowerCAmelCase:       Optional[int]          =			ids_tensor([self.batch_size, self.seq_length]     ,  self.vocab_size)

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

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

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

						__lowerCAmelCase:       Any          =			self.get_config()

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

				def 							lowercase_     (	self							:				str)->   Dict:



						'''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=UpperCamelCase__     ,  initializer_range=self.initializer_range     ,  use_stable_embedding=UpperCamelCase__     ,  )

				def 							lowercase_     (	self							:				Optional[Any]     ,  UpperCamelCase__							:				int     ,  UpperCamelCase__							:				Tuple     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				Tuple     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				Optional[Any])->   str:



						'''simple docstring'''


						__lowerCAmelCase:       List[Any]          =			OpenLlamaModel(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:       Union[str, Any]          =			model(UpperCamelCase__     ,  attention_mask=UpperCamelCase__)
						__lowerCAmelCase:       int          =			model(UpperCamelCase__)
						self.parent.assertEqual(result.last_hidden_state.shape     ,  (self.batch_size, self.seq_length, self.hidden_size))

				def 							lowercase_     (	self							:				Optional[Any]     ,  UpperCamelCase__							:				List[str]     ,  UpperCamelCase__							:				Dict     ,  UpperCamelCase__							:				Union[str, Any]     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				Any     ,  UpperCamelCase__							:				Any     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				Union[str, Any]     ,  UpperCamelCase__							:				Tuple     ,  )->   Optional[int]:



						'''simple docstring'''


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

				def 							lowercase_     (	self							:				Dict     ,  UpperCamelCase__							:				Optional[int]     ,  UpperCamelCase__							:				Optional[int]     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				Tuple     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				Union[str, Any]     ,  UpperCamelCase__							:				Union[str, Any]     ,  UpperCamelCase__							:				List[str]     ,  )->   Optional[int]:



						'''simple docstring'''


						__lowerCAmelCase:       str          =			OpenLlamaForCausalLM(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:       Optional[int]          =			model(UpperCamelCase__     ,  attention_mask=UpperCamelCase__     ,  labels=UpperCamelCase__)
						self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.seq_length, self.vocab_size))

				def 							lowercase_     (	self							:				Tuple     ,  UpperCamelCase__							:				Union[str, Any]     ,  UpperCamelCase__							:				str     ,  UpperCamelCase__							:				Optional[int]     ,  UpperCamelCase__							:				Any     ,  UpperCamelCase__							:				int     ,  UpperCamelCase__							:				int     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				int     ,  )->   Tuple:



						'''simple docstring'''


						__lowerCAmelCase:       Union[str, Any]          =			True
						__lowerCAmelCase:       Dict          =			True
						__lowerCAmelCase:       Union[str, Any]          =			OpenLlamaForCausalLM(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()

						# first forward pass
						__lowerCAmelCase:       Optional[int]          =			model(
						    UpperCamelCase__     ,  attention_mask=UpperCamelCase__     ,  encoder_hidden_states=UpperCamelCase__     ,  encoder_attention_mask=UpperCamelCase__     ,  use_cache=UpperCamelCase__     ,  )
						__lowerCAmelCase:       Optional[int]          =			outputs.past_key_values

						# create hypothetical multiple next token and extent to next_input_ids
						__lowerCAmelCase:       Optional[int]          =			ids_tensor((self.batch_size, 3)     ,  config.vocab_size)
						__lowerCAmelCase:       Optional[int]          =			ids_tensor((self.batch_size, 3)     ,  vocab_size=2)

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

						__lowerCAmelCase:       Union[str, Any]          =			model(
						    UpperCamelCase__     ,  attention_mask=UpperCamelCase__     ,  encoder_hidden_states=UpperCamelCase__     ,  encoder_attention_mask=UpperCamelCase__     ,  output_hidden_states=UpperCamelCase__     ,  )["hidden_states"][0]
						__lowerCAmelCase:       List[str]          =			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
						__lowerCAmelCase:       List[str]          =			ids_tensor((1,)     ,  output_from_past.shape[-1]).item()
						__lowerCAmelCase:       List[Any]          =			output_from_no_past[:, -3:, random_slice_idx].detach()
						__lowerCAmelCase:       Union[str, Any]          =			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 							lowercase_     (	self							:				Tuple)->   str:



						'''simple docstring'''


						__lowerCAmelCase:       Union[str, Any]          =			self.prepare_config_and_inputs()
						(
						    (
						    __lowerCAmelCase
						)      , (
						    __lowerCAmelCase
						)      , (
						    __lowerCAmelCase
						)      , (
						    __lowerCAmelCase
						)      , (
						    __lowerCAmelCase
						)      , (
						    __lowerCAmelCase
						)      , (
						    __lowerCAmelCase
						)      , 
						):       List[Any]          =			config_and_inputs
						__lowerCAmelCase:       Any          =			{"input_ids": input_ids, "attention_mask": input_mask}
						return config, inputs_dict





@require_torch
class     snake_case	(				__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 							lowercase_     (	self							:				Dict)->   Union[str, Any]:



						'''simple docstring'''


						__lowerCAmelCase:       int          =			OpenLlamaModelTester(self)
						__lowerCAmelCase:       Optional[Any]          =			ConfigTester(self     ,  config_class=UpperCamelCase__     ,  hidden_size=3_7)

				def 							lowercase_     (	self							:				List[str])->   Union[str, Any]:



						'''simple docstring'''


						self.config_tester.run_common_tests()

				def 							lowercase_     (	self							:				Union[str, Any])->   List[Any]:



						'''simple docstring'''


						__lowerCAmelCase:       Any          =			self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*UpperCamelCase__)

				def 							lowercase_     (	self							:				int)->   str:



						'''simple docstring'''


						__lowerCAmelCase:       Optional[Any]          =			self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
								__lowerCAmelCase:       Union[str, Any]          =			type
								self.model_tester.create_and_check_model(*UpperCamelCase__)

				def 							lowercase_     (	self							:				Tuple)->   Optional[Any]:



						'''simple docstring'''


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

				def 							lowercase_     (	self							:				Dict)->   Tuple:



						'''simple docstring'''


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

				def 							lowercase_     (	self							:				Optional[int])->   Any:



						'''simple docstring'''


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



						'''simple docstring'''


						pass




				@parameterized.expand([("linear",), ("dynamic",)])
				def 							lowercase_     (	self							:				Any     ,  UpperCamelCase__							:				List[str])->   Dict:



						'''simple docstring'''


						__lowerCAmelCase      , __lowerCAmelCase:       Optional[Any]          =			self.model_tester.prepare_config_and_inputs_for_common()
						__lowerCAmelCase:       Any          =			ids_tensor([1, 1_0]     ,  config.vocab_size)
						__lowerCAmelCase:       Union[str, Any]          =			ids_tensor([1, int(config.max_position_embeddings * 1.5)]     ,  config.vocab_size)

						set_seed(4_2)  # Fixed seed at init time so the two models get the same random weights
						__lowerCAmelCase:       List[Any]          =			OpenLlamaModel(UpperCamelCase__)
						original_model.to(UpperCamelCase__)
						original_model.eval()
						__lowerCAmelCase:       int          =			original_model(UpperCamelCase__).last_hidden_state
						__lowerCAmelCase:       str          =			original_model(UpperCamelCase__).last_hidden_state

						set_seed(4_2)  # Fixed seed at init time so the two models get the same random weights
						__lowerCAmelCase:       Dict          =			{"type": scaling_type, "factor": 10.0}
						__lowerCAmelCase:       List[str]          =			OpenLlamaModel(UpperCamelCase__)
						scaled_model.to(UpperCamelCase__)
						scaled_model.eval()
						__lowerCAmelCase:       Dict          =			scaled_model(UpperCamelCase__).last_hidden_state
						__lowerCAmelCase:       Any          =			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))


