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 luke model import unittest from transformers import lukeconfig istorchavailable from transformers testingutils import requiretorch slow torchdevice from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin idstensor randomattentionmask from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import lukeforentityclassification lukeforentitypairclassification lukeforentityspanclassification lukeformaskedlm lukeformultiplechoice lukeforquestionanswering lukeforsequenceclassification lukefortokenclassification lukemodel luketokenizer from transformers models luke modelingluke import lukepretrainedmodelarchivelist class lukemodeltester def init self parent batchsize13 seqlength7 istrainingtrue entitylength3 mentionlength5 useattentionmasktrue usetokentypeidstrue useentityidstrue useentityattentionmasktrue useentitytokentypeidstrue useentitypositionidstrue uselabelstrue vocabsize99 entityvocabsize10 entityembsize6 hiddensize32 numhiddenlayers2 numattentionheads4 intermediatesize37 hiddenactgelu hiddendropoutprob0 1 attentionprobsdropoutprob0 1 maxpositionembeddings512 typevocabsize16 typesequencelabelsize2 initializerrange0 02 numlabels3 numchoices4 numentityclassificationlabels9 numentitypairclassificationlabels6 numentityspanclassificationlabels4 useentityawareattentiontrue scopenone self parent parent self batchsize batchsize self seqlength seqlength self istraining istraining self entitylength entitylength self mentionlength mentionlength self useattentionmask useattentionmask self usetokentypeids usetokentypeids self useentityids useentityids self useentityattentionmask useentityattentionmask self useentitytokentypeids useentitytokentypeids self useentitypositionids useentitypositionids self uselabels uselabels self vocabsize vocabsize self entityvocabsize entityvocabsize self entityembsize entityembsize 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 numentityclassificationlabels numentityclassificationlabels self numentitypairclassificationlabels numentitypairclassificationlabels self numentityspanclassificationlabels numentityspanclassificationlabels self scope scope self useentityawareattention useentityawareattention self encoderseqlength seqlength self keylength seqlength self numhiddenstatestypes 2 hiddenstates and entityhiddenstates def prepareconfigandinputsself prepare words inputids idstensorself batchsize self seqlength self vocabsize attentionmask none if self useattentionmask attentionmask randomattentionmaskself batchsize self seqlength tokentypeids none if self usetokentypeids tokentypeids idstensorself batchsize self seqlength self typevocabsize prepare entities entityids idstensorself batchsize self entitylength self entityvocabsize entityattentionmask none if self useentityattentionmask entityattentionmask randomattentionmaskself batchsize self entitylength entitytokentypeids none if self usetokentypeids entitytokentypeids idstensorself batchsize self entitylength self typevocabsize entitypositionids none if self useentitypositionids entitypositionids idstensor self batchsize self entitylength self mentionlength self mentionlength sequencelabels none tokenlabels none choicelabels none entitylabels none entityclassificationlabels none entitypairclassificationlabels none entityspanclassificationlabels none if self uselabels sequencelabels idstensorself batchsize self typesequencelabelsize tokenlabels idstensorself batchsize self seqlength self numlabels choicelabels idstensorself batchsize self numchoices entitylabels idstensorself batchsize self entitylength self entityvocabsize entityclassificationlabels idstensorself batchsize self numentityclassificationlabels entitypairclassificationlabels idstensor self batchsize self numentitypairclassificationlabels entityspanclassificationlabels idstensor self batchsize self entitylength self numentityspanclassificationlabels config self getconfig return config inputids attentionmask tokentypeids entityids entityattentionmask entitytokentypeids entitypositionids sequencelabels tokenlabels choicelabels entitylabels entityclassificationlabels entitypairclassificationlabels entityspanclassificationlabels def getconfigself return lukeconfig vocabsizeself vocabsize entityvocabsizeself entityvocabsize entityembsizeself entityembsize hiddensizeself hiddensize numhiddenlayersself numhiddenlayers numattentionheadsself numattentionheads intermediatesizeself intermediatesize hiddenactself hiddenact hiddendropoutprobself hiddendropoutprob attentionprobsdropoutprobself attentionprobsdropoutprob maxpositionembeddingsself maxpositionembeddings typevocabsizeself typevocabsize isdecoderfalse initializerrangeself initializerrange useentityawareattentionself useentityawareattention def createandcheckmodel self config inputids attentionmask tokentypeids entityids entityattentionmask entitytokentypeids entitypositionids sequencelabels tokenlabels choicelabels entitylabels entityclassificationlabels entitypairclassificationlabels entityspanclassificationlabels model lukemodelconfigconfig model totorchdevice model eval test with words entities result model inputids attentionmaskattentionmask tokentypeidstokentypeids entityidsentityids entityattentionmaskentityattentionmask entitytokentypeidsentitytokentypeids entitypositionidsentitypositionids self parent assertequalresult lasthiddenstate shape self batchsize self seqlength self hiddensize self parent assertequal result entitylasthiddenstate shape self batchsize self entitylength self hiddensize test with words only result modelinputids tokentypeidstokentypeids result modelinputids self parent assertequalresult lasthiddenstate shape self batchsize self seqlength self hiddensize def createandcheckformaskedlm self config inputids attentionmask tokentypeids entityids entityattentionmask entitytokentypeids entitypositionids sequencelabels tokenlabels choicelabels entitylabels entityclassificationlabels entitypairclassificationlabels entityspanclassificationlabels config numlabels self numentityclassificationlabels model lukeformaskedlmconfig model totorchdevice model eval result model inputids attentionmaskattentionmask tokentypeidstokentypeids entityidsentityids entityattentionmaskentityattentionmask entitytokentypeidsentitytokentypeids entitypositionidsentitypositionids labelstokenlabels entitylabelsentitylabels self parent assertequalresult logits shape self batchsize self seqlength self vocabsize if entityids is not none self parent assertequal result entitylogits shape self batchsize self entitylength self entityvocabsize else self parent assertisnoneresult entitylogits def createandcheckforentityclassification self config inputids attentionmask tokentypeids entityids entityattentionmask entitytokentypeids entitypositionids sequencelabels tokenlabels choicelabels entitylabels entityclassificationlabels entitypairclassificationlabels entityspanclassificationlabels config numlabels self numentityclassificationlabels model lukeforentityclassificationconfig model totorchdevice model eval result model inputids attentionmaskattentionmask tokentypeidstokentypeids entityidsentityids entityattentionmaskentityattentionmask entitytokentypeidsentitytokentypeids entitypositionidsentitypositionids labelsentityclassificationlabels self parent assertequalresult logits shape self batchsize self numentityclassificationlabels def createandcheckforentitypairclassification self config inputids attentionmask tokentypeids entityids entityattentionmask entitytokentypeids entitypositionids sequencelabels tokenlabels choicelabels entitylabels entityclassificationlabels entitypairclassificationlabels entityspanclassificationlabels config numlabels self numentitypairclassificationlabels model lukeforentityclassificationconfig model totorchdevice model eval result model inputids attentionmaskattentionmask tokentypeidstokentypeids entityidsentityids entityattentionmaskentityattentionmask entitytokentypeidsentitytokentypeids entitypositionidsentitypositionids labelsentitypairclassificationlabels self parent assertequalresult logits shape self batchsize self numentitypairclassificationlabels def createandcheckforentityspanclassification self config inputids attentionmask tokentypeids entityids entityattentionmask entitytokentypeids entitypositionids sequencelabels tokenlabels choicelabels entitylabels entityclassificationlabels entitypairclassificationlabels entityspanclassificationlabels config numlabels self numentityspanclassificationlabels model lukeforentityspanclassificationconfig model totorchdevice model eval entitystartpositions idstensorself batchsize self entitylength self seqlength entityendpositions idstensorself batchsize self entitylength self seqlength result model inputids attentionmaskattentionmask tokentypeidstokentypeids entityidsentityids entityattentionmaskentityattentionmask entitytokentypeidsentitytokentypeids entitypositionidsentitypositionids entitystartpositionsentitystartpositions entityendpositionsentityendpositions labelsentityspanclassificationlabels self parent assertequal result logits shape self batchsize self entitylength self numentityspanclassificationlabels def createandcheckforquestionanswering self config inputids attentionmask tokentypeids entityids entityattentionmask entitytokentypeids entitypositionids sequencelabels tokenlabels choicelabels entitylabels entityclassificationlabels entitypairclassificationlabels entityspanclassificationlabels model lukeforquestionansweringconfigconfig model totorchdevice model eval result model inputids attentionmaskattentionmask tokentypeidstokentypeids entityidsentityids entityattentionmaskentityattentionmask entitytokentypeidsentitytokentypeids entitypositionidsentitypositionids startpositionssequencelabels endpositionssequencelabels self parent assertequalresult startlogits shape self batchsize self seqlength self parent assertequalresult endlogits shape self batchsize self seqlength def createandcheckforsequenceclassification self config inputids attentionmask tokentypeids entityids entityattentionmask entitytokentypeids entitypositionids sequencelabels tokenlabels choicelabels entitylabels entityclassificationlabels entitypairclassificationlabels entityspanclassificationlabels config numlabels self numlabels model lukeforsequenceclassificationconfig model totorchdevice model eval result model inputids attentionmaskattentionmask tokentypeidstokentypeids entityidsentityids entityattentionmaskentityattentionmask entitytokentypeidsentitytokentypeids entitypositionidsentitypositionids labelssequencelabels self parent assertequalresult logits shape self batchsize self numlabels def createandcheckfortokenclassification self config inputids attentionmask tokentypeids entityids entityattentionmask entitytokentypeids entitypositionids sequencelabels tokenlabels choicelabels entitylabels entityclassificationlabels entitypairclassificationlabels entityspanclassificationlabels config numlabels self numlabels model lukefortokenclassificationconfigconfig model totorchdevice model eval result model inputids attentionmaskattentionmask tokentypeidstokentypeids entityidsentityids entityattentionmaskentityattentionmask entitytokentypeidsentitytokentypeids entitypositionidsentitypositionids labelstokenlabels self parent assertequalresult logits shape self batchsize self seqlength self numlabels def createandcheckformultiplechoice self config inputids attentionmask tokentypeids entityids entityattentionmask entitytokentypeids entitypositionids sequencelabels tokenlabels choicelabels entitylabels entityclassificationlabels entitypairclassificationlabels entityspanclassificationlabels config numchoices self numchoices model lukeformultiplechoiceconfigconfig model totorchdevice model eval multiplechoiceinputsids inputids unsqueeze1 expand1 self numchoices 1 contiguous multiplechoicetokentypeids tokentypeids unsqueeze1 expand1 self numchoices 1 contiguous multiplechoiceattentionmask attentionmask unsqueeze1 expand1 self numchoices 1 contiguous multiplechoiceentityids entityids unsqueeze1 expand1 self numchoices 1 contiguous multiplechoiceentitytokentypeids entitytokentypeids unsqueeze1 expand1 self numchoices 1 contiguous multiplechoiceentityattentionmask entityattentionmask unsqueeze1 expand1 self numchoices 1 contiguous multiplechoiceentitypositionids entitypositionids unsqueeze1 expand1 self numchoices 1 1 contiguous result model multiplechoiceinputsids attentionmaskmultiplechoiceattentionmask tokentypeidsmultiplechoicetokentypeids entityidsmultiplechoiceentityids entityattentionmaskmultiplechoiceentityattentionmask entitytokentypeidsmultiplechoiceentitytokentypeids entitypositionidsmultiplechoiceentitypositionids labelschoicelabels self parent assertequalresult logits shape self batchsize self numchoices def prepareconfigandinputsforcommonself configandinputs self prepareconfigandinputs config inputids attentionmask tokentypeids entityids entityattentionmask entitytokentypeids entitypositionids sequencelabels tokenlabels choicelabels entitylabels entityclassificationlabels entitypairclassificationlabels entityspanclassificationlabels configandinputs inputsdict inputids inputids tokentypeids tokentypeids attentionmask attentionmask entityids entityids entitytokentypeids entitytokentypeids entityattentionmask entityattentionmask entitypositionids entitypositionids return config inputsdict requiretorch class lukemodeltestmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses lukemodel lukeformaskedlm lukeforentityclassification lukeforentitypairclassification lukeforentityspanclassification lukeforquestionanswering lukeforsequenceclassification lukefortokenclassification lukeformultiplechoice if istorchavailable else pipelinemodelmapping featureextraction lukemodel fillmask lukeformaskedlm questionanswering lukeforquestionanswering textclassification lukeforsequenceclassification tokenclassification lukefortokenclassification zeroshot lukeforsequenceclassification if istorchavailable else testpruning false testtorchscript false testresizeembeddings true testheadmasking true todo fix the failed tests def ispipelinetesttoskip self pipelinetestcassename configclass modelarchitecture tokenizername processorname if pipelinetestcassename in qapipelinetests zeroshotclassificationpipelinetests return true return false def prepareforclassself inputsdict modelclass returnlabelsfalse entityinputsdict k v for k v in inputsdict items if k startswithentity inputsdict k v for k v in inputsdict items if not k startswithentity inputsdict super prepareforclassinputsdict modelclass returnlabelsreturnlabels if modelclass lukeformultiplechoice entityinputsdict k v unsqueeze1 expand1 self modeltester numchoices 1 contiguous if v ndim 2 else v unsqueeze1 expand1 self modeltester numchoices 1 1 contiguous for k v in entityinputsdict items inputsdict updateentityinputsdict if modelclass lukeforentityspanclassification inputsdictentitystartpositions torch zeros self modeltester batchsize self modeltester entitylength dtypetorch long devicetorchdevice inputsdictentityendpositions torch ones self modeltester batchsize self modeltester entitylength dtypetorch long devicetorchdevice if returnlabels if modelclass in lukeforentityclassification lukeforentitypairclassification lukeforsequenceclassification lukeformultiplechoice inputsdictlabels torch zeros self modeltester batchsize dtypetorch long devicetorchdevice elif modelclass lukeforentityspanclassification inputsdictlabels torch zeros self modeltester batchsize self modeltester entitylength dtypetorch long devicetorchdevice elif modelclass lukefortokenclassification inputsdictlabels torch zeros self modeltester batchsize self modeltester seqlength dtypetorch long devicetorchdevice elif modelclass lukeformaskedlm inputsdictlabels torch zeros self modeltester batchsize self modeltester seqlength dtypetorch long devicetorchdevice inputsdictentitylabels torch zeros self modeltester batchsize self modeltester entitylength dtypetorch long devicetorchdevice return inputsdict def setupself self modeltester lukemodeltesterself self configtester configtesterself configclasslukeconfig hiddensize37 def testconfigself self configtester runcommontests def testmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfigandinputs slow def testmodelfrompretrainedself for modelname in lukepretrainedmodelarchivelist model lukemodel frompretrainedmodelname self assertisnotnonemodel def testformaskedlmself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckformaskedlmconfigandinputs def testformaskedlmwithwordonlyself configandinputs self modeltester prepareconfigandinputs configandinputs configandinputs 4 none lenconfigandinputs4 self modeltester createandcheckformaskedlmconfigandinputs def testforquestionansweringself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckforquestionansweringconfigandinputs def testforsequenceclassificationself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckforsequenceclassificationconfigandinputs def testfortokenclassificationself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckfortokenclassificationconfigandinputs def testformultiplechoiceself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckformultiplechoiceconfigandinputs def testforentityclassificationself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckforentityclassificationconfigandinputs def testforentitypairclassificationself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckforentitypairclassificationconfigandinputs def testforentityspanclassificationself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckforentityspanclassificationconfigandinputs def testattentionoutputsself config inputsdict self modeltester prepareconfigandinputsforcommon config returndict true seqlength self modeltester seqlength entitylength self modeltester entitylength keylength seqlength entitylength for modelclass in self allmodelclasses inputsdictoutputattentions true inputsdictoutputhiddenstates false config returndict true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass attentions outputs attentions self assertequallenattentions self modeltester numhiddenlayers check that outputattentions also work using config del inputsdictoutputattentions config outputattentions true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass attentions outputs attentions self assertequallenattentions self modeltester numhiddenlayers self assertlistequal listattentions0 shape3 self modeltester numattentionheads seqlength entitylength keylength outlen lenoutputs check attention is always last and order is fine inputsdictoutputattentions true inputsdictoutputhiddenstates true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass addedhiddenstates self modeltester numhiddenstatestypes self assertequaloutlen addedhiddenstates lenoutputs selfattentions outputs attentions self assertequallenselfattentions self modeltester numhiddenlayers self assertlistequal listselfattentions0 shape3 self modeltester numattentionheads seqlength entitylength keylength def testentityhiddenstatesoutputself def checkhiddenstatesoutputinputsdict config modelclass model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass entityhiddenstates outputs entityhiddenstates expectednumlayers getattr self modeltester expectednumhiddenlayers self modeltester numhiddenlayers 1 self assertequallenentityhiddenstates expectednumlayers entitylength self modeltester entitylength self assertlistequal listentityhiddenstates0 shape2 entitylength self modeltester hiddensize config inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses inputsdictoutputhiddenstates true checkhiddenstatesoutputinputsdict config modelclass check that outputhiddenstates also work using config del inputsdictoutputhiddenstates config outputhiddenstates true checkhiddenstatesoutputinputsdict config modelclass def testretaingradentityhiddenstatesself config inputsdict self modeltester prepareconfigandinputsforcommon config outputhiddenstates true config outputattentions true no need to test all models as different heads yield the same functionality modelclass self allmodelclasses0 model modelclassconfig model totorchdevice inputs self prepareforclassinputsdict modelclass outputs modelinputs output outputs0 entityhiddenstates outputs entityhiddenstates0 entityhiddenstates retaingrad output flatten0 backwardretaingraphtrue self assertisnotnoneentityhiddenstates grad unittest skip reasonthis architecure seem to not compute gradients properly when using gc check https github comhuggingfacetransformerspull27124 def testtraininggradientcheckpointingself pass unittest skip reasonthis architecure seem to not compute gradients properly when using gc check https github comhuggingfacetransformerspull27124 def testtraininggradientcheckpointingusereentrantself pass unittest skip reasonthis architecure seem to not compute gradients properly when using gc check https github comhuggingfacetransformerspull27124 def testtraininggradientcheckpointingusereentrantfalseself pass requiretorch class lukemodelintegrationtestsunittest testcase slow def testinferencebasemodelself model lukemodel frompretrainedstudioousialukebase eval model totorchdevice tokenizer luketokenizer frompretrainedstudioousialukebase taskentityclassification text top seed ana ivanovic said on thursday she could hardly believe her luck as a fortuitous netcord helped the new world number one avoid a humiliating second round exit at wimbledon span 39 42 encoding tokenizertext entityspansspan addprefixspacetrue returntensorspt move all values to device for key value in encoding items encodingkey encodingkey totorchdevice outputs modelencoding verify word hidden states expectedshape torch size1 42 768 self assertequaloutputs lasthiddenstate shape expectedshape expectedslice torch tensor 0 0037 0 1368 0 0091 0 1099 0 3329 0 1095 0 0765 0 5335 0 1179 totorchdevice self asserttruetorch allcloseoutputs lasthiddenstate0 3 3 expectedslice atol1e4 verify entity hidden states expectedshape torch size1 1 768 self assertequaloutputs entitylasthiddenstate shape expectedshape expectedslice torch tensor0 1457 0 1044 0 0174 totorchdevice self asserttruetorch allcloseoutputs entitylasthiddenstate0 3 3 expectedslice atol1e4 slow def testinferencelargemodelself model lukemodel frompretrainedstudioousialukelarge eval model totorchdevice tokenizer luketokenizer frompretrainedstudioousialukelarge taskentityclassification text top seed ana ivanovic said on thursday she could hardly believe her luck as a fortuitous netcord helped the new world number one avoid a humiliating second round exit at wimbledon span 39 42 encoding tokenizertext entityspansspan addprefixspacetrue returntensorspt move all values to device for key value in encoding items encodingkey encodingkey totorchdevice outputs modelencoding verify word hidden states expectedshape torch size1 42 1024 self assertequaloutputs lasthiddenstate shape expectedshape expectedslice torch tensor 0 0133 0 0865 0 0095 0 3093 0 2576 0 7418 0 1720 0 2117 0 2869 totorchdevice self asserttruetorch allcloseoutputs lasthiddenstate0 3 3 expectedslice atol1e4 verify entity hidden states expectedshape torch size1 1 1024 self assertequaloutputs entitylasthiddenstate shape expectedshape expectedslice torch tensor0 0466 0 0106 0 0179 totorchdevice self asserttruetorch allcloseoutputs entitylasthiddenstate0 3 3 expectedslice atol1e4 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 luke model hidden_states and entity_hidden_states prepare words prepare entities test with words entities test with words only todo fix the failed tests check that output_attentions also work using config check attention is always last and order is fine check that output_hidden_states also work using config no need to test all models as different heads yield the same functionality move all values to device verify word hidden states verify entity hidden states move all values to device verify word hidden states verify entity hidden states
import unittest

