

"""simple docstring"""


import dataclasses
import re
import string
from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple

import numpy as np

from . import residue_constants


__A						   =			Mapping[str, np.ndarray]
__A						   =			Mapping[str, Any]  # Is a nested dict.
__A						   =			0.01





@dataclasses.dataclass(frozen=__snake_case  )
class     snake_case	:
				SCREAMING_SNAKE_CASE_   :      np.ndarray  # [num_res, num_atom_type, 3]

				# Amino-acid type for each residue represented as an integer between 0 and
				# 20, where 20 is 'X'.
				SCREAMING_SNAKE_CASE_   :      np.ndarray  # [num_res]

				# Binary float mask to indicate presence of a particular atom. 1.0 if an atom
				# is present and 0.0 if not. This should be used for loss masking.
				SCREAMING_SNAKE_CASE_   :      np.ndarray  # [num_res, num_atom_type]

				# Residue index as used in PDB. It is not necessarily continuous or 0-indexed.
				SCREAMING_SNAKE_CASE_   :      np.ndarray  # [num_res]

				# B-factors, or temperature factors, of each residue (in sq. angstroms units),
				# representing the displacement of the residue from its ground truth mean
				# value.
				SCREAMING_SNAKE_CASE_   :      np.ndarray  # [num_res, num_atom_type]

				# Chain indices for multi-chain predictions
				SCREAMING_SNAKE_CASE_   :      Optional[np.ndarray]         =  None

				# Optional remark about the protein. Included as a comment in output PDB
				# files
				SCREAMING_SNAKE_CASE_   :      Optional[str]         =  None

				# Templates used to generate this protein (prediction-only)
				SCREAMING_SNAKE_CASE_   :      Optional[Sequence[str]]         =  None

				# Chain corresponding to each parent
				SCREAMING_SNAKE_CASE_   :      Optional[Sequence[int]]         =  None






