



"""simple docstring"""






from abc import ABC, abstractmethod
from typing import List, Optional






class        __snake_case       (	_lowercase):




  def __init__( self     :						str		):






         """simple docstring"""






         self.test()




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






         """simple docstring"""






         _lowerCamelCase						:				Union[str, Any]    			= 0
         _lowerCamelCase						:				Optional[int]    			= False
         while not completed:
                if counter == 1:
                       self.reset()
                _lowerCamelCase						:				Optional[int]    			= self.advance()
                if not self.does_advance(__lowerCAmelCase		):
                       raise Exception(
                           '''Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true.'''		)

                _lowerCamelCase			,    _lowerCamelCase			,    _lowerCamelCase						:				int    			= self.update(__lowerCAmelCase		)
                counter += 1

                if counter > 1_0_0_0_0:
                       raise Exception('''update() does not fulfill the constraint.'''		)

         if self.remaining() != 0:
                raise Exception('''Custom Constraint is not defined correctly.'''		)




  @abstractmethod
  def 						SCREAMING_SNAKE_CASE				( self     :						Optional[int]		):






         """simple docstring"""






         raise NotImplementedError(
             f'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)




  @abstractmethod
  def 						SCREAMING_SNAKE_CASE				( self     :						Tuple						,					__lowerCAmelCase     :						int		):






         """simple docstring"""






         raise NotImplementedError(
             f'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)




  @abstractmethod
  def 						SCREAMING_SNAKE_CASE				( self     :						str						,					__lowerCAmelCase     :						int		):






         """simple docstring"""






         raise NotImplementedError(
             f'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)




  @abstractmethod
  def 						SCREAMING_SNAKE_CASE				( self     :						List[str]		):






         """simple docstring"""






         raise NotImplementedError(
             f'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)




  @abstractmethod
  def 						SCREAMING_SNAKE_CASE				( self     :						Any		):






         """simple docstring"""






         raise NotImplementedError(
             f'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)




  @abstractmethod
  def 						SCREAMING_SNAKE_CASE				( self     :						Optional[Any]						,					__lowerCAmelCase     :						Tuple=False		):






         """simple docstring"""






         raise NotImplementedError(
             f'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''		)






class        __snake_case       (	_lowercase):




  def __init__( self     :						Optional[int]						,					__lowerCAmelCase     :						List[int]		):






         """simple docstring"""






         super(__lowerCAmelCase						,					self		).__init__()

         if not isinstance(__lowerCAmelCase						,					__lowerCAmelCase		) or len(__lowerCAmelCase		) == 0:
                raise ValueError(f'''`token_ids` has to be a non-empty list, but is {token_ids}.'''		)
         if any((not isinstance(__lowerCAmelCase						,					__lowerCAmelCase		) or token_id < 0) for token_id in token_ids		):
                raise ValueError(f'''Each list in `token_ids` has to be a list of positive integers, but is {token_ids}.'''		)

         _lowerCamelCase						:				Union[str, Any]    			= token_ids

         _lowerCamelCase						:				Dict    			= len(self.token_ids		)
         _lowerCamelCase						:				List[str]    			= -1  # the index of the currently fulfilled step
         _lowerCamelCase						:				List[Any]    			= False




  def 						SCREAMING_SNAKE_CASE				( self     :						Optional[Any]		):






         """simple docstring"""






         if self.completed:
                return None
         return self.token_ids[self.fulfilled_idx + 1]




  def 						SCREAMING_SNAKE_CASE				( self     :						int						,					__lowerCAmelCase     :						int		):






         """simple docstring"""






         if not isinstance(__lowerCAmelCase						,					__lowerCAmelCase		):
                raise ValueError(f'''`token_id` has to be an `int`, but is {token_id} of type {type(__lowerCAmelCase		)}'''		)

         if self.completed:
                return False

         return token_id == self.token_ids[self.fulfilled_idx + 1]




  def 						SCREAMING_SNAKE_CASE				( self     :						Union[str, Any]						,					__lowerCAmelCase     :						int		):






         """simple docstring"""






         if not isinstance(__lowerCAmelCase						,					__lowerCAmelCase		):
                raise ValueError(f'''`token_id` has to be an `int`, but is {token_id} of type {type(__lowerCAmelCase		)}'''		)

         _lowerCamelCase						:				List[str]    			= False
         _lowerCamelCase						:				Union[str, Any]    			= False
         _lowerCamelCase						:				Optional[int]    			= False

         if self.does_advance(__lowerCAmelCase		):
                self.fulfilled_idx += 1
                _lowerCamelCase						:				List[Any]    			= True
                if self.fulfilled_idx == (self.seqlen - 1):
                       _lowerCamelCase						:				int    			= True
                _lowerCamelCase						:				Optional[Any]    			= completed
         else:
                # failed to make progress.
                _lowerCamelCase						:				str    			= True
                self.reset()
         return stepped, completed, reset




  def 						SCREAMING_SNAKE_CASE				( self     :						Dict		):






         """simple docstring"""






         _lowerCamelCase						:				str    			= False
         _lowerCamelCase						:				Tuple    			= 0




  def 						SCREAMING_SNAKE_CASE				( self     :						List[str]		):






         """simple docstring"""






         return self.seqlen - (self.fulfilled_idx + 1)




  def 						SCREAMING_SNAKE_CASE				( self     :						Dict						,					__lowerCAmelCase     :						Dict=False		):






         """simple docstring"""






         _lowerCamelCase						:				Optional[int]    			= PhrasalConstraint(self.token_ids		)

         if stateful:
                _lowerCamelCase						:				Any    			= self.seqlen
                _lowerCamelCase						:				List[Any]    			= self.fulfilled_idx
                _lowerCamelCase						:				Dict    			= self.completed

         return new_constraint






class        __snake_case       :




  def __init__( self     :						Dict						,					__lowerCAmelCase     :						List[List[int]]						,					__lowerCAmelCase     :						Dict=True		):






         """simple docstring"""






         _lowerCamelCase						:				Any    			= max([len(__lowerCAmelCase		) for one in nested_token_ids]		)

         _lowerCamelCase						:				Any    			= {}
         for token_ids in nested_token_ids:
                _lowerCamelCase						:				Dict    			= root
                for tidx, token_id in enumerate(__lowerCAmelCase		):
                       if token_id not in level:
                              _lowerCamelCase						:				List[Any]    			= {}

                       _lowerCamelCase						:				Union[str, Any]    			= level[token_id]

         if no_subsets and self.has_subsets(__lowerCAmelCase						,					__lowerCAmelCase		):
                raise ValueError(
                    '''Each list in `nested_token_ids` can\'t be a complete subset of another list, but is'''
                    f''' {nested_token_ids}.'''		)

         _lowerCamelCase						:				Dict    			= root




  def 						SCREAMING_SNAKE_CASE				( self     :						List[Any]						,					__lowerCAmelCase     :						int		):






         """simple docstring"""






         _lowerCamelCase						:				List[Any]    			= self.trie

         for current_token in current_seq:
                _lowerCamelCase						:				str    			= start[current_token]

         _lowerCamelCase						:				Any    			= list(start.keys()		)

         return next_tokens




  def 						SCREAMING_SNAKE_CASE				( self     :						Any						,					__lowerCAmelCase     :						Union[str, Any]		):






         """simple docstring"""






         _lowerCamelCase						:				Dict    			= self.next_tokens(__lowerCAmelCase		)

         return len(__lowerCAmelCase		) == 0




  def 						SCREAMING_SNAKE_CASE				( self     :						List[Any]						,					__lowerCAmelCase     :						int		):






         """simple docstring"""






         _lowerCamelCase						:				Dict    			= list(root.values()		)
         if len(__lowerCAmelCase		) == 0:
                return 1
         else:
                return sum([self.count_leaves(__lowerCAmelCase		) for nn in next_nodes]		)




  def 						SCREAMING_SNAKE_CASE				( self     :						Any						,					__lowerCAmelCase     :						int						,					__lowerCAmelCase     :						str		):






         """simple docstring"""






         _lowerCamelCase						:				Optional[Any]    			= self.count_leaves(__lowerCAmelCase		)
         return len(__lowerCAmelCase		) != leaf_count






class        __snake_case       (	_lowercase):




  def __init__( self     :						Dict						,					__lowerCAmelCase     :						List[List[int]]		):






         """simple docstring"""






         super(__lowerCAmelCase						,					self		).__init__()

         if not isinstance(__lowerCAmelCase						,					__lowerCAmelCase		) or len(__lowerCAmelCase		) == 0:
                raise ValueError(f'''`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}.'''		)
         if any(not isinstance(__lowerCAmelCase						,					__lowerCAmelCase		) for token_ids in nested_token_ids		):
                raise ValueError(f'''`nested_token_ids` has to be a list of lists, but is {nested_token_ids}.'''		)
         if any(
             any((not isinstance(__lowerCAmelCase						,					__lowerCAmelCase		) or token_id < 0) for token_id in token_ids		)
             for token_ids in nested_token_ids		):
                raise ValueError(
                    f'''Each list in `nested_token_ids` has to be a list of positive integers, but is {nested_token_ids}.'''		)

         _lowerCamelCase						:				str    			= DisjunctiveTrie(__lowerCAmelCase		)
         _lowerCamelCase						:				Any    			= nested_token_ids

         _lowerCamelCase						:				List[Any]    			= self.trie.max_height
         _lowerCamelCase						:				Optional[int]    			= []
         _lowerCamelCase						:				List[str]    			= False




  def 						SCREAMING_SNAKE_CASE				( self     :						int		):






         """simple docstring"""






         _lowerCamelCase						:				Any    			= self.trie.next_tokens(self.current_seq		)

         if len(__lowerCAmelCase		) == 0:
                return None
         else:
                return token_list




  def 						SCREAMING_SNAKE_CASE				( self     :						List[Any]						,					__lowerCAmelCase     :						int		):






         """simple docstring"""






         if not isinstance(__lowerCAmelCase						,					__lowerCAmelCase		):
                raise ValueError(f'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(__lowerCAmelCase		)}'''		)

         _lowerCamelCase						:				int    			= self.trie.next_tokens(self.current_seq		)

         return token_id in next_tokens




  def 						SCREAMING_SNAKE_CASE				( self     :						int						,					__lowerCAmelCase     :						int		):






         """simple docstring"""






         if not isinstance(__lowerCAmelCase						,					__lowerCAmelCase		):
                raise ValueError(f'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(__lowerCAmelCase		)}'''		)

         _lowerCamelCase						:				int    			= False
         _lowerCamelCase						:				str    			= False
         _lowerCamelCase						:				Optional[Any]    			= False

         if self.does_advance(__lowerCAmelCase		):
                self.current_seq.append(__lowerCAmelCase		)
                _lowerCamelCase						:				List[str]    			= True
         else:
                _lowerCamelCase						:				List[Any]    			= True
                self.reset()

         _lowerCamelCase						:				List[str]    			= self.trie.reached_leaf(self.current_seq		)
         _lowerCamelCase						:				int    			= completed

         return stepped, completed, reset




  def 						SCREAMING_SNAKE_CASE				( self     :						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:				Dict    			= False
         _lowerCamelCase						:				List[str]    			= []




  def 						SCREAMING_SNAKE_CASE				( self     :						List[Any]		):






         """simple docstring"""






         if self.completed:
                # since this can be completed without reaching max height
                return 0
         else:
                return self.seqlen - len(self.current_seq		)




  def 						SCREAMING_SNAKE_CASE				( self     :						str						,					__lowerCAmelCase     :						str=False		):






         """simple docstring"""






         _lowerCamelCase						:				Optional[Any]    			= DisjunctiveConstraint(self.token_ids		)

         if stateful:
                _lowerCamelCase						:				int    			= self.seqlen
                _lowerCamelCase						:				List[str]    			= self.current_seq
                _lowerCamelCase						:				Union[str, Any]    			= self.completed

         return new_constraint






class        __snake_case       :




  def __init__( self     :						str						,					__lowerCAmelCase     :						List[Constraint]		):






         """simple docstring"""






         _lowerCamelCase						:				Any    			= constraints

         # max # of steps required to fulfill a given constraint
         _lowerCamelCase						:				str    			= max([c.seqlen for c in constraints]		)
         _lowerCamelCase						:				List[Any]    			= len(__lowerCAmelCase		)
         _lowerCamelCase						:				str    			= False

         self.init_state()




  def 						SCREAMING_SNAKE_CASE				( self     :						List[Any]		):






         """simple docstring"""






         _lowerCamelCase						:				Union[str, Any]    			= []
         _lowerCamelCase						:				List[str]    			= None
         _lowerCamelCase						:				Any    			= [constraint.copy(stateful=__lowerCAmelCase		) for constraint in self.constraints]




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






         """simple docstring"""






         _lowerCamelCase						:				str    			= 0
         if self.inprogress_constraint:
                # extra points for having a constraint mid-fulfilled
                add += self.max_seqlen - self.inprogress_constraint.remaining()

         return (len(self.complete_constraints		) * self.max_seqlen) + add




  def 						SCREAMING_SNAKE_CASE				( self     :						str		):






         """simple docstring"""






         _lowerCamelCase						:				Union[str, Any]    			= []
         if self.inprogress_constraint is None:
                for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
                       _lowerCamelCase						:				Union[str, Any]    			= constraint.advance()
                       if isinstance(__lowerCAmelCase						,					__lowerCAmelCase		):
                              token_list.append(__lowerCAmelCase		)
                       elif isinstance(__lowerCAmelCase						,					__lowerCAmelCase		):
                              token_list.extend(__lowerCAmelCase		)
         else:
                _lowerCamelCase						:				Any    			= self.inprogress_constraint.advance()
                if isinstance(__lowerCAmelCase						,					__lowerCAmelCase		):
                       token_list.append(__lowerCAmelCase		)
                elif isinstance(__lowerCAmelCase						,					__lowerCAmelCase		):
                       token_list.extend(__lowerCAmelCase		)

         if len(__lowerCAmelCase		) == 0:
                return None
         else:
                return token_list




  def 						SCREAMING_SNAKE_CASE				( self     :						Dict						,					__lowerCAmelCase     :						Optional[List[int]]		):






         """simple docstring"""






         self.init_state()

         if token_ids is not None:
                for token in token_ids:
                       # completes or steps **one** constraint
                       _lowerCamelCase			,    _lowerCamelCase						:				str    			= self.add(__lowerCAmelCase		)

                       # the entire list of constraints are fulfilled
                       if self.completed:
                              break




  def 						SCREAMING_SNAKE_CASE				( self     :						List[str]						,					__lowerCAmelCase     :						int		):






         """simple docstring"""






         if not isinstance(__lowerCAmelCase						,					__lowerCAmelCase		):
                raise ValueError(f'''`token_id` should be an `int`, but is `{token_id}`.'''		)

         _lowerCamelCase			,    _lowerCamelCase						:				List[str]    			= False, False

         if self.completed:
                _lowerCamelCase						:				int    			= True
                _lowerCamelCase						:				Optional[Any]    			= False
                return complete, stepped

         if self.inprogress_constraint is not None:
                # In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current
                # job, simply update the state

                _lowerCamelCase			,    _lowerCamelCase			,    _lowerCamelCase						:				Tuple    			= self.inprogress_constraint.update(__lowerCAmelCase		)
                if reset:
                       # 1. If the next token breaks the progress, then we must restart.
                       #     e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books".

                       #     But that doesn't mean we self.init_state(), since we only reset the state for this particular
                       #     constraint, not the full list of constraints.

                       self.pending_constraints.append(self.inprogress_constraint.copy(stateful=__lowerCAmelCase		)		)
                       _lowerCamelCase						:				List[str]    			= None

                if complete:
                       # 2. If the next token completes the constraint, move it to completed list, set
                       #     inprogress to None. If there are no pending constraints either, then this full list of constraints
                       #     is complete.

                       self.complete_constraints.append(self.inprogress_constraint		)
                       _lowerCamelCase						:				str    			= None

                       if len(self.pending_constraints		) == 0:
                              # we're done!
                              _lowerCamelCase						:				Any    			= True

         else:
                # Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list
                # of constraints?

                for cidx, pending_constraint in enumerate(self.pending_constraints		):
                       if pending_constraint.does_advance(__lowerCAmelCase		):
                              _lowerCamelCase			,    _lowerCamelCase			,    _lowerCamelCase						:				Any    			= pending_constraint.update(__lowerCAmelCase		)

                              if not stepped:
                                     raise Exception(
                                         '''`constraint.update(token_id)` is not yielding incremental progress, '''
                                         '''even though `constraint.does_advance(token_id)` is true.'''		)

                              if complete:
                                     self.complete_constraints.append(__lowerCAmelCase		)
                                     _lowerCamelCase						:				Optional[int]    			= None

                              if not complete and stepped:
                                     _lowerCamelCase						:				Union[str, Any]    			= pending_constraint

                              if complete or stepped:
                                     # If we made any progress at all, then it's at least not a "pending constraint".

                                     _lowerCamelCase						:				Union[str, Any]    			= (
                                         self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :]
                                     )

                                     if len(self.pending_constraints		) == 0 and self.inprogress_constraint is None:
                                            # If there's no longer any pending after this and no inprogress either, then we must be
                                            # complete.

                                            _lowerCamelCase						:				List[Any]    			= True

                                     break  # prevent accidentally stepping through multiple constraints with just one token.

         return complete, stepped




  def 						SCREAMING_SNAKE_CASE				( self     :						Dict						,					__lowerCAmelCase     :						Dict=True		):






         """simple docstring"""






         _lowerCamelCase						:				Union[str, Any]    			= ConstraintListState(self.constraints		)  # we actually never though self.constraints objects
         # throughout this process. So it's at initialization state.

         if stateful:
                _lowerCamelCase						:				Optional[Any]    			= [
                    constraint.copy(stateful=__lowerCAmelCase		) for constraint in self.complete_constraints
                ]
                if self.inprogress_constraint is not None:
                       _lowerCamelCase						:				List[str]    			= self.inprogress_constraint.copy(stateful=__lowerCAmelCase		)
                _lowerCamelCase						:				Optional[int]    			= [constraint.copy() for constraint in self.pending_constraints]

         return new_state





"""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))


