codingutf8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch falcon model import unittest from parameterized import parameterized from transformers import automodelforcausallm autotokenizer falconconfig istorchavailable setseed from transformers testingutils import requirebitsandbytes requiretorch slow torchdevice from generation testutils import generationtestermixin from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin idstensor randomattentionmask from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import falconforcausallm falconforquestionanswering falconforsequenceclassification falconfortokenclassification falconmodel class falconmodeltester def init self parent batchsize3 seqlength7 istrainingtrue useinputmasktrue usetokentypeidsfalse uselabelstrue vocabsize99 hiddensize32 numhiddenlayers2 numattentionheads4 intermediatesize37 hiddenactgelu hiddendropoutprob0 1 attentionprobsdropoutprob0 1 maxpositionembeddings512 typevocabsize16 typesequencelabelsize2 initializerrange0 02 numlabels3 numchoices4 scopenone self parent parent self batchsize batchsize self seqlength seqlength self istraining istraining self useinputmask useinputmask self usetokentypeids usetokentypeids self uselabels uselabels self vocabsize vocabsize self hiddensize hiddensize self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads self intermediatesize intermediatesize self hiddenact hiddenact self hiddendropoutprob hiddendropoutprob self attentionprobsdropoutprob attentionprobsdropoutprob self maxpositionembeddings maxpositionembeddings self typevocabsize typevocabsize self typesequencelabelsize typesequencelabelsize self initializerrange initializerrange self numlabels numlabels self numchoices numchoices self scope scope def prepareconfigandinputsself inputids idstensorself batchsize self seqlength self vocabsize inputmask none if self useinputmask inputmask randomattentionmaskself batchsize self seqlength tokentypeids none sequencelabels none tokenlabels none choicelabels none if self uselabels sequencelabels idstensorself batchsize self typesequencelabelsize tokenlabels idstensorself batchsize self seqlength self numlabels choicelabels idstensorself batchsize self numchoices config self getconfig return config inputids tokentypeids inputmask sequencelabels tokenlabels choicelabels def getconfigself return falconconfig vocabsizeself vocabsize hiddensizeself hiddensize numhiddenlayersself numhiddenlayers numattentionheadsself numattentionheads intermediatesizeself intermediatesize hiddenactself hiddenact hiddendropoutprobself hiddendropoutprob attentionprobsdropoutprobself attentionprobsdropoutprob maxpositionembeddingsself maxpositionembeddings typevocabsizeself typevocabsize isdecoderfalse initializerrangeself initializerrange padtokenid1 newdecoderarchitecturetrue def createandcheckmodel self config inputids tokentypeids inputmask sequencelabels tokenlabels choicelabels model falconmodelconfigconfig model totorchdevice model eval result modelinputids attentionmaskinputmask result modelinputids self parent assertequalresult lasthiddenstate shape self batchsize self seqlength self hiddensize def createandcheckmodelasdecoder self config inputids tokentypeids inputmask sequencelabels tokenlabels choicelabels encoderhiddenstates encoderattentionmask config addcrossattention true model falconmodelconfig model totorchdevice model eval result model inputids attentionmaskinputmask encoderhiddenstatesencoderhiddenstates encoderattentionmaskencoderattentionmask result model inputids attentionmaskinputmask encoderhiddenstatesencoderhiddenstates result modelinputids attentionmaskinputmask self parent assertequalresult lasthiddenstate shape self batchsize self seqlength self hiddensize def createandcheckforcausallm self config inputids tokentypeids inputmask sequencelabels tokenlabels choicelabels encoderhiddenstates encoderattentionmask model falconforcausallmconfigconfig model totorchdevice model eval result modelinputids attentionmaskinputmask labelstokenlabels self parent assertequalresult logits shape self batchsize self seqlength self vocabsize def createandcheckdecodermodelpastlargeinputs self config inputids tokentypeids inputmask sequencelabels tokenlabels choicelabels encoderhiddenstates encoderattentionmask config isdecoder true config addcrossattention true model falconforcausallmconfigconfig model totorchdevice model eval first forward pass outputs model inputids attentionmaskinputmask encoderhiddenstatesencoderhiddenstates encoderattentionmaskencoderattentionmask usecachetrue pastkeyvalues outputs pastkeyvalues create hypothetical multiple next token and extent to nextinputids nexttokens idstensorself batchsize 3 config vocabsize nextmask idstensorself batchsize 3 vocabsize2 append to next inputids and nextinputids torch catinputids nexttokens dim1 nextattentionmask torch catinputmask nextmask dim1 outputfromnopast model nextinputids attentionmasknextattentionmask encoderhiddenstatesencoderhiddenstates encoderattentionmaskencoderattentionmask outputhiddenstatestrue hiddenstates0 outputfrompast model nexttokens attentionmasknextattentionmask encoderhiddenstatesencoderhiddenstates encoderattentionmaskencoderattentionmask pastkeyvaluespastkeyvalues outputhiddenstatestrue hiddenstates0 select random slice randomsliceidx idstensor1 outputfrompast shape1 item outputfromnopastslice outputfromnopast 3 randomsliceidx detach outputfrompastslice outputfrompast randomsliceidx detach self parent asserttrueoutputfrompastslice shape1 nexttokens shape1 test that outputs are equal for slice self parent asserttruetorch allcloseoutputfrompastslice outputfromnopastslice atol1e3 def prepareconfigandinputsforcommonself configandinputs self prepareconfigandinputs config inputids tokentypeids inputmask sequencelabels tokenlabels choicelabels configandinputs inputsdict inputids inputids attentionmask inputmask return config inputsdict requiretorch class falconmodeltestmodeltestermixin generationtestermixin pipelinetestermixin unittest testcase allmodelclasses falconmodel falconforcausallm falconforsequenceclassification falconfortokenclassification falconforquestionanswering if istorchavailable else allgenerativemodelclasses falconforcausallm if istorchavailable else pipelinemodelmapping featureextraction falconmodel questionanswering falconforquestionanswering textclassification falconforsequenceclassification textgeneration falconforcausallm tokenclassification falconfortokenclassification zeroshot falconforsequenceclassification if istorchavailable else testheadmasking false testpruning false todo ydshieh check this see https app circleci compipelinesgithubhuggingfacetransformers79245workflows9490ef5879c2410d8f51e3495156cf9cjobs1012146 def ispipelinetesttoskip self pipelinetestcassename configclass modelarchitecture tokenizername processorname return true def setupself self modeltester falconmodeltesterself self configtester configtesterself configclassfalconconfig hiddensize37 def testconfigself self configtester runcommontests def testmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfigandinputs def testpositionembeddingtypesself config inputs self modeltester prepareconfigandinputs for alibi in true false config alibi alibi self modeltester createandcheckmodelconfig inputs def testfalconsequenceclassificationmodelself config inputdict self modeltester prepareconfigandinputsforcommon config numlabels 3 inputids inputdictinputids attentionmask inputids ne1 totorchdevice sequencelabels idstensorself modeltester batchsize self modeltester typesequencelabelsize model falconforsequenceclassificationconfig model totorchdevice model eval result modelinputids attentionmaskattentionmask labelssequencelabels self assertequalresult logits shape self modeltester batchsize self modeltester numlabels def testfalconsequenceclassificationmodelforsinglelabelself config inputdict self modeltester prepareconfigandinputsforcommon config numlabels 3 config problemtype singlelabelclassification inputids inputdictinputids attentionmask inputids ne1 totorchdevice sequencelabels idstensorself modeltester batchsize self modeltester typesequencelabelsize model falconforsequenceclassificationconfig model totorchdevice model eval result modelinputids attentionmaskattentionmask labelssequencelabels self assertequalresult logits shape self modeltester batchsize self modeltester numlabels def testfalconsequenceclassificationmodelformultilabelself config inputdict self modeltester prepareconfigandinputsforcommon config numlabels 3 config problemtype multilabelclassification inputids inputdictinputids attentionmask inputids ne1 totorchdevice sequencelabels idstensor self modeltester batchsize config numlabels self modeltester typesequencelabelsize totorch float model falconforsequenceclassificationconfig model totorchdevice model eval result modelinputids attentionmaskattentionmask labelssequencelabels self assertequalresult logits shape self modeltester batchsize self modeltester numlabels def testpastkeyvaluesformatself falcon can have different numbers of kvheads than the number of query heads so we need to override this test to use the right head counts for modelclass in self allgenerativemodelclasses config inputs self modeltester prepareconfigandinputsforcommon if it doesn t support cache pass the test if not hasattrconfig usecache return model modelclassconfig totorchdevice if usecache not in inputs inputsusecache true outputs modelinputs if pastkeyvalues is not returned pass the test e g rwkv uses a different cache name and format if pastkeyvalues not in outputs return numhiddenlayers getattrconfig decoderlayers none or getattrconfig numdecoderlayers none or config numhiddenlayers numattentionheads getattrconfig numkvheads config numattentionheads embeddim getattrconfig dmodel config hiddensize perheadembeddim embeddim numattentionheads pastkv outputspastkeyvalues self assertequallenpastkv numhiddenlayers batchsize seqlength inputsinputids shape for i in rangenumhiddenlayers if config newdecoderarchitecture numattentionheads config numattentionheads elif config multiquery numattentionheads 1 self assertequallenpastkv0 2 k v for the decoder 2 self assertequal pastkvi0 shape batchsize numattentionheads seqlength perheadembeddim self assertequal pastkvi1 shape batchsize numattentionheads seqlength perheadembeddim parameterized expandlinear dynamic def testmodelropescalingself scalingtype config self modeltester prepareconfigandinputsforcommon shortinput idstensor1 10 config vocabsize longinput idstensor1 intconfig maxpositionembeddings 1 5 config vocabsize setseed42 fixed seed at init time so the two models get the same random weights originalmodel falconmodelconfig originalmodel totorchdevice originalmodel eval originalshortoutput originalmodelshortinput lasthiddenstate originallongoutput originalmodellonginput lasthiddenstate setseed42 fixed seed at init time so the two models get the same random weights config ropescaling type scalingtype factor 10 0 scaledmodel falconmodelconfig scaledmodel totorchdevice scaledmodel eval scaledshortoutput scaledmodelshortinput lasthiddenstate scaledlongoutput scaledmodellonginput lasthiddenstate dynamic scaling does not change the rope embeddings until it receives an input longer than the original maximum sequence length so the outputs for the short input should match if scalingtype dynamic self asserttruetorch allcloseoriginalshortoutput scaledshortoutput atol1e5 else self assertfalsetorch allcloseoriginalshortoutput scaledshortoutput atol1e5 the output should be different for long inputs self assertfalsetorch allcloseoriginallongoutput scaledlongoutput atol1e5 requiretorch class falconlanguagegenerationtestunittest testcase slow def testlmgeneratefalconself tokenizer autotokenizer frompretrainedrocketknight1falconrw1b model falconforcausallm frompretrainedrocketknight1falconrw1b model eval model totorchdevice inputs tokenizermy favorite food is returntensorspt totorchdevice expectedoutput my favorite food is pizza i love it so much that i have a pizza party every year for my birthday outputids model generateinputs dosamplefalse maxnewtokens19 outputstr tokenizer batchdecodeoutputids0 self assertequaloutputstr expectedoutput slow def testlmgenerationbigmodelsself 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 rocketknight1tinyrandomfalcon7b rocketknight1tinyrandomfalcon40b tokenizer autotokenizer frompretrainedrepo model falconforcausallm frompretrainedrepo model eval model totorchdevice inputs tokenizermy favorite food is returntensorspt totorchdevice we just test that these run without errors the models are randomly initialized and so the actual text outputs will be garbage model generateinputs dosamplefalse maxnewtokens4 model generateinputs dosampletrue maxnewtokens4 model generateinputs numbeams2 maxnewtokens4 slow def testlmgenerationusecacheself 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 nograd for repo in rocketknight1falconrw1b rocketknight1tinyrandomfalcon7b rocketknight1tinyrandomfalcon40b tokenizer autotokenizer frompretrainedrepo model falconforcausallm frompretrainedrepo model eval model todevicetorchdevice inputs tokenizermy favorite food is returntensorspt totorchdevice test results are the same with and without cache outputsnocache model generateinputs dosamplefalse maxnewtokens20 usecachefalse outputscache model generateinputs dosamplefalse maxnewtokens20 usecachetrue self asserttrueoutputscache outputsnocache sum item 0 requirebitsandbytes slow def testbatchedgenerationself tokenizer autotokenizer frompretrainedtiiuaefalcon7b paddingsideleft tokenizer padtoken tokenizer eostoken model automodelforcausallm frompretrained tiiuaefalcon7b devicemapauto loadin4bittrue testtext a sequence 1 2 should generate the rest of the sequence unpaddedinputs tokenizertesttext returntensorspt tocuda 0 unpaddedgenout model generateunpaddedinputs maxnewtokens20 unpaddedgentext tokenizer batchdecodeunpaddedgenout skipspecialtokenstrue dummytext this is a longer text 2 forces leftpadding on testtext paddedinputs tokenizertesttext dummytext returntensorspt paddingtrue tocuda 0 paddedgenout model generatepaddedinputs maxnewtokens20 paddedgentext tokenizer batchdecodepaddedgenout skipspecialtokenstrue expectedoutput a sequence 1 2 3 4 5 6 7 8 self assertlessunpaddedinputs inputids shape1 paddedinputs inputids shape1 leftpadding exists self assertequalunpaddedgentext0 expectedoutput self assertequalpaddedgentext0 expectedoutput coding utf 8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch falcon model first forward pass create hypothetical multiple next token and extent to next_input_ids append to next input_ids and select random slice test that outputs are equal for slice todo ydshieh check this see https app circleci com pipelines github huggingface transformers 79245 workflows 9490ef58 79c2 410d 8f51 e3495156cf9c jobs 1012146 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 if it doesn t support cache pass the test if past_key_values is not returned pass the test e g rwkv uses a different cache name and format k v for the decoder 2 fixed seed at init time so the two models get the same random weights fixed seed at init time so the two models get the same random weights dynamic scaling does not change the rope embeddings until it receives an input longer than the original maximum sequence length so the outputs for the short input should match the output should be different for long inputs 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 we just test that these run without errors the models are randomly initialized and so the actual text outputs will be garbage 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 test results are the same with and without cache should generate the rest of the sequence forces left padding on test_text left padding exists
import unittest

