






import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, 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 (
	    FalconForCausalLM,
	    FalconForQuestionAnswering,
	    FalconForSequenceClassification,
	    FalconForTokenClassification,
	    FalconModel,
	)


class    _lowerCAmelCase			:



				"""simple docstring"""



				def __init__(      self    :    List[Any]			,		UpperCamelCase__    :    List[str]			,		UpperCamelCase__    :    str=3			,		UpperCamelCase__    :    Tuple=7			,		UpperCamelCase__    :    int=True			,		UpperCamelCase__    :    List[Any]=True			,		UpperCamelCase__    :    Optional[Any]=False			,		UpperCamelCase__    :    Any=True			,		UpperCamelCase__    :    Dict=9_9			,		UpperCamelCase__    :    List[Any]=3_2			,		UpperCamelCase__    :    Optional[Any]=5			,		UpperCamelCase__    :    Optional[Any]=4			,		UpperCamelCase__    :    Dict=3_7			,		UpperCamelCase__    :    Dict="gelu"			,		UpperCamelCase__    :    List[str]=0.1			,		UpperCamelCase__    :    int=0.1			,		UpperCamelCase__    :    int=5_1_2			,		UpperCamelCase__    :    Any=1_6			,		UpperCamelCase__    :    List[Any]=2			,		UpperCamelCase__    :    str=0.02			,		UpperCamelCase__    :    int=3			,		UpperCamelCase__    :    Any=4			,		UpperCamelCase__    :    Tuple=None			,		):






											'''simple docstring'''







											snake_case__	  =       parent
											snake_case__	  =       batch_size
											snake_case__	  =       seq_length
											snake_case__	  =       is_training
											snake_case__	  =       use_input_mask
											snake_case__	  =       use_token_type_ids
											snake_case__	  =       use_labels
											snake_case__	  =       vocab_size
											snake_case__	  =       hidden_size
											snake_case__	  =       num_hidden_layers
											snake_case__	  =       num_attention_heads
											snake_case__	  =       intermediate_size
											snake_case__	  =       hidden_act
											snake_case__	  =       hidden_dropout_prob
											snake_case__	  =       attention_probs_dropout_prob
											snake_case__	  =       max_position_embeddings
											snake_case__	  =       type_vocab_size
											snake_case__	  =       type_sequence_label_size
											snake_case__	  =       initializer_range
											snake_case__	  =       num_labels
											snake_case__	  =       num_choices
											snake_case__	  =       scope



				def 					__magic_name__							(      self    :    List[Any]):






											'''simple docstring'''







											snake_case__	  =       ids_tensor([self.batch_size, self.seq_length]			,		self.vocab_size)

											snake_case__	  =       None
											if self.use_input_mask:
																		snake_case__	  =       random_attention_mask([self.batch_size, self.seq_length])

											snake_case__	  =       None

											snake_case__	  =       None
											snake_case__	  =       None
											snake_case__	  =       None
											if self.use_labels:
																		snake_case__	  =       ids_tensor([self.batch_size]			,		self.type_sequence_label_size)
																		snake_case__	  =       ids_tensor([self.batch_size, self.seq_length]			,		self.num_labels)
																		snake_case__	  =       ids_tensor([self.batch_size]			,		self.num_choices)

											snake_case__	  =       self.get_config()

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



				def 					__magic_name__							(      self    :    Optional[int]):






											'''simple docstring'''







											return FalconConfig(
											    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			,		pad_token_id=1			,		new_decoder_architecture=_a			,		)



				def 					__magic_name__							(      self    :    Union[str, Any]			,		UpperCamelCase__    :    Dict			,		UpperCamelCase__    :    Optional[Any]			,		UpperCamelCase__    :    List[str]			,		UpperCamelCase__    :    Any			,		UpperCamelCase__    :    Optional[Any]			,		UpperCamelCase__    :    List[Any]			,		UpperCamelCase__    :    Tuple):






											'''simple docstring'''







											snake_case__	  =       FalconModel(config=_a)
											model.to(_a)
											model.eval()
											snake_case__	  =       model(_a			,		attention_mask=_a)
											snake_case__	  =       model(_a)
											self.parent.assertEqual(result.last_hidden_state.shape			,		(self.batch_size, self.seq_length, self.hidden_size))



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






											'''simple docstring'''







											snake_case__	  =       True
											snake_case__	  =       FalconModel(_a)
											model.to(_a)
											model.eval()
											snake_case__	  =       model(
											    _a			,		attention_mask=_a			,		encoder_hidden_states=_a			,		encoder_attention_mask=_a			,		)
											snake_case__	  =       model(
											    _a			,		attention_mask=_a			,		encoder_hidden_states=_a			,		)
											snake_case__	  =       model(_a			,		attention_mask=_a)
											self.parent.assertEqual(result.last_hidden_state.shape			,		(self.batch_size, self.seq_length, self.hidden_size))



				def 					__magic_name__							(      self    :    Union[str, Any]			,		UpperCamelCase__    :    Union[str, Any]			,		UpperCamelCase__    :    List[Any]			,		UpperCamelCase__    :    str			,		UpperCamelCase__    :    List[Any]			,		UpperCamelCase__    :    Optional[int]			,		UpperCamelCase__    :    Union[str, Any]			,		UpperCamelCase__    :    List[Any]			,		UpperCamelCase__    :    Optional[int]			,		UpperCamelCase__    :    Dict			,		):






											'''simple docstring'''







											snake_case__	  =       FalconForCausalLM(config=_a)
											model.to(_a)
											model.eval()
											snake_case__	  =       model(_a			,		attention_mask=_a			,		labels=_a)
											self.parent.assertEqual(result.logits.shape			,		(self.batch_size, self.seq_length, self.vocab_size))



				def 					__magic_name__							(      self    :    int			,		UpperCamelCase__    :    int			,		UpperCamelCase__    :    Dict			,		UpperCamelCase__    :    int			,		UpperCamelCase__    :    Any			,		UpperCamelCase__    :    Optional[int]			,		UpperCamelCase__    :    str			,		UpperCamelCase__    :    List[str]			,		UpperCamelCase__    :    Any			,		UpperCamelCase__    :    Tuple			,		):






											'''simple docstring'''







											snake_case__	  =       True
											snake_case__	  =       True
											snake_case__	  =       FalconForCausalLM(config=_a)
											model.to(_a)
											model.eval()

											# first forward pass
											snake_case__	  =       model(
											    _a			,		attention_mask=_a			,		encoder_hidden_states=_a			,		encoder_attention_mask=_a			,		use_cache=_a			,		)
											snake_case__	  =       outputs.past_key_values

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

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

											snake_case__	  =       model(
											    _a			,		attention_mask=_a			,		encoder_hidden_states=_a			,		encoder_attention_mask=_a			,		output_hidden_states=_a			,		)["""hidden_states"""][0]
											snake_case__	  =       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
											snake_case__	  =       ids_tensor((1,)			,		output_from_past.shape[-1]).item()
											snake_case__	  =       output_from_no_past[:, -3:, random_slice_idx].detach()
											snake_case__	  =       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 					__magic_name__							(      self    :    Any):






											'''simple docstring'''







											snake_case__	  =       self.prepare_config_and_inputs()
											(
											    (
											    snake_case__  
											)  , (
											    snake_case__  
											)  , (
											    snake_case__  
											)  , (
											    snake_case__  
											)  , (
											    snake_case__  
											)  , (
											    snake_case__  
											)  , (
											    snake_case__  
											)  ,   
											)	  =       config_and_inputs
											snake_case__	  =       {"""input_ids""": input_ids, """attention_mask""": input_mask}
											return config, inputs_dict


