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



    def __init__(						self   :							Optional[Any],__A   :							Optional[int],__A   :							Optional[int]=1_3,__A   :							Optional[Any]=7,__A   :							str=True,__A   :							Dict=True,__A   :							int=False,__A   :							Optional[Any]=True,__A   :							Optional[Any]=9_9,__A   :							Dict=3_2,__A   :							List[Any]=5,__A   :							int=4,__A   :							List[str]=3_7,__A   :							Optional[Any]="gelu",__A   :							str=0.1,__A   :							Optional[int]=0.1,__A   :							int=5_1_2,__A   :							Union[str, Any]=1_6,__A   :							Optional[Any]=2,__A   :							Tuple=0.02,__A   :							Dict=3,__A   :							List[str]=4,__A   :							List[str]=None,):
          _lowerCamelCase					: List[str]		    = parent
          _lowerCamelCase					: Tuple		    = batch_size
          _lowerCamelCase					: str		    = seq_length
          _lowerCamelCase					: str		    = is_training
          _lowerCamelCase					: Optional[Any]		    = use_input_mask
          _lowerCamelCase					: Tuple		    = use_token_type_ids
          _lowerCamelCase					: Tuple		    = use_labels
          _lowerCamelCase					: Optional[int]		    = vocab_size
          _lowerCamelCase					: List[Any]		    = hidden_size
          _lowerCamelCase					: Optional[Any]		    = num_hidden_layers
          _lowerCamelCase					: str		    = num_attention_heads
          _lowerCamelCase					: Optional[int]		    = intermediate_size
          _lowerCamelCase					: str		    = hidden_act
          _lowerCamelCase					: Union[str, Any]		    = hidden_dropout_prob
          _lowerCamelCase					: Dict		    = attention_probs_dropout_prob
          _lowerCamelCase					: Optional[Any]		    = max_position_embeddings
          _lowerCamelCase					: Union[str, Any]		    = type_vocab_size
          _lowerCamelCase					: List[str]		    = type_sequence_label_size
          _lowerCamelCase					: List[str]		    = initializer_range
          _lowerCamelCase					: Tuple		    = num_labels
          _lowerCamelCase					: Union[str, Any]		    = num_choices
          _lowerCamelCase					: Tuple		    = scope



    def 	lowerCamelCase_	(						self   :							Dict  ):
          _lowerCamelCase					: Any		    = ids_tensor([self.batch_size, self.seq_length],self.vocab_size  )

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

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

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

          _lowerCamelCase					: Optional[Any]		    = self.get_config()

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



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



    def 	lowerCamelCase_	(						self   :							List[str],__A   :							List[str],__A   :							List[Any],__A   :							List[str],__A   :							int,__A   :							Tuple,__A   :							List[str],__A   :							Optional[Any]  ):
          _lowerCamelCase					: int		    = OpenLlamaModel(config=a_  )
          model.to(a_  )
          model.eval()
          _lowerCamelCase					: Dict		    = model(a_,attention_mask=a_  )
          _lowerCamelCase					: int		    = model(a_  )
          self.parent.assertEqual(result.last_hidden_state.shape,(self.batch_size, self.seq_length, self.hidden_size)  )



    def 	lowerCamelCase_	(						self   :							Tuple,__A   :							Dict,__A   :							Optional[Any],__A   :							Optional[Any],__A   :							Union[str, Any],__A   :							Optional[int],__A   :							Tuple,__A   :							Union[str, Any],__A   :							Union[str, Any],__A   :							Optional[int],):
          _lowerCamelCase					: Optional[int]		    = True
          _lowerCamelCase					: Tuple		    = OpenLlamaModel(a_  )
          model.to(a_  )
          model.eval()
          _lowerCamelCase					: Optional[int]		    = model(
              a_,attention_mask=a_,encoder_hidden_states=a_,encoder_attention_mask=a_,)
          _lowerCamelCase					: List[Any]		    = model(
              a_,attention_mask=a_,encoder_hidden_states=a_,)
          _lowerCamelCase					: Tuple		    = model(a_,attention_mask=a_  )
          self.parent.assertEqual(result.last_hidden_state.shape,(self.batch_size, self.seq_length, self.hidden_size)  )



    def 	lowerCamelCase_	(						self   :							Optional[int],__A   :							Optional[int],__A   :							Any,__A   :							str,__A   :							Tuple,__A   :							List[str],__A   :							Optional[int],__A   :							str,__A   :							str,__A   :							Tuple,):
          _lowerCamelCase					: str		    = OpenLlamaForCausalLM(config=a_  )
          model.to(a_  )
          model.eval()
          _lowerCamelCase					: int		    = model(a_,attention_mask=a_,labels=a_  )
          self.parent.assertEqual(result.logits.shape,(self.batch_size, self.seq_length, self.vocab_size)  )



    def 	lowerCamelCase_	(						self   :							Optional[Any],__A   :							Union[str, Any],__A   :							Dict,__A   :							Dict,__A   :							int,__A   :							Any,__A   :							Union[str, Any],__A   :							Union[str, Any],__A   :							Optional[Any],__A   :							Tuple,):
          _lowerCamelCase					: Any		    = True
          _lowerCamelCase					: Optional[int]		    = True
          _lowerCamelCase					: str		    = OpenLlamaForCausalLM(config=a_  )
          model.to(a_  )
          model.eval()

          # first forward pass
          _lowerCamelCase					: Optional[int]		    = model(
              a_,attention_mask=a_,encoder_hidden_states=a_,encoder_attention_mask=a_,use_cache=a_,)
          _lowerCamelCase					: Union[str, Any]		    = outputs.past_key_values

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

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

          _lowerCamelCase					: Dict		    = model(
              a_,attention_mask=a_,encoder_hidden_states=a_,encoder_attention_mask=a_,output_hidden_states=a_,)["hidden_states"][0]
          _lowerCamelCase					: List[str]		    = 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
          _lowerCamelCase					: Any		    = ids_tensor((1,),output_from_past.shape[-1]  ).item()
          _lowerCamelCase					: Any		    = output_from_no_past[:, -3:, random_slice_idx].detach()
          _lowerCamelCase					: 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(a_,a_,atol=1e-3  )  )



    def 	lowerCamelCase_	(						self   :							List[str]  ):
          _lowerCamelCase					: Any		    = self.prepare_config_and_inputs()
          (
              _lowerCamelCase	
          )					: int		    = config_and_inputs
          _lowerCamelCase					: List[str]		    = {"input_ids": input_ids, "attention_mask": input_mask}
          return config, inputs_dict


