





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     :


				def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_=3   , UpperCAmelCase_=7   , UpperCAmelCase_=True   , UpperCAmelCase_=True   , UpperCAmelCase_=False   , UpperCAmelCase_=True   , UpperCAmelCase_=99   , UpperCAmelCase_=32   , UpperCAmelCase_=5   , UpperCAmelCase_=4   , UpperCAmelCase_=37   , UpperCAmelCase_="gelu"   , UpperCAmelCase_=0.1   , UpperCAmelCase_=0.1   , UpperCAmelCase_=512   , UpperCAmelCase_=16   , UpperCAmelCase_=2   , UpperCAmelCase_=0.0_2   , UpperCAmelCase_=3   , UpperCAmelCase_=4   , UpperCAmelCase_=None   , ):
								lowerCamelCase	 =parent
								lowerCamelCase	 =batch_size
								lowerCamelCase	 =seq_length
								lowerCamelCase	 =is_training
								lowerCamelCase	 =use_input_mask
								lowerCamelCase	 =use_token_type_ids
								lowerCamelCase	 =use_labels
								lowerCamelCase	 =vocab_size
								lowerCamelCase	 =hidden_size
								lowerCamelCase	 =num_hidden_layers
								lowerCamelCase	 =num_attention_heads
								lowerCamelCase	 =intermediate_size
								lowerCamelCase	 =hidden_act
								lowerCamelCase	 =hidden_dropout_prob
								lowerCamelCase	 =attention_probs_dropout_prob
								lowerCamelCase	 =max_position_embeddings
								lowerCamelCase	 =type_vocab_size
								lowerCamelCase	 =type_sequence_label_size
								lowerCamelCase	 =initializer_range
								lowerCamelCase	 =num_labels
								lowerCamelCase	 =num_choices
								lowerCamelCase	 =scope


				def _snake_case       (		self							):
								lowerCamelCase	 =ids_tensor([self.batch_size, self.seq_length]   , self.vocab_size							)

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

								lowerCamelCase	 =None

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

								lowerCamelCase	 =self.get_config()

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


				def _snake_case       (		self							):
								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=UpperCAmelCase_   , initializer_range=self.initializer_range   , pad_token_id=1   , new_decoder_architecture=UpperCAmelCase_   , )


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_							):
								lowerCamelCase	 =FalconModel(config=UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()
								lowerCamelCase	 =model(UpperCAmelCase_   , attention_mask=UpperCAmelCase_							)
								lowerCamelCase	 =model(UpperCAmelCase_							)
								self.parent.assertEqual(result.last_hidden_state.shape   , (self.batch_size, self.seq_length, self.hidden_size)							)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , ):
								lowerCamelCase	 =True
								lowerCamelCase	 =FalconModel(UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()
								lowerCamelCase	 =model(
								    UpperCAmelCase_   , attention_mask=UpperCAmelCase_   , encoder_hidden_states=UpperCAmelCase_   , encoder_attention_mask=UpperCAmelCase_   , )
								lowerCamelCase	 =model(
								    UpperCAmelCase_   , attention_mask=UpperCAmelCase_   , encoder_hidden_states=UpperCAmelCase_   , )
								lowerCamelCase	 =model(UpperCAmelCase_   , attention_mask=UpperCAmelCase_							)
								self.parent.assertEqual(result.last_hidden_state.shape   , (self.batch_size, self.seq_length, self.hidden_size)							)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , ):
								lowerCamelCase	 =FalconForCausalLM(config=UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()
								lowerCamelCase	 =model(UpperCAmelCase_   , attention_mask=UpperCAmelCase_   , labels=UpperCAmelCase_							)
								self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.seq_length, self.vocab_size)							)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , ):
								lowerCamelCase	 =True
								lowerCamelCase	 =True
								lowerCamelCase	 =FalconForCausalLM(config=UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()

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

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

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

								lowerCamelCase	 =model(
								    UpperCAmelCase_   , attention_mask=UpperCAmelCase_   , encoder_hidden_states=UpperCAmelCase_   , encoder_attention_mask=UpperCAmelCase_   , output_hidden_states=UpperCAmelCase_   , )["""hidden_states"""][0]
								lowerCamelCase	 =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	 =ids_tensor((1,)   , output_from_past.shape[-1]							).item()
								lowerCamelCase	 =output_from_no_past[:, -3:, random_slice_idx].detach()
								lowerCamelCase	 =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 _snake_case       (		self							):
								lowerCamelCase	 =self.prepare_config_and_inputs()
								(
								    (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  (
								    lowerCamelCase
								) ,  
								)	 =config_and_inputs
								lowerCamelCase	 ={"""input_ids""": input_ids, """attention_mask""": input_mask}
								return config, inputs_dict




@require_torch
class       __A     (						a    ,					a    ,					a    ,					unittest.TestCase    ):
				__A      =				(
				    (
				        FalconModel,
				        FalconForCausalLM,
				        FalconForSequenceClassification,
				        FalconForTokenClassification,
				        FalconForQuestionAnswering,
				    )
				    if is_torch_available()
				    else ()
				)
				__A      =				(FalconForCausalLM,) if is_torch_available() else ()
				__A      =				(
				    {
				        """feature-extraction""": FalconModel,
				        """text-classification""": FalconForSequenceClassification,
				        """text-generation""": FalconForCausalLM,
				        """question-answering""": FalconForQuestionAnswering,
				        """token-classification""": FalconForTokenClassification,
				        """zero-shot""": FalconForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				__A      =				False
				__A      =				False


				def _snake_case       (		self							):
								lowerCamelCase	 =FalconModelTester(self							)
								lowerCamelCase	 =ConfigTester(self   , config_class=UpperCAmelCase_   , hidden_size=37							)


				def _snake_case       (		self							):
								self.config_tester.run_common_tests()


				def _snake_case       (		self							):
								lowerCamelCase	 =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*UpperCAmelCase_							)


				def _snake_case       (		self							):
								lowerCamelCase ,  *lowerCamelCase	 =self.model_tester.prepare_config_and_inputs()
								for alibi in [True, False]:
												lowerCamelCase	 =alibi
												self.model_tester.create_and_check_model(UpperCAmelCase_   , *UpperCAmelCase_							)


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


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


				def _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase	 =self.model_tester.prepare_config_and_inputs_for_common()
								lowerCamelCase	 =input_dict["""input_ids"""]
								lowerCamelCase	 =FalconForCausalLM(UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()
								lowerCamelCase	 =model(UpperCAmelCase_   , use_cache=UpperCAmelCase_							)
								lowerCamelCase	 =input_ids.shape[0]
								lowerCamelCase	 =model._convert_to_rw_cache(result.past_key_values							)
								lowerCamelCase	 =model._convert_cache_to_standard_format(UpperCAmelCase_   , UpperCAmelCase_							)
								for layer in range(len(UpperCAmelCase_							)							):
												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 _snake_case       (		self							):
								lowerCamelCase ,  lowerCamelCase	 =self.model_tester.prepare_config_and_inputs_for_common()
								lowerCamelCase	 =3
								lowerCamelCase	 ="""multi_label_classification"""
								lowerCamelCase	 =input_dict["""input_ids"""]
								lowerCamelCase	 =input_ids.ne(1							).to(UpperCAmelCase_							)
								lowerCamelCase	 =ids_tensor(
								    [self.model_tester.batch_size, config.num_labels]   , self.model_tester.type_sequence_label_size							).to(torch.float							)
								lowerCamelCase	 =FalconForSequenceClassification(UpperCAmelCase_							)
								model.to(UpperCAmelCase_							)
								model.eval()
								lowerCamelCase	 =model(UpperCAmelCase_   , attention_mask=UpperCAmelCase_   , labels=UpperCAmelCase_							)
								self.assertEqual(result.logits.shape   , (self.model_tester.batch_size, self.model_tester.num_labels)							)


				def _snake_case       (		self							):
								# Falcon can have different numbers of KV-heads than the number of query heads, so we need
								# to override this test to use the right head counts.
								for model_class in self.all_generative_model_classes:
												lowerCamelCase ,  lowerCamelCase	 =self.model_tester.prepare_config_and_inputs_for_common()

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

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

												# 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	 =(
												    getattr(UpperCAmelCase_   , """decoder_layers"""   , UpperCAmelCase_							)
												    or getattr(UpperCAmelCase_   , """num_decoder_layers"""   , UpperCAmelCase_							)
												    or config.num_hidden_layers
												)
												lowerCamelCase	 =getattr(UpperCAmelCase_   , """num_kv_heads"""   , config.num_attention_heads							)
												lowerCamelCase	 =getattr(UpperCAmelCase_   , """d_model"""   , config.hidden_size							)
												lowerCamelCase	 =embed_dim // num_attention_heads

												lowerCamelCase	 =outputs["""past_key_values"""]
												self.assertEqual(len(UpperCAmelCase_							)   , UpperCAmelCase_							)

												lowerCamelCase ,  lowerCamelCase	 =inputs["""input_ids"""].shape
												for i in range(UpperCAmelCase_							):
																if config.new_decoder_architecture:
																				lowerCamelCase	 =config.num_attention_heads
																elif config.multi_query:
																				lowerCamelCase	 =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    ):


				@slow
				def _snake_case       (		self							):
								lowerCamelCase	 =AutoTokenizer.from_pretrained("""Rocketknight1/falcon-rw-1b"""							)
								lowerCamelCase	 =FalconForCausalLM.from_pretrained("""Rocketknight1/falcon-rw-1b"""							)
								model.eval()
								model.to(UpperCAmelCase_							)
								lowerCamelCase	 =tokenizer("""My favorite food is"""   , return_tensors="""pt"""							).to(UpperCAmelCase_							)

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

								lowerCamelCase	 =model.generate(**UpperCAmelCase_   , do_sample=UpperCAmelCase_   , max_new_tokens=19							)
								lowerCamelCase	 =tokenizer.batch_decode(UpperCAmelCase_							)[0]

								self.assertEqual(UpperCAmelCase_   , UpperCAmelCase_							)


				@slow
				def _snake_case       (		self							):
								# The big models are way too big for the CI, so we use tiny random models that resemble their
								# architectures but with much smaller and fewer layers
								for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
												lowerCamelCase	 =AutoTokenizer.from_pretrained(UpperCAmelCase_							)
												lowerCamelCase	 =FalconForCausalLM.from_pretrained(UpperCAmelCase_							)
												model.eval()
												model.to(UpperCAmelCase_							)
												lowerCamelCase	 =tokenizer("""My favorite food is"""   , return_tensors="""pt"""							).to(UpperCAmelCase_							)

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


				@slow
				def _snake_case       (		self							):
								# The big models are way too big for the CI, so we use tiny random models that resemble their
								# architectures but with much smaller and fewer layers
								with torch.no_grad():
												for repo in [
												    "Rocketknight1/falcon-rw-1b",
												    "Rocketknight1/tiny-random-falcon-7b",
												    "Rocketknight1/tiny-random-falcon-40b",
												]:
																lowerCamelCase	 =AutoTokenizer.from_pretrained(UpperCAmelCase_							)
																lowerCamelCase	 =FalconForCausalLM.from_pretrained(UpperCAmelCase_							)
																model.eval()
																model.to(device=UpperCAmelCase_							)
																lowerCamelCase	 =tokenizer("""My favorite food is"""   , return_tensors="""pt"""							).to(UpperCAmelCase_							)

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












def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase				)						->  list[str]:
				return [sentence[i : i + ngram_size] for i in range(len(_UpperCAmelCase				) - ngram_size + 1				)]


if __name__ == "__main__":
							from doctest import testmod

							testmod()