from parameterized import parameterized

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    FalconConfig,
    is_torch_available,
    set_seed,
)
from transformers.testing_utils import require_bitsandbytes, 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 FalconModelTester:
    def __init__(
        self,
        parent,
        batch_size=3,
        seq_length=7,
        is_training=True,
        use_input_mask=True,
        use_token_type_ids=False,
        use_labels=True,
        vocab_size=99,
        hidden_size=32,
        num_hidden_layers=2,
        num_attention_heads=4,
        intermediate_size=37,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        max_position_embeddings=512,
        type_vocab_size=16,
        type_sequence_label_size=2,
        initializer_range=0.02,
        num_labels=3,
        num_choices=4,
        scope=None,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        self.use_input_mask = use_input_mask
        self.use_token_type_ids = use_token_type_ids
        self.use_labels = use_labels
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.type_vocab_size = type_vocab_size
        self.type_sequence_label_size = type_sequence_label_size
        self.initializer_range = initializer_range
        self.num_labels = num_labels
        self.num_choices = num_choices
        self.scope = scope

    def prepare_config_and_inputs(self):
        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)

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

        token_type_ids = None

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

        config = self.get_config()

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

    def get_config(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=False,
            initializer_range=self.initializer_range,
            pad_token_id=1,
            new_decoder_architecture=True,
        )

    def create_and_check_model(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        model = FalconModel(config=config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=input_mask)
        result = model(input_ids)
        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))

    def create_and_check_model_as_decoder(
        self,
        config,
        input_ids,
        token_type_ids,
        input_mask,
        sequence_labels,
        token_labels,
        choice_labels,
        encoder_hidden_states,
        encoder_attention_mask,
    ):
        config.add_cross_attention = True
        model = FalconModel(config)
        model.to(torch_device)
        model.eval()
        result = model(
            input_ids,
            attention_mask=input_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
        )
        result = model(
            input_ids,
            attention_mask=input_mask,
            encoder_hidden_states=encoder_hidden_states,
        )
        result = model(input_ids, attention_mask=input_mask)
        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))

    def create_and_check_for_causal_lm(
        self,
        config,
        input_ids,
        token_type_ids,
        input_mask,
        sequence_labels,
        token_labels,
        choice_labels,
        encoder_hidden_states,
        encoder_attention_mask,
    ):
        model = FalconForCausalLM(config=config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=input_mask, labels=token_labels)
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))

    def create_and_check_decoder_model_past_large_inputs(
        self,
        config,
        input_ids,
        token_type_ids,
        input_mask,
        sequence_labels,
        token_labels,
        choice_labels,
        encoder_hidden_states,
        encoder_attention_mask,
    ):
        config.is_decoder = True
        config.add_cross_attention = True
        model = FalconForCausalLM(config=config)
        model.to(torch_device)
        model.eval()

        
        outputs = model(
            input_ids,
            attention_mask=input_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            use_cache=True,
        )
        past_key_values = outputs.past_key_values

        
        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
        next_mask = ids_tensor((self.batch_size, 3), vocab_size=2)

        
        next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
        next_attention_mask = torch.cat([input_mask, next_mask], dim=-1)

        output_from_no_past = model(
            next_input_ids,
            attention_mask=next_attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_hidden_states=True,
        )["hidden_states"][0]
        output_from_past = model(
            next_tokens,
            attention_mask=next_attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            past_key_values=past_key_values,
            output_hidden_states=True,
        )["hidden_states"][0]

        
        random_slice_idx = ids_tensor((1,), output_from_past.shape[-1]).item()
        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx].detach()
        output_from_past_slice = output_from_past[:, :, random_slice_idx].detach()

        self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1])

        
        self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        (
            config,
            input_ids,
            token_type_ids,
            input_mask,
            sequence_labels,
            token_labels,
            choice_labels,
        ) = config_and_inputs
        inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask}
        return config, inputs_dict