from transformers import LukeConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

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 (
        LukeForEntityClassification,
        LukeForEntityPairClassification,
        LukeForEntitySpanClassification,
        LukeForMaskedLM,
        LukeForMultipleChoice,
        LukeForQuestionAnswering,
        LukeForSequenceClassification,
        LukeForTokenClassification,
        LukeModel,
        LukeTokenizer,
    )
    from transformers.models.luke.modeling_luke import LUKE_PRETRAINED_MODEL_ARCHIVE_LIST


class LukeModelTester:
    def __init__(
        self,
        parent,
        batch_size=13,
        seq_length=7,
        is_training=True,
        entity_length=3,
        mention_length=5,
        use_attention_mask=True,
        use_token_type_ids=True,
        use_entity_ids=True,
        use_entity_attention_mask=True,
        use_entity_token_type_ids=True,
        use_entity_position_ids=True,
        use_labels=True,
        vocab_size=99,
        entity_vocab_size=10,
        entity_emb_size=6,
        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,
        num_entity_classification_labels=9,
        num_entity_pair_classification_labels=6,
        num_entity_span_classification_labels=4,
        use_entity_aware_attention=True,
        scope=None,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        self.entity_length = entity_length
        self.mention_length = mention_length
        self.use_attention_mask = use_attention_mask
        self.use_token_type_ids = use_token_type_ids
        self.use_entity_ids = use_entity_ids
        self.use_entity_attention_mask = use_entity_attention_mask
        self.use_entity_token_type_ids = use_entity_token_type_ids
        self.use_entity_position_ids = use_entity_position_ids
        self.use_labels = use_labels
        self.vocab_size = vocab_size
        self.entity_vocab_size = entity_vocab_size
        self.entity_emb_size = entity_emb_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.num_entity_classification_labels = num_entity_classification_labels
        self.num_entity_pair_classification_labels = num_entity_pair_classification_labels
        self.num_entity_span_classification_labels = num_entity_span_classification_labels
        self.scope = scope
        self.use_entity_aware_attention = use_entity_aware_attention

        self.encoder_seq_length = seq_length
        self.key_length = seq_length
        self.num_hidden_states_types = 2  

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

        attention_mask = None
        if self.use_attention_mask:
            attention_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)

        
        entity_ids = ids_tensor([self.batch_size, self.entity_length], self.entity_vocab_size)

        entity_attention_mask = None
        if self.use_entity_attention_mask:
            entity_attention_mask = random_attention_mask([self.batch_size, self.entity_length])

        entity_token_type_ids = None
        if self.use_token_type_ids:
            entity_token_type_ids = ids_tensor([self.batch_size, self.entity_length], self.type_vocab_size)

        entity_position_ids = None
        if self.use_entity_position_ids:
            entity_position_ids = ids_tensor(
                [self.batch_size, self.entity_length, self.mention_length], self.mention_length
            )

        sequence_labels = None
        token_labels = None
        choice_labels = None
        entity_labels = None
        entity_classification_labels = None
        entity_pair_classification_labels = None
        entity_span_classification_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)

            entity_labels = ids_tensor([self.batch_size, self.entity_length], self.entity_vocab_size)

            entity_classification_labels = ids_tensor([self.batch_size], self.num_entity_classification_labels)
            entity_pair_classification_labels = ids_tensor(
                [self.batch_size], self.num_entity_pair_classification_labels
            )
            entity_span_classification_labels = ids_tensor(
                [self.batch_size, self.entity_length], self.num_entity_span_classification_labels
            )

        config = self.get_config()

        return (
            config,
            input_ids,
            attention_mask,
            token_type_ids,
            entity_ids,
            entity_attention_mask,
            entity_token_type_ids,
            entity_position_ids,
            sequence_labels,
            token_labels,
            choice_labels,
            entity_labels,
            entity_classification_labels,
            entity_pair_classification_labels,
            entity_span_classification_labels,
        )

    def get_config(self):
        return LukeConfig(
            vocab_size=self.vocab_size,
            entity_vocab_size=self.entity_vocab_size,
            entity_emb_size=self.entity_emb_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,
            use_entity_aware_attention=self.use_entity_aware_attention,
        )

    def create_and_check_model(
        self,
        config,
        input_ids,
        attention_mask,
        token_type_ids,
        entity_ids,
        entity_attention_mask,
        entity_token_type_ids,
        entity_position_ids,
        sequence_labels,
        token_labels,
        choice_labels,
        entity_labels,
        entity_classification_labels,
        entity_pair_classification_labels,
        entity_span_classification_labels,
    ):
        model = LukeModel(config=config)
        model.to(torch_device)
        model.eval()
        
        result = model(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            entity_ids=entity_ids,
            entity_attention_mask=entity_attention_mask,
            entity_token_type_ids=entity_token_type_ids,
            entity_position_ids=entity_position_ids,
        )
        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))
        self.parent.assertEqual(
            result.entity_last_hidden_state.shape, (self.batch_size, self.entity_length, self.hidden_size)
        )

        
        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_for_masked_lm(
        self,
        config,
        input_ids,
        attention_mask,
        token_type_ids,
        entity_ids,
        entity_attention_mask,
        entity_token_type_ids,
        entity_position_ids,
        sequence_labels,
        token_labels,
        choice_labels,
        entity_labels,
        entity_classification_labels,
        entity_pair_classification_labels,
        entity_span_classification_labels,
    ):
        config.num_labels = self.num_entity_classification_labels
        model = LukeForMaskedLM(config)
        model.to(torch_device)
        model.eval()

        result = model(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            entity_ids=entity_ids,
            entity_attention_mask=entity_attention_mask,
            entity_token_type_ids=entity_token_type_ids,
            entity_position_ids=entity_position_ids,
            labels=token_labels,
            entity_labels=entity_labels,
        )
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
        if entity_ids is not None:
            self.parent.assertEqual(
                result.entity_logits.shape, (self.batch_size, self.entity_length, self.entity_vocab_size)
            )
        else:
            self.parent.assertIsNone(result.entity_logits)

    def create_and_check_for_entity_classification(
        self,
        config,
        input_ids,
        attention_mask,
        token_type_ids,
        entity_ids,
        entity_attention_mask,
        entity_token_type_ids,
        entity_position_ids,
        sequence_labels,
        token_labels,
        choice_labels,
        entity_labels,
        entity_classification_labels,
        entity_pair_classification_labels,
        entity_span_classification_labels,
    ):
        config.num_labels = self.num_entity_classification_labels
        model = LukeForEntityClassification(config)
        model.to(torch_device)
        model.eval()

        result = model(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            entity_ids=entity_ids,
            entity_attention_mask=entity_attention_mask,
            entity_token_type_ids=entity_token_type_ids,
            entity_position_ids=entity_position_ids,
            labels=entity_classification_labels,
        )
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_entity_classification_labels))

    def create_and_check_for_entity_pair_classification(
        self,
        config,
        input_ids,
        attention_mask,
        token_type_ids,
        entity_ids,
        entity_attention_mask,
        entity_token_type_ids,
        entity_position_ids,
        sequence_labels,
        token_labels,
        choice_labels,
        entity_labels,
        entity_classification_labels,
        entity_pair_classification_labels,
        entity_span_classification_labels,
    ):
        config.num_labels = self.num_entity_pair_classification_labels
        model = LukeForEntityClassification(config)
        model.to(torch_device)
        model.eval()

        result = model(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            entity_ids=entity_ids,
            entity_attention_mask=entity_attention_mask,
            entity_token_type_ids=entity_token_type_ids,
            entity_position_ids=entity_position_ids,
            labels=entity_pair_classification_labels,
        )
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_entity_pair_classification_labels))

    def create_and_check_for_entity_span_classification(
        self,
        config,
        input_ids,
        attention_mask,
        token_type_ids,
        entity_ids,
        entity_attention_mask,
        entity_token_type_ids,
        entity_position_ids,
        sequence_labels,
        token_labels,
        choice_labels,
        entity_labels,
        entity_classification_labels,
        entity_pair_classification_labels,
        entity_span_classification_labels,
    ):
        config.num_labels = self.num_entity_span_classification_labels
        model = LukeForEntitySpanClassification(config)
        model.to(torch_device)
        model.eval()

        entity_start_positions = ids_tensor([self.batch_size, self.entity_length], self.seq_length)
        entity_end_positions = ids_tensor([self.batch_size, self.entity_length], self.seq_length)

        result = model(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            entity_ids=entity_ids,
            entity_attention_mask=entity_attention_mask,
            entity_token_type_ids=entity_token_type_ids,
            entity_position_ids=entity_position_ids,
            entity_start_positions=entity_start_positions,
            entity_end_positions=entity_end_positions,
            labels=entity_span_classification_labels,
        )
        self.parent.assertEqual(
            result.logits.shape, (self.batch_size, self.entity_length, self.num_entity_span_classification_labels)
        )

    def create_and_check_for_question_answering(
        self,
        config,
        input_ids,
        attention_mask,
        token_type_ids,
        entity_ids,
        entity_attention_mask,
        entity_token_type_ids,
        entity_position_ids,
        sequence_labels,
        token_labels,
        choice_labels,
        entity_labels,
        entity_classification_labels,
        entity_pair_classification_labels,
        entity_span_classification_labels,
    ):
        model = LukeForQuestionAnswering(config=config)
        model.to(torch_device)
        model.eval()
        result = model(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            entity_ids=entity_ids,
            entity_attention_mask=entity_attention_mask,
            entity_token_type_ids=entity_token_type_ids,
            entity_position_ids=entity_position_ids,
            start_positions=sequence_labels,
            end_positions=sequence_labels,
        )
        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_for_sequence_classification(
        self,
        config,
        input_ids,
        attention_mask,
        token_type_ids,
        entity_ids,
        entity_attention_mask,
        entity_token_type_ids,
        entity_position_ids,
        sequence_labels,
        token_labels,
        choice_labels,
        entity_labels,
        entity_classification_labels,
        entity_pair_classification_labels,
        entity_span_classification_labels,
    ):
        config.num_labels = self.num_labels
        model = LukeForSequenceClassification(config)
        model.to(torch_device)
        model.eval()
        result = model(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            entity_ids=entity_ids,
            entity_attention_mask=entity_attention_mask,
            entity_token_type_ids=entity_token_type_ids,
            entity_position_ids=entity_position_ids,
            labels=sequence_labels,
        )
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_labels))

    def create_and_check_for_token_classification(
        self,
        config,
        input_ids,
        attention_mask,
        token_type_ids,
        entity_ids,
        entity_attention_mask,
        entity_token_type_ids,
        entity_position_ids,
        sequence_labels,
        token_labels,
        choice_labels,
        entity_labels,
        entity_classification_labels,
        entity_pair_classification_labels,
        entity_span_classification_labels,
    ):
        config.num_labels = self.num_labels
        model = LukeForTokenClassification(config=config)
        model.to(torch_device)
        model.eval()
        result = model(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            entity_ids=entity_ids,
            entity_attention_mask=entity_attention_mask,
            entity_token_type_ids=entity_token_type_ids,
            entity_position_ids=entity_position_ids,
            labels=token_labels,
        )
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.num_labels))

    def create_and_check_for_multiple_choice(
        self,
        config,
        input_ids,
        attention_mask,
        token_type_ids,
        entity_ids,
        entity_attention_mask,
        entity_token_type_ids,
        entity_position_ids,
        sequence_labels,
        token_labels,
        choice_labels,
        entity_labels,
        entity_classification_labels,
        entity_pair_classification_labels,
        entity_span_classification_labels,
    ):
        config.num_choices = self.num_choices
        model = LukeForMultipleChoice(config=config)
        model.to(torch_device)
        model.eval()
        multiple_choice_inputs_ids = input_ids.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
        multiple_choice_token_type_ids = token_type_ids.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
        multiple_choice_attention_mask = attention_mask.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
        multiple_choice_entity_ids = entity_ids.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
        multiple_choice_entity_token_type_ids = (
            entity_token_type_ids.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
        )
        multiple_choice_entity_attention_mask = (
            entity_attention_mask.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
        )
        multiple_choice_entity_position_ids = (
            entity_position_ids.unsqueeze(1).expand(-1, self.num_choices, -1, -1).contiguous()
        )
        result = model(
            multiple_choice_inputs_ids,
            attention_mask=multiple_choice_attention_mask,
            token_type_ids=multiple_choice_token_type_ids,
            entity_ids=multiple_choice_entity_ids,
            entity_attention_mask=multiple_choice_entity_attention_mask,
            entity_token_type_ids=multiple_choice_entity_token_type_ids,
            entity_position_ids=multiple_choice_entity_position_ids,
            labels=choice_labels,
        )
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_choices))

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        (
            config,
            input_ids,
            attention_mask,
            token_type_ids,
            entity_ids,
            entity_attention_mask,
            entity_token_type_ids,
            entity_position_ids,
            sequence_labels,
            token_labels,
            choice_labels,
            entity_labels,
            entity_classification_labels,
            entity_pair_classification_labels,
            entity_span_classification_labels,
        ) = config_and_inputs
        inputs_dict = {
            "input_ids": input_ids,
            "token_type_ids": token_type_ids,
            "attention_mask": attention_mask,
            "entity_ids": entity_ids,
            "entity_token_type_ids": entity_token_type_ids,
            "entity_attention_mask": entity_attention_mask,
            "entity_position_ids": entity_position_ids,
        }
        return config, inputs_dict


