



"""simple docstring"""






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 			A__   :




	'''simple docstring'''





	def __init__(     self:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		List[str]=3       ,			_SCREAMING_SNAKE_CASE:		int=7       ,			_SCREAMING_SNAKE_CASE:		Optional[Any]=True       ,			_SCREAMING_SNAKE_CASE:		List[Any]=True       ,			_SCREAMING_SNAKE_CASE:		Tuple=False       ,			_SCREAMING_SNAKE_CASE:		str=True       ,			_SCREAMING_SNAKE_CASE:		str=99       ,			_SCREAMING_SNAKE_CASE:		int=32       ,			_SCREAMING_SNAKE_CASE:		List[Any]=5       ,			_SCREAMING_SNAKE_CASE:		Union[str, Any]=4       ,			_SCREAMING_SNAKE_CASE:		Optional[Any]=37       ,			_SCREAMING_SNAKE_CASE:		Optional[Any]="gelu"       ,			_SCREAMING_SNAKE_CASE:		Tuple=0.1       ,			_SCREAMING_SNAKE_CASE:		int=0.1       ,			_SCREAMING_SNAKE_CASE:		List[Any]=512       ,			_SCREAMING_SNAKE_CASE:		Any=16       ,			_SCREAMING_SNAKE_CASE:		Any=2       ,			_SCREAMING_SNAKE_CASE:		Optional[Any]=0.02       ,			_SCREAMING_SNAKE_CASE:		int=3       ,			_SCREAMING_SNAKE_CASE:		Optional[int]=4       ,			_SCREAMING_SNAKE_CASE:		Any=None       ,			)				->       int:

				"""simple docstring"""

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

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

				"""simple docstring"""

				__lowerCAmelCase      :  Union[str, Any]												=       ids_tensor([self.batch_size, self.seq_length]       ,			self.vocab_size)

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

				__lowerCAmelCase      :  int												=       None

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

				__lowerCAmelCase      :  List[str]												=       self.get_config()

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

	def 							_SCREAMING_SNAKE_CASE   (     self:		Any)				->       Any:

				"""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=_SCREAMING_SNAKE_CASE       ,			initializer_range=self.initializer_range       ,			pad_token_id=1       ,			new_decoder_architecture=_SCREAMING_SNAKE_CASE       ,			)

	def 							_SCREAMING_SNAKE_CASE   (     self:		int       ,			_SCREAMING_SNAKE_CASE:		Any       ,			_SCREAMING_SNAKE_CASE:		Any       ,			_SCREAMING_SNAKE_CASE:		List[str]       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		List[str]       ,			_SCREAMING_SNAKE_CASE:		List[str]       ,			_SCREAMING_SNAKE_CASE:		Optional[Any])				->       List[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  str												=       FalconModel(config=_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()
				__lowerCAmelCase      :  Dict												=       model(_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  List[str]												=       model(_SCREAMING_SNAKE_CASE)
				self.parent.assertEqual(result.last_hidden_state.shape       ,			(self.batch_size, self.seq_length, self.hidden_size))

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

				"""simple docstring"""

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

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

				"""simple docstring"""

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

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

				"""simple docstring"""

				__lowerCAmelCase      :  Dict												=       True
				__lowerCAmelCase      :  Tuple												=       True
				__lowerCAmelCase      :  Any												=       FalconForCausalLM(config=_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()

				# first forward pass
				__lowerCAmelCase      :  Any												=       model(
				    _SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE       ,			encoder_hidden_states=_SCREAMING_SNAKE_CASE       ,			encoder_attention_mask=_SCREAMING_SNAKE_CASE       ,			use_cache=_SCREAMING_SNAKE_CASE       ,			)
				__lowerCAmelCase      :  Any												=       outputs.past_key_values

				# create hypothetical multiple next token and extent to next_input_ids
				__lowerCAmelCase      :  List[Any]												=       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      :  Dict												=       torch.cat([input_ids, next_tokens]       ,			dim=-1)
				__lowerCAmelCase      :  Optional[int]												=       torch.cat([input_mask, next_mask]       ,			dim=-1)

				__lowerCAmelCase      :  Any												=       model(
				    _SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE       ,			encoder_hidden_states=_SCREAMING_SNAKE_CASE       ,			encoder_attention_mask=_SCREAMING_SNAKE_CASE       ,			output_hidden_states=_SCREAMING_SNAKE_CASE       ,			)["hidden_states"][0]
				__lowerCAmelCase      :  Tuple												=       model(
				    _SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE       ,			encoder_hidden_states=_SCREAMING_SNAKE_CASE       ,			encoder_attention_mask=_SCREAMING_SNAKE_CASE       ,			past_key_values=_SCREAMING_SNAKE_CASE       ,			output_hidden_states=_SCREAMING_SNAKE_CASE       ,			)["hidden_states"][0]

				# select random slice
				__lowerCAmelCase      :  List[str]												=       ids_tensor((1,)       ,			output_from_past.shape[-1]).item()
				__lowerCAmelCase      :  Optional[int]												=       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(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			atol=1e-3))







	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str])				->       List[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  str												=       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 			A__   (  __SCREAMING_SNAKE_CASE  ,		__SCREAMING_SNAKE_CASE  ,		__SCREAMING_SNAKE_CASE  ,		unittest.TestCase							):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE          =	(
	    (
	        FalconModel,
	        FalconForCausalLM,
	        FalconForSequenceClassification,
	        FalconForTokenClassification,
	        FalconForQuestionAnswering,
	    )
	    if is_torch_available()
	    else ()
	)
	SCREAMING_SNAKE_CASE          =	(FalconForCausalLM,) if is_torch_available() else ()
	SCREAMING_SNAKE_CASE          =	(
	    {
	        'feature-extraction': FalconModel,
	        'text-classification': FalconForSequenceClassification,
	        'text-generation': FalconForCausalLM,
	        'question-answering': FalconForQuestionAnswering,
	        'token-classification': FalconForTokenClassification,
	        'zero-shot': FalconForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	SCREAMING_SNAKE_CASE          =	False
	SCREAMING_SNAKE_CASE          =	False

	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[Any])				->       Dict:

				"""simple docstring"""

				__lowerCAmelCase      :  Tuple												=       FalconModelTester(self)
				__lowerCAmelCase      :  Tuple												=       ConfigTester(self       ,			config_class=_SCREAMING_SNAKE_CASE       ,			hidden_size=37)

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str])				->       Dict:

				"""simple docstring"""

				self.config_tester.run_common_tests()

	def 							_SCREAMING_SNAKE_CASE   (     self:		str)				->       Dict:

				"""simple docstring"""

				__lowerCAmelCase      :  Optional[Any]												=       self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[int])				->       Tuple:

				"""simple docstring"""

				__lowerCAmelCase						,	*__lowerCAmelCase      :  str												=       self.model_tester.prepare_config_and_inputs()
				for alibi in [True, False]:
							__lowerCAmelCase      :  str												=       alibi
							self.model_tester.create_and_check_model(_SCREAMING_SNAKE_CASE       ,			*_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		int)				->       Optional[int]:

				"""simple docstring"""

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

	def 							_SCREAMING_SNAKE_CASE   (     self:		Tuple)				->       Any:

				"""simple docstring"""

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

	def 							_SCREAMING_SNAKE_CASE   (     self:		int)				->       Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase						,	__lowerCAmelCase      :  str												=       self.model_tester.prepare_config_and_inputs_for_common()
				__lowerCAmelCase      :  str												=       input_dict["input_ids"]
				__lowerCAmelCase      :  Dict												=       FalconForCausalLM(_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()
				__lowerCAmelCase      :  List[str]												=       model(_SCREAMING_SNAKE_CASE       ,			use_cache=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Optional[Any]												=       input_ids.shape[0]
				__lowerCAmelCase      :  int												=       model._convert_to_rw_cache(result.past_key_values)
				__lowerCAmelCase      :  Union[str, Any]												=       model._convert_cache_to_standard_format(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)
				for layer in range(len(_SCREAMING_SNAKE_CASE)):
							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 							_SCREAMING_SNAKE_CASE   (     self:		List[Any])				->       Dict:

				"""simple docstring"""

				__lowerCAmelCase						,	__lowerCAmelCase      :  Tuple												=       self.model_tester.prepare_config_and_inputs_for_common()
				__lowerCAmelCase      :  Any												=       3
				__lowerCAmelCase      :  Optional[Any]												=       "multi_label_classification"
				__lowerCAmelCase      :  Dict												=       input_dict["input_ids"]
				__lowerCAmelCase      :  Optional[int]												=       input_ids.ne(1).to(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  str												=       ids_tensor(
				    [self.model_tester.batch_size, config.num_labels]       ,			self.model_tester.type_sequence_label_size).to(torch.float)
				__lowerCAmelCase      :  int												=       FalconForSequenceClassification(_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()
				__lowerCAmelCase      :  int												=       model(_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE       ,			labels=_SCREAMING_SNAKE_CASE)
				self.assertEqual(result.logits.shape       ,			(self.model_tester.batch_size, self.model_tester.num_labels))







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

				"""simple docstring"""

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

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

							__lowerCAmelCase      :  Dict												=       model_class(_SCREAMING_SNAKE_CASE).to(_SCREAMING_SNAKE_CASE)
							if "use_cache" not in inputs:
										__lowerCAmelCase      :  str												=       True
							__lowerCAmelCase      :  Tuple												=       model(**_SCREAMING_SNAKE_CASE)

							# 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

							__lowerCAmelCase      :  int												=       (
							    getattr(_SCREAMING_SNAKE_CASE       ,			"decoder_layers"       ,			_SCREAMING_SNAKE_CASE)
							    or getattr(_SCREAMING_SNAKE_CASE       ,			"num_decoder_layers"       ,			_SCREAMING_SNAKE_CASE)
							    or config.num_hidden_layers
							)
							__lowerCAmelCase      :  List[Any]												=       getattr(_SCREAMING_SNAKE_CASE       ,			"num_kv_heads"       ,			config.num_attention_heads)
							__lowerCAmelCase      :  List[Any]												=       getattr(_SCREAMING_SNAKE_CASE       ,			"d_model"       ,			config.hidden_size)
							__lowerCAmelCase      :  Optional[int]												=       embed_dim // num_attention_heads

							__lowerCAmelCase      :  List[Any]												=       outputs["past_key_values"]
							self.assertEqual(len(_SCREAMING_SNAKE_CASE)       ,			_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase						,	__lowerCAmelCase      :  int												=       inputs["input_ids"].shape
							for i in range(_SCREAMING_SNAKE_CASE):
										if config.new_decoder_architecture:
													__lowerCAmelCase      :  List[str]												=       config.num_attention_heads
										elif config.multi_query:
													__lowerCAmelCase      :  Optional[int]												=       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 			A__   (  unittest.TestCase							):




	'''simple docstring'''





	@slow
	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any])				->       Union[str, Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  Tuple												=       AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b")
				__lowerCAmelCase      :  int												=       FalconForCausalLM.from_pretrained("Rocketknight1/falcon-rw-1b")
				model.eval()
				model.to(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Any												=       tokenizer("My favorite food is"       ,			return_tensors="pt").to(_SCREAMING_SNAKE_CASE)

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

				__lowerCAmelCase      :  int												=       model.generate(**_SCREAMING_SNAKE_CASE       ,			do_sample=_SCREAMING_SNAKE_CASE       ,			max_new_tokens=19)
				__lowerCAmelCase      :  Dict												=       tokenizer.batch_decode(_SCREAMING_SNAKE_CASE)[0]

				self.assertEqual(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)

	@slow
	def 							_SCREAMING_SNAKE_CASE   (     self:		Tuple)				->       Any:

				"""simple docstring"""

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

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







	@slow
	def 							_SCREAMING_SNAKE_CASE   (     self:		Union[str, Any])				->       Optional[Any]:

				"""simple docstring"""

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

										# Test results are the same with and without cache
										__lowerCAmelCase      :  int												=       model.generate(**_SCREAMING_SNAKE_CASE       ,			do_sample=_SCREAMING_SNAKE_CASE       ,			max_new_tokens=20       ,			use_cache=_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase      :  str												=       model.generate(**_SCREAMING_SNAKE_CASE       ,			do_sample=_SCREAMING_SNAKE_CASE       ,			max_new_tokens=20       ,			use_cache=_SCREAMING_SNAKE_CASE)
										self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0)




"""simple docstring"""






import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
	import optax
	from flax.training.common_utils import onehot

	from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
	from transformers.models.ta.modeling_flax_ta import shift_tokens_right



@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class 			A__   (  unittest.TestCase							):




	'''simple docstring'''





	@slow
	def 							_SCREAMING_SNAKE_CASE   (     self:		Union[str, Any])				->       str:

				"""simple docstring"""

				__lowerCAmelCase      :  Optional[Any]												=       FlaxMTaForConditionalGeneration.from_pretrained("google/mt5-small")
				__lowerCAmelCase      :  Optional[Any]												=       AutoTokenizer.from_pretrained("google/mt5-small")

				__lowerCAmelCase      :  Tuple												=       tokenizer("Hello there"       ,			return_tensors="np").input_ids
				__lowerCAmelCase      :  Dict												=       tokenizer("Hi I am"       ,			return_tensors="np").input_ids

				__lowerCAmelCase      :  str												=       shift_tokens_right(_SCREAMING_SNAKE_CASE       ,			model.config.pad_token_id       ,			model.config.decoder_start_token_id)

				__lowerCAmelCase      :  Optional[int]												=       model(_SCREAMING_SNAKE_CASE       ,			decoder_input_ids=_SCREAMING_SNAKE_CASE).logits
				__lowerCAmelCase      :  int												=       optax.softmax_cross_entropy(_SCREAMING_SNAKE_CASE       ,			onehot(_SCREAMING_SNAKE_CASE       ,			logits.shape[-1])).mean()

				__lowerCAmelCase      :  List[str]												=       -(labels.shape[-1] * loss.item())

				__lowerCAmelCase      :  str												=       -84.9127
				self.assertTrue(abs(mtf_score - EXPECTED_SCORE) < 1e-4)