@require_torch
class    _lowerCAmelCase			(		__SCREAMING_SNAKE_CASE     ,      __SCREAMING_SNAKE_CASE     ,      __SCREAMING_SNAKE_CASE     ,      unittest.TestCase      ):



				"""simple docstring"""



				_lowercase				:   List[str]      =  (
				    (
				        FalconModel,
				        FalconForCausalLM,
				        FalconForSequenceClassification,
				        FalconForTokenClassification,
				        FalconForQuestionAnswering,
				    )
				    if is_torch_available()
				    else ()
				)
				_lowercase				:   Optional[int]      =  (FalconForCausalLM,) if is_torch_available() else ()
				_lowercase				:   Dict      =  (
				    {
				        'feature-extraction': FalconModel,
				        'text-classification': FalconForSequenceClassification,
				        'text-generation': FalconForCausalLM,
				        'question-answering': FalconForQuestionAnswering,
				        'token-classification': FalconForTokenClassification,
				        'zero-shot': FalconForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				_lowercase				:   Optional[int]      =  False
				_lowercase				:   List[str]      =  False



				def 					__magic_name__							(      self    :    List[Any]):






											'''simple docstring'''







											snake_case__	  =       FalconModelTester(self)
											snake_case__	  =       ConfigTester(self			,		config_class=_a			,		hidden_size=3_7)



				def 					__magic_name__							(      self    :    Optional[Any]):






											'''simple docstring'''







											self.config_tester.run_common_tests()



				def 					__magic_name__							(      self    :    str):






											'''simple docstring'''







											snake_case__	  =       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_a)



				def 					__magic_name__							(      self    :    Tuple):






											'''simple docstring'''







											snake_case__  , *snake_case__	  =       self.model_tester.prepare_config_and_inputs()
											for alibi in [True, False]:
																		snake_case__	  =       alibi
																		self.model_tester.create_and_check_model(_a			,		*_a)



				def 					__magic_name__							(      self    :    Optional[Any]):






											'''simple docstring'''







											snake_case__  , snake_case__	  =       self.model_tester.prepare_config_and_inputs_for_common()
											snake_case__	  =       3
											snake_case__	  =       input_dict["""input_ids"""]
											snake_case__	  =       input_ids.ne(1).to(_a)
											snake_case__	  =       ids_tensor([self.model_tester.batch_size]			,		self.model_tester.type_sequence_label_size)
											snake_case__	  =       FalconForSequenceClassification(_a)
											model.to(_a)
											model.eval()
											snake_case__	  =       model(_a			,		attention_mask=_a			,		labels=_a)
											self.assertEqual(result.logits.shape			,		(self.model_tester.batch_size, self.model_tester.num_labels))



				def 					__magic_name__							(      self    :    str):






											'''simple docstring'''







											snake_case__  , snake_case__	  =       self.model_tester.prepare_config_and_inputs_for_common()
											snake_case__	  =       3
											snake_case__	  =       """single_label_classification"""
											snake_case__	  =       input_dict["""input_ids"""]
											snake_case__	  =       input_ids.ne(1).to(_a)
											snake_case__	  =       ids_tensor([self.model_tester.batch_size]			,		self.model_tester.type_sequence_label_size)
											snake_case__	  =       FalconForSequenceClassification(_a)
											model.to(_a)
											model.eval()
											snake_case__	  =       model(_a			,		attention_mask=_a			,		labels=_a)
											self.assertEqual(result.logits.shape			,		(self.model_tester.batch_size, self.model_tester.num_labels))



				def 					__magic_name__							(      self    :    List[str]):






											'''simple docstring'''







											snake_case__  , snake_case__	  =       self.model_tester.prepare_config_and_inputs_for_common()
											snake_case__	  =       input_dict["""input_ids"""]
											snake_case__	  =       FalconForCausalLM(_a)
											model.to(_a)
											model.eval()
											snake_case__	  =       model(_a			,		use_cache=_a)
											snake_case__	  =       input_ids.shape[0]
											snake_case__	  =       model._convert_to_rw_cache(result.past_key_values)
											snake_case__	  =       model._convert_cache_to_standard_format(_a			,		_a)
											for layer in range(len(_a)):
																		for tensor_idx in range(2):
																									self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3)
																									self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4)
																									self.assertTrue(
																									    torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]))



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






											'''simple docstring'''







											snake_case__  , snake_case__	  =       self.model_tester.prepare_config_and_inputs_for_common()
											snake_case__	  =       3
											snake_case__	  =       """multi_label_classification"""
											snake_case__	  =       input_dict["""input_ids"""]
											snake_case__	  =       input_ids.ne(1).to(_a)
											snake_case__	  =       ids_tensor(
											    [self.model_tester.batch_size, config.num_labels]			,		self.model_tester.type_sequence_label_size).to(torch.float)
											snake_case__	  =       FalconForSequenceClassification(_a)
											model.to(_a)
											model.eval()
											snake_case__	  =       model(_a			,		attention_mask=_a			,		labels=_a)
											self.assertEqual(result.logits.shape			,		(self.model_tester.batch_size, self.model_tester.num_labels))



				def 					__magic_name__							(      self    :    List[str]):






											'''simple docstring'''







											for model_class in self.all_generative_model_classes:
																		snake_case__  , snake_case__	  =       self.model_tester.prepare_config_and_inputs_for_common()

																		# If it doesn't support cache, pass the test
																		if not hasattr(_a			,		"""use_cache"""):
																									return

																		snake_case__	  =       model_class(_a).to(_a)
																		if "use_cache" not in inputs:
																									snake_case__	  =       True
																		snake_case__	  =       model(**_a)

																		# If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
																		if "past_key_values" not in outputs:
																									return

																		snake_case__	  =       (
																		    getattr(_a			,		"""decoder_layers"""			,		_a)
																		    or getattr(_a			,		"""num_decoder_layers"""			,		_a)
																		    or config.num_hidden_layers
																		)
																		snake_case__	  =       getattr(_a			,		"""num_kv_heads"""			,		config.num_attention_heads)
																		snake_case__	  =       getattr(_a			,		"""d_model"""			,		config.hidden_size)
																		snake_case__	  =       embed_dim // num_attention_heads

																		snake_case__	  =       outputs["""past_key_values"""]
																		self.assertEqual(len(_a)			,		_a)

																		snake_case__  , snake_case__	  =       inputs["""input_ids"""].shape
																		for i in range(_a):
																									if config.new_decoder_architecture:
																																snake_case__	  =       config.num_attention_heads
																									elif config.multi_query:
																																snake_case__	  =       1
																									self.assertEqual(len(past_kv[0])			,		2)  # K V for the decoder = 2
																									self.assertEqual(
																									    past_kv[i][0].shape			,		(batch_size, num_attention_heads, seq_length, per_head_embed_dim))
																									self.assertEqual(
																									    past_kv[i][1].shape			,		(batch_size, num_attention_heads, seq_length, per_head_embed_dim))


