codingutf8 2021 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 gpt neo model import unittest from transformers import gptneoconfig istorchavailable from transformers testingutils import requiretorch slow torchdevice from transformers utils import cachedproperty from generation testutils import generationtestermixin from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin floatstensor idstensor randomattentionmask from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import gptneopretrainedmodelarchivelist gpt2tokenizer gptneoforcausallm gptneoforquestionanswering gptneoforsequenceclassification gptneofortokenclassification gptneomodel class gptneomodeltester def init self parent batchsize14 seqlength7 istrainingtrue usetokentypeidstrue useinputmasktrue uselabelstrue usemctokenidstrue vocabsize99 hiddensize32 numhiddenlayers2 attentiontypesglobal local 1 numattentionheads4 intermediatesize37 hiddenactgelu hiddendropoutprob0 1 attentionprobsdropoutprob0 1 maxpositionembeddings512 windowsize7 typevocabsize16 typesequencelabelsize2 initializerrange0 02 numlabels3 numchoices4 self parent parent self batchsize batchsize self seqlength seqlength self istraining istraining self usetokentypeids usetokentypeids self useinputmask useinputmask self uselabels uselabels self usemctokenids usemctokenids 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 windowsize windowsize self typevocabsize typevocabsize self typesequencelabelsize typesequencelabelsize self initializerrange initializerrange self numlabels numlabels self numchoices numchoices self bostokenid vocabsize 1 self eostokenid vocabsize 1 self padtokenid vocabsize 1 self attentiontypes attentiontypes def getlargemodelconfigself return gptneoconfig frompretrainedgptneo125m def prepareconfigandinputsself inputids idstensorself batchsize self seqlength self vocabsize inputmask none if self useinputmask inputmask randomattentionmaskself batchsize self seqlength tokentypeids none if self usetokentypeids tokentypeids idstensorself batchsize self seqlength self typevocabsize mctokenids none if self usemctokenids mctokenids idstensorself batchsize self numchoices self seqlength 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 headmask idstensorself numhiddenlayers self numattentionheads 2 return config inputids inputmask headmask tokentypeids mctokenids sequencelabels tokenlabels choicelabels def getconfigself return gptneoconfig vocabsizeself vocabsize hiddensizeself hiddensize numlayersself numhiddenlayers numheadsself numattentionheads maxpositionembeddingsself maxpositionembeddings usecachetrue bostokenidself bostokenid eostokenidself eostokenid padtokenidself padtokenid windowsizeself windowsize attentiontypesself attentiontypes def getpipelineconfigself config self getconfig config vocabsize 300 return config def prepareconfigandinputsfordecoderself config inputids inputmask headmask tokentypeids mctokenids sequencelabels tokenlabels choicelabels self prepareconfigandinputs encoderhiddenstates floatstensorself batchsize self seqlength self hiddensize encoderattentionmask idstensorself batchsize self seqlength vocabsize2 return config inputids inputmask headmask tokentypeids sequencelabels tokenlabels choicelabels encoderhiddenstates encoderattentionmask def createandcheckgptneomodelself config inputids inputmask headmask tokentypeids args model gptneomodelconfigconfig model totorchdevice model eval result modelinputids tokentypeidstokentypeids headmaskheadmask result modelinputids tokentypeidstokentypeids result modelinputids self parent assertequalresult lasthiddenstate shape self batchsize self seqlength self hiddensize pastkeyvalues is not implemented self parent assertequallenresult pastkeyvalues config nlayer def createandcheckgptneomodelpastself config inputids inputmask headmask tokentypeids args model gptneomodelconfigconfig model totorchdevice model eval first forward pass outputs modelinputids tokentypeidstokentypeids usecachetrue outputsusecacheconf modelinputids tokentypeidstokentypeids outputsnopast modelinputids tokentypeidstokentypeids usecachefalse self parent asserttruelenoutputs lenoutputsusecacheconf self parent asserttruelenoutputs lenoutputsnopast 1 output past outputs totuple create hypothetical next token and extent to nextinputids nexttokens idstensorself batchsize 1 config vocabsize nexttokentypes idstensorself batchsize 1 self typevocabsize append to next inputids and tokentypeids nextinputids torch catinputids nexttokens dim1 nexttokentypeids torch cattokentypeids nexttokentypes dim1 outputfromnopast modelnextinputids tokentypeidsnexttokentypeidslasthiddenstate outputfrompast modelnexttokens tokentypeidsnexttokentypes pastkeyvaluespast lasthiddenstate select random slice randomsliceidx idstensor1 outputfrompast shape1 item outputfromnopastslice outputfromnopast 1 randomsliceidx detach outputfrompastslice outputfrompast 0 randomsliceidx detach test that outputs are equal for slice self parent asserttruetorch allcloseoutputfrompastslice outputfromnopastslice atol1e3 def createandcheckgptneomodelattentionmaskpast self config inputids inputmask headmask tokentypeids args model gptneomodelconfigconfig model totorchdevice model eval create attention mask attnmask torch onesinputids shape dtypetorch long devicetorchdevice halfseqlength self seqlength 2 attnmask halfseqlength 0 first forward pass output past modelinputids attentionmaskattnmask totuple create hypothetical next token and extent to nextinputids nexttokens idstensorself batchsize 1 config vocabsize change a random masked slice from inputids randomseqidxtochange idstensor1 halfseqlength item 1 randomothernexttokens idstensorself batchsize 1 config vocabsize squeeze1 inputids randomseqidxtochange randomothernexttokens append to next inputids and attnmask nextinputids torch catinputids nexttokens dim1 attnmask torch cat attnmask torch onesattnmask shape0 1 dtypetorch long devicetorchdevice dim1 get two different outputs outputfromnopast modelnextinputids attentionmaskattnmasklasthiddenstate outputfrompast modelnexttokens pastkeyvaluespast attentionmaskattnmasklasthiddenstate select random slice randomsliceidx idstensor1 outputfrompast shape1 item outputfromnopastslice outputfromnopast 1 randomsliceidx detach outputfrompastslice outputfrompast 0 randomsliceidx detach test that outputs are equal for slice self parent asserttruetorch allcloseoutputfrompastslice outputfromnopastslice atol1e3 def createandcheckgptneomodelpastlargeinputs self config inputids inputmask headmask tokentypeids args model gptneomodelconfigconfig model totorchdevice model eval first forward pass outputs modelinputids tokentypeidstokentypeids attentionmaskinputmask usecachetrue output past outputs totuple create hypothetical next token and extent to nextinputids nexttokens idstensorself batchsize 3 config vocabsize nexttokentypes idstensorself batchsize 3 self typevocabsize nextmask idstensorself batchsize 3 vocabsize2 append to next inputids and tokentypeids nextinputids torch catinputids nexttokens dim1 nexttokentypeids torch cattokentypeids nexttokentypes dim1 nextattentionmask torch catinputmask nextmask dim1 outputfromnopast model nextinputids tokentypeidsnexttokentypeids attentionmasknextattentionmask lasthiddenstate outputfrompast model nexttokens tokentypeidsnexttokentypes attentionmasknextattentionmask pastkeyvaluespast lasthiddenstate self parent asserttrueoutputfrompast shape1 nexttokens shape1 select random slice randomsliceidx idstensor1 outputfrompast shape1 item outputfromnopastslice outputfromnopast 3 randomsliceidx detach outputfrompastslice outputfrompast randomsliceidx detach test that outputs are equal for slice self parent asserttruetorch allcloseoutputfrompastslice outputfromnopastslice atol1e3 def createandchecklmheadmodelself config inputids inputmask headmask tokentypeids args model gptneoforcausallmconfig model totorchdevice model eval result modelinputids tokentypeidstokentypeids labelsinputids self parent assertequalresult loss shape self parent assertequalresult logits shape self batchsize self seqlength self vocabsize def createandcheckgptneoforquestionanswering self config inputids inputmask headmask tokentypeids mctokenids sequencelabels args config numlabels self numlabels model gptneoforquestionansweringconfig model totorchdevice model eval result modelinputids attentionmaskinputmask tokentypeidstokentypeids self parent assertequalresult startlogits shape self batchsize self seqlength self parent assertequalresult endlogits shape self batchsize self seqlength def createandcheckgptneoforsequenceclassification self config inputids inputmask headmask tokentypeids mctokenids sequencelabels args config numlabels self numlabels model gptneoforsequenceclassificationconfig model totorchdevice model eval result modelinputids attentionmaskinputmask tokentypeidstokentypeids labelssequencelabels self parent assertequalresult logits shape self batchsize self numlabels def createandcheckgptneofortokenclassification self config inputids inputmask headmask tokentypeids mctokenids sequencelabels args config numlabels self numlabels model gptneofortokenclassificationconfig model totorchdevice model eval result modelinputids attentionmaskinputmask tokentypeidstokentypeids self parent assertequalresult logits shape self batchsize self seqlength self numlabels def createandcheckforwardandbackwards self config inputids inputmask headmask tokentypeids args gradientcheckpointingfalse model gptneoforcausallmconfig if gradientcheckpointing model gradientcheckpointingenable model totorchdevice result modelinputids tokentypeidstokentypeids labelsinputids self parent assertequalresult loss shape self parent assertequalresult logits shape self batchsize self seqlength self vocabsize result loss backward def prepareconfigandinputsforcommonself configandinputs self prepareconfigandinputs config inputids inputmask headmask tokentypeids mctokenids sequencelabels tokenlabels choicelabels configandinputs inputsdict inputids inputids tokentypeids tokentypeids headmask headmask return config inputsdict requiretorch class gptneomodeltestmodeltestermixin generationtestermixin pipelinetestermixin unittest testcase allmodelclasses gptneomodel gptneoforcausallm gptneoforquestionanswering gptneoforsequenceclassification gptneofortokenclassification if istorchavailable else allgenerativemodelclasses gptneoforcausallm if istorchavailable else pipelinemodelmapping featureextraction gptneomodel questionanswering gptneoforquestionanswering textclassification gptneoforsequenceclassification textgeneration gptneoforcausallm tokenclassification gptneofortokenclassification zeroshot gptneoforsequenceclassification if istorchavailable else fxcompatible true testmissingkeys false testpruning false testmodelparallel false special case for doubleheads model def prepareforclassself inputsdict modelclass returnlabelsfalse inputsdict super prepareforclassinputsdict modelclass returnlabelsreturnlabels return inputsdict def setupself self modeltester gptneomodeltesterself self configtester configtesterself configclassgptneoconfig nembd37 def testconfigself self configtester runcommontests def testgptneomodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckgptneomodelconfigandinputs def testgptneomodelpastself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckgptneomodelpastconfigandinputs def testgptneomodelattmaskpastself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckgptneomodelattentionmaskpastconfigandinputs def testgptneomodelpastlargeinputsself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckgptneomodelpastlargeinputsconfigandinputs def testgptneolmheadmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandchecklmheadmodelconfigandinputs def testgptneoquestionansweringmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckgptneoforquestionansweringconfigandinputs def testgptneosequenceclassificationmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckgptneoforsequenceclassificationconfigandinputs def testgptneotokenclassificationmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckgptneofortokenclassificationconfigandinputs def testgptneogradientcheckpointingself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckforwardandbackwardsconfigandinputs gradientcheckpointingtrue def gethiddenstatesself return torch tensor 0 4983 0 7584 1 6944 0 5440 2 6918 0 4206 0 4176 0 2055 0 0071 0 0405 1 4920 0 3630 1 0492 0 1599 1 7648 0 2419 1 8348 2 0514 0 1946 0 3203 0 7672 1 1600 1 7118 0 9056 0 2986 0 5372 0 7729 0 1927 0 0285 0 2629 1 1156 1 1992 dtypetorch float32 devicetorchdevice def testlocalattnprobsself model gptneomodel frompretrainedvalhallagptneorandomtiny eval layer model h1 attn attention totorchdevice hiddenstates self gethiddenstates hiddenstates torch cathiddenstates hiddenstates 0 5 dim2 batchsize seqlength hiddenstates shape masktokens 2 attentionmask torch onesbatchsize seqlength devicetorchdevice dtypetorch long attentionmask masktokens 0 dont attend last masktokens attentionmask attentionmask viewbatchsize 1 attentionmask attentionmask none none attentionmask 1 0 attentionmask 10000 0 attnprobs layerhiddenstates attentionmaskattentionmask outputattentionstrue1 the last 2 tokens are masked and should have 0 attnprobs self asserttruetorch allattnprobs masktokens masktokens 0 in loacal attention each token can only attend to the previous windowsize tokens inlcuding itself here windowsize is 4 so a token at index 5 can only attend to indcies 2 3 4 5 and the attnprobs should be 0 for token 0 1 self asserttruetorch allattnprobs 5 2 6 0 self asserttruetorch allattnprobs 5 2 0 requiretorch class gptneomodellanguagegenerationtestunittest testcase cachedproperty def modelself return gptneoforcausallm frompretrainedeleutheraigptneo1 3b totorchdevice cachedproperty def tokenizerself return gpt2tokenizer frompretrainedeleutheraigptneo1 3b slow def testlmgenerategptneoself for checkpointing in true false model self model if checkpointing model gradientcheckpointingenable else model gradientcheckpointingdisable inputids torch tensor464 3290 dtypetorch long devicetorchdevice the dog the dogeared copy of the book which is a collection of essays by the late expectedoutputids 464 3290 12 3380 4866 286 262 1492 11 543 318 257 4947 286 27126 416 262 2739 1772 11 fmt skip outputids model generateinputids dosamplefalse self assertlistequaloutputids0 tolist expectedoutputids slow def testgptneosampleself model self model tokenizer self tokenizer torch manualseed0 tokenized tokenizertoday is a nice day and returntensorspt returntokentypeidstrue inputids tokenized inputids totorchdevice outputids model generateinputids dosampletrue outputstr tokenizer decodeoutputids0 skipspecialtokenstrue expectedoutputstr today is a nice day and if you dont get the memo here is what you can self assertequaloutputstr expectedoutputstr slow def testbatchgenerationself model self model tokenizer self tokenizer tokenizer paddingside left define pad token eos token 50256 tokenizer padtoken tokenizer eostoken model config padtokenid model config eostokenid use different length sentences to test batching sentences hello my dog is a little today i am inputs tokenizersentences returntensorspt paddingtrue inputids inputsinputids totorchdevice outputs model generate inputidsinputids attentionmaskinputsattentionmask totorchdevice inputsnonpadded tokenizersentences0 returntensorspt inputids totorchdevice outputnonpadded model generateinputidsinputsnonpadded numpaddings inputsnonpadded shape1 inputsattentionmask1 long sum cpu item inputspadded tokenizersentences1 returntensorspt inputids totorchdevice outputpadded model generateinputidsinputspadded maxlengthmodel config maxlength numpaddings batchoutsentence tokenizer batchdecodeoutputs skipspecialtokenstrue nonpaddedsentence tokenizer decodeoutputnonpadded0 skipspecialtokenstrue paddedsentence tokenizer decodeoutputpadded0 skipspecialtokenstrue expectedoutputsentence hello my dog is a little bit of a kitty she is a very sweet and loving today i am going to talk about the best way to get a job in the self assertlistequalexpectedoutputsentence batchoutsentence self assertlistequalexpectedoutputsentence nonpaddedsentence paddedsentence slow def testmodelfrompretrainedself for modelname in gptneopretrainedmodelarchivelist 1 model gptneomodel frompretrainedmodelname self assertisnotnonemodel coding utf 8 2021 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 gpt neo model past_key_values is not implemented self parent assertequal len result past_key_values config n_layer first forward pass create hypothetical next token and extent to next_input_ids append to next input_ids and token_type_ids select random slice test that outputs are equal for slice create attention mask first forward pass create hypothetical next token and extent to next_input_ids change a random masked slice from input_ids append to next input_ids and attn_mask get two different outputs select random slice test that outputs are equal for slice first forward pass create hypothetical next token and extent to next_input_ids append to next input_ids and token_type_ids select random slice test that outputs are equal for slice special case for doubleheads model dont attend last mask_tokens the last 2 tokens are masked and should have 0 attn_probs in loacal attention each token can only attend to the previous window_size tokens inlcuding itself here window_size is 4 so a token at index 5 can only attend to indcies 2 3 4 5 and the attn_probs should be 0 for token 0 1 the dog the dog eared copy of the book which is a collection of essays by the late fmt skip define pad token eos token 50256 use different length sentences to test batching
import unittest