@require_torch
class FalconModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (
        (
            FalconModel,
            FalconForCausalLM,
            FalconForSequenceClassification,
            FalconForTokenClassification,
            FalconForQuestionAnswering,
        )
        if is_torch_available()
        else ()
    )
    all_generative_model_classes = (FalconForCausalLM,) if is_torch_available() else ()
    pipeline_model_mapping = (
        {
            "feature-extraction": FalconModel,
            "question-answering": FalconForQuestionAnswering,
            "text-classification": FalconForSequenceClassification,
            "text-generation": FalconForCausalLM,
            "token-classification": FalconForTokenClassification,
            "zero-shot": FalconForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    test_headmasking = False
    test_pruning = False

    
    def is_pipeline_test_to_skip(
        self, pipeline_test_casse_name, config_class, model_architecture, tokenizer_name, processor_name
    ):
        return True

    def setUp(self):
        self.model_tester = FalconModelTester(self)
        self.config_tester = ConfigTester(self, config_class=FalconConfig, hidden_size=37)

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

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)

    def test_position_embedding_types(self):
        config, *inputs = self.model_tester.prepare_config_and_inputs()
        for alibi in [True, False]:
            config.alibi = alibi
            self.model_tester.create_and_check_model(config, *inputs)

    def test_falcon_sequence_classification_model(self):
        config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.num_labels = 3
        input_ids = input_dict["input_ids"]
        attention_mask = input_ids.ne(1).to(torch_device)
        sequence_labels = ids_tensor([self.model_tester.batch_size], self.model_tester.type_sequence_label_size)
        model = FalconForSequenceClassification(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels)
        self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels))

    def test_falcon_sequence_classification_model_for_single_label(self):
        config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.num_labels = 3
        config.problem_type = "single_label_classification"
        input_ids = input_dict["input_ids"]
        attention_mask = input_ids.ne(1).to(torch_device)
        sequence_labels = ids_tensor([self.model_tester.batch_size], self.model_tester.type_sequence_label_size)
        model = FalconForSequenceClassification(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels)
        self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels))

    def test_falcon_sequence_classification_model_for_multi_label(self):
        config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.num_labels = 3
        config.problem_type = "multi_label_classification"
        input_ids = input_dict["input_ids"]
        attention_mask = input_ids.ne(1).to(torch_device)
        sequence_labels = ids_tensor(
            [self.model_tester.batch_size, config.num_labels], self.model_tester.type_sequence_label_size
        ).to(torch.float)
        model = FalconForSequenceClassification(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels)
        self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels))

    def test_past_key_values_format(self):
        
        
        for model_class in self.all_generative_model_classes:
            config, inputs = self.model_tester.prepare_config_and_inputs_for_common()

            
            if not hasattr(config, "use_cache"):
                return

            model = model_class(config).to(torch_device)
            if "use_cache" not in inputs:
                inputs["use_cache"] = True
            outputs = model(**inputs)

            
            if "past_key_values" not in outputs:
                return

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

            past_kv = outputs["past_key_values"]
            self.assertEqual(len(past_kv), num_hidden_layers)

            batch_size, seq_length = inputs["input_ids"].shape
            for i in range(num_hidden_layers):
                if config.new_decoder_architecture:
                    num_attention_heads = config.num_attention_heads
                elif config.multi_query:
                    num_attention_heads = 1
                self.assertEqual(len(past_kv[0]), 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)
                )

    @parameterized.expand([("linear",), ("dynamic",)])
    def test_model_rope_scaling(self, scaling_type):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        short_input = ids_tensor([1, 10], config.vocab_size)
        long_input = ids_tensor([1, int(config.max_position_embeddings * 1.5)], config.vocab_size)

        set_seed(42)  
        original_model = FalconModel(config)
        original_model.to(torch_device)
        original_model.eval()
        original_short_output = original_model(short_input).last_hidden_state
        original_long_output = original_model(long_input).last_hidden_state

        set_seed(42)  
        config.rope_scaling = {"type": scaling_type, "factor": 10.0}
        scaled_model = FalconModel(config)
        scaled_model.to(torch_device)
        scaled_model.eval()
        scaled_short_output = scaled_model(short_input).last_hidden_state
        scaled_long_output = scaled_model(long_input).last_hidden_state

        
        
        if scaling_type == "dynamic":
            self.assertTrue(torch.allclose(original_short_output, scaled_short_output, atol=1e-5))
        else:
            self.assertFalse(torch.allclose(original_short_output, scaled_short_output, atol=1e-5))

        
        self.assertFalse(torch.allclose(original_long_output, scaled_long_output, atol=1e-5))