@require_torch
class 			UpperCAmelCase__    (							_UpperCamelCase    ,						_UpperCamelCase    ,						_UpperCamelCase    ,						unittest.TestCase			):
    lowerCAmelCase_					     =		(
        (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
    )
    lowerCAmelCase_					     =		(OpenLlamaForCausalLM,) if is_torch_available() else ()
    lowerCAmelCase_					     =		(
        {
            "feature-extraction": OpenLlamaModel,
            "text-classification": OpenLlamaForSequenceClassification,
            "text-generation": OpenLlamaForCausalLM,
            "zero-shot": OpenLlamaForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    lowerCAmelCase_					     =		False
    lowerCAmelCase_					     =		False



    def 	lowerCamelCase_	(						self   :							Union[str, Any]  ):
          _lowerCamelCase					: Optional[Any]		    = OpenLlamaModelTester(self  )
          _lowerCamelCase					: Dict		    = ConfigTester(self,config_class=a_,hidden_size=3_7  )



    def 	lowerCamelCase_	(						self   :							Optional[int]  ):
          self.config_tester.run_common_tests()



    def 	lowerCamelCase_	(						self   :							Any  ):
          _lowerCamelCase					: Optional[Any]		    = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*a_  )



    def 	lowerCamelCase_	(						self   :							Dict  ):
          _lowerCamelCase					: Union[str, Any]		    = self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
                _lowerCamelCase					: int		    = type
                self.model_tester.create_and_check_model(*a_  )



    def 	lowerCamelCase_	(						self   :							Dict  ):
          _lowerCamelCase					: int		    = self.model_tester.prepare_config_and_inputs_for_common()
          _lowerCamelCase					: Tuple		    = 3
          _lowerCamelCase					: Tuple		    = input_dict["input_ids"]
          _lowerCamelCase					: Dict		    = input_ids.ne(1  ).to(a_  )
          _lowerCamelCase					: Tuple		    = ids_tensor([self.model_tester.batch_size],self.model_tester.type_sequence_label_size  )
          _lowerCamelCase					: Optional[Any]		    = OpenLlamaForSequenceClassification(a_  )
          model.to(a_  )
          model.eval()
          _lowerCamelCase					: Any		    = model(a_,attention_mask=a_,labels=a_  )
          self.assertEqual(result.logits.shape,(self.model_tester.batch_size, self.model_tester.num_labels)  )



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



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



    @parameterized.expand([("linear",), ("dynamic",)]  )
    def 	lowerCamelCase_	(						self   :							Optional[int],__A   :							int  ):
          _lowerCamelCase					: Optional[int]		    = self.model_tester.prepare_config_and_inputs_for_common()
          _lowerCamelCase					: List[Any]		    = ids_tensor([1, 1_0],config.vocab_size  )
          _lowerCamelCase					: List[str]		    = 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					: Tuple		    = OpenLlamaModel(a_  )
          original_model.to(a_  )
          original_model.eval()
          _lowerCamelCase					: int		    = original_model(a_  ).last_hidden_state
          _lowerCamelCase					: Tuple		    = original_model(a_  ).last_hidden_state

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






"""simple docstring"""






from typing import List, Union

import numpy as np

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, logging
from .base import PIPELINE_INIT_ARGS, ArgumentHandler, ChunkPipeline


__UpperCAmelCase						=			logging.get_logger(__name__)



class 					__UpperCAmelCase							(						_UpperCamelCase  ):
				def       UpperCAmelCase      (		self    :			int	, a_    :			List[Any]    )	->	Optional[Any]:



							'''simple docstring'''



							if isinstance(a_	, a_    ):
										a__  :	Any		 =			[label.strip() for label in labels.split(","    ) if label.strip()]
							return labels






				def __call__(		self    :			Union[str, Any]	, a_    :			Tuple	, a_    :			Optional[Any]	, a_    :			List[str]    )	->	Optional[Any]:



							'''simple docstring'''



							if len(a_    ) == 0 or len(a_    ) == 0:
										raise ValueError("You must include at least one label and at least one sequence."    )
							if hypothesis_template.format(labels[0]    ) == hypothesis_template:
										raise ValueError(
										    (
										        "The provided hypothesis_template \"{}\" was not able to be formatted with the target labels. "
										        "Make sure the passed template includes formatting syntax such as {{}} where the label should go."
										    ).format(a_    )    )

							if isinstance(a_	, a_    ):
										a__  :	str		 =			[sequences]

							a__  :	Optional[int]		 =			[]
							for sequence in sequences:
										sequence_pairs.extend([[sequence, hypothesis_template.format(a_    )] for label in labels]    )

							return sequence_pairs, sequences



@add_end_docstrings(_UpperCamelCase  )
class 					__UpperCAmelCase							(						_UpperCamelCase  ):
				def __init__(		self    :			str	, a_    :			Optional[Any]=ZeroShotClassificationArgumentHandler()	, *a_    :			Tuple	, **a_    :			str    )	->	Optional[int]:



							'''simple docstring'''



							a__  :	List[Any]		 =			args_parser
							super().__init__(*a_	, **a_    )
							if self.entailment_id == -1:
										logger.warning(
										    "Failed to determine 'entailment' label id from the label2id mapping in the model config. Setting to "
										    "-1. Define a descriptive label2id mapping in the model config to ensure correct outputs."    )
				@property
				def       UpperCAmelCase      (		self    :			Tuple    )	->	str:



							'''simple docstring'''



							for label, ind in self.model.config.labelaid.items():
										if label.lower().startswith("entail"    ):
													return ind
							return -1
				def       UpperCAmelCase      (		self    :			Optional[int]	, a_    :			List[Any]	, a_    :			int=True	, a_    :			Tuple=True	, a_    :			Tuple=TruncationStrategy.ONLY_FIRST	, **a_    :			Dict    )	->	List[Any]:



							'''simple docstring'''



							a__  :	str		 =			self.framework
							if self.tokenizer.pad_token is None:
										# Override for tokenizers not supporting padding
										logger.error(
										    "Tokenizer was not supporting padding necessary for zero-shot, attempting to use "
										    " `pad_token=eos_token`"    )
										a__  :	List[str]		 =			self.tokenizer.eos_token
							try:
										a__  :	List[str]		 =			self.tokenizer(
										    a_	, add_special_tokens=a_	, return_tensors=a_	, padding=a_	, truncation=a_	, )
							except Exception as e:
										if "too short" in str(a_    ):
													# tokenizers might yell that we want to truncate
													# to a value that is not even reached by the input.
													# In that case we don't want to truncate.
													# It seems there's not a really better way to catch that
													# exception.

													a__  :	List[str]		 =			self.tokenizer(
													    a_	, add_special_tokens=a_	, return_tensors=a_	, padding=a_	, truncation=TruncationStrategy.DO_NOT_TRUNCATE	, )
										else:
													raise e

							return inputs
				def       UpperCAmelCase      (		self    :			Tuple	, **a_    :			Tuple    )	->	Optional[int]:



							'''simple docstring'''



							if kwargs.get("multi_class"	, a_    ) is not None:
										a__  :	str		 =			kwargs["multi_class"]
										logger.warning(
										    "The `multi_class` argument has been deprecated and renamed to `multi_label`. "
										    "`multi_class` will be removed in a future version of Transformers."    )
							a__  :	Tuple		 =			{}
							if "candidate_labels" in kwargs:
										a__  :	Any		 =			self._args_parser._parse_labels(kwargs["candidate_labels"]    )
							if "hypothesis_template" in kwargs:
										a__  :	str		 =			kwargs["hypothesis_template"]

							a__  :	Tuple		 =			{}
							if "multi_label" in kwargs:
										a__  :	Dict		 =			kwargs["multi_label"]
							return preprocess_params, {}, postprocess_params
				def __call__(		self    :			str	, a_    :			Union[str, List[str]]	, *a_    :			List[str]	, **a_    :			List[Any]	, )	->	Tuple:



							'''simple docstring'''



							if len(a_    ) == 0:
										pass
							elif len(a_    ) == 1 and "candidate_labels" not in kwargs:
										a__  :	Any		 =			args[0]
							else:
										raise ValueError(F"Unable to understand extra arguments {args}"    )

							return super().__call__(a_	, **a_    )
				def       UpperCAmelCase      (		self    :			Optional[int]	, a_    :			Tuple	, a_    :			Any=None	, a_    :			Dict="This example is {}."    )	->	Optional[int]:



							'''simple docstring'''



							a__ ,	a__  :	Optional[Any]		 =			self._args_parser(a_	, a_	, a_    )

							for i, (candidate_label, sequence_pair) in enumerate(zip(a_	, a_    )    ):
										a__  :	Union[str, Any]		 =			self._parse_and_tokenize([sequence_pair]    )

										yield {
										    "candidate_label": candidate_label,
										    "sequence": sequences[0],
										    "is_last": i == len(a_    ) - 1,
										    **model_input,
										}
				def       UpperCAmelCase      (		self    :			Optional[int]	, a_    :			Optional[Any]    )	->	List[Any]:



							'''simple docstring'''



							a__  :	Dict		 =			inputs["candidate_label"]
							a__  :	Optional[int]		 =			inputs["sequence"]
							a__  :	Optional[int]		 =			{k: inputs[k] for k in self.tokenizer.model_input_names}
							a__  :	int		 =			self.model(**a_    )

							a__  :	Optional[int]		 =			{
							    "candidate_label": candidate_label,
							    "sequence": sequence,
							    "is_last": inputs["is_last"],
							    **outputs,
							}
							return model_outputs






				def       UpperCAmelCase      (		self    :			Dict	, a_    :			Any	, a_    :			List[str]=False    )	->	Union[str, Any]:



							'''simple docstring'''



							a__  :	int		 =			[outputs["candidate_label"] for outputs in model_outputs]
							a__  :	Optional[int]		 =			[outputs["sequence"] for outputs in model_outputs]
							a__  :	Union[str, Any]		 =			np.concatenate([output["logits"].numpy() for output in model_outputs]    )
							a__  :	List[str]		 =			logits.shape[0]
							a__  :	Optional[int]		 =			len(a_    )
							a__  :	List[str]		 =			N // n
							a__  :	int		 =			logits.reshape((num_sequences, n, -1)    )

							if multi_label or len(a_    ) == 1:
										# softmax over the entailment vs. contradiction dim for each label independently
										a__  :	str		 =			self.entailment_id
										a__  :	str		 =			-1 if entailment_id == 0 else 0
										a__  :	str		 =			reshaped_outputs[..., [contradiction_id, entailment_id]]
										a__  :	List[Any]		 =			np.exp(a_    ) / np.exp(a_    ).sum(-1	, keepdims=a_    )
										a__  :	str		 =			scores[..., 1]
							else:
										# softmax the "entailment" logits over all candidate labels
										a__  :	str		 =			reshaped_outputs[..., self.entailment_id]
										a__  :	Optional[int]		 =			np.exp(a_    ) / np.exp(a_    ).sum(-1	, keepdims=a_    )

							a__  :	List[str]		 =			list(reversed(scores[0].argsort()    )    )
							return {
							    "sequence": sequences[0],
							    "labels": [candidate_labels[i] for i in top_inds],
							    "scores": scores[0, top_inds].tolist(),
							}