from transformers import GPTNeoConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import cached_property

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


if is_torch_available():
    import torch

    from transformers import (
        GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST,
        GPT2Tokenizer,
        GPTNeoForCausalLM,
        GPTNeoForQuestionAnswering,
        GPTNeoForSequenceClassification,
        GPTNeoForTokenClassification,
        GPTNeoModel,
    )


class GPTNeoModelTester:
    def __init__(
        self,
        parent,
        batch_size=14,
        seq_length=7,
        is_training=True,
        use_token_type_ids=True,
        use_input_mask=True,
        use_labels=True,
        use_mc_token_ids=True,
        vocab_size=99,
        hidden_size=32,
        num_hidden_layers=2,
        attention_types=[[["global", "local"], 1]],
        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,
        window_size=7,
        type_vocab_size=16,
        type_sequence_label_size=2,
        initializer_range=0.02,
        num_labels=3,
        num_choices=4,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        self.use_token_type_ids = use_token_type_ids
        self.use_input_mask = use_input_mask
        self.use_labels = use_labels
        self.use_mc_token_ids = use_mc_token_ids
        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.window_size = window_size
        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.bos_token_id = vocab_size - 1
        self.eos_token_id = vocab_size - 1
        self.pad_token_id = vocab_size - 1
        self.attention_types = attention_types

    def get_large_model_config(self):
        return GPTNeoConfig.from_pretrained("gpt-neo-125M")

    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
        if self.use_token_type_ids:
            token_type_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size)

        mc_token_ids = None
        if self.use_mc_token_ids:
            mc_token_ids = ids_tensor([self.batch_size, self.num_choices], self.seq_length)

        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()

        head_mask = ids_tensor([self.num_hidden_layers, self.num_attention_heads], 2)

        return (
            config,
            input_ids,
            input_mask,
            head_mask,
            token_type_ids,
            mc_token_ids,
            sequence_labels,
            token_labels,
            choice_labels,
        )

    def get_config(self):
        return GPTNeoConfig(
            vocab_size=self.vocab_size,
            hidden_size=self.hidden_size,
            num_layers=self.num_hidden_layers,
            num_heads=self.num_attention_heads,
            max_position_embeddings=self.max_position_embeddings,
            use_cache=True,
            bos_token_id=self.bos_token_id,
            eos_token_id=self.eos_token_id,
            pad_token_id=self.pad_token_id,
            window_size=self.window_size,
            attention_types=self.attention_types,
        )

    def get_pipeline_config(self):
        config = self.get_config()
        config.vocab_size = 300
        return config

    def prepare_config_and_inputs_for_decoder(self):
        (
            config,
            input_ids,
            input_mask,
            head_mask,
            token_type_ids,
            mc_token_ids,
            sequence_labels,
            token_labels,
            choice_labels,
        ) = self.prepare_config_and_inputs()

        encoder_hidden_states = floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
        encoder_attention_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

        return (
            config,
            input_ids,
            input_mask,
            head_mask,
            token_type_ids,
            sequence_labels,
            token_labels,
            choice_labels,
            encoder_hidden_states,
            encoder_attention_mask,
        )

    def create_and_check_gpt_neo_model(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
        model = GPTNeoModel(config=config)
        model.to(torch_device)
        model.eval()

        result = model(input_ids, token_type_ids=token_type_ids, head_mask=head_mask)
        result = model(input_ids, token_type_ids=token_type_ids)
        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_gpt_neo_model_past(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
        model = GPTNeoModel(config=config)
        model.to(torch_device)
        model.eval()

        
        outputs = model(input_ids, token_type_ids=token_type_ids, use_cache=True)
        outputs_use_cache_conf = model(input_ids, token_type_ids=token_type_ids)
        outputs_no_past = model(input_ids, token_type_ids=token_type_ids, use_cache=False)

        self.parent.assertTrue(len(outputs) == len(outputs_use_cache_conf))
        self.parent.assertTrue(len(outputs) == len(outputs_no_past) + 1)

        output, past = outputs.to_tuple()

        
        next_tokens = ids_tensor((self.batch_size, 1), config.vocab_size)
        next_token_types = ids_tensor([self.batch_size, 1], self.type_vocab_size)

        
        next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
        next_token_type_ids = torch.cat([token_type_ids, next_token_types], dim=-1)

        output_from_no_past = model(next_input_ids, token_type_ids=next_token_type_ids)["last_hidden_state"]
        output_from_past = model(next_tokens, token_type_ids=next_token_types, past_key_values=past)[
            "last_hidden_state"
        ]

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

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

    def create_and_check_gpt_neo_model_attention_mask_past(
        self, config, input_ids, input_mask, head_mask, token_type_ids, *args
    ):
        model = GPTNeoModel(config=config)
        model.to(torch_device)
        model.eval()

        
        attn_mask = torch.ones(input_ids.shape, dtype=torch.long, device=torch_device)
        half_seq_length = self.seq_length // 2
        attn_mask[:, half_seq_length:] = 0

        
        output, past = model(input_ids, attention_mask=attn_mask).to_tuple()

        
        next_tokens = ids_tensor((self.batch_size, 1), config.vocab_size)

        
        random_seq_idx_to_change = ids_tensor((1,), half_seq_length).item() + 1
        random_other_next_tokens = ids_tensor((self.batch_size, 1), config.vocab_size).squeeze(-1)
        input_ids[:, -random_seq_idx_to_change] = random_other_next_tokens

        
        next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
        attn_mask = torch.cat(
            [attn_mask, torch.ones((attn_mask.shape[0], 1), dtype=torch.long, device=torch_device)],
            dim=1,
        )

        
        output_from_no_past = model(next_input_ids, attention_mask=attn_mask)["last_hidden_state"]
        output_from_past = model(next_tokens, past_key_values=past, attention_mask=attn_mask)["last_hidden_state"]

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

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

    def create_and_check_gpt_neo_model_past_large_inputs(
        self, config, input_ids, input_mask, head_mask, token_type_ids, *args
    ):
        model = GPTNeoModel(config=config)
        model.to(torch_device)
        model.eval()

        
        outputs = model(input_ids, token_type_ids=token_type_ids, attention_mask=input_mask, use_cache=True)

        output, past = outputs.to_tuple()

        
        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
        next_token_types = ids_tensor([self.batch_size, 3], self.type_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_token_type_ids = torch.cat([token_type_ids, next_token_types], dim=-1)
        next_attention_mask = torch.cat([input_mask, next_mask], dim=-1)

        output_from_no_past = model(
            next_input_ids, token_type_ids=next_token_type_ids, attention_mask=next_attention_mask
        )["last_hidden_state"]
        output_from_past = model(
            next_tokens, token_type_ids=next_token_types, attention_mask=next_attention_mask, past_key_values=past
        )["last_hidden_state"]
        self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1])

        
        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(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))

    def create_and_check_lm_head_model(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
        model = GPTNeoForCausalLM(config)
        model.to(torch_device)
        model.eval()

        result = model(input_ids, token_type_ids=token_type_ids, labels=input_ids)
        self.parent.assertEqual(result.loss.shape, ())
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))

    def create_and_check_gpt_neo_for_question_answering(
        self, config, input_ids, input_mask, head_mask, token_type_ids, mc_token_ids, sequence_labels, *args
    ):
        config.num_labels = self.num_labels
        model = GPTNeoForQuestionAnswering(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids)
        self.parent.assertEqual(result.start_logits.shape, (self.batch_size, self.seq_length))
        self.parent.assertEqual(result.end_logits.shape, (self.batch_size, self.seq_length))

    def create_and_check_gpt_neo_for_sequence_classification(
        self, config, input_ids, input_mask, head_mask, token_type_ids, mc_token_ids, sequence_labels, *args
    ):
        config.num_labels = self.num_labels
        model = GPTNeoForSequenceClassification(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=sequence_labels)
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_labels))

    def create_and_check_gpt_neo_for_token_classification(
        self, config, input_ids, input_mask, head_mask, token_type_ids, mc_token_ids, sequence_labels, *args
    ):
        config.num_labels = self.num_labels
        model = GPTNeoForTokenClassification(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids)
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.num_labels))

    def create_and_check_forward_and_backwards(
        self, config, input_ids, input_mask, head_mask, token_type_ids, *args, gradient_checkpointing=False
    ):
        model = GPTNeoForCausalLM(config)
        if gradient_checkpointing:
            model.gradient_checkpointing_enable()
        model.to(torch_device)

        result = model(input_ids, token_type_ids=token_type_ids, labels=input_ids)
        self.parent.assertEqual(result.loss.shape, ())
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
        result.loss.backward()

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()

        (
            config,
            input_ids,
            input_mask,
            head_mask,
            token_type_ids,
            mc_token_ids,
            sequence_labels,
            token_labels,
            choice_labels,
        ) = config_and_inputs

        inputs_dict = {
            "input_ids": input_ids,
            "token_type_ids": token_type_ids,
            "head_mask": head_mask,
        }

        return config, inputs_dict


