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






    '''simple docstring'''




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




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

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

       lowerCAmelCase_:							Optional[Any]       =	None

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

       lowerCAmelCase_:							Any       =	self.get_config()

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




    def      _a    (    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      _a    (    self       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__ ):
       lowerCAmelCase_:							Any       =	FalconModel(config=UpperCamelCase_ )
       model.to(UpperCamelCase_ )
       model.eval()
       lowerCAmelCase_:							Dict       =	model(UpperCamelCase_       ,						attention_mask=UpperCamelCase_ )
       lowerCAmelCase_:							Tuple       =	model(UpperCamelCase_ )
       self.parent.assertEqual(result.last_hidden_state.shape       ,						(self.batch_size, self.seq_length, self.hidden_size) )




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




    def      _a    (    self       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						):
       lowerCAmelCase_:							Tuple       =	FalconForCausalLM(config=UpperCamelCase_ )
       model.to(UpperCamelCase_ )
       model.eval()
       lowerCAmelCase_:							Any       =	model(UpperCamelCase_       ,						attention_mask=UpperCamelCase_       ,						labels=UpperCamelCase_ )
       self.parent.assertEqual(result.logits.shape       ,						(self.batch_size, self.seq_length, self.vocab_size) )




    def      _a    (    self       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						lowerCamelCase__       ,						):
       lowerCAmelCase_:							str       =	True
       lowerCAmelCase_:							Tuple       =	True
       lowerCAmelCase_:							Any       =	FalconForCausalLM(config=UpperCamelCase_ )
       model.to(UpperCamelCase_ )
       model.eval()

       # first forward pass
       lowerCAmelCase_:							Dict       =	model(
           UpperCamelCase_       ,						attention_mask=UpperCamelCase_       ,						encoder_hidden_states=UpperCamelCase_       ,						encoder_attention_mask=UpperCamelCase_       ,						use_cache=UpperCamelCase_       ,						)
       lowerCAmelCase_:							List[Any]       =	outputs.past_key_values

       # create hypothetical multiple next token and extent to next_input_ids
       lowerCAmelCase_:							Any       =	ids_tensor((self.batch_size, 3)       ,						config.vocab_size )
       lowerCAmelCase_:							Optional[Any]       =	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_:							Any       =	torch.cat([input_mask, next_mask]       ,						dim=-1 )

       lowerCAmelCase_:							Any       =	model(
           UpperCamelCase_       ,						attention_mask=UpperCamelCase_       ,						encoder_hidden_states=UpperCamelCase_       ,						encoder_attention_mask=UpperCamelCase_       ,						output_hidden_states=UpperCamelCase_       ,						)['''hidden_states'''][0]
       lowerCAmelCase_:							Optional[Any]       =	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_:							List[Any]       =	ids_tensor((1,)       ,						output_from_past.shape[-1] ).item()
       lowerCAmelCase_:							Optional[int]       =	output_from_no_past[:, -3:, random_slice_idx].detach()
       lowerCAmelCase_:							Tuple       =	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      _a    (    self ):
       lowerCAmelCase_:							Dict       =	self.prepare_config_and_inputs()
       (
           lowerCAmelCase_							
       ):							Optional[int]       =	config_and_inputs
       lowerCAmelCase_:							Optional[int]       =	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict





@require_torch
class 					_lowercase    (     A__ ,  A__ ,  A__ ,  unittest.TestCase			):






    '''simple docstring'''
    SCREAMING_SNAKE_CASE:				str          =    (
        (
            FalconModel,
            FalconForCausalLM,
            FalconForSequenceClassification,
            FalconForTokenClassification,
            FalconForQuestionAnswering,
        )
        if is_torch_available()
        else ()
    )
    SCREAMING_SNAKE_CASE:				Any          =    (FalconForCausalLM,) if is_torch_available() else ()
    SCREAMING_SNAKE_CASE:				Dict          =    (
        {
            '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:				Any          =    False
    SCREAMING_SNAKE_CASE:				Union[str, Any]          =    False




    def      _a    (    self ):
       lowerCAmelCase_:							Union[str, Any]       =	FalconModelTester(self )
       lowerCAmelCase_:							List[Any]       =	ConfigTester(self       ,						config_class=UpperCamelCase_       ,						hidden_size=37 )




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




    def      _a    (    self ):
       lowerCAmelCase_:							Tuple       =	self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*UpperCamelCase_ )




    def      _a    (    self ):
       lowerCAmelCase_:							Union[str, Any]       =	self.model_tester.prepare_config_and_inputs()
       for alibi in [True, False]:
          lowerCAmelCase_:							Any       =	alibi
          self.model_tester.create_and_check_model(UpperCamelCase_       ,						*UpperCamelCase_ )




    def      _a    (    self ):
       lowerCAmelCase_:							List[str]       =	self.model_tester.prepare_config_and_inputs_for_common()
       lowerCAmelCase_:							str       =	3
       lowerCAmelCase_:							Any       =	input_dict['''input_ids''']
       lowerCAmelCase_:							Tuple       =	input_ids.ne(1 ).to(UpperCamelCase_ )
       lowerCAmelCase_:							Optional[Any]       =	ids_tensor([self.model_tester.batch_size]       ,						self.model_tester.type_sequence_label_size )
       lowerCAmelCase_:							List[Any]       =	FalconForSequenceClassification(UpperCamelCase_ )
       model.to(UpperCamelCase_ )
       model.eval()
       lowerCAmelCase_:							Dict       =	model(UpperCamelCase_       ,						attention_mask=UpperCamelCase_       ,						labels=UpperCamelCase_ )
       self.assertEqual(result.logits.shape       ,						(self.model_tester.batch_size, self.model_tester.num_labels) )




    def      _a    (    self ):
       lowerCAmelCase_:							List[Any]       =	self.model_tester.prepare_config_and_inputs_for_common()
       lowerCAmelCase_:							List[str]       =	3
       lowerCAmelCase_:							str       =	'''single_label_classification'''
       lowerCAmelCase_:							Union[str, Any]       =	input_dict['''input_ids''']
       lowerCAmelCase_:							List[str]       =	input_ids.ne(1 ).to(UpperCamelCase_ )
       lowerCAmelCase_:							Dict       =	ids_tensor([self.model_tester.batch_size]       ,						self.model_tester.type_sequence_label_size )
       lowerCAmelCase_:							Dict       =	FalconForSequenceClassification(UpperCamelCase_ )
       model.to(UpperCamelCase_ )
       model.eval()
       lowerCAmelCase_:							str       =	model(UpperCamelCase_       ,						attention_mask=UpperCamelCase_       ,						labels=UpperCamelCase_ )
       self.assertEqual(result.logits.shape       ,						(self.model_tester.batch_size, self.model_tester.num_labels) )




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




    def      _a    (    self ):


       for model_class in self.all_generative_model_classes:
          lowerCAmelCase_:							List[Any]       =	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_:							Tuple       =	model_class(UpperCamelCase_ ).to(UpperCamelCase_ )
          if "use_cache" not in inputs:
             lowerCAmelCase_:							Tuple       =	True
          lowerCAmelCase_:							List[str]       =	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_:							str       =	(
              getattr(UpperCamelCase_       ,						"decoder_layers"       ,						UpperCamelCase_ )
              or getattr(UpperCamelCase_       ,						"num_decoder_layers"       ,						UpperCamelCase_ )
              or config.num_hidden_layers
          )
          lowerCAmelCase_:							List[Any]       =	getattr(UpperCamelCase_       ,						"num_kv_heads"       ,						config.num_attention_heads )
          lowerCAmelCase_:							str       =	getattr(UpperCamelCase_       ,						"d_model"       ,						config.hidden_size )
          lowerCAmelCase_:							Optional[Any]       =	embed_dim // num_attention_heads

          lowerCAmelCase_:							Optional[Any]       =	outputs['''past_key_values''']
          self.assertEqual(len(UpperCamelCase_ )       ,						UpperCamelCase_ )

          lowerCAmelCase_:							Tuple       =	inputs['''input_ids'''].shape
          for i in range(UpperCamelCase_ ):
             if config.new_decoder_architecture:
                lowerCAmelCase_:							List[Any]       =	config.num_attention_heads
             elif config.multi_query:
                lowerCAmelCase_:							str       =	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 					_lowercase    (     unittest.TestCase			):






    '''simple docstring'''




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

       lowerCAmelCase_:							Optional[Any]       =	(
           '''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(**UpperCamelCase_       ,						do_sample=UpperCamelCase_       ,						max_new_tokens=19 )
       lowerCAmelCase_:							str       =	tokenizer.batch_decode(UpperCamelCase_ )[0]

       self.assertEqual(UpperCamelCase_       ,						UpperCamelCase_ )




    @slow
    def      _a    (    self ):


       for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
          lowerCAmelCase_:							Any       =	AutoTokenizer.from_pretrained(UpperCamelCase_ )
          lowerCAmelCase_:							Dict       =	FalconForCausalLM.from_pretrained(UpperCamelCase_ )
          model.eval()
          model.to(UpperCamelCase_ )
          lowerCAmelCase_:							str       =	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      _a    (    self ):


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

             # Test results are the same with and without cache
             lowerCAmelCase_:							Union[str, Any]       =	model.generate(**UpperCamelCase_       ,						do_sample=UpperCamelCase_       ,						max_new_tokens=20       ,						use_cache=UpperCamelCase_ )
             lowerCAmelCase_:							Any       =	model.generate(**UpperCamelCase_       ,						do_sample=UpperCamelCase_       ,						max_new_tokens=20       ,						use_cache=UpperCamelCase_ )
             self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0 )



'''simple docstring'''
from __future__ import annotations

from collections import namedtuple
from dataclasses import dataclass
@dataclass
class   lowercase      :

		"""simple docstring"""

		_a       					=							42
		_a       					=							None
		_a       					=							None



__snake_case				 =	namedtuple('''CoinsDistribResult''', '''moves excess''')


def      a (     __a  )    ->      int:



					'''simple docstring'''






					if root is None:
										return 0

					# Validation
					def count_nodes(__a  ) -> int:
										if node is None:
															return 0

										return count_nodes(node.left  ) + count_nodes(node.right  ) + 1

					def count_coins(__a  ) -> int:
										if node is None:
															return 0

										return count_coins(node.left  ) + count_coins(node.right  ) + node.data

					if count_nodes(__a  ) != count_coins(__a  ):
										raise ValueError('''The nodes number should be same as the number of coins'''  )

					# Main calculation
					def get_distrib(__a  ) -> CoinsDistribResult:

										if node is None:
															return CoinsDistribResult(0					,   1  )

										UpperCamelCase__   ,							UpperCamelCase__  :Optional[Any]				   =					get_distrib(node.left  )
										UpperCamelCase__   ,							UpperCamelCase__  :Union[str, Any]				   =					get_distrib(node.right  )

										UpperCamelCase__  :Tuple				   =					1 - left_distrib_excess
										UpperCamelCase__  :Optional[int]				   =					1 - right_distrib_excess

										UpperCamelCase__  :int				   =					(
										    left_distrib_moves
										    + right_distrib_moves
										    + abs(__a  )
										    + abs(__a  )
										)
										UpperCamelCase__  :Optional[Any]				   =					node.data - coins_to_left - coins_to_right

										return CoinsDistribResult(__a					,   __a  )

					return get_distrib(__a  )[0]


if __name__ == "__main__":
		import doctest

		doctest.testmod()