
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	:





     def __init__(self   :					str       ,					_A   :					Dict       ,					_A   :					Union[str, Any]=3       ,					_A   :					int=7       ,					_A   :					Union[str, Any]=True       ,					_A   :					int=True       ,					_A   :					Optional[int]=False       ,					_A   :					Optional[Any]=True       ,					_A   :					Tuple=9_9       ,					_A   :					List[str]=3_2       ,					_A   :					Optional[int]=5       ,					_A   :					Optional[Any]=4       ,					_A   :					Any=3_7       ,					_A   :					List[Any]="gelu"       ,					_A   :					List[Any]=0.1       ,					_A   :					str=0.1       ,					_A   :					List[str]=5_1_2       ,					_A   :					Union[str, Any]=1_6       ,					_A   :					Optional[Any]=2       ,					_A   :					Optional[int]=0.02       ,					_A   :					List[Any]=3       ,					_A   :					Dict=4       ,					_A   :					Optional[Any]=None       ,					)		->		Union[str, Any]:
         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 	UpperCAmelCase(self   :					List[str]     )		->		List[str]:
         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 	UpperCAmelCase(self   :					str     )		->		Optional[Any]:
         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=_lowerCamelCase       ,					initializer_range=self.initializer_range       ,					pad_token_id=1       ,					new_decoder_architecture=_lowerCamelCase       ,					)





     def 	UpperCAmelCase(self   :					Any       ,					_A   :					str       ,					_A   :					Optional[Any]       ,					_A   :					Any       ,					_A   :					int       ,					_A   :					Dict       ,					_A   :					Dict       ,					_A   :					Dict     )		->		List[str]:
         snake_case				  =      FalconModel(config=_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         snake_case				  =      model(_lowerCamelCase       ,					attention_mask=_lowerCamelCase     )
         snake_case				  =      model(_lowerCamelCase     )
         self.parent.assertEqual(result.last_hidden_state.shape       ,					(self.batch_size, self.seq_length, self.hidden_size)     )





     def 	UpperCAmelCase(self   :					Union[str, Any]       ,					_A   :					int       ,					_A   :					Optional[int]       ,					_A   :					int       ,					_A   :					List[str]       ,					_A   :					List[Any]       ,					_A   :					Dict       ,					_A   :					Dict       ,					_A   :					Union[str, Any]       ,					_A   :					Optional[int]       ,					)		->		Tuple:
         snake_case				  =      True
         snake_case				  =      FalconModel(_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         snake_case				  =      model(
             _lowerCamelCase       ,					attention_mask=_lowerCamelCase       ,					encoder_hidden_states=_lowerCamelCase       ,					encoder_attention_mask=_lowerCamelCase       ,					)
         snake_case				  =      model(
             _lowerCamelCase       ,					attention_mask=_lowerCamelCase       ,					encoder_hidden_states=_lowerCamelCase       ,					)
         snake_case				  =      model(_lowerCamelCase       ,					attention_mask=_lowerCamelCase     )
         self.parent.assertEqual(result.last_hidden_state.shape       ,					(self.batch_size, self.seq_length, self.hidden_size)     )





     def 	UpperCAmelCase(self   :					List[str]       ,					_A   :					str       ,					_A   :					List[str]       ,					_A   :					Union[str, Any]       ,					_A   :					Union[str, Any]       ,					_A   :					Optional[int]       ,					_A   :					Optional[Any]       ,					_A   :					Tuple       ,					_A   :					Dict       ,					_A   :					Dict       ,					)		->		Dict:
         snake_case				  =      FalconForCausalLM(config=_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         snake_case				  =      model(_lowerCamelCase       ,					attention_mask=_lowerCamelCase       ,					labels=_lowerCamelCase     )
         self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.seq_length, self.vocab_size)     )





     def 	UpperCAmelCase(self   :					List[str]       ,					_A   :					Tuple       ,					_A   :					Optional[Any]       ,					_A   :					Tuple       ,					_A   :					int       ,					_A   :					List[Any]       ,					_A   :					int       ,					_A   :					Dict       ,					_A   :					int       ,					_A   :					Union[str, Any]       ,					)		->		Optional[Any]:
         snake_case				  =      True
         snake_case				  =      True
         snake_case				  =      FalconForCausalLM(config=_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()

         # first forward pass
         snake_case				  =      model(
             _lowerCamelCase       ,					attention_mask=_lowerCamelCase       ,					encoder_hidden_states=_lowerCamelCase       ,					encoder_attention_mask=_lowerCamelCase       ,					use_cache=_lowerCamelCase       ,					)
         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(
             _lowerCamelCase       ,					attention_mask=_lowerCamelCase       ,					encoder_hidden_states=_lowerCamelCase       ,					encoder_attention_mask=_lowerCamelCase       ,					output_hidden_states=_lowerCamelCase       ,					)["hidden_states"][0]
         snake_case				  =      model(
             _lowerCamelCase       ,					attention_mask=_lowerCamelCase       ,					encoder_hidden_states=_lowerCamelCase       ,					encoder_attention_mask=_lowerCamelCase       ,					past_key_values=_lowerCamelCase       ,					output_hidden_states=_lowerCamelCase       ,					)["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(_lowerCamelCase       ,					_lowerCamelCase       ,					atol=1E-3     )     )





     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		List[str]:
         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	(						__UpperCAmelCase							,				__UpperCAmelCase							,				__UpperCAmelCase							,				unittest.TestCase						):
     UpperCAmelCase__       :							Dict     =	(
         (
             FalconModel,
             FalconForCausalLM,
             FalconForSequenceClassification,
             FalconForTokenClassification,
             FalconForQuestionAnswering,
         )
         if is_torch_available()
         else ()
     )
     UpperCAmelCase__       :							Union[str, Any]     =	(FalconForCausalLM,) if is_torch_available() else ()
     UpperCAmelCase__       :							str     =	(
         {
             "feature-extraction": FalconModel,
             "text-classification": FalconForSequenceClassification,
             "text-generation": FalconForCausalLM,
             "question-answering": FalconForQuestionAnswering,
             "token-classification": FalconForTokenClassification,
             "zero-shot": FalconForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     UpperCAmelCase__       :							List[str]     =	False
     UpperCAmelCase__       :							List[Any]     =	False





     def 	UpperCAmelCase(self   :					Optional[int]     )		->		Union[str, Any]:
         snake_case				  =      FalconModelTester(self     )
         snake_case				  =      ConfigTester(self       ,					config_class=_lowerCamelCase       ,					hidden_size=3_7     )





     def 	UpperCAmelCase(self   :					int     )		->		List[Any]:
         self.config_tester.run_common_tests()





     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		str:
         snake_case				  =      self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_lowerCamelCase     )





     def 	UpperCAmelCase(self   :					str     )		->		Any:
         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(_lowerCamelCase       ,					*_lowerCamelCase     )





     def 	UpperCAmelCase(self   :					List[str]     )		->		List[str]:
         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(_lowerCamelCase     )
         snake_case				  =      ids_tensor([self.model_tester.batch_size]       ,					self.model_tester.type_sequence_label_size     )
         snake_case				  =      FalconForSequenceClassification(_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         snake_case				  =      model(_lowerCamelCase       ,					attention_mask=_lowerCamelCase       ,					labels=_lowerCamelCase     )
         self.assertEqual(result.logits.shape       ,					(self.model_tester.batch_size, self.model_tester.num_labels)     )





     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		Any:
         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(_lowerCamelCase     )
         snake_case				  =      ids_tensor([self.model_tester.batch_size]       ,					self.model_tester.type_sequence_label_size     )
         snake_case				  =      FalconForSequenceClassification(_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         snake_case				  =      model(_lowerCamelCase       ,					attention_mask=_lowerCamelCase       ,					labels=_lowerCamelCase     )
         self.assertEqual(result.logits.shape       ,					(self.model_tester.batch_size, self.model_tester.num_labels)     )





     def 	UpperCAmelCase(self   :					Any     )		->		int:
         snake_case  ,							snake_case				  =      self.model_tester.prepare_config_and_inputs_for_common()
         snake_case				  =      input_dict["input_ids"]
         snake_case				  =      FalconForCausalLM(_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         snake_case				  =      model(_lowerCamelCase       ,					use_cache=_lowerCamelCase     )
         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(_lowerCamelCase       ,					_lowerCamelCase     )
         for layer in range(len(_lowerCamelCase     )     ):
             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 	UpperCAmelCase(self   :					Any     )		->		int:
         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(_lowerCamelCase     )
         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(_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         snake_case				  =      model(_lowerCamelCase       ,					attention_mask=_lowerCamelCase       ,					labels=_lowerCamelCase     )
         self.assertEqual(result.logits.shape       ,					(self.model_tester.batch_size, self.model_tester.num_labels)     )





     def 	UpperCAmelCase(self   :					int     )		->		Tuple:
         # 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:
             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(_lowerCamelCase       ,					"use_cache"     ):
                 return

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

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

             snake_case				  =      outputs["past_key_values"]
             self.assertEqual(len(_lowerCamelCase     )       ,					_lowerCamelCase     )

             snake_case  ,							snake_case				  =      inputs["input_ids"].shape
             for i in range(_lowerCamelCase     ):
                 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						):





     @slow
     def 	UpperCAmelCase(self   :					List[str]     )		->		Optional[Any]:
         snake_case				  =      AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b"     )
         snake_case				  =      FalconForCausalLM.from_pretrained("Rocketknight1/falcon-rw-1b"     )
         model.eval()
         model.to(_lowerCamelCase     )
         snake_case				  =      tokenizer("My favorite food is"       ,					return_tensors="pt"     ).to(_lowerCamelCase     )

         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(**_lowerCamelCase       ,					do_sample=_lowerCamelCase       ,					max_new_tokens=1_9     )
         snake_case				  =      tokenizer.batch_decode(_lowerCamelCase     )[0]

         self.assertEqual(_lowerCamelCase       ,					_lowerCamelCase     )





     @slow
     def 	UpperCAmelCase(self   :					Any     )		->		str:
         # 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"]:
             snake_case				  =      AutoTokenizer.from_pretrained(_lowerCamelCase     )
             snake_case				  =      FalconForCausalLM.from_pretrained(_lowerCamelCase     )
             model.eval()
             model.to(_lowerCamelCase     )
             snake_case				  =      tokenizer("My favorite food is"       ,					return_tensors="pt"     ).to(_lowerCamelCase     )

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





     @slow
     def 	UpperCAmelCase(self   :					Dict     )		->		Optional[Any]:
         # 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",
             ]:
                 snake_case				  =      AutoTokenizer.from_pretrained(_lowerCamelCase     )
                 snake_case				  =      FalconForCausalLM.from_pretrained(_lowerCamelCase     )
                 model.eval()
                 model.to(device=_lowerCamelCase     )
                 snake_case				  =      tokenizer("My favorite food is"       ,					return_tensors="pt"     ).to(_lowerCamelCase     )

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








import os

# Precomputes a list of the 100 first triangular numbers
_A						      =				[int(0.5 * n * (n + 1)) for n in range(1, 1_01)]



def 		lowercase_    (     )							->							Tuple:
    """simple docstring"""

    snake_case				  =      os.path.dirname(os.path.realpath(A__	)	)
    snake_case				  =      os.path.join(A__				,	"words.txt"	)

    snake_case				  =      ""
    with open(A__	) as f:
        snake_case				  =      f.readline()

    snake_case				  =      [word.strip("\""	) for word in words.strip("\r\n"	).split(","	)]
    snake_case				  =      [
        word
        for word in [sum(ord(A__	) - 64 for x in word	) for word in words]
        if word in TRIANGULAR_NUMBERS
    ]
    return len(A__	)


if __name__ == "__main__":
   print(solution())