@require_torch
class GPTNeoModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (
        (
            GPTNeoModel,
            GPTNeoForCausalLM,
            GPTNeoForQuestionAnswering,
            GPTNeoForSequenceClassification,
            GPTNeoForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    all_generative_model_classes = (GPTNeoForCausalLM,) if is_torch_available() else ()
    pipeline_model_mapping = (
        {
            "feature-extraction": GPTNeoModel,
            "question-answering": GPTNeoForQuestionAnswering,
            "text-classification": GPTNeoForSequenceClassification,
            "text-generation": GPTNeoForCausalLM,
            "token-classification": GPTNeoForTokenClassification,
            "zero-shot": GPTNeoForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    fx_compatible = True
    test_missing_keys = False
    test_pruning = False
    test_model_parallel = False

    
    def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
        inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels)
        return inputs_dict

    def setUp(self):
        self.model_tester = GPTNeoModelTester(self)
        self.config_tester = ConfigTester(self, config_class=GPTNeoConfig, n_embd=37)

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

    def test_gpt_neo_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_gpt_neo_model(*config_and_inputs)

    def test_gpt_neo_model_past(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_gpt_neo_model_past(*config_and_inputs)

    def test_gpt_neo_model_att_mask_past(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_gpt_neo_model_attention_mask_past(*config_and_inputs)

    def test_gpt_neo_model_past_large_inputs(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_gpt_neo_model_past_large_inputs(*config_and_inputs)

    def test_gpt_neo_lm_head_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_lm_head_model(*config_and_inputs)

    def test_gpt_neo_question_answering_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_gpt_neo_for_question_answering(*config_and_inputs)

    def test_gpt_neo_sequence_classification_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_gpt_neo_for_sequence_classification(*config_and_inputs)

    def test_gpt_neo_token_classification_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_gpt_neo_for_token_classification(*config_and_inputs)

    def test_gpt_neo_gradient_checkpointing(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_forward_and_backwards(*config_and_inputs, gradient_checkpointing=True)

    def _get_hidden_states(self):
        return torch.tensor(
            [
                [
                    [0.4983, -0.7584, -1.6944, 0.5440],
                    [2.6918, 0.4206, 0.4176, 0.2055],
                    [-0.0071, -0.0405, -1.4920, -0.3630],
                    [1.0492, 0.1599, -1.7648, 0.2419],
                    [-1.8348, 2.0514, -0.1946, 0.3203],
                    [0.7672, -1.1600, -1.7118, -0.9056],
                    [0.2986, 0.5372, 0.7729, -0.1927],
                    [0.0285, 0.2629, -1.1156, -1.1992],
                ]
            ],
            dtype=torch.float32,
            device=torch_device,
        )

    def test_local_attn_probs(self):
        model = GPTNeoModel.from_pretrained("valhalla/gpt-neo-random-tiny").eval()
        layer = model.h[1].attn.attention.to(torch_device)
        hidden_states = self._get_hidden_states()
        hidden_states = torch.cat([hidden_states, hidden_states - 0.5], dim=2)

        batch_size, seq_length, _ = hidden_states.shape
        mask_tokens = 2
        attention_mask = torch.ones(batch_size, seq_length, device=torch_device, dtype=torch.long)
        attention_mask[:, -mask_tokens:] = 0  

        attention_mask = attention_mask.view(batch_size, -1)
        attention_mask = attention_mask[:, None, None, :]
        attention_mask = (1.0 - attention_mask) * -10000.0

        attn_probs = layer(hidden_states, attention_mask=attention_mask, output_attentions=True)[-1]

        
        self.assertTrue(torch.all(attn_probs[:, :, -mask_tokens:, -mask_tokens:] == 0))

        
        
        
        self.assertTrue(torch.all(attn_probs[:, :, 5, 2:6] != 0))
        self.assertTrue(torch.all(attn_probs[:, :, 5, :2] == 0))


@require_torch
class GPTNeoModelLanguageGenerationTest(unittest.TestCase):
    @cached_property
    def model(self):
        return GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B").to(torch_device)

    @cached_property
    def tokenizer(self):
        return GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")

    @slow
    def test_lm_generate_gpt_neo(self):
        for checkpointing in [True, False]:
            model = self.model
            if checkpointing:
                model.gradient_checkpointing_enable()
            else:
                model.gradient_checkpointing_disable()
            input_ids = torch.tensor([[464, 3290]], dtype=torch.long, device=torch_device)  
            
            expected_output_ids = [464, 3290, 12, 3380, 4866, 286, 262, 1492, 11, 543, 318, 257, 4947, 286, 27126, 416, 262, 2739, 1772, 11]  
            output_ids = model.generate(input_ids, do_sample=False)
            self.assertListEqual(output_ids[0].tolist(), expected_output_ids)

    @slow
    def test_gpt_neo_sample(self):
        model = self.model
        tokenizer = self.tokenizer

        torch.manual_seed(0)
        tokenized = tokenizer("Today is a nice day and", return_tensors="pt", return_token_type_ids=True)
        input_ids = tokenized.input_ids.to(torch_device)
        output_ids = model.generate(input_ids, do_sample=True)
        output_str = tokenizer.decode(output_ids[0], skip_special_tokens=True)

        EXPECTED_OUTPUT_STR = "Today is a nice day and if you don’t get the memo here is what you can"
        self.assertEqual(output_str, EXPECTED_OUTPUT_STR)

    @slow
    def test_batch_generation(self):
        model = self.model
        tokenizer = self.tokenizer

        tokenizer.padding_side = "left"

        
        tokenizer.pad_token = tokenizer.eos_token
        model.config.pad_token_id = model.config.eos_token_id

        
        sentences = [
            "Hello, my dog is a little",
            "Today, I am",
        ]

        inputs = tokenizer(sentences, return_tensors="pt", padding=True)
        input_ids = inputs["input_ids"].to(torch_device)

        outputs = model.generate(
            input_ids=input_ids,
            attention_mask=inputs["attention_mask"].to(torch_device),
        )

        inputs_non_padded = tokenizer(sentences[0], return_tensors="pt").input_ids.to(torch_device)
        output_non_padded = model.generate(input_ids=inputs_non_padded)

        num_paddings = inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].long().sum().cpu().item()
        inputs_padded = tokenizer(sentences[1], return_tensors="pt").input_ids.to(torch_device)
        output_padded = model.generate(input_ids=inputs_padded, max_length=model.config.max_length - num_paddings)

        batch_out_sentence = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        non_padded_sentence = tokenizer.decode(output_non_padded[0], skip_special_tokens=True)
        padded_sentence = tokenizer.decode(output_padded[0], skip_special_tokens=True)

        expected_output_sentence = [
            "Hello, my dog is a little bit of a kitty. She is a very sweet and loving",
            "Today, I am going to talk about the best way to get a job in the",
        ]
        self.assertListEqual(expected_output_sentence, batch_out_sentence)
        self.assertListEqual(expected_output_sentence, [non_padded_sentence, padded_sentence])

    @slow
    def test_model_from_pretrained(self):
        for model_name in GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            model = GPTNeoModel.from_pretrained(model_name)
            self.assertIsNotNone(model)