def     a__    (    __SCREAMING_SNAKE_CASE )	->	Protein:
		__lowerCAmelCase:       str          =			R"(\[[A-Z]+\]\n)"
		__lowerCAmelCase:       List[str]          =			[tag.strip() for tag in re.split(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE ) if len(__SCREAMING_SNAKE_CASE ) > 0]
		__lowerCAmelCase:       Iterator[Tuple[str, List[str]]]          =			zip(tags[0::2]					, [l.split("\n" ) for l in tags[1::2]] )

		__lowerCAmelCase:       List[str]          =			["N", "CA", "C"]
		__lowerCAmelCase:       List[str]          =			None
		__lowerCAmelCase:       Tuple          =			None
		__lowerCAmelCase:       Dict          =			None
		for g in groups:
				if "[PRIMARY]" == g[0]:
						__lowerCAmelCase:       Optional[int]          =			g[1][0].strip()
						for i in range(len(__SCREAMING_SNAKE_CASE ) ):
								if seq[i] not in residue_constants.restypes:
										__lowerCAmelCase:       Any          =			"X"  # FIXME: strings are immutable
						__lowerCAmelCase:       int          =			np.array(
						    [residue_constants.restype_order.get(__SCREAMING_SNAKE_CASE					, residue_constants.restype_num ) for res_symbol in seq] )
				elif "[TERTIARY]" == g[0]:
						__lowerCAmelCase:       List[List[float]]          =			[]
						for axis in range(3 ):
								tertiary.append(list(map(__SCREAMING_SNAKE_CASE					, g[1][axis].split() ) ) )
						__lowerCAmelCase:       Any          =			np.array(__SCREAMING_SNAKE_CASE )
						__lowerCAmelCase:       List[str]          =			np.zeros((len(tertiary[0] ) // 3, residue_constants.atom_type_num, 3) ).astype(np.floataa )
						for i, atom in enumerate(__SCREAMING_SNAKE_CASE ):
								__lowerCAmelCase:       Union[str, Any]          =			np.transpose(tertiary_np[:, i::3] )
						atom_positions *= PICO_TO_ANGSTROM
				elif "[MASK]" == g[0]:
						__lowerCAmelCase:       List[Any]          =			np.array(list(map({"-": 0, "+": 1}.get					, g[1][0].strip() ) ) )
						__lowerCAmelCase:       Dict          =			np.zeros(
						    (
						        len(__SCREAMING_SNAKE_CASE ),
						        residue_constants.atom_type_num,
						    ) ).astype(np.floataa )
						for i, atom in enumerate(__SCREAMING_SNAKE_CASE ):
								__lowerCAmelCase:       Dict          =			1
						atom_mask *= mask[..., None]

		assert aatype is not None

		return Protein(
		    atom_positions=__SCREAMING_SNAKE_CASE					, atom_mask=__SCREAMING_SNAKE_CASE					, aatype=__SCREAMING_SNAKE_CASE					, residue_index=np.arange(len(__SCREAMING_SNAKE_CASE ) )					, b_factors=__SCREAMING_SNAKE_CASE					, )






def     a__    (    __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE = 0 )	->	List[str]:
		__lowerCAmelCase:       List[str]          =			[]

		__lowerCAmelCase:       Any          =			prot.remark
		if remark is not None:
				pdb_headers.append(F"REMARK {remark}" )

		__lowerCAmelCase:       Optional[Any]          =			prot.parents
		__lowerCAmelCase:       Optional[int]          =			prot.parents_chain_index
		if parents is not None and parents_chain_index is not None:
				__lowerCAmelCase:       Optional[Any]          =			[p for i, p in zip(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE ) if i == chain_id]

		if parents is None or len(__SCREAMING_SNAKE_CASE ) == 0:
				__lowerCAmelCase:       List[str]          =			["N/A"]

		pdb_headers.append(F"PARENT {' '.join(__SCREAMING_SNAKE_CASE )}" )

		return pdb_headers






def     a__    (    __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )	->	str:
		__lowerCAmelCase:       List[str]          =			[]
		__lowerCAmelCase:       str          =			pdb_str.split("\n" )

		__lowerCAmelCase:       List[Any]          =			prot.remark
		if remark is not None:
				out_pdb_lines.append(F"REMARK {remark}" )

		__lowerCAmelCase:       List[List[str]]
		if prot.parents is not None and len(prot.parents ) > 0:
				__lowerCAmelCase:       Dict          =			[]
				if prot.parents_chain_index is not None:
						__lowerCAmelCase:       Dict[str, List[str]]          =			{}
						for p, i in zip(prot.parents					, prot.parents_chain_index ):
								parent_dict.setdefault(str(__SCREAMING_SNAKE_CASE )					, [] )
								parent_dict[str(__SCREAMING_SNAKE_CASE )].append(__SCREAMING_SNAKE_CASE )

						__lowerCAmelCase:       List[str]          =			max([int(__SCREAMING_SNAKE_CASE ) for chain_idx in parent_dict] )
						for i in range(max_idx + 1 ):
								__lowerCAmelCase:       str          =			parent_dict.get(str(__SCREAMING_SNAKE_CASE )					, ["N/A"] )
								parents_per_chain.append(__SCREAMING_SNAKE_CASE )
				else:
						parents_per_chain.append(list(prot.parents ) )
		else:
				__lowerCAmelCase:       Dict          =			[["N/A"]]

		def make_parent_line(__SCREAMING_SNAKE_CASE ) -> str:
				return F"PARENT {' '.join(__SCREAMING_SNAKE_CASE )}"

		out_pdb_lines.append(make_parent_line(parents_per_chain[0] ) )

		__lowerCAmelCase:       Optional[Any]          =			0
		for i, l in enumerate(__SCREAMING_SNAKE_CASE ):
				if "PARENT" not in l and "REMARK" not in l:
						out_pdb_lines.append(__SCREAMING_SNAKE_CASE )
				if "TER" in l and "END" not in lines[i + 1]:
						chain_counter += 1
						if not chain_counter >= len(__SCREAMING_SNAKE_CASE ):
								__lowerCAmelCase:       List[str]          =			parents_per_chain[chain_counter]
						else:
								__lowerCAmelCase:       Tuple          =			["N/A"]

						out_pdb_lines.append(make_parent_line(__SCREAMING_SNAKE_CASE ) )

		return "\n".join(__SCREAMING_SNAKE_CASE )






def     a__    (    __SCREAMING_SNAKE_CASE )	->	str:
		__lowerCAmelCase:       List[Any]          =			residue_constants.restypes + ["X"]

		def res_atoa(__SCREAMING_SNAKE_CASE ) -> str:
				return residue_constants.restype_atoa.get(restypes[r]					, "UNK" )

		__lowerCAmelCase:       int          =			residue_constants.atom_types

		__lowerCAmelCase:       List[str]          =			[]

		__lowerCAmelCase:       str          =			prot.atom_mask
		__lowerCAmelCase:       Any          =			prot.aatype
		__lowerCAmelCase:       List[Any]          =			prot.atom_positions
		__lowerCAmelCase:       List[str]          =			prot.residue_index.astype(np.intaa )
		__lowerCAmelCase:       List[str]          =			prot.b_factors
		__lowerCAmelCase:       Optional[Any]          =			prot.chain_index

		if np.any(aatype > residue_constants.restype_num ):
				raise ValueError("Invalid aatypes." )

		__lowerCAmelCase:       str          =			get_pdb_headers(__SCREAMING_SNAKE_CASE )
		if len(__SCREAMING_SNAKE_CASE ) > 0:
				pdb_lines.extend(__SCREAMING_SNAKE_CASE )

		__lowerCAmelCase:       List[str]          =			aatype.shape[0]
		__lowerCAmelCase:       Optional[int]          =			1
		__lowerCAmelCase:       List[Any]          =			0
		__lowerCAmelCase:       int          =			string.ascii_uppercase
		__lowerCAmelCase:       Tuple          =			None
		# Add all atom sites.
		for i in range(__SCREAMING_SNAKE_CASE ):
				__lowerCAmelCase:       Tuple          =			res_atoa(aatype[i] )
				for atom_name, pos, mask, b_factor in zip(__SCREAMING_SNAKE_CASE					, atom_positions[i]					, atom_mask[i]					, b_factors[i] ):
						if mask < 0.5:
								continue

						__lowerCAmelCase:       Optional[Any]          =			"ATOM"
						__lowerCAmelCase:       List[Any]          =			atom_name if len(__SCREAMING_SNAKE_CASE ) == 4 else F" {atom_name}"
						__lowerCAmelCase:       List[Any]          =			""
						__lowerCAmelCase:       Any          =			""
						__lowerCAmelCase:       Union[str, Any]          =			1.00
						__lowerCAmelCase:       List[str]          =			atom_name[0]  # Protein supports only C, N, O, S, this works.
						__lowerCAmelCase:       Any          =			""

						__lowerCAmelCase:       int          =			"A"
						if chain_index is not None:
								__lowerCAmelCase:       Dict          =			chain_tags[chain_index[i]]

						# PDB is a columnar format, every space matters here!
						__lowerCAmelCase:       str          =			(
						    F"{record_type:<6}{atom_index:>5} {name:<4}{alt_loc:>1}"
						    F"{res_name_a:>3} {chain_tag:>1}"
						    F"{residue_index[i]:>4}{insertion_code:>1}   "
						    F"{pos[0]:>8.3f}{pos[1]:>8.3f}{pos[2]:>8.3f}"
						    F"{occupancy:>6.2f}{b_factor:>6.2f}          "
						    F"{element:>2}{charge:>2}"
						)
						pdb_lines.append(__SCREAMING_SNAKE_CASE )
						atom_index += 1

				__lowerCAmelCase:       Any          =			i == n - 1
				if chain_index is not None:
						if i != n - 1 and chain_index[i + 1] != prev_chain_index:
								__lowerCAmelCase:       Dict          =			True
								__lowerCAmelCase:       Optional[Any]          =			chain_index[i + 1]

				if should_terminate:
						# Close the chain.
						__lowerCAmelCase:       List[str]          =			"TER"
						__lowerCAmelCase:       List[str]          =			(
						    F"{chain_end:<6}{atom_index:>5}      {res_atoa(aatype[i] ):>3} {chain_tag:>1}{residue_index[i]:>4}"
						)
						pdb_lines.append(__SCREAMING_SNAKE_CASE )
						atom_index += 1

						if i != n - 1:
								# "prev" is a misnomer here. This happens at the beginning of
								# each new chain.
								pdb_lines.extend(get_pdb_headers(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE ) )

		pdb_lines.append("END" )
		pdb_lines.append("" )
		return "\n".join(__SCREAMING_SNAKE_CASE )






def     a__    (    __SCREAMING_SNAKE_CASE )	->	np.ndarray:
		return residue_constants.STANDARD_ATOM_MASK[prot.aatype]






def     a__    (    __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE = None					, __SCREAMING_SNAKE_CASE = None					, __SCREAMING_SNAKE_CASE = None					, __SCREAMING_SNAKE_CASE = None					, __SCREAMING_SNAKE_CASE = None					, )	->	Protein:
		return Protein(
		    aatype=features["aatype"]					, atom_positions=result["final_atom_positions"]					, atom_mask=result["final_atom_mask"]					, residue_index=features["residue_index"] + 1					, b_factors=b_factors if b_factors is not None else np.zeros_like(result["final_atom_mask"] )					, chain_index=__SCREAMING_SNAKE_CASE					, remark=__SCREAMING_SNAKE_CASE					, parents=__SCREAMING_SNAKE_CASE					, parents_chain_index=__SCREAMING_SNAKE_CASE					, )





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


