



'''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    lowerCamelCase__				:






      """simple docstring"""




      def __init__(       self							:     int			,a__							:     Optional[Any]			,a__							:     Optional[Any]=3			,a__							:     Optional[int]=7			,a__							:     Optional[int]=True			,a__							:     Optional[int]=True			,a__							:     Union[str, Any]=False			,a__							:     int=True			,a__							:     Any=99			,a__							:     Tuple=32			,a__							:     str=5			,a__							:     List[Any]=4			,a__							:     List[str]=37			,a__							:     List[Any]="gelu"			,a__							:     Tuple=0.1			,a__							:     str=0.1			,a__							:     int=5_12			,a__							:     Dict=16			,a__							:     str=2			,a__							:     Optional[int]=0.02			,a__							:     Union[str, Any]=3			,a__							:     Dict=4			,a__							:     Union[str, Any]=None			,):
             a__								= parent
             a__								= batch_size
             a__								= seq_length
             a__								= is_training
             a__								= use_input_mask
             a__								= use_token_type_ids
             a__								= use_labels
             a__								= vocab_size
             a__								= hidden_size
             a__								= num_hidden_layers
             a__								= num_attention_heads
             a__								= intermediate_size
             a__								= hidden_act
             a__								= hidden_dropout_prob
             a__								= attention_probs_dropout_prob
             a__								= max_position_embeddings
             a__								= type_vocab_size
             a__								= type_sequence_label_size
             a__								= initializer_range
             a__								= num_labels
             a__								= num_choices
             a__								= scope




      def lowerCAmelCase_					(       self							:     List[Any]   ):
             a__								= ids_tensor([self.batch_size, self.seq_length]			,self.vocab_size   )

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

             a__								= None

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

             a__								= self.get_config()

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




      def lowerCAmelCase_					(       self							:     List[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 lowerCAmelCase_					(       self							:     List[str]			,a__							:     Tuple			,a__							:     Dict			,a__							:     int			,a__							:     Union[str, Any]			,a__							:     Optional[Any]			,a__							:     List[Any]			,a__							:     List[str]   ):
             a__								= FalconModel(config=__lowerCamelCase   )
             model.to(__lowerCamelCase   )
             model.eval()
             a__								= model(__lowerCamelCase			,attention_mask=__lowerCamelCase   )
             a__								= model(__lowerCamelCase   )
             self.parent.assertEqual(result.last_hidden_state.shape			,(self.batch_size, self.seq_length, self.hidden_size)   )




      def lowerCAmelCase_					(       self							:     Dict			,a__							:     str			,a__							:     str			,a__							:     Any			,a__							:     Tuple			,a__							:     Any			,a__							:     List[Any]			,a__							:     Union[str, Any]			,a__							:     int			,a__							:     str			,):
             a__								= True
             a__								= FalconModel(__lowerCamelCase   )
             model.to(__lowerCamelCase   )
             model.eval()
             a__								= model(
                 __lowerCamelCase			,attention_mask=__lowerCamelCase			,encoder_hidden_states=__lowerCamelCase			,encoder_attention_mask=__lowerCamelCase			,)
             a__								= model(
                 __lowerCamelCase			,attention_mask=__lowerCamelCase			,encoder_hidden_states=__lowerCamelCase			,)
             a__								= model(__lowerCamelCase			,attention_mask=__lowerCamelCase   )
             self.parent.assertEqual(result.last_hidden_state.shape			,(self.batch_size, self.seq_length, self.hidden_size)   )




      def lowerCAmelCase_					(       self							:     List[Any]			,a__							:     str			,a__							:     List[str]			,a__							:     Tuple			,a__							:     int			,a__							:     int			,a__							:     List[str]			,a__							:     Optional[int]			,a__							:     Tuple			,a__							:     str			,):
             a__								= FalconForCausalLM(config=__lowerCamelCase   )
             model.to(__lowerCamelCase   )
             model.eval()
             a__								= model(__lowerCamelCase			,attention_mask=__lowerCamelCase			,labels=__lowerCamelCase   )
             self.parent.assertEqual(result.logits.shape			,(self.batch_size, self.seq_length, self.vocab_size)   )




      def lowerCAmelCase_					(       self							:     Dict			,a__							:     Dict			,a__							:     List[str]			,a__							:     List[Any]			,a__							:     str			,a__							:     Dict			,a__							:     Tuple			,a__							:     List[str]			,a__							:     int			,a__							:     str			,):
             a__								= True
             a__								= True
             a__								= FalconForCausalLM(config=__lowerCamelCase   )
             model.to(__lowerCamelCase   )
             model.eval()

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

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

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

             a__								= model(
                 __lowerCamelCase			,attention_mask=__lowerCamelCase			,encoder_hidden_states=__lowerCamelCase			,encoder_attention_mask=__lowerCamelCase			,output_hidden_states=__lowerCamelCase			,)["hidden_states"][0]
             a__								= 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
             a__								= ids_tensor((1,)			,output_from_past.shape[-1]   ).item()
             a__								= output_from_no_past[:, -3:, random_slice_idx].detach()
             a__								= 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 lowerCAmelCase_					(       self							:     Union[str, Any]   ):
             a__								= self.prepare_config_and_inputs()
             (
                 (
                 a__     
             )						,					(
                 a__     
             )						,					(
                 a__     
             )						,					(
                 a__     
             )						,					(
                 a__     
             )						,					(
                 a__     
             )						,					(
                 a__     
             )						,					     
             )								= config_and_inputs
             a__								= {"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"""
      UpperCamelCase__			 =     (
          (
              FalconModel,
              FalconForCausalLM,
              FalconForSequenceClassification,
              FalconForTokenClassification,
              FalconForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      UpperCamelCase__			 =     (FalconForCausalLM,) if is_torch_available() else ()
      UpperCamelCase__			 =     (
          {
              '''feature-extraction''': FalconModel,
              '''text-classification''': FalconForSequenceClassification,
              '''text-generation''': FalconForCausalLM,
              '''question-answering''': FalconForQuestionAnswering,
              '''token-classification''': FalconForTokenClassification,
              '''zero-shot''': FalconForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      UpperCamelCase__			 =     False
      UpperCamelCase__			 =     False




      def lowerCAmelCase_					(       self							:     Optional[int]   ):
             a__								= FalconModelTester(self   )
             a__								= ConfigTester(self			,config_class=__lowerCamelCase			,hidden_size=37   )




      def lowerCAmelCase_					(       self							:     Any   ):
             self.config_tester.run_common_tests()




      def lowerCAmelCase_					(       self							:     Any   ):
             a__								= self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*__lowerCamelCase   )




      def lowerCAmelCase_					(       self							:     Optional[int]   ):
             a__						,					*a__								= self.model_tester.prepare_config_and_inputs()
             for alibi in [True, False]:
                    a__								= alibi
                    self.model_tester.create_and_check_model(__lowerCamelCase			,*__lowerCamelCase   )




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




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




      def lowerCAmelCase_					(       self							:     List[str]   ):
             a__						,					a__								= self.model_tester.prepare_config_and_inputs_for_common()
             a__								= input_dict["input_ids"]
             a__								= FalconForCausalLM(__lowerCamelCase   )
             model.to(__lowerCamelCase   )
             model.eval()
             a__								= model(__lowerCamelCase			,use_cache=__lowerCamelCase   )
             a__								= input_ids.shape[0]
             a__								= model._convert_to_rw_cache(result.past_key_values   )
             a__								= 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 lowerCAmelCase_					(       self							:     Optional[Any]   ):
             a__						,					a__								= self.model_tester.prepare_config_and_inputs_for_common()
             a__								= 3
             a__								= "multi_label_classification"
             a__								= input_dict["input_ids"]
             a__								= input_ids.ne(1   ).to(__lowerCamelCase   )
             a__								= ids_tensor(
                 [self.model_tester.batch_size, config.num_labels]			,self.model_tester.type_sequence_label_size   ).to(torch.float   )
             a__								= FalconForSequenceClassification(__lowerCamelCase   )
             model.to(__lowerCamelCase   )
             model.eval()
             a__								= model(__lowerCamelCase			,attention_mask=__lowerCamelCase			,labels=__lowerCamelCase   )
             self.assertEqual(result.logits.shape			,(self.model_tester.batch_size, self.model_tester.num_labels)   )




      def lowerCAmelCase_					(       self							:     Union[str, Any]   ):
             # 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:
                    a__						,					a__								= 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

                    a__								= model_class(__lowerCamelCase   ).to(__lowerCamelCase   )
                    if "use_cache" not in inputs:
                           a__								= True
                    a__								= 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

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

                    a__								= outputs["past_key_values"]
                    self.assertEqual(len(__lowerCamelCase   )			,__lowerCamelCase   )

                    a__						,					a__								= inputs["input_ids"].shape
                    for i in range(__lowerCamelCase   ):
                           if config.new_decoder_architecture:
                                  a__								= config.num_attention_heads
                           elif config.multi_query:
                                  a__								= 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 lowerCAmelCase_					(       self							:     int   ):
             a__								= AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b"   )
             a__								= FalconForCausalLM.from_pretrained("Rocketknight1/falcon-rw-1b"   )
             model.eval()
             model.to(__lowerCamelCase   )
             a__								= tokenizer("My favorite food is"			,return_tensors="pt"   ).to(__lowerCamelCase   )

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

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

             self.assertEqual(__lowerCamelCase			,__lowerCamelCase   )




      @slow
      def lowerCAmelCase_					(       self							:     Union[str, 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
             for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
                    a__								= AutoTokenizer.from_pretrained(__lowerCamelCase   )
                    a__								= FalconForCausalLM.from_pretrained(__lowerCamelCase   )
                    model.eval()
                    model.to(__lowerCamelCase   )
                    a__								= 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 lowerCAmelCase_					(       self							:     Union[str, 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",
                    ]:
                           a__								= AutoTokenizer.from_pretrained(__lowerCamelCase   )
                           a__								= FalconForCausalLM.from_pretrained(__lowerCamelCase   )
                           model.eval()
                           model.to(device=__lowerCamelCase   )
                           a__								= tokenizer("My favorite food is"			,return_tensors="pt"   ).to(__lowerCamelCase   )

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





"""simple docstring"""






from copy import deepcopy



class     UpperCAmelCase    :



   def __init__(       self    :	Optional[Any]							,       __lowerCamelCase    :	list[int] | None = None							,       __lowerCamelCase    :	int | None = None				):



     """simple docstring"""

     if arr is None and size is not None:
       _snake_case   	=      size
       _snake_case   	=      [0] * size
     elif arr is not None:
       self.init(__lowerCamelCase				)
     else:
       raise ValueError('''Either arr or size must be specified'''				)



   def    __UpperCAmelCase    (       self    :	Any							,       __lowerCamelCase    :	list[int]				):



     """simple docstring"""
     _snake_case   	=      len(__lowerCamelCase				)
     _snake_case   	=      deepcopy(__lowerCamelCase				)
     for i in range(1							,       self.size				):
       _snake_case   	=      self.next_(__lowerCamelCase				)
       if j < self.size:
         self.tree[j] += self.tree[i]



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



     """simple docstring"""
     _snake_case   	=      self.tree[:]
     for i in range(self.size - 1							,       0							,       -1				):
       _snake_case   	=      self.next_(__lowerCamelCase				)
       if j < self.size:
         arr[j] -= arr[i]
     return arr



   @staticmethod
   def    __UpperCAmelCase    (       __lowerCamelCase    :	int				):



     """simple docstring"""
     return index + (index & (-index))



   @staticmethod
   def    __UpperCAmelCase    (       __lowerCamelCase    :	int				):



     """simple docstring"""
     return index - (index & (-index))



   def    __UpperCAmelCase    (       self    :	Dict							,       __lowerCamelCase    :	int							,       __lowerCamelCase    :	int				):



     """simple docstring"""
     if index == 0:
       self.tree[0] += value
       return
     while index < self.size:
       self.tree[index] += value
       _snake_case   	=      self.next_(__lowerCamelCase				)



   def    __UpperCAmelCase    (       self    :	List[str]							,       __lowerCamelCase    :	int							,       __lowerCamelCase    :	int				):



     """simple docstring"""
     self.add(__lowerCamelCase							,       value - self.get(__lowerCamelCase				)				)



   def    __UpperCAmelCase    (       self    :	Optional[Any]							,       __lowerCamelCase    :	int				):



     """simple docstring"""
     if right == 0:
       return 0
     _snake_case   	=      self.tree[0]
     right -= 1  # make right inclusive
     while right > 0:
       result += self.tree[right]
       _snake_case   	=      self.prev(__lowerCamelCase				)
     return result



   def    __UpperCAmelCase    (       self    :	str							,       __lowerCamelCase    :	int							,       __lowerCamelCase    :	int				):



     """simple docstring"""
     return self.prefix(__lowerCamelCase				) - self.prefix(__lowerCamelCase				)



   def    __UpperCAmelCase    (       self    :	List[str]							,       __lowerCamelCase    :	int				):



     """simple docstring"""
     return self.query(__lowerCamelCase							,       index + 1				)





   def    __UpperCAmelCase    (       self    :	Optional[int]							,       __lowerCamelCase    :	int				):



     """simple docstring"""
     value -= self.tree[0]
     if value < 0:
       return -1

     _snake_case   	=      1  # Largest power of 2 <= size
     while j * 2 < self.size:
       j *= 2

     _snake_case   	=      0

     while j > 0:
       if i + j < self.size and self.tree[i + j] <= value:
         value -= self.tree[i + j]
         i += j
       j //= 2
     return i


if __name__ == "__main__":
     import doctest

     doctest.testmod()