@require_torch
class LukeModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (
        (
            LukeModel,
            LukeForMaskedLM,
            LukeForEntityClassification,
            LukeForEntityPairClassification,
            LukeForEntitySpanClassification,
            LukeForQuestionAnswering,
            LukeForSequenceClassification,
            LukeForTokenClassification,
            LukeForMultipleChoice,
        )
        if is_torch_available()
        else ()
    )
    pipeline_model_mapping = (
        {
            "feature-extraction": LukeModel,
            "fill-mask": LukeForMaskedLM,
            "question-answering": LukeForQuestionAnswering,
            "text-classification": LukeForSequenceClassification,
            "token-classification": LukeForTokenClassification,
            "zero-shot": LukeForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    test_pruning = False
    test_torchscript = False
    test_resize_embeddings = True
    test_head_masking = True

    
    def is_pipeline_test_to_skip(
        self, pipeline_test_casse_name, config_class, model_architecture, tokenizer_name, processor_name
    ):
        if pipeline_test_casse_name in ["QAPipelineTests", "ZeroShotClassificationPipelineTests"]:
            return True

        return False

    def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
        entity_inputs_dict = {k: v for k, v in inputs_dict.items() if k.startswith("entity")}
        inputs_dict = {k: v for k, v in inputs_dict.items() if not k.startswith("entity")}

        inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels)
        if model_class == LukeForMultipleChoice:
            entity_inputs_dict = {
                k: v.unsqueeze(1).expand(-1, self.model_tester.num_choices, -1).contiguous()
                if v.ndim == 2
                else v.unsqueeze(1).expand(-1, self.model_tester.num_choices, -1, -1).contiguous()
                for k, v in entity_inputs_dict.items()
            }
        inputs_dict.update(entity_inputs_dict)

        if model_class == LukeForEntitySpanClassification:
            inputs_dict["entity_start_positions"] = torch.zeros(
                (self.model_tester.batch_size, self.model_tester.entity_length), dtype=torch.long, device=torch_device
            )
            inputs_dict["entity_end_positions"] = torch.ones(
                (self.model_tester.batch_size, self.model_tester.entity_length), dtype=torch.long, device=torch_device
            )

        if return_labels:
            if model_class in (
                LukeForEntityClassification,
                LukeForEntityPairClassification,
                LukeForSequenceClassification,
                LukeForMultipleChoice,
            ):
                inputs_dict["labels"] = torch.zeros(
                    self.model_tester.batch_size, dtype=torch.long, device=torch_device
                )
            elif model_class == LukeForEntitySpanClassification:
                inputs_dict["labels"] = torch.zeros(
                    (self.model_tester.batch_size, self.model_tester.entity_length),
                    dtype=torch.long,
                    device=torch_device,
                )
            elif model_class == LukeForTokenClassification:
                inputs_dict["labels"] = torch.zeros(
                    (self.model_tester.batch_size, self.model_tester.seq_length),
                    dtype=torch.long,
                    device=torch_device,
                )
            elif model_class == LukeForMaskedLM:
                inputs_dict["labels"] = torch.zeros(
                    (self.model_tester.batch_size, self.model_tester.seq_length),
                    dtype=torch.long,
                    device=torch_device,
                )
                inputs_dict["entity_labels"] = torch.zeros(
                    (self.model_tester.batch_size, self.model_tester.entity_length),
                    dtype=torch.long,
                    device=torch_device,
                )

        return inputs_dict

    def setUp(self):
        self.model_tester = LukeModelTester(self)
        self.config_tester = ConfigTester(self, config_class=LukeConfig, 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)

    @slow
    def test_model_from_pretrained(self):
        for model_name in LUKE_PRETRAINED_MODEL_ARCHIVE_LIST:
            model = LukeModel.from_pretrained(model_name)
            self.assertIsNotNone(model)

    def test_for_masked_lm(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*config_and_inputs)

    def test_for_masked_lm_with_word_only(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        config_and_inputs = (*config_and_inputs[:4], *((None,) * len(config_and_inputs[4:])))
        self.model_tester.create_and_check_for_masked_lm(*config_and_inputs)

    def test_for_question_answering(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*config_and_inputs)

    def test_for_sequence_classification(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*config_and_inputs)

    def test_for_token_classification(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*config_and_inputs)

    def test_for_multiple_choice(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*config_and_inputs)

    def test_for_entity_classification(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_entity_classification(*config_and_inputs)

    def test_for_entity_pair_classification(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_entity_pair_classification(*config_and_inputs)

    def test_for_entity_span_classification(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_entity_span_classification(*config_and_inputs)

    def test_attention_outputs(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.return_dict = True

        seq_length = self.model_tester.seq_length
        entity_length = self.model_tester.entity_length
        key_length = seq_length + entity_length

        for model_class in self.all_model_classes:
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = False
            config.return_dict = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))
            attentions = outputs.attentions
            self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)

            
            del inputs_dict["output_attentions"]
            config.output_attentions = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))
            attentions = outputs.attentions
            self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)

            self.assertListEqual(
                list(attentions[0].shape[-3:]),
                [self.model_tester.num_attention_heads, seq_length + entity_length, key_length],
            )
            out_len = len(outputs)

            
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))

            added_hidden_states = self.model_tester.num_hidden_states_types
            self.assertEqual(out_len + added_hidden_states, len(outputs))

            self_attentions = outputs.attentions

            self.assertEqual(len(self_attentions), self.model_tester.num_hidden_layers)
            self.assertListEqual(
                list(self_attentions[0].shape[-3:]),
                [self.model_tester.num_attention_heads, seq_length + entity_length, key_length],
            )

    def test_entity_hidden_states_output(self):
        def check_hidden_states_output(inputs_dict, config, model_class):
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))

            entity_hidden_states = outputs.entity_hidden_states

            expected_num_layers = getattr(
                self.model_tester, "expected_num_hidden_layers", self.model_tester.num_hidden_layers + 1
            )
            self.assertEqual(len(entity_hidden_states), expected_num_layers)

            entity_length = self.model_tester.entity_length

            self.assertListEqual(
                list(entity_hidden_states[0].shape[-2:]),
                [entity_length, self.model_tester.hidden_size],
            )

        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            inputs_dict["output_hidden_states"] = True
            check_hidden_states_output(inputs_dict, config, model_class)

            
            del inputs_dict["output_hidden_states"]
            config.output_hidden_states = True

            check_hidden_states_output(inputs_dict, config, model_class)

    def test_retain_grad_entity_hidden_states(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.output_hidden_states = True
        config.output_attentions = True

        
        model_class = self.all_model_classes[0]
        model = model_class(config)
        model.to(torch_device)

        inputs = self._prepare_for_class(inputs_dict, model_class)

        outputs = model(**inputs)

        output = outputs[0]

        entity_hidden_states = outputs.entity_hidden_states[0]
        entity_hidden_states.retain_grad()

        output.flatten()[0].backward(retain_graph=True)

        self.assertIsNotNone(entity_hidden_states.grad)

    @unittest.skip(
        reason="This architecure seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
    )
    def test_training_gradient_checkpointing(self):
        pass

    @unittest.skip(
        reason="This architecure seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
    )
    def test_training_gradient_checkpointing_use_reentrant(self):
        pass

    @unittest.skip(
        reason="This architecure seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
    )
    def test_training_gradient_checkpointing_use_reentrant_false(self):
        pass


@require_torch
class LukeModelIntegrationTests(unittest.TestCase):
    @slow
    def test_inference_base_model(self):
        model = LukeModel.from_pretrained("studio-ousia/luke-base").eval()
        model.to(torch_device)

        tokenizer = LukeTokenizer.from_pretrained("studio-ousia/luke-base", task="entity_classification")
        text = (
            "Top seed Ana Ivanovic said on Thursday she could hardly believe her luck as a fortuitous netcord helped"
            " the new world number one avoid a humiliating second- round exit at Wimbledon ."
        )
        span = (39, 42)
        encoding = tokenizer(text, entity_spans=[span], add_prefix_space=True, return_tensors="pt")

        
        for key, value in encoding.items():
            encoding[key] = encoding[key].to(torch_device)

        outputs = model(**encoding)

        
        expected_shape = torch.Size((1, 42, 768))
        self.assertEqual(outputs.last_hidden_state.shape, expected_shape)

        expected_slice = torch.tensor(
            [[0.0037, 0.1368, -0.0091], [0.1099, 0.3329, -0.1095], [0.0765, 0.5335, 0.1179]]
        ).to(torch_device)
        self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3], expected_slice, atol=1e-4))

        
        expected_shape = torch.Size((1, 1, 768))
        self.assertEqual(outputs.entity_last_hidden_state.shape, expected_shape)

        expected_slice = torch.tensor([[0.1457, 0.1044, 0.0174]]).to(torch_device)
        self.assertTrue(torch.allclose(outputs.entity_last_hidden_state[0, :3, :3], expected_slice, atol=1e-4))

    @slow
    def test_inference_large_model(self):
        model = LukeModel.from_pretrained("studio-ousia/luke-large").eval()
        model.to(torch_device)

        tokenizer = LukeTokenizer.from_pretrained("studio-ousia/luke-large", task="entity_classification")
        text = (
            "Top seed Ana Ivanovic said on Thursday she could hardly believe her luck as a fortuitous netcord helped"
            " the new world number one avoid a humiliating second- round exit at Wimbledon ."
        )
        span = (39, 42)
        encoding = tokenizer(text, entity_spans=[span], add_prefix_space=True, return_tensors="pt")

        
        for key, value in encoding.items():
            encoding[key] = encoding[key].to(torch_device)

        outputs = model(**encoding)

        
        expected_shape = torch.Size((1, 42, 1024))
        self.assertEqual(outputs.last_hidden_state.shape, expected_shape)

        expected_slice = torch.tensor(
            [[0.0133, 0.0865, 0.0095], [0.3093, -0.2576, -0.7418], [-0.1720, -0.2117, -0.2869]]
        ).to(torch_device)
        self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3], expected_slice, atol=1e-4))

        
        expected_shape = torch.Size((1, 1, 1024))
        self.assertEqual(outputs.entity_last_hidden_state.shape, expected_shape)

        expected_slice = torch.tensor([[0.0466, -0.0106, -0.0179]]).to(torch_device)
        self.assertTrue(torch.allclose(outputs.entity_last_hidden_state[0, :3, :3], expected_slice, atol=1e-4))