


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     ,   __snake_case     ,   __snake_case=3     ,   __snake_case=7     ,   __snake_case=True     ,   __snake_case=True     ,   __snake_case=False     ,   __snake_case=True     ,   __snake_case=9_9     ,   __snake_case=3_2     ,   __snake_case=5     ,   __snake_case=4     ,   __snake_case=3_7     ,   __snake_case="gelu"     ,   __snake_case=0.1     ,   __snake_case=0.1     ,   __snake_case=5_1_2     ,   __snake_case=1_6     ,   __snake_case=2     ,   __snake_case=0.02     ,   __snake_case=3     ,   __snake_case=4     ,   __snake_case=None     ,   ):
         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 						a_     (			self       ):
         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 						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=__snake_case     ,   initializer_range=self.initializer_range     ,   pad_token_id=1     ,   new_decoder_architecture=__snake_case     ,   )
    def 						a_     (			self     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case       ):
         snake_case											=							FalconModel(config=__snake_case       )
         model.to(__snake_case       )
         model.eval()
         snake_case											=							model(__snake_case     ,   attention_mask=__snake_case       )
         snake_case											=							model(__snake_case       )
         self.parent.assertEqual(result.last_hidden_state.shape     ,   (self.batch_size, self.seq_length, self.hidden_size)       )
    def 						a_     (			self     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   ):
         snake_case											=							True
         snake_case											=							FalconModel(__snake_case       )
         model.to(__snake_case       )
         model.eval()
         snake_case											=							model(
             __snake_case     ,   attention_mask=__snake_case     ,   encoder_hidden_states=__snake_case     ,   encoder_attention_mask=__snake_case     ,   )
         snake_case											=							model(
             __snake_case     ,   attention_mask=__snake_case     ,   encoder_hidden_states=__snake_case     ,   )
         snake_case											=							model(__snake_case     ,   attention_mask=__snake_case       )
         self.parent.assertEqual(result.last_hidden_state.shape     ,   (self.batch_size, self.seq_length, self.hidden_size)       )
    def 						a_     (			self     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   ):
         snake_case											=							FalconForCausalLM(config=__snake_case       )
         model.to(__snake_case       )
         model.eval()
         snake_case											=							model(__snake_case     ,   attention_mask=__snake_case     ,   labels=__snake_case       )
         self.parent.assertEqual(result.logits.shape     ,   (self.batch_size, self.seq_length, self.vocab_size)       )
    def 						a_     (			self     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   __snake_case     ,   ):
         snake_case											=							True
         snake_case											=							True
         snake_case											=							FalconForCausalLM(config=__snake_case       )
         model.to(__snake_case       )
         model.eval()

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




    def 						a_     (			self       ):
         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     A__   (       snake_case__			,   snake_case__			,   snake_case__			,   unittest.TestCase      ):







    """simple docstring"""

    __magic_name__         =						(
        (
            FalconModel,
            FalconForCausalLM,
            FalconForSequenceClassification,
            FalconForTokenClassification,
            FalconForQuestionAnswering,
        )
        if is_torch_available()
        else ()
    )
    __magic_name__         =						(FalconForCausalLM,) if is_torch_available() else ()
    __magic_name__         =						(
        {
            'feature-extraction': FalconModel,
            'text-classification': FalconForSequenceClassification,
            'text-generation': FalconForCausalLM,
            'question-answering': FalconForQuestionAnswering,
            'token-classification': FalconForTokenClassification,
            'zero-shot': FalconForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    __magic_name__         =						False
    __magic_name__         =						False
    def 						a_     (			self       ):
         snake_case											=							FalconModelTester(self       )
         snake_case											=							ConfigTester(self     ,   config_class=__snake_case     ,   hidden_size=3_7       )
    def 						a_     (			self       ):
         self.config_tester.run_common_tests()
    def 						a_     (			self       ):
         snake_case											=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__snake_case       )
    def 						a_     (			self       ):
         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(__snake_case     ,   *__snake_case       )
    def 						a_     (			self       ):
         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(__snake_case       )
         snake_case											=							ids_tensor([self.model_tester.batch_size]     ,   self.model_tester.type_sequence_label_size       )
         snake_case											=							FalconForSequenceClassification(__snake_case       )
         model.to(__snake_case       )
         model.eval()
         snake_case											=							model(__snake_case     ,   attention_mask=__snake_case     ,   labels=__snake_case       )
         self.assertEqual(result.logits.shape     ,   (self.model_tester.batch_size, self.model_tester.num_labels)       )
    def 						a_     (			self       ):
         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(__snake_case       )
         snake_case											=							ids_tensor([self.model_tester.batch_size]     ,   self.model_tester.type_sequence_label_size       )
         snake_case											=							FalconForSequenceClassification(__snake_case       )
         model.to(__snake_case       )
         model.eval()
         snake_case											=							model(__snake_case     ,   attention_mask=__snake_case     ,   labels=__snake_case       )
         self.assertEqual(result.logits.shape     ,   (self.model_tester.batch_size, self.model_tester.num_labels)       )
    def 						a_     (			self       ):
         snake_case					,    snake_case											=							self.model_tester.prepare_config_and_inputs_for_common()
         snake_case											=							input_dict['''input_ids''']
         snake_case											=							FalconForCausalLM(__snake_case       )
         model.to(__snake_case       )
         model.eval()
         snake_case											=							model(__snake_case     ,   use_cache=__snake_case       )
         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(__snake_case     ,   __snake_case       )
         for layer in range(len(__snake_case       )       ):
              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       ):
         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(__snake_case       )
         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(__snake_case       )
         model.to(__snake_case       )
         model.eval()
         snake_case											=							model(__snake_case     ,   attention_mask=__snake_case     ,   labels=__snake_case       )
         self.assertEqual(result.logits.shape     ,   (self.model_tester.batch_size, self.model_tester.num_labels)       )




    def 						a_     (			self       ):
         # Falcon can have different numbers of KV-heads than the number of query heads, so we need
         # to override this test to use the right head counts.
         for model_class in self.all_generative_model_classes:
              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(__snake_case     ,   '''use_cache'''       ):
                   return

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

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

              snake_case											=							outputs['''past_key_values''']
              self.assertEqual(len(__snake_case       )     ,   __snake_case       )

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







    """simple docstring"""

    @slow
    def 						a_     (			self       ):
         snake_case											=							AutoTokenizer.from_pretrained('''Rocketknight1/falcon-rw-1b'''       )
         snake_case											=							FalconForCausalLM.from_pretrained('''Rocketknight1/falcon-rw-1b'''       )
         model.eval()
         model.to(__snake_case       )
         snake_case											=							tokenizer('''My favorite food is'''     ,   return_tensors='''pt'''       ).to(__snake_case       )

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

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

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




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

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








from ...utils import is_torch_available, is_transformers_available


if is_transformers_available() and is_torch_available():
 from .pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings, VQDiffusionPipeline




