

import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         FalconForCausalLM,
         FalconForQuestionAnswering,
         FalconForSequenceClassification,
         FalconForTokenClassification,
         FalconModel,
     )


class 			A						:







    """simple docstring"""





    def __init__(				self			:	Optional[Any],lowercase_			:	List[str],lowercase_			:	List[str]=3,lowercase_			:	Union[str, Any]=7,lowercase_			:	int=True,lowercase_			:	int=True,lowercase_			:	Tuple=False,lowercase_			:	Tuple=True,lowercase_			:	Any=9_9,lowercase_			:	Any=3_2,lowercase_			:	Tuple=5,lowercase_			:	int=4,lowercase_			:	Tuple=3_7,lowercase_			:	Any="gelu",lowercase_			:	Any=0.1,lowercase_			:	Tuple=0.1,lowercase_			:	List[str]=5_1_2,lowercase_			:	Tuple=1_6,lowercase_			:	Dict=2,lowercase_			:	Optional[Any]=0.02,lowercase_			:	Tuple=3,lowercase_			:	Union[str, Any]=4,lowercase_			:	Union[str, Any]=None,)->   List[str]:





        '''simple docstring'''







        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 						snake_case__					(				self			:	Union[str, Any]		)->   int:





        '''simple docstring'''







        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 						snake_case__					(				self			:	Tuple		)->   Union[str, Any]:





        '''simple docstring'''







        return FalconConfig(
            vocab_size=self.vocab_size,hidden_size=self.hidden_size,num_hidden_layers=self.num_hidden_layers,num_attention_heads=self.num_attention_heads,intermediate_size=self.intermediate_size,hidden_act=self.hidden_act,hidden_dropout_prob=self.hidden_dropout_prob,attention_probs_dropout_prob=self.attention_probs_dropout_prob,max_position_embeddings=self.max_position_embeddings,type_vocab_size=self.type_vocab_size,is_decoder=lowercase_,initializer_range=self.initializer_range,pad_token_id=1,new_decoder_architecture=lowercase_,)


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





        '''simple docstring'''







        A__   				=       FalconModel(config=lowercase_		)
        model.to(lowercase_		)
        model.eval()
        A__   				=       model(lowercase_,attention_mask=lowercase_		)
        A__   				=       model(lowercase_		)
        self.parent.assertEqual(result.last_hidden_state.shape,(self.batch_size, self.seq_length, self.hidden_size)		)


    def 						snake_case__					(				self			:	List[str],lowercase_			:	Union[str, Any],lowercase_			:	List[str],lowercase_			:	Any,lowercase_			:	Optional[Any],lowercase_			:	str,lowercase_			:	Union[str, Any],lowercase_			:	Any,lowercase_			:	List[Any],lowercase_			:	Optional[Any],)->   str:





        '''simple docstring'''







        A__   				=       True
        A__   				=       FalconModel(lowercase_		)
        model.to(lowercase_		)
        model.eval()
        A__   				=       model(
            lowercase_,attention_mask=lowercase_,encoder_hidden_states=lowercase_,encoder_attention_mask=lowercase_,)
        A__   				=       model(
            lowercase_,attention_mask=lowercase_,encoder_hidden_states=lowercase_,)
        A__   				=       model(lowercase_,attention_mask=lowercase_		)
        self.parent.assertEqual(result.last_hidden_state.shape,(self.batch_size, self.seq_length, self.hidden_size)		)


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





        '''simple docstring'''







        A__   				=       FalconForCausalLM(config=lowercase_		)
        model.to(lowercase_		)
        model.eval()
        A__   				=       model(lowercase_,attention_mask=lowercase_,labels=lowercase_		)
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.seq_length, self.vocab_size)		)


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





        '''simple docstring'''







        A__   				=       True
        A__   				=       True
        A__   				=       FalconForCausalLM(config=lowercase_		)
        model.to(lowercase_		)
        model.eval()

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


    def 						snake_case__					(				self			:	str		)->   Tuple:





        '''simple docstring'''







        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 			A						(							_UpperCAmelCase   ,							_UpperCAmelCase   ,							_UpperCAmelCase   ,							unittest.TestCase			):







    """simple docstring"""





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


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





        '''simple docstring'''







        A__   				=       FalconModelTester(self		)
        A__   				=       ConfigTester(self,config_class=lowercase_,hidden_size=3_7		)


    def 						snake_case__					(				self			:	Any		)->   Dict:





        '''simple docstring'''







        self.config_tester.run_common_tests()


    def 						snake_case__					(				self			:	Any		)->   Any:





        '''simple docstring'''







        A__   				=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase_		)


    def 						snake_case__					(				self			:	Tuple		)->   List[str]:





        '''simple docstring'''







        A__				,						*A__   				=       self.model_tester.prepare_config_and_inputs()
        for alibi in [True, False]:
            A__   				=       alibi
            self.model_tester.create_and_check_model(lowercase_,*lowercase_		)


    def 						snake_case__					(				self			:	str		)->   Optional[int]:





        '''simple docstring'''







        A__				,						A__   				=       self.model_tester.prepare_config_and_inputs_for_common()
        A__   				=       3
        A__   				=       input_dict['input_ids']
        A__   				=       input_ids.ne(1		).to(lowercase_		)
        A__   				=       ids_tensor([self.model_tester.batch_size],self.model_tester.type_sequence_label_size		)
        A__   				=       FalconForSequenceClassification(lowercase_		)
        model.to(lowercase_		)
        model.eval()
        A__   				=       model(lowercase_,attention_mask=lowercase_,labels=lowercase_		)
        self.assertEqual(result.logits.shape,(self.model_tester.batch_size, self.model_tester.num_labels)		)


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





        '''simple docstring'''







        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(lowercase_		)
        A__   				=       ids_tensor([self.model_tester.batch_size],self.model_tester.type_sequence_label_size		)
        A__   				=       FalconForSequenceClassification(lowercase_		)
        model.to(lowercase_		)
        model.eval()
        A__   				=       model(lowercase_,attention_mask=lowercase_,labels=lowercase_		)
        self.assertEqual(result.logits.shape,(self.model_tester.batch_size, self.model_tester.num_labels)		)


    def 						snake_case__					(				self			:	int		)->   str:





        '''simple docstring'''







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





        '''simple docstring'''







        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(lowercase_		)
        A__   				=       ids_tensor(
            [self.model_tester.batch_size, config.num_labels],self.model_tester.type_sequence_label_size		).to(torch.float		)
        A__   				=       FalconForSequenceClassification(lowercase_		)
        model.to(lowercase_		)
        model.eval()
        A__   				=       model(lowercase_,attention_mask=lowercase_,labels=lowercase_		)
        self.assertEqual(result.logits.shape,(self.model_tester.batch_size, self.model_tester.num_labels)		)


    def 						snake_case__					(				self			:	str		)->   Tuple:





        '''simple docstring'''







        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(lowercase_,'use_cache'		):
                return

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

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

            A__   				=       outputs['past_key_values']
            self.assertEqual(len(lowercase_		),lowercase_		)

            A__				,						A__   				=       inputs['input_ids'].shape
            for i in range(lowercase_		):
                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 			A						(							unittest.TestCase			):







    """simple docstring"""





    @slow
    def 						snake_case__					(				self			:	Optional[Any]		)->   Optional[Any]:





        '''simple docstring'''







        A__   				=       AutoTokenizer.from_pretrained('Rocketknight1/falcon-rw-1b'		)
        A__   				=       FalconForCausalLM.from_pretrained('Rocketknight1/falcon-rw-1b'		)
        model.eval()
        model.to(lowercase_		)
        A__   				=       tokenizer('My favorite food is',return_tensors='pt'		).to(lowercase_		)

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

        self.assertEqual(lowercase_,lowercase_		)


    @slow
    def 						snake_case__					(				self			:	Optional[int]		)->   Optional[int]:





        '''simple docstring'''







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

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


    @slow
    def 						snake_case__					(				self			:	Tuple		)->   List[str]:





        '''simple docstring'''







        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(lowercase_		)
                A__   				=       FalconForCausalLM.from_pretrained(lowercase_		)
                model.eval()
                model.to(device=lowercase_		)
                A__   				=       tokenizer('My favorite food is',return_tensors='pt'		).to(lowercase_		)

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