@require_torch
class    _lowerCAmelCase			(		unittest.TestCase      ):



				"""simple docstring"""



				@slow
				def 					__magic_name__							(      self    :    int):






											'''simple docstring'''







											snake_case__	  =       AutoTokenizer.from_pretrained("""Rocketknight1/falcon-rw-1b""")
											snake_case__	  =       FalconForCausalLM.from_pretrained("""Rocketknight1/falcon-rw-1b""")
											model.eval()
											model.to(_a)
											snake_case__	  =       tokenizer("""My favorite food is"""			,		return_tensors="""pt""").to(_a)

											snake_case__	  =       (
											    """My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday."""
											)

											snake_case__	  =       model.generate(**_a			,		do_sample=_a			,		max_new_tokens=1_9)
											snake_case__	  =       tokenizer.batch_decode(_a)[0]

											self.assertEqual(_a			,		_a)



				@slow
				def 					__magic_name__							(      self    :    str):






											'''simple docstring'''







											for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
																		snake_case__	  =       AutoTokenizer.from_pretrained(_a)
																		snake_case__	  =       FalconForCausalLM.from_pretrained(_a)
																		model.eval()
																		model.to(_a)
																		snake_case__	  =       tokenizer("""My favorite food is"""			,		return_tensors="""pt""").to(_a)

																		# We just test that these run without errors - the models are randomly initialized
																		# and so the actual text outputs will be garbage
																		model.generate(**_a			,		do_sample=_a			,		max_new_tokens=4)
																		model.generate(**_a			,		do_sample=_a			,		max_new_tokens=4)
																		model.generate(**_a			,		num_beams=2			,		max_new_tokens=4)



				@slow
				def 					__magic_name__							(      self    :    Tuple):






											'''simple docstring'''







											with torch.no_grad():
																		for repo in [
																		    "Rocketknight1/falcon-rw-1b",
																		    "Rocketknight1/tiny-random-falcon-7b",
																		    "Rocketknight1/tiny-random-falcon-40b",
																		]:
																									snake_case__	  =       AutoTokenizer.from_pretrained(_a)
																									snake_case__	  =       FalconForCausalLM.from_pretrained(_a)
																									model.eval()
																									model.to(device=_a)
																									snake_case__	  =       tokenizer("""My favorite food is"""			,		return_tensors="""pt""").to(_a)

																									# Test results are the same with and without cache
																									snake_case__	  =       model.generate(**_a			,		do_sample=_a			,		max_new_tokens=2_0			,		use_cache=_a)
																									snake_case__	  =       model.generate(**_a			,		do_sample=_a			,		max_new_tokens=2_0			,		use_cache=_a)
																									self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0)








from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


a__       =     {
    """configuration_autoformer""": [
        """AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """AutoformerConfig""",
    ],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	a__       =     [
	    """AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """AutoformerForPrediction""",
	    """AutoformerModel""",
	    """AutoformerPreTrainedModel""",
	]


if TYPE_CHECKING:
	from .configuration_autoformer import (
	    AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    AutoformerConfig,
	)

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_autoformer import (
		    AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    AutoformerForPrediction,
		    AutoformerModel,
		    AutoformerPreTrainedModel,
		)

else:
	import sys

	a__       =     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