@require_torch
class FalconLanguageGenerationTest(unittest.TestCase):
    @slow
    def test_lm_generate_falcon(self):
        tokenizer = AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b")
        model = FalconForCausalLM.from_pretrained("Rocketknight1/falcon-rw-1b")
        model.eval()
        model.to(torch_device)
        inputs = tokenizer("My favorite food is", return_tensors="pt").to(torch_device)

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

        output_ids = model.generate(**inputs, do_sample=False, max_new_tokens=19)
        output_str = tokenizer.batch_decode(output_ids)[0]

        self.assertEqual(output_str, EXPECTED_OUTPUT)

    @slow
    def test_lm_generation_big_models(self):
        
        
        for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
            tokenizer = AutoTokenizer.from_pretrained(repo)
            model = FalconForCausalLM.from_pretrained(repo)
            model.eval()
            model.to(torch_device)
            inputs = tokenizer("My favorite food is", return_tensors="pt").to(torch_device)

            
            
            model.generate(**inputs, do_sample=False, max_new_tokens=4)
            model.generate(**inputs, do_sample=True, max_new_tokens=4)
            model.generate(**inputs, num_beams=2, max_new_tokens=4)

    @slow
    def test_lm_generation_use_cache(self):
        
        
        with torch.no_grad():
            for repo in [
                "Rocketknight1/falcon-rw-1b",
                "Rocketknight1/tiny-random-falcon-7b",
                "Rocketknight1/tiny-random-falcon-40b",
            ]:
                tokenizer = AutoTokenizer.from_pretrained(repo)
                model = FalconForCausalLM.from_pretrained(repo)
                model.eval()
                model.to(device=torch_device)
                inputs = tokenizer("My favorite food is", return_tensors="pt").to(torch_device)

                
                outputs_no_cache = model.generate(**inputs, do_sample=False, max_new_tokens=20, use_cache=False)
                outputs_cache = model.generate(**inputs, do_sample=False, max_new_tokens=20, use_cache=True)
                self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0)

    @require_bitsandbytes
    @slow
    def test_batched_generation(self):
        tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b", padding_side="left")
        tokenizer.pad_token = tokenizer.eos_token
        model = AutoModelForCausalLM.from_pretrained(
            "tiiuae/falcon-7b",
            device_map="auto",
            load_in_4bit=True,
        )

        test_text = "A sequence: 1, 2"  

        unpadded_inputs = tokenizer([test_text], return_tensors="pt").to("cuda:0")
        unpadded_gen_out = model.generate(**unpadded_inputs, max_new_tokens=20)
        unpadded_gen_text = tokenizer.batch_decode(unpadded_gen_out, skip_special_tokens=True)

        dummy_text = "This is a longer text " * 2  
        padded_inputs = tokenizer([test_text, dummy_text], return_tensors="pt", padding=True).to("cuda:0")
        padded_gen_out = model.generate(**padded_inputs, max_new_tokens=20)
        padded_gen_text = tokenizer.batch_decode(padded_gen_out, skip_special_tokens=True)

        expected_output = "A sequence: 1, 2, 3, 4, 5, 6, 7, 8, "
        self.assertLess(unpadded_inputs.input_ids.shape[-1], padded_inputs.input_ids.shape[-1])  
        self.assertEqual(unpadded_gen_text[0], expected_output)
        self.assertEqual(padded_gen_text[0], expected_output)