import numpy as np

from transformers import Pipeline
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     int				)							->				int:





    '''simple docstring'''






    A__   				=       np.max(SCREAMING_SNAKE_CASE__   ,      axis=-1   ,      keepdims=SCREAMING_SNAKE_CASE__				)
    A__   				=       np.exp(outputs - maxes				)
    return shifted_exp / shifted_exp.sum(axis=-1   ,      keepdims=SCREAMING_SNAKE_CASE__				)


class 			A						(							_UpperCAmelCase			):







    """simple docstring"""





    def 						snake_case__					(				self			:	Dict,**lowercase_			:	Tuple		)->   Tuple:





        '''simple docstring'''







        A__   				=       {}
        if "second_text" in kwargs:
            A__   				=       kwargs['second_text']
        return preprocess_kwargs, {}, {}


    def 						snake_case__					(				self			:	List[Any],lowercase_			:	int,lowercase_			:	Optional[int]=None		)->   List[str]:





        '''simple docstring'''







        return self.tokenizer(lowercase_,text_pair=lowercase_,return_tensors=self.framework		)


    def 						snake_case__					(				self			:	str,lowercase_			:	Dict		)->   List[str]:





        '''simple docstring'''







        return self.model(**lowercase_		)


    def 						snake_case__					(				self			:	Dict,lowercase_			:	Optional[int]		)->   Dict:





        '''simple docstring'''







        A__   				=       model_outputs.logits[0].numpy()
        A__   				=       softmax(lowercase_		)

        A__   				=       np.argmax(lowercase_		)
        A__   				=       self.model.config.idalabel[best_class]
        A__   				=       probabilities[best_class].item()
        A__   				=       logits.tolist()
        return {"label": label, "score": score, "logits": logits}

