codingutf8 2023 toshiyuki sakamototanreinama and 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 pytorch gptsanjapanese model import copy from typing import list optional tuple union import torch import torch nn as nn from activations import act2fn from modelingoutputs import moecausallmoutputwithpast moemodeloutputwithpastandcrossattentions from modelingutils import pretrainedmodel from utils import dummyinputs dummymask addstartdocstrings addstartdocstringstomodelforward istorchfxproxy logging from configurationgptsanjapanese import gptsanjapaneseconfig logger logging getloggername configfordoc gptsanjapaneseconfig checkpointfordoc tanreigptsanjapanese this dict contains ids and associated url for the pretrained weights provided with the models gptsanjapanesepretrainedmodelarchivelist tanreigptsanjapanese see all gptsanjapanese models at https huggingface comodels filtergptsanjapanese copied from transformers models switchtransformers modelingswitchtransformers routerzlossfunc def routerzlossfuncrouterlogits torch tensor float r compute the router zloss implemented in pytorch the router zloss was introduced in designing effective sparse expert modelshttps arxiv orgabs2202 08906 it encourages router logits to remain small in an effort to improve stability args routerlogits float input logits of shape batchsize sequencelength numexperts returns scalar router zloss copied from transformers models switchtransformers modelingswitchtransformers loadbalancinglossfunc numexperts routerprobs shape1 cast the expert indices to int64 otherwise onehot encoding will fail if expertindices dtype torch int64 expertindices expertindices totorch int64 if lenexpertindices shape 2 expertindices expertindices unsqueeze2 expertmask torch nn functional onehotexpertindices numexperts for a given token determine if it was routed to a given expert expertmask torch maxexpertmask axis2 values cast to float32 otherwise mean will fail expertmask expertmask totorch float32 tokenspergroupandexpert torch meanexpertmask axis2 routerprobpergroupandexpert torch meanrouterprobs axis2 return torch meantokenspergroupandexpert routerprobpergroupandexpert numexperts2 class gptsanjapanesedenseactdensenn module def initself config gptsanjapaneseconfig extlayerfalse super init dinter config dext if extlayer else config dff self wi nn linearconfig dmodel dinter biasextlayer self wo nn lineardinter config dmodel biasextlayer self dropout nn identity if extlayer else nn dropoutconfig dropoutrate self act act2fnswish if extlayer else relu def forwardself hiddenstates r args hiddenstates torch tensor numgroups tokenspergroup hiddendim inputs to send to experts returns torch tensornumgroups tokenspergroup hiddendim copied from transformers models switchtransformers modelingswitchtransformers switchtransformerstop1router with switchtransformersgptsanjapanese router using tokens choose top1 experts assignment this router uses the same mechanism as in switch transformer https arxiv orgabs2101 03961 and vmoe https arxiv orgabs2106 05974 tokens choose their top experts items are sorted by routerprobs and then routed to their choice of expert until the expert s expertcapacity is reached there is no guarantee that each token is processed by an expert or that each expert receives at least one token float32 is used to ensure stability see the discussion of selective precision in https arxiv orgabs2101 03961 we also store the previous dtype to cast back the output to the previous dtype self inputdtype hiddenstates dtype hiddenstates hiddenstates toself dtype if self jitternoise 0 multiply the token inputs by the uniform distribution adding some noise hiddenstates torch emptylikehiddenstates uniform1 0 self jitternoise 1 0 self jitternoise shape numgroups tokenspergroup numexperts self castclassifier routerlogits self classifierhiddenstates apply softmax and cast back to the original dtype routerprobabilities nn functional softmaxrouterlogits dim1 dtypeself dtype toself inputdtype return routerprobabilities routerlogits def castclassifierself r bitsandbytes linear8bitlt layers does not support manual casting therefore we need to check if they are an instance of the linear8bitlt class by checking special attributes routerprobs routerlogits self computerouterprobabilitieshiddenstates expertindex torch argmaxrouterprobs dim1 expertindex torch nn functional onehotexpertindex numclassesself numexperts mask tokens outside expert capacity sum over each sequence tokenpriority torch cumsumexpertindex dim2 mask if the token routed to to the expert will overflow expertcapacitymask tokenpriority self expertcapacity expertindex expertindex expertcapacitymask routerprobs torch maxrouterprobs dim1 values unsqueeze1 return expertindex routerprobs routerlogits copied from transformers models switchtransformers modelingswitchtransformers switchtransformerssparsemlp with switchtransformersgptsanjapanese class gptsanjapanesesparsemlpnn module r implementation of the switch transformers sparse mlp module step 1 get the correct router according to its class step 2 get the experts step 1 get the routermask from the router as wel as the probabilities routermask routerprobs routerlogits self routerhiddenstates expertindex torch argmaxroutermask dim1 the routers introduced might not always map all the tokens to a router which means that some hidden states can be unchanged from one layer to another that is why the hidden states are cloned before updating only the seleced ones nextstates hiddenstates clone for idx expert in enumerateself experts values tokenindices routermask idx bool nextstatestokenindices experthiddenstatestokenindices tonextstates dtype hiddenstates routerprobs nextstates return hiddenstates routerlogits expertindex class gptsanjapaneselayersparseffnn module r switch transformers feed forward layer module this is a wrapper around the mixture of experts module parameters config gptsanjapaneseconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel frompretrained method to load the model weights forwardedstates routertuple self mlphiddenstates forwardedstates torch tanhself softbypassmlphiddenstates output hiddenstates self normforwardedstates if outputrouterlogits and routertuple is not none return output routertuple else return output class gptsanjapaneselayerdenseffnn module r extra transformers feed forward layer module parameters config gptsanjapaneseconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel frompretrained method to load the model weights check if it is a sparse layer if not then it is a dense layer forwardedstates self mlphiddenstates output hiddenstates self normforwardedstates return output copied from transformers models bart modelingbart bartattention with bartgptsanjapanese class gptsanjapaneseattentionnn module input shape batch x time x channel if keyvaluestates are provided this layer is used as a crossattention layer for the decoder iscrossattention keyvaluestates is not none bsz tgtlen hiddenstates size get query proj querystates self qprojhiddenstates self scaling get key value proj pastkeyvalue0 shape2 keyvaluestates shape1 is checking that the sequencelength of the pastkeyvalue is the same as the provided keyvaluestates to support prefix tuning if iscrossattention and pastkeyvalue is not none and pastkeyvalue0 shape2 keyvaluestates shape1 reuse k v crossattentions keystates pastkeyvalue0 valuestates pastkeyvalue1 elif iscrossattention crossattentions keystates self shapeself kprojkeyvaluestates 1 bsz valuestates self shapeself vprojkeyvaluestates 1 bsz elif pastkeyvalue is not none reuse k v selfattention keystates self shapeself kprojhiddenstates 1 bsz valuestates self shapeself vprojhiddenstates 1 bsz keystates torch catpastkeyvalue0 keystates dim2 valuestates torch catpastkeyvalue1 valuestates dim2 else selfattention keystates self shapeself kprojhiddenstates 1 bsz valuestates self shapeself vprojhiddenstates 1 bsz if self isdecoder if crossattention save tupletorch tensor torch tensor of all cross attention keyvaluestates further calls to crossattention layer can then reuse all crossattention keyvaluestates first if case if unidirectional selfattention decoder save tupletorch tensor torch tensor of all previous decoder keyvaluestates further calls to unidirectional selfattention can concat previous decoder keyvaluestates to current projected keyvaluestates third elif case if encoder bidirectional selfattention pastkeyvalue is always none pastkeyvalue keystates valuestates projshape bsz self numheads 1 self headdim querystates self shapequerystates tgtlen bsz viewprojshape keystates keystates reshapeprojshape valuestates valuestates reshapeprojshape srclen keystates size1 attnweights torch bmmquerystates keystates transpose1 2 if attnweights size bsz self numheads tgtlen srclen raise valueerror fattention weights should be of size bsz self numheads tgtlen srclen but is f attnweights size if attentionmask is not none if attentionmask size bsz 1 tgtlen srclen raise valueerror fattention mask should be of size bsz 1 tgtlen srclen but is attentionmask size attnweights attnweights viewbsz self numheads tgtlen srclen attentionmask attnweights attnweights viewbsz self numheads tgtlen srclen attnweights nn functional softmaxattnweights dim1 if layerheadmask is not none if layerheadmask size self numheads raise valueerror fhead mask for a single layer should be of size self numheads but is f layerheadmask size attnweights layerheadmask view1 1 1 1 attnweights viewbsz self numheads tgtlen srclen attnweights attnweights viewbsz self numheads tgtlen srclen if outputattentions this operation is a bit awkward but it s required to make sure that attnweights keeps its gradient in order to do so attnweights have to be reshaped twice and have to be reused in the following attnweightsreshaped attnweights viewbsz self numheads tgtlen srclen attnweights attnweightsreshaped viewbsz self numheads tgtlen srclen else attnweightsreshaped none attnprobs nn functional dropoutattnweights pself dropout trainingself training attnoutput torch bmmattnprobs valuestates if attnoutput size bsz self numheads tgtlen self headdim raise valueerror fattnoutput should be of size bsz self numheads tgtlen self headdim but is f attnoutput size attnoutput attnoutput viewbsz self numheads tgtlen self headdim attnoutput attnoutput transpose1 2 use the embeddim from the config stored in the class rather than hiddenstate because attnoutput can be partitioned across gpus when using tensorparallelism attnoutput attnoutput reshapebsz tgtlen self embeddim attnoutput self outprojattnoutput return attnoutput attnweightsreshaped pastkeyvalue class gptsanjapaneselayerselfattentionnn module def initself config hasrelativeattentionbiasfalse super init self selfattn gptsanjapaneseattention embeddimconfig dmodel numheadsconfig numheads isdecodertrue biashasrelativeattentionbias self norm nn layernormconfig dmodel epsconfig layernormepsilon def forward self hiddenstates optionaltupletorch floattensor pastkeyvalue optionaltupletorch tensor none attentionmask optionaltorch floattensor none headmask optionaltorch floattensor none usecache optionalbool false outputattentions optionalbool false tupleuniontorch tensor tupletorch tensor r selfattention and normalize block args hiddenstates torch floattensor of shape batchsize sequencelength hiddensize optional sequence of hiddenstates at the output of the last layer of the encoder used in the crossattention if the model is configured as a decoder pastkeyvalues tupletupletorch floattensor of length config nlayers with each tuple having 4 tensors of shape batchsize numheads sequencelength 1 embedsizeperhead contains precomputed key and value hidden states of the attention blocks can be used to speed up decoding if pastkeyvalues are used the user can optionally input only the last decoderinputids those that don t have their past key value states given to this model of shape batchsize 1 instead of all decoderinputids of shape batchsize sequencelength attentionmask torch floattensor of shape batchsize sequencelength optional mask to avoid performing attention on the padding token indices of the encoder input this mask is used in the crossattention if the model is configured as a decoder mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked headmask numpy ndarray of shape 0 optional mask to nullify selected heads of the attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked usecache bool optional if set to true pastkeyvalues key value states are returned and can be used to speed up decoding see pastkeyvalues outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail returns tupletorch tensornumgroups tokenspergroup hiddendim self attention decoder unidirectional selfattention cached keyvalues tuple is at positions 1 2 add present selfattn cache to positions 1 2 of presentkeyvalue tuple self attention and ffn unit attenout self selfattn hiddenstateshiddenstates pastkeyvaluepastkeyvalue attentionmaskattentionmask headmaskheadmask usecacheusecache outputattentionsoutputattentions attentionoutput attenout0 if isinstanceself feedforward gptsanjapaneselayersparseff sparseout self feedforwardattentionoutput outputroutertuple if outputroutertuple hidden routertuple sparseout else hidden sparseout else hidden self feedforwardattentionoutput outputs hidden attenout1 if isinstanceself feedforward gptsanjapaneselayersparseff and outputroutertuple outputs routertuple return outputs class gptsanjapanesepretrainedmodelpretrainedmodel configclass gptsanjapaneseconfig basemodelprefix gptsanjapanese supportsgradientcheckpointing false nosplitmodules gptsanjapaneseblock skipkeysdeviceplacement pastkeyvalues property def dummyinputsself inputids torch tensordummyinputs inputmask torch tensordummymask dummyinputs inputids inputids attentionmask inputmask return dummyinputs def initweightsself module mesh tensorflow embeddings initialization see https github comtensorflowmeshblobfa19d69eafc9a482aff0b59ddd96b025c0cb207dmeshtensorflowlayers pyl1624 mesh tensorflow embeddings initialization see https github comtensorflowmeshblobfa19d69eafc9a482aff0b59ddd96b025c0cb207dmeshtensorflowlayers pyl1624 mesh tensorflow ff initialization see https github comtensorflowmeshblobmastermeshtensorflowtransformertransformerlayers pyl56 and https github comtensorflowmeshblobfa19d69eafc9a482aff0b59ddd96b025c0cb207dmeshtensorflowlayers pyl89 multiheaded attention mesh tensorflow attention initialization to avoid scaling before softmax see https github comtensorflowmeshblobfa19d69eafc9a482aff0b59ddd96b025c0cb207dmeshtensorflowtransformerattention pyl136 copied from transformers models t5 modelingt5 t5pretrainedmodel shiftright shift inputs to the right item assignment is not supported natively for proxies replace possible 100 values in labels by padtokenid gptsanjapaneseinputsdocstring r args inputids torch longtensor of shape batchsize sequencelength indices of input sequence tokens in the vocabulary gptsanjapanese is a model that generates sentence continuations or predicts tokens at mask positions special tokens required for inputs to the model are automatically appended attentionmask torch floattensor of shape batchsize sequencelength optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask tokentypeids torch floattensor of shape batchsize sequencelength optional an input that masks the prefix part in the prefixlm input mask values selected in 0 1 1 for tokens that are prefix input 0 for tokens that are notprefix input spout torch tensor of shape batchsize config dspout this vector is transformed through an 8layer ffn and can be used instead of pastkeyvalues pastkeyvalues tupletupletorch floattensor of length config nlayers with each tuple having 4 tensors of shape batchsize numheads sequencelength 1 embedsizeperhead contains precomputed key and value hidden states of the attention blocks can be used to speed up decoding if pastkeyvalues are used the user can optionally input only the last decoderinputids those that don t have their past key value states given to this model of shape batchsize 1 instead of all decoderinputids of shape batchsize sequencelength headmask torch floattensor of shape numheads or numlayers numheads optional mask to nullify selected heads of the selfattention modules mask values selected in 0 1 usecache bool optional if set to true pastkeyvalues key value states are returned and can be used to speed up decoding see pastkeyvalues inputsembeds torch floattensor of shape batchsize sequencelength hiddensize optional optionally instead of passing inputids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert inputids indices into associated vectors than the model s internal embedding lookup matrix decoderinputsembeds torch floattensor of shape batchsize targetsequencelength hiddensize optional optionally instead of passing decoderinputids you can choose to directly pass an embedded representation if pastkeyvalues is used optionally only the last decoderinputsembeds have to be input see pastkeyvalues this is useful if you want more control over how to convert decoderinputids indices into associated vectors than the model s internal embedding lookup matrix outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple routerlogits tupletorch floattensor optional returned when outputrouterlogitstrue is passed or when config addrouterprobstrue tuple of torch floattensor one for each layer of shape batchsize sequencelength numexperts router logits of the decoder model useful to compute the auxiliary loss for mixture of experts models returndict returndict if returndict is not none else self config usereturndict device self positionembeddings weight device if inputids is none inputids torch zeros1 1 int todevice dummy for inputids was none numpastscontexts 0 numbatch inputids shape0 pastsorspoutvalue none if pastkeyvalues is not none numpastscontexts pastkeyvalues00 shape2 elif self config dspout and spout is not none spout is a special input vector specific to gptsan this controls the output by projecting embedded information such as the class of sentences during learning it should passed instead of the first pastkeyvalue see the original gptsan repository for details numpastscontexts 1 if there is an attentionmask increase first one for spout if self config dspout and spout is not none and attentionmask is not none attentionmaskwithspout torch onesnumbatch attentionmask shape1 1 devicedevice attentionmaskwithspout 1 1 attentionmask 1st token should be spout attentionmask attentionmaskwithspout update attentionmask if numprecontext is not none numprecontext is the number of tokens that refer to each other in prefixlm created per batch so dimension of numprecontext should be batch 1 if not lennumprecontext shape 2 and numprecontext shape1 1 numprecontext should be batch 1 raise valueerrornumprecontext should be batch 1 size numprecontext torch reshapenumprecontext 1 else numprecontext torch zerosnumbatch int todevice numinputcontexts inputids shape1 numoutputcontexts numinputcontexts numpastscontexts hiddenstates self embedtokensinputids if pastkeyvalues is not none pastsorspoutvalue pastkeyvalues elif self config dspout and spout is not none make vector from spout of gptsan to the same shape as pastkeyvalues pastsorspoutvalue self spoutspout projecting spout vector pastsorspoutvalue torch reshape pastsorspoutvalue numbatch self config numlayers 2 self config numheads numpastscontexts self config dmodel self config numheads pastsorspoutvalue torch splitpastsorspoutvalue 1 self config numlayers dim1 make same shape as pastkeyvalues pastsorspoutvalue tuple tupleb squeeze1 for b in torch splita squeeze1 1 1 dim1 for a in pastsorspoutvalue else pastsorspoutvalue none self config numlayers token position considering spout and pasts tokenposition torch arangenuminputcontexts todevice numpastscontexts if attentionmask is none attentionmask torch onesnumbatch numinputcontexts devicedevice positions for get positionembeddings gatherposition torch zerosnumbatch self config dmodel numinputcontexts todevice tokenposition unsqueeze0 transpose1 2 long when padding with paddingsideleft zeros line up on the left side of attentionmask so positionembeddings is shifted accordingly gatherposition 1 attentionmask argmindim1 unsqueeze1 unsqueeze2 gatherposition torch clipgatherposition numpastscontexts self config maxpositionembeddings 1 attentionmask is applied per batch for i in rangenumbatch hiddenstatesi torch gatherself positionembeddings weight dim0 indexgatherpositioni create a mask to be used when making the prefix input length of prefixlm variable causalmask torch triltorch onesnumoutputcontexts numoutputcontexts dtypetorch uint8 view1 1 numoutputcontexts numoutputcontexts todevice prefixlmmask causalmask numinputcontexts if tokentypeids is not none tokentypeids tokentypeids unsqueeze1 unsqueeze2 prefixlmmask prefixlmmask tokentypeids 0 float marge prefixlmmask and attentionmask extendedattentionmask prefixlmmask attentionmask unsqueeze1 unsqueeze2 prepare head mask if needed if headmask is not none headmask self getheadmask headmask self config numswitchlayers self config numextlayers nlayer x batch x nheads x n x n outputs presentkeyvaluestates if self config usecache or usecache else none allhiddenstates if self config outputhiddenstates or outputhiddenstates else none allattentions if self config outputattentions or outputattentions else none allrouterprobs if self config outputrouterlogits or outputrouterlogits else none for layer past in enumeratepastsorspoutvalue if layer self config numswitchlayers if self config numextlayers 0 extrapositionembeddings are extra position embeddings that are only created when extending the model with code from the original gptsan repository not used in the default model however it is created when you create an additional layer and partially train only that location therefore convertgptsantfcheckpointtopytorch py is used when converting and loading models created in the original gptsan repository for i in rangenumbatch hiddenstatesi torch gather self extrapositionembeddings weight dim0 indexgatherpositioni outputroutertuple self config outputrouterlogits or outputrouterlogits and layer self config numswitchlayers blockoutput self blockslayer hiddenstateshiddenstates pastkeyvaluepast attentionmaskextendedattentionmask headmaskheadmask usecacheself config usecache or usecache outputattentionsself config outputattentions or outputattentions outputroutertupleoutputroutertuple outpos 0 hiddenstates blockoutputoutpos if self config outputhiddenstates or outputhiddenstates allhiddenstates hiddenstates if self config usecache or usecache outpos 1 present blockoutputoutpos presentkeyvaluestates present if self config outputattentions or outputattentions outpos 1 attentionprobs blockoutputoutpos allattentions attentionprobs if outputroutertuple outpos 1 routertuple blockoutputoutpos allrouterprobs appendroutertuple0 hiddenstates self lastprojecthiddenstates hiddenstates self acthiddenstates if self config outputhiddenstates or outputhiddenstates allhiddenstates allhiddenstates hiddenstates if not returndict return tuple v for v in hiddenstates presentkeyvaluestates allhiddenstates allattentions allrouterprobs if v is not none return moemodeloutputwithpastandcrossattentions lasthiddenstatehiddenstates pastkeyvaluespresentkeyvaluestates hiddenstatesallhiddenstates attentionsallattentions routerprobsallrouterprobs addstartdocstrings the bare gptsanjapanese model with a language modeling head gptsanjapanesestartdocstring class gptsanjapaneseforconditionalgenerationgptsanjapanesepretrainedmodel tiedweightskeys lmhead weight def initself config gptsanjapaneseconfig super initconfig self model gptsanjapanesemodelconfig self registerbufferfinallogitsbias torch zeros1 config vocabsize self lmhead nn linearconfig dmodel config vocabsize biasfalse if not self config torchscript self lmhead weight self model embedtokens weight addstartdocstringstomodelforwardgptsanjapaneseinputsdocstring def forward self inputids optionaltorch longtensor none attentionmask optionaltorch floattensor none tokentypeids optionaltorch floattensor none spout optionaltorch floattensor none pastkeyvalues optionaltupletupletorch floattensor none headmask optionaltorch floattensor none usecache optionalbool false inputsembeds optionaltorch floattensor none decoderinputsembeds optionaltorch floattensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none outputrouterlogits optionalbool none labels optionaltorch longtensor none uniontupletorch floattensor moecausallmoutputwithpast r labels torch longtensor of shape batchsize optional labels for computing the sequence classification loss indices should be in 100 0 config vocabsize 1 all labels set to 100 are ignored masked the loss is only computed for labels in 0 config vocabsize returns moecausallmoutputwithpast or tuple if returndict returns moecausallmoutputwithpast insted of tuple example text generation with regular lm model python from transformers import automodel autotokenizer trainerutils device cuda model automodel frompretrainedtanreigptsanjapanese todevice tokenizer autotokenizer frompretrainedtanreigptsanjapanese xtoken tokenizer returntensorspt trainerutils setseed30 inputids xtoken inputids todevice gentoken model generateinputids maxnewtokens50 tokenizer decodegentoken0 text generation with prefixlm model python from transformers import automodel autotokenizer trainerutils device cuda model automodel frompretrainedtanreigptsanjapanese todevice tokenizer autotokenizer frompretrainedtanreigptsanjapanese xtoken tokenizer prefixtext returntensorspt trainerutils setseed30 inputids xtoken inputids todevice tokentypeids xtoken tokentypeids todevice gentoken model generateinputids tokentypeidstokentypeids maxnewtokens50 tokenizer decodegentoken0 1568 simultaneously text generation and masked language model python from transformers import automodel autotokenizer trainerutils device cuda model automodel frompretrainedtanreigptsanjapanese todevice tokenizer autotokenizer frompretrainedtanreigptsanjapanese maskedsentence inputmaskinputmask xtoken tokenizer prefixtextmaskedsentence returntensorspt trainerutils setseed30 inputids xtoken inputids todevice tokentypeids xtoken tokentypeids todevice outlmtoken model generateinputids tokentypeidstokentypeids maxnewtokens50 outmlmtoken modelinputids tokentypeidstokentypeids logits argmaxaxis1 tokenizer decodeoutmlmtoken0 tokenizer decodeoutlmtoken0inputids shape1 n segtoken self config separatortokenid usecache usecache or self config usecache returndict returndict if returndict is not none else self config usereturndict modelreturndict true numprecontext none if inputids is not none numbatch inputids shape0 numprecontext torch zerosnumbatch int toinputids device whereseparators torch whereinputids segtoken numprecontextwhereseparators0 whereseparators1 numprecontext numprecontext unsqueeze1 outputs self model inputids attentionmask tokentypeids spout pastkeyvalues headmask usecache inputsembeds decoderinputsembeds outputattentions outputhiddenstates modelreturndict outputrouterlogits numprecontext lmlogits self lmheadoutputs0 if lmlogits shape1 self finallogitsbias shape1 lmlogits lmlogits self finallogitsbias loss none zloss none routerprobs none auxloss none if labels is not none move labels to correct device to enable model parallelism labels labels tolmlogits device lossfct nn crossentropylossignoreindex100 if outputrouterlogits compute the router loss zloss auxiliary loss for each router in the encoder and decoder routerlogits expertindexes self unpackrouterlogitsoutputs routerprobs zloss routerzlossfuncrouterlogits routerprobs nn softmaxdim1routerlogits auxloss loadbalancinglossfuncrouterprobs expertindexes loss lossfctlmlogits view1 lmlogits size1 labels view1 if not returndict return tuple v for v in loss lmlogits outputs pastkeyvalues outputs hiddenstates outputs routerprobs zloss auxloss if v is not none return moecausallmoutputwithpast lossloss logitslmlogits pastkeyvaluesoutputs pastkeyvalues hiddenstatesoutputs hiddenstates attentionsoutputs attentions routerlogitsoutputs routerprobs zlosszloss auxlossauxloss def prepareinputsforgeneration self inputids torch longtensor attentionmask torch floattensor tokentypeids optionaltorch floattensor none spout optionalunionlist torch floattensor none pastkeyvalues optionaltupletupletorch floattensor none kwargs if isinstancespout list spout torch tensorspout float if inputids is not none spout spout toinputids device if pastkeyvalues is not none return inputids inputids 1 if inputids is not none else none attentionmask attentionmask tokentypeids tokentypeids 1 if tokentypeids is not none else none spout spout pastkeyvalues pastkeyvalues return inputids inputids attentionmask attentionmask tokentypeids tokentypeids spout spout pastkeyvalues none copied from transformers models switchtransformers modelingswitchtransformers switchtransformersforconditionalgeneration preparedecoderinputidsfromlabels with switchtransformersgptsanjapanese def preparedecoderinputidsfromlabelsself labels torch tensor return self shiftrightlabels copied from transformers models mbart modelingmbart mbartforconditionalgeneration resizetokenembeddings with mbartgptsanjapanese def resizetokenembeddingsself newnumtokens int padtomultipleof optionalint none nn embedding newembeddings super resizetokenembeddingsnewnumtokens padtomultipleof self resizefinallogitsbiasnewembeddings weight shape0 return newembeddings copied from transformers models mbart modelingmbart mbartforconditionalgeneration resizefinallogitsbias with mbartgptsanjapanese def resizefinallogitsbiasself newnumtokens int none oldnumtokens self finallogitsbias shape1 if newnumtokens oldnumtokens newbias self finallogitsbias newnumtokens else extrabias torch zeros1 newnumtokens oldnumtokens deviceself finallogitsbias device newbias torch catself finallogitsbias extrabias dim1 self registerbufferfinallogitsbias newbias def getinputembeddingsself return self model getinputembeddings def setinputembeddingsself newembeddings self model setinputembeddingsnewembeddings copied from transformers models switchtransformers modelingswitchtransformers switchtransformersforconditionalgeneration setoutputembeddings with switchtransformersgptsanjapanese def setoutputembeddingsself newembeddings self lmhead newembeddings copied from transformers models switchtransformers modelingswitchtransformers switchtransformersforconditionalgeneration getoutputembeddings with switchtransformersgptsanjapanese def getoutputembeddingsself return self lmhead copied from transformers models switchtransformers modelingswitchtransformers switchtransformersforconditionalgeneration unpackrouterlogits with switchtransformersgptsanjapanese def unpackrouterlogitsself routeroutputs totalrouterlogits totalexpertindexes for routeroutput in routeroutputs if lenrouteroutput0 shape 1 routerlogits expertindexes routeroutput totalrouterlogits appendrouterlogits totalexpertindexes appendexpertindexes return torch cattotalrouterlogits dim1 torch cattotalexpertindexes dim1 coding utf 8 2023 toshiyuki sakamoto tanreinama and 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 pytorch gptsanjapanese model this dict contains ids and associated url for the pretrained weights provided with the models see all gptsan japanese models at https huggingface co models filter gptsan japanese copied from transformers models switch_transformers modeling_switch_transformers router_z_loss_func compute the router z loss implemented in pytorch the router z loss was introduced in designing effective sparse expert models https arxiv org abs 2202 08906 it encourages router logits to remain small in an effort to improve stability args router_logits float input logits of shape batch_size sequence_length num_experts returns scalar router z loss copied from transformers models switch_transformers modeling_switch_transformers load_balancing_loss_func computes auxiliary load balancing loss as in switch transformer implemented in pytorch see switch transformer https arxiv org abs 2101 03961 for more details this function implements the loss function presented in equations 4 6 of the paper it aims at penalizing cases where the routing between experts is too unbalanced args router_probs torch tensor probability assigned to each expert per token shape batch_size seqeunce_length num_experts expert_indices torch tensor indices tensor of shape batch_size seqeunce_length identifying the selected expert for a given token returns the auxiliary loss cast the expert indices to int64 otherwise one hot encoding will fail for a given token determine if it was routed to a given expert cast to float32 otherwise mean will fail ffn layer for switch transformer and extra layers gptsan can mix switch transformer layers and normal transformer layers this class is used as expert in switch transformer layers and as ffn in regular transformer layers relu is used in the switch transformer layer and swish is used in the normal transformer layer so there is a choice of which is used in the argument args hidden_states torch tensor num_groups tokens_per_group hidden_dim inputs to send to experts returns torch tensor num_groups tokens_per_group hidden_dim copied from transformers models switch_transformers modeling_switch_transformers switchtransformerstop1router with switchtransformers gptsanjapanese router using tokens choose top 1 experts assignment this router uses the same mechanism as in switch transformer https arxiv org abs 2101 03961 and v moe https arxiv org abs 2106 05974 tokens choose their top experts items are sorted by router_probs and then routed to their choice of expert until the expert s expert_capacity is reached there is no guarantee that each token is processed by an expert or that each expert receives at least one token computes router probabilities from input hidden states args hidden_states torch tensor batch_size sequence_length hidden_dim from which router probabilities are computed returns router_probabilities torch tensor tensor of shape batch_size sequence_length num_experts corresponding to the probabilities for each token and expert used for routing tokens to experts router_logits torch tensor logits tensor of shape batch_size sequence_length num_experts corresponding to raw router logits this is used later for computing router z loss float32 is used to ensure stability see the discussion of selective precision in https arxiv org abs 2101 03961 we also store the previous dtype to cast back the output to the previous dtype multiply the token inputs by the uniform distribution adding some noise shape num_groups tokens_per_group num_experts apply softmax and cast back to the original dtype bitsandbytes linear8bitlt layers does not support manual casting therefore we need to check if they are an instance of the linear8bitlt class by checking special attributes generic forward function for every router class each router expects to have the same input hidden states hidden_states corresponding to the hidden states for each token the expert_capacity corresponding to the number of tokens the router will send to each expert some routers can send up to few tokens to each expert each router works as the following it expects the hidden states for each token gets the router_probs and router_logits from the router_weights this will assign for each token the raw probability to be assigned to an expert then each router class will have to define its own _compute_routing_instructions args hidden_states torch tensor num_groups tokens_per_group hidden_dim inputs to send to experts returns tuple torch tensor torch tensor torch tensor tuple containing the expert index the router probs and the router logits the router probabilities and logits are required to compute the loss mask tokens outside expert capacity sum over each sequence mask if the token routed to to the expert will overflow copied from transformers models switch_transformers modeling_switch_transformers switchtransformerssparsemlp with switchtransformers gptsanjapanese implementation of the switch transformers sparse mlp module step 1 get the correct router according to its class step 2 get the experts hold on this will be slightly tricky to understand in the correct order a moe layer does the following 1 gets the router_mask from the router the shape of the mask is batch_size sequence_length num_expert and corresponds to the argmax of the router_probs the probabilities are needed in the computation of the hidden states they are broadcasted to the hidden states values can be interpreted as a scaling factor 2 dispatch the tokens to its associated experts we do a classic for loop over the experts and assign for each expert the corresponding hidden states step 1 get the router_mask from the router as wel as the probabilities the routers introduced might not always map all the tokens to a router which means that some hidden states can be unchanged from one layer to another that is why the hidden states are cloned before updating only the seleced ones switch transformers feed forward layer module this is a wrapper around the mixture of experts module parameters config gptsanjapaneseconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel from_pretrained method to load the model weights args hidden_states torch tensor num_groups tokens_per_group hidden_dim inputs to send to experts output_router_logits bool output experts router output returns torch tensor num_groups tokens_per_group hidden_dim extra transformers feed forward layer module parameters config gptsanjapaneseconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel from_pretrained method to load the model weights check if it is a sparse layer if not then it is a dense layer args hidden_states torch tensor num_groups tokens_per_group hidden_dim inputs to send to experts returns torch tensor num_groups tokens_per_group hidden_dim copied from transformers models bart modeling_bart bartattention with bart gptsanjapanese multi headed attention from attention is all you need paper input shape batch x time x channel if key_value_states are provided this layer is used as a cross attention layer for the decoder get query proj get key value proj past_key_value 0 shape 2 key_value_states shape 1 is checking that the sequence_length of the past_key_value is the same as the provided key_value_states to support prefix tuning reuse k v cross_attentions cross_attentions reuse k v self_attention self_attention if cross_attention save tuple torch tensor torch tensor of all cross attention key value_states further calls to cross_attention layer can then reuse all cross attention key value_states first if case if uni directional self attention decoder save tuple torch tensor torch tensor of all previous decoder key value_states further calls to uni directional self attention can concat previous decoder key value_states to current projected key value_states third elif case if encoder bi directional self attention past_key_value is always none this operation is a bit awkward but it s required to make sure that attn_weights keeps its gradient in order to do so attn_weights have to be reshaped twice and have to be reused in the following use the embed_dim from the config stored in the class rather than hidden_state because attn_output can be partitioned across gpus when using tensor parallelism self attention and normalization unit self attention and normalize block args hidden_states torch floattensor of shape batch_size sequence_length hidden_size optional sequence of hidden states at the output of the last layer of the encoder used in the cross attention if the model is configured as a decoder past_key_values tuple tuple torch floattensor of length config n_layers with each tuple having 4 tensors of shape batch_size num_heads sequence_length 1 embed_size_per_head contains precomputed key and value hidden states of the attention blocks can be used to speed up decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape batch_size sequence_length attention_mask torch floattensor of shape batch_size sequence_length optional mask to avoid performing attention on the padding token indices of the encoder input this mask is used in the cross attention if the model is configured as a decoder mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked head_mask numpy ndarray of shape 0 optional mask to nullify selected heads of the attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked use_cache bool optional if set to true past_key_values key value states are returned and can be used to speed up decoding see past_key_values output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail returns tuple torch tensor num_groups tokens_per_group hidden_dim self attention decoder uni directional self attention cached key values tuple is at positions 1 2 add present self attn cache to positions 1 2 of present_key_value tuple hidden present attentions hidden attentions self attention and ffn unit gptsan transformer block args hidden_states torch floattensor of shape batch_size sequence_length hidden_size optional sequence of hidden states at the output of the last layer of the encoder used in the cross attention if the model is configured as a decoder past_key_values tuple tuple torch floattensor of length config n_layers with each tuple having 4 tensors of shape batch_size num_heads sequence_length 1 embed_size_per_head contains precomputed key and value hidden states of the attention blocks can be used to speed up decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape batch_size sequence_length attention_mask torch floattensor of shape batch_size sequence_length optional mask to avoid performing attention on the padding token indices of the encoder input this mask is used in the cross attention if the model is configured as a decoder mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked head_mask numpy ndarray of shape 0 optional mask to nullify selected heads of the attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked use_cache bool optional if set to true past_key_values key value states are returned and can be used to speed up decoding see past_key_values output_attentions bool output attention probabirities output_router_tuple output experts router logits and expert id returns tuple torch tensor num_groups tokens_per_group hidden_dim an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights used for testing weights initialization mesh tensorflow embeddings initialization see https github com tensorflow mesh blob fa19d69eafc9a482aff0b59ddd96b025c0cb207d mesh_tensorflow layers py l1624 mesh tensorflow embeddings initialization see https github com tensorflow mesh blob fa19d69eafc9a482aff0b59ddd96b025c0cb207d mesh_tensorflow layers py l1624 mesh tensorflow ff initialization see https github com tensorflow mesh blob master mesh_tensorflow transformer transformer_layers py l56 and https github com tensorflow mesh blob fa19d69eafc9a482aff0b59ddd96b025c0cb207d mesh_tensorflow layers py l89 multi headed attention mesh tensorflow attention initialization to avoid scaling before softmax see https github com tensorflow mesh blob fa19d69eafc9a482aff0b59ddd96b025c0cb207d mesh_tensorflow transformer attention py l136 copied from transformers models t5 modeling_t5 t5pretrainedmodel _shift_right shift inputs to the right item assignment is not supported natively for proxies replace possible 100 values in labels by pad_token_id the gptsan japanese https github com tanreinama gptsan model was proposed in general purpose swich transformer based japanese language model this model is also a pytorch torch nn module https pytorch org docs stable nn html torch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config gptsanjapaneseconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel from_pretrained method to load the model weights args input_ids torch longtensor of shape batch_size sequence_length indices of input sequence tokens in the vocabulary gptsan japanese is a model that generates sentence continuations or predicts tokens at mask positions special tokens required for inputs to the model are automatically appended attention_mask torch floattensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask token_type_ids torch floattensor of shape batch_size sequence_length optional an input that masks the prefix part in the prefix lm input mask values selected in 0 1 1 for tokens that are prefix input 0 for tokens that are not prefix input spout torch tensor of shape batch_size config d_spout this vector is transformed through an 8 layer ffn and can be used instead of past_key_values past_key_values tuple tuple torch floattensor of length config n_layers with each tuple having 4 tensors of shape batch_size num_heads sequence_length 1 embed_size_per_head contains precomputed key and value hidden states of the attention blocks can be used to speed up decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape batch_size sequence_length head_mask torch floattensor of shape num_heads or num_layers num_heads optional mask to nullify selected heads of the self attention modules mask values selected in 0 1 use_cache bool optional if set to true past_key_values key value states are returned and can be used to speed up decoding see past_key_values inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size optional optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix decoder_inputs_embeds torch floattensor of shape batch_size target_sequence_length hidden_size optional optionally instead of passing decoder_input_ids you can choose to directly pass an embedded representation if past_key_values is used optionally only the last decoder_inputs_embeds have to be input see past_key_values this is useful if you want more control over how to convert decoder_input_ids indices into associated vectors than the model s internal embedding lookup matrix output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple router_logits tuple torch floattensor optional returned when output_router_logits true is passed or when config add_router_probs true tuple of torch floattensor one for each layer of shape batch_size sequence_length num_experts router logits of the decoder model useful to compute the auxiliary loss for mixture of experts models num_precontext torch longtensor of shape batch_size 1 length of hybrid input tokens in the input tokens up to this length refer to both front and back like bert tokens after that refer only to front like gpt see also https github com tanreinama gptsan blob main report model md returns moemodeloutputwithpastandcrossattentions or tuple if return_dict returns moemodeloutputwithpastandcrossattentions insted of tuple dummy for input_ids was none spout is a special input vector specific to gptsan this controls the output by projecting embedded information such as the class of sentences during learning it should passed instead of the first past_key_value see the original gptsan repository for details if there is an attention_mask increase first one for spout 1st token should be spout update attention_mask num_precontext is the number of tokens that refer to each other in prefix lm created per batch so dimension of num_precontext should be batch 1 num_precontext should be batch 1 make vector from spout of gptsan to the same shape as past_key_values projecting spout vector make same shape as past_key_values token position considering spout and pasts positions for get position_embeddings when padding with padding_side left zeros line up on the left side of attention_mask so position_embeddings is shifted accordingly attention_mask is applied per batch create a mask to be used when making the prefix input length of prefix lm variable marge prefix_lm_mask and attention_mask prepare head mask if needed n_layer x batch x n_heads x n x n outputs extra_position_embeddings are extra position embeddings that are only created when extending the model with code from the original gptsan repository not used in the default model however it is created when you create an additional layer and partially train only that location therefore convert_gptsan_tf_checkpoint_to_pytorch py is used when converting and loading models created in the original gptsan repository labels torch longtensor of shape batch_size optional labels for computing the sequence classification loss indices should be in 100 0 config vocab_size 1 all labels set to 100 are ignored masked the loss is only computed for labels in 0 config vocab_size returns moecausallmoutputwithpast or tuple if return_dict returns moecausallmoutputwithpast insted of tuple example text generation with regular lm model python from transformers import automodel autotokenizer trainer_utils device cuda model automodel from_pretrained tanrei gptsan japanese to device tokenizer autotokenizer from_pretrained tanrei gptsan japanese x_token tokenizer 織田信長は return_tensors pt trainer_utils set_seed 30 input_ids x_token input_ids to device gen_token model generate input_ids max_new_tokens 50 tokenizer decode gen_token 0 織田信長は 政治 軍事の中枢まで掌握した政治家であり 日本史上類を見ない驚異的な軍事侵攻を続け text generation with prefix lm model python from transformers import automodel autotokenizer trainer_utils device cuda model automodel from_pretrained tanrei gptsan japanese to device tokenizer autotokenizer from_pretrained tanrei gptsan japanese x_token tokenizer prefix_text 織田信長は return_tensors pt trainer_utils set_seed 30 input_ids x_token input_ids to device token_type_ids x_token token_type_ids to device gen_token model generate input_ids token_type_ids token_type_ids max_new_tokens 50 tokenizer decode gen_token 0 織田信長は 政治 外交で数々の戦果を上げるが 1568年からは いわゆる本能寺の変で細川晴元に暗殺される simultaneously text generation and masked language model python from transformers import automodel autotokenizer trainer_utils device cuda model automodel from_pretrained tanrei gptsan japanese to device tokenizer autotokenizer from_pretrained tanrei gptsan japanese masked_sentence 武田信玄は inputmask 時代ファンならぜひ押さえ inputmask きたい名将の一人 x_token tokenizer prefix_text masked_sentence return_tensors pt trainer_utils set_seed 30 input_ids x_token input_ids to device token_type_ids x_token token_type_ids to device out_lm_token model generate input_ids token_type_ids token_type_ids max_new_tokens 50 out_mlm_token model input_ids token_type_ids token_type_ids logits argmax axis 1 tokenizer decode out_mlm_token 0 武田信玄は 戦国時代ファンならぜひ押さえておきたい名将の一人 tokenizer decode out_lm_token 0 input_ids shape 1 武田氏の三代に渡った武田家のひとり n甲斐市に住む 日本史上最大の戦国大名 move labels to correct device to enable model parallelism compute the router loss z_loss auxiliary loss for each router in the encoder and decoder copied from transformers models switch_transformers modeling_switch_transformers switchtransformersforconditionalgeneration prepare_decoder_input_ids_from_labels with switchtransformers gptsanjapanese copied from transformers models mbart modeling_mbart mbartforconditionalgeneration resize_token_embeddings with mbart gptsanjapanese copied from transformers models mbart modeling_mbart mbartforconditionalgeneration _resize_final_logits_bias with mbart gptsanjapanese copied from transformers models switch_transformers modeling_switch_transformers switchtransformersforconditionalgeneration set_output_embeddings with switchtransformers gptsanjapanese copied from transformers models switch_transformers modeling_switch_transformers switchtransformersforconditionalgeneration get_output_embeddings with switchtransformers gptsanjapanese copied from transformers models switch_transformers modeling_switch_transformers switchtransformersforconditionalgeneration _unpack_router_logits with switchtransformers gptsanjapanese
import copy
from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn

from ...activations import ACT2FN
from ...modeling_outputs import MoECausalLMOutputWithPast, MoEModelOutputWithPastAndCrossAttentions
from ...modeling_utils import PreTrainedModel
from ...utils import (
    DUMMY_INPUTS,
    DUMMY_MASK,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    is_torch_fx_proxy,
    logging,
)
from .configuration_gptsan_japanese import GPTSanJapaneseConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "GPTSanJapaneseConfig"
_CHECKPOINT_FOR_DOC = "Tanrei/GPTSAN-japanese"





GPTSAN_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "Tanrei/GPTSAN-japanese",
    
]



def router_z_loss_func(router_logits: torch.Tensor) -> float:
    r
    num_groups, tokens_per_group, _ = router_logits.shape
    log_z = torch.logsumexp(router_logits, dim=-1)
    z_loss = log_z**2
    return torch.sum(z_loss) / (num_groups * tokens_per_group)



def load_balancing_loss_func(router_probs: torch.Tensor, expert_indices: torch.Tensor) -> float:
    r
    num_experts = router_probs.shape[-1]

    
    if expert_indices.dtype != torch.int64:
        expert_indices = expert_indices.to(torch.int64)

    if len(expert_indices.shape) == 2:
        expert_indices = expert_indices.unsqueeze(2)

    expert_mask = torch.nn.functional.one_hot(expert_indices, num_experts)

    
    expert_mask = torch.max(expert_mask, axis=-2).values

    
    expert_mask = expert_mask.to(torch.float32)
    tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)

    router_prob_per_group_and_expert = torch.mean(router_probs, axis=-2)
    return torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert) * (num_experts**2)


class GPTSanJapaneseDenseActDense(nn.Module):
    

    def __init__(self, config: GPTSanJapaneseConfig, ext_layer=False):
        super().__init__()
        d_inter = config.d_ext if ext_layer else config.d_ff
        self.wi = nn.Linear(config.d_model, d_inter, bias=ext_layer)
        self.wo = nn.Linear(d_inter, config.d_model, bias=ext_layer)
        self.dropout = nn.Identity() if ext_layer else nn.Dropout(config.dropout_rate)
        self.act = ACT2FN["swish" if ext_layer else "relu"]

    def forward(self, hidden_states):
        r
        hidden_states = self.wi(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.wo(hidden_states)
        return hidden_states



class GPTSanJapaneseTop1Router(nn.Module):
    

    def __init__(self, config: GPTSanJapaneseConfig):
        super().__init__()
        self.num_experts = config.num_experts
        self.expert_capacity = config.expert_capacity
        self.classifier = nn.Linear(config.hidden_size, self.num_experts, bias=config.router_bias)
        self.jitter_noise = config.router_jitter_noise
        self.ignore_padding_tokens = config.router_ignore_padding_tokens
        self.dtype = getattr(torch, config.router_dtype)

    def _compute_router_probabilities(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        r
        
        
        
        self.input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(self.dtype)

        if self.jitter_noise > 0:
            
            hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise)

        
        self._cast_classifier()
        router_logits = self.classifier(hidden_states)

        
        router_probabilities = nn.functional.softmax(router_logits, dim=-1, dtype=self.dtype).to(self.input_dtype)
        return router_probabilities, router_logits

    def _cast_classifier(self):
        r
        if not (hasattr(self.classifier, "SCB") or hasattr(self.classifier, "CB")):
            self.classifier = self.classifier.to(self.dtype)

    def forward(self, hidden_states: torch.Tensor) -> Tuple:
        r
        router_probs, router_logits = self._compute_router_probabilities(hidden_states)

        expert_index = torch.argmax(router_probs, dim=-1)
        expert_index = torch.nn.functional.one_hot(expert_index, num_classes=self.num_experts)

        
        token_priority = torch.cumsum(expert_index, dim=-2)
        
        expert_capacity_mask = token_priority <= self.expert_capacity
        expert_index = expert_index * expert_capacity_mask

        router_probs = torch.max(router_probs, dim=-1).values.unsqueeze(-1)
        return expert_index, router_probs, router_logits



class GPTSanJapaneseSparseMLP(nn.Module):
    r

    def __init__(self, config: GPTSanJapaneseConfig, expert_class: nn.Module = GPTSanJapaneseDenseActDense):
        super().__init__()
        
        self.router = GPTSanJapaneseTop1Router(config)

        
        self.experts = nn.ModuleDict()
        for idx in range(config.num_experts):
            self.experts[f"expert_{idx}"] = expert_class(config)

    def forward(self, hidden_states):
        r
        
        router_mask, router_probs, router_logits = self.router(hidden_states)
        expert_index = torch.argmax(router_mask, dim=-1)

        
        

        next_states = hidden_states.clone()
        for idx, expert in enumerate(self.experts.values()):
            token_indices = router_mask[:, :, idx].bool()
            next_states[token_indices] = expert(hidden_states[token_indices]).to(next_states.dtype)

        hidden_states = router_probs * next_states
        return hidden_states, (router_logits, expert_index)


class GPTSanJapaneseLayerSparseFF(nn.Module):
    r

    def __init__(self, config: GPTSanJapaneseConfig):
        super().__init__()
        self.mlp = GPTSanJapaneseSparseMLP(config)
        self.soft_bypass_mlp = nn.Linear(config.d_model, config.d_model, bias=False)
        self.norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_epsilon)

    def forward(self, hidden_states, output_router_logits):
        r
        forwarded_states, router_tuple = self.mlp(hidden_states)
        forwarded_states += torch.tanh(self.soft_bypass_mlp(hidden_states))
        output = hidden_states + self.norm(forwarded_states)

        if output_router_logits and router_tuple is not None:
            return output, router_tuple
        else:
            return output


class GPTSanJapaneseLayerDenseFF(nn.Module):
    r

    def __init__(self, config: GPTSanJapaneseConfig):
        super().__init__()
        
        self.mlp = GPTSanJapaneseDenseActDense(config, ext_layer=True)
        self.norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_epsilon)

    def forward(self, hidden_states):
        r
        forwarded_states = self.mlp(hidden_states)
        output = hidden_states + self.norm(forwarded_states)
        return output



class GPTSanJapaneseAttention(nn.Module):
    

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        is_decoder: bool = False,
        bias: bool = True,
        is_causal: bool = False,
        config: Optional[GPTSanJapaneseConfig] = None,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        self.config = config

        if (self.head_dim * num_heads) != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
                f" and `num_heads`: {num_heads})."
            )
        self.scaling = self.head_dim**-0.5
        self.is_decoder = is_decoder
        self.is_causal = is_causal

        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

    def forward(
        self,
        hidden_states: torch.Tensor,
        key_value_states: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        layer_head_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        

        
        
        is_cross_attention = key_value_states is not None

        bsz, tgt_len, _ = hidden_states.size()

        
        query_states = self.q_proj(hidden_states) * self.scaling
        
        
        
        
        if (
            is_cross_attention
            and past_key_value is not None
            and past_key_value[0].shape[2] == key_value_states.shape[1]
        ):
            
            key_states = past_key_value[0]
            value_states = past_key_value[1]
        elif is_cross_attention:
            
            key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
            value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
        elif past_key_value is not None:
            
            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
            key_states = torch.cat([past_key_value[0], key_states], dim=2)
            value_states = torch.cat([past_key_value[1], value_states], dim=2)
        else:
            
            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

        if self.is_decoder:
            
            
            
            
            
            
            
            past_key_value = (key_states, value_states)

        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
        query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
        key_states = key_states.reshape(*proj_shape)
        value_states = value_states.reshape(*proj_shape)

        src_len = key_states.size(1)
        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))

        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
            raise ValueError(
                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
                f" {attn_weights.size()}"
            )

        if attention_mask is not None:
            if attention_mask.size() != (bsz, 1, tgt_len, src_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
                )
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        attn_weights = nn.functional.softmax(attn_weights, dim=-1)

        if layer_head_mask is not None:
            if layer_head_mask.size() != (self.num_heads,):
                raise ValueError(
                    f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
                    f" {layer_head_mask.size()}"
                )
            attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        if output_attentions:
            
            
            
            
            attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
        else:
            attn_weights_reshaped = None

        attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)

        attn_output = torch.bmm(attn_probs, value_states)

        if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
            raise ValueError(
                f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is"
                f" {attn_output.size()}"
            )

        attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
        attn_output = attn_output.transpose(1, 2)

        
        
        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)

        attn_output = self.out_proj(attn_output)

        return attn_output, attn_weights_reshaped, past_key_value


class GPTSanJapaneseLayerSelfAttention(nn.Module):
    

    def __init__(self, config, has_relative_attention_bias=False):
        super().__init__()
        self.self_attn = GPTSanJapaneseAttention(
            embed_dim=config.d_model,
            num_heads=config.num_heads,
            is_decoder=True,
            bias=has_relative_attention_bias,
        )
        self.norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_epsilon)

    def forward(
        self,
        hidden_states: Optional[Tuple[torch.FloatTensor]],
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = False,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]:
        r
        
        
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        
        atten_out = self.self_attn(
            hidden_states=hidden_states,
            past_key_value=self_attn_past_key_value,
            attention_mask=(1 - attention_mask) * torch.finfo(hidden_states.dtype).min,
            layer_head_mask=head_mask,
            output_attentions=output_attentions,
        )
        if output_attentions:
            attn_weights = (atten_out[1],)
        else:
            attn_weights = ()

        attention_output = atten_out[0]

        hidden = hidden_states + self.norm(attention_output)

        if use_cache:
            outputs = (hidden, atten_out[2])  
        else:
            outputs = (hidden,)  

        return outputs + attn_weights


class GPTSanJapaneseBlock(nn.Module):
    

    def __init__(self, config, ext_layer=False):
        super().__init__()
        self.self_attn = GPTSanJapaneseLayerSelfAttention(config)
        self.feed_forward = GPTSanJapaneseLayerDenseFF(config) if ext_layer else GPTSanJapaneseLayerSparseFF(config)

    def forward(
        self,
        hidden_states: Optional[Tuple[torch.FloatTensor]],
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = False,
        output_attentions: Optional[bool] = False,
        output_router_tuple: Optional[bool] = False,
    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]:
        r
        atten_out = self.self_attn(
            hidden_states=hidden_states,
            past_key_value=past_key_value,
            attention_mask=attention_mask,
            head_mask=head_mask,
            use_cache=use_cache,
            output_attentions=output_attentions,
        )
        attention_output = atten_out[0]

        if isinstance(self.feed_forward, GPTSanJapaneseLayerSparseFF):
            sparse_out = self.feed_forward(attention_output, output_router_tuple)
            if output_router_tuple:
                hidden, router_tuple = sparse_out
            else:
                hidden = sparse_out
        else:
            hidden = self.feed_forward(attention_output)

        outputs = (hidden,) + atten_out[1:]

        if isinstance(self.feed_forward, GPTSanJapaneseLayerSparseFF) and output_router_tuple:
            outputs += (router_tuple,)

        return outputs


class GPTSanJapanesePreTrainedModel(PreTrainedModel):
    

    config_class = GPTSanJapaneseConfig
    base_model_prefix = "gptsan_japanese"
    supports_gradient_checkpointing = False
    _no_split_modules = ["GPTSanJapaneseBlock"]
    _skip_keys_device_placement = "past_key_values"

    @property
    def dummy_inputs(self):
        input_ids = torch.tensor(DUMMY_INPUTS)
        input_mask = torch.tensor(DUMMY_MASK)
        dummy_inputs = {
            "input_ids": input_ids,
            "attention_mask": input_mask,
        }
        return dummy_inputs

    def _init_weights(self, module):
        
        factor = self.config.initializer_factor  
        if isinstance(module, nn.LayerNorm):
            module.weight.data.fill_(factor * 1.0)
            module.bias.data.zero_()
        elif isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5))
            if hasattr(module, "bias") and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=factor * 1.0)
        elif isinstance(module, GPTSanJapaneseModel):
            
            
            module.embed_tokens.weight.data.normal_(mean=0.0, std=factor * 1.0)
            module.position_embeddings.weight.data.normal_(mean=0.0, std=factor * 1.0)
            if hasattr(module, "extra_position_embeddings") and module.extra_position_embeddings is not None:
                module.extra_position_embeddings.weight.data.normal_(mean=0.0, std=factor * 1.0)
        elif isinstance(module, (GPTSanJapaneseModel, GPTSanJapaneseForConditionalGeneration)):
            
            
            module.final_logits_bias.data.normal_(mean=0.0, std=factor * 1.0)
            if hasattr(module, "lm_head") and not self.config.tie_word_embeddings:
                module.lm_head.weight.data.normal_(mean=0.0, std=factor * 1.0)
        elif isinstance(module, GPTSanJapaneseDenseActDense):
            
            
            
            module.wi.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5))
            if hasattr(module.wi, "bias") and module.wi.bias is not None:
                module.wi.bias.data.zero_()
            module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5))
            if hasattr(module.wo, "bias") and module.wo.bias is not None:
                module.wo.bias.data.zero_()
        elif isinstance(module, GPTSanJapaneseAttention):
            
            d_model = self.config.d_model
            key_value_proj_dim = self.config.d_model
            n_heads = self.config.num_heads
            module.k_proj.weight.data.normal_(mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5))
            module.v_proj.weight.data.normal_(mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5))
            module.q_proj.weight.data.normal_(mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5))
            module.out_proj.weight.data.normal_(mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5))
        elif isinstance(module, GPTSanJapaneseSparseMLP):
            
            
            d_model = self.config.d_model
            key_value_proj_dim = self.config.d_model
            n_heads = self.config.num_heads
            module.router.classifier.weight.data.normal_(mean=0.0, std=factor * 1)
            for idx in range(self.config.num_experts):
                module.experts[f"expert_{idx}"].wi.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5))
                module.experts[f"expert_{idx}"].wo.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5))

    
    def _shift_right(self, input_ids):
        decoder_start_token_id = self.config.decoder_start_token_id
        pad_token_id = self.config.pad_token_id

        if decoder_start_token_id is None:
            raise ValueError(
                "self.model.config.decoder_start_token_id has to be defined. In T5 it is usually set to the pad_token_id. "
                "See T5 docs for more information."
            )

        
        if is_torch_fx_proxy(input_ids):
            
            shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), decoder_start_token_id)
            shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1)
        else:
            shifted_input_ids = input_ids.new_zeros(input_ids.shape)
            shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()
            shifted_input_ids[..., 0] = decoder_start_token_id

        if pad_token_id is None:
            raise ValueError("self.model.config.pad_token_id has to be defined.")
        
        shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)

        return shifted_input_ids


GPTSAN_JAPANESE_START_DOCSTRING = r

GPTSAN_JAPANESE_INPUTS_DOCSTRING = r


@add_start_docstrings(
    "The bare GPTSAN-japanese Model transformer outputting raw hidden-states without any specific head on top.",
    GPTSAN_JAPANESE_START_DOCSTRING,
)
class GPTSanJapaneseModel(GPTSanJapanesePreTrainedModel):
    def __init__(self, config: GPTSanJapaneseConfig):
        super().__init__(config)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.d_model)
        self.config = copy.deepcopy(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
        self.last_project = nn.Linear(config.d_model, config.d_model, bias=True)
        self.act = ACT2FN["swish"]

        self.blocks = torch.nn.ModuleList([])
        for _ in range(config.num_switch_layers):
            self.blocks.append(GPTSanJapaneseBlock(config))
        for _ in range(config.num_ext_layers):
            self.blocks.append(GPTSanJapaneseBlock(config, ext_layer=True))

        if config.num_ext_layers > 0:
            self.extra_position_embeddings = nn.Embedding(config.max_position_embeddings, config.d_model)

        if config.d_spout:
            spouts = []
            for _ in range(8):
                spouts.append(nn.Linear(config.d_spout, config.d_spout, bias=False))
                spouts.append(nn.Tanh())
            spouts.append(nn.Linear(config.d_spout, config.num_layers * 2 * config.d_model, bias=False))
            self.spout = nn.Sequential(*spouts)

        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, new_embeddings):
        self.embed_tokens = new_embeddings

    @add_start_docstrings_to_model_forward(GPTSAN_JAPANESE_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.FloatTensor] = None,
        spout: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = False,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        output_router_logits: Optional[bool] = None,
        num_precontext: Optional[torch.LongTensor] = None,
    ) -> Union[MoEModelOutputWithPastAndCrossAttentions, Tuple[torch.FloatTensor]]:
        r
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        device = self.position_embeddings.weight.device
        if input_ids is None:
            input_ids = torch.zeros([1, 1]).int().to(device)  
        num_pasts_contexts = 0
        num_batch = input_ids.shape[0]
        pasts_or_spout_value = None
        if past_key_values is not None:
            num_pasts_contexts = past_key_values[0][0].shape[2]
        elif self.config.d_spout and spout is not None:
            
            
            
            
            num_pasts_contexts += 1

        
        if self.config.d_spout and spout is not None and attention_mask is not None:
            attention_mask_with_spout = torch.ones(num_batch, attention_mask.shape[1] + 1, device=device)
            attention_mask_with_spout[:, 1:] -= 1 - attention_mask  
            attention_mask = attention_mask_with_spout  

        if num_precontext is not None:
            
            
            if not (
                len(num_precontext.shape) == 2 and num_precontext.shape[1] == 1
            ):  
                raise ValueError("num_precontext should be [batch, 1] size.")
            num_precontext = torch.reshape(num_precontext, [-1])
        else:
            num_precontext = torch.zeros([num_batch]).int().to(device)

        num_input_contexts = input_ids.shape[1]
        num_output_contexts = num_input_contexts + num_pasts_contexts

        hidden_states = self.embed_tokens(input_ids)

        if past_key_values is not None:
            pasts_or_spout_value = past_key_values
        elif self.config.d_spout and spout is not None:
            
            pasts_or_spout_value = self.spout(spout)  
            pasts_or_spout_value = torch.reshape(
                pasts_or_spout_value,
                [
                    num_batch,
                    self.config.num_layers,
                    2,
                    self.config.num_heads,
                    num_pasts_contexts,
                    self.config.d_model // self.config.num_heads,
                ],
            )
            pasts_or_spout_value = torch.split(pasts_or_spout_value, [1] * self.config.num_layers, dim=1)
            
            pasts_or_spout_value = tuple(
                tuple([b.squeeze(1) for b in torch.split(a.squeeze(1), [1, 1], dim=1)]) for a in pasts_or_spout_value
            )
        else:
            pasts_or_spout_value = [None] * self.config.num_layers

        
        token_position = torch.arange(num_input_contexts).to(device) + num_pasts_contexts

        if attention_mask is None:
            attention_mask = torch.ones(num_batch, num_input_contexts, device=device)

        
        gather_position = (
            (
                torch.zeros((num_batch, self.config.d_model, num_input_contexts)).to(device)
                + token_position.unsqueeze(0)
            )
            .transpose(1, 2)
            .long()
        )
        
        gather_position -= (1 - attention_mask).argmin(dim=-1).unsqueeze(1).unsqueeze(2)
        gather_position = torch.clip(gather_position, num_pasts_contexts, self.config.max_position_embeddings - 1)

        
        for i in range(num_batch):
            hidden_states[i] += torch.gather(self.position_embeddings.weight, dim=0, index=gather_position[i])

        
        causal_mask = (
            torch.tril(torch.ones((num_output_contexts, num_output_contexts), dtype=torch.uint8))
            .view(1, 1, num_output_contexts, num_output_contexts)
            .to(device)
        )
        prefix_lm_mask = causal_mask[:, :, -num_input_contexts:, :]
        if token_type_ids is not None:
            token_type_ids = token_type_ids.unsqueeze(1).unsqueeze(2)
            prefix_lm_mask = ((prefix_lm_mask + token_type_ids) > 0).float()
        
        extended_attention_mask = prefix_lm_mask * attention_mask.unsqueeze(1).unsqueeze(2)

        
        if head_mask is not None:
            head_mask = self.get_head_mask(
                head_mask, self.config.num_switch_layers + self.config.num_ext_layers
            )  

        
        present_key_value_states = () if self.config.use_cache or use_cache else None
        all_hidden_states = () if self.config.output_hidden_states or output_hidden_states else None
        all_attentions = () if self.config.output_attentions or output_attentions else None
        all_router_probs = () if self.config.output_router_logits or output_router_logits else None

        for layer, past in enumerate(pasts_or_spout_value):
            if layer == self.config.num_switch_layers:
                if self.config.num_ext_layers > 0:
                    
                    
                    
                    for i in range(num_batch):
                        hidden_states[i] += torch.gather(
                            self.extra_position_embeddings.weight, dim=0, index=gather_position[i]
                        )

            output_router_tuple = (
                self.config.output_router_logits or output_router_logits
            ) and layer < self.config.num_switch_layers
            block_output = self.blocks[layer](
                hidden_states=hidden_states,
                past_key_value=past,
                attention_mask=extended_attention_mask,
                head_mask=head_mask,
                use_cache=self.config.use_cache or use_cache,
                output_attentions=self.config.output_attentions or output_attentions,
                output_router_tuple=output_router_tuple,
            )

            outpos = 0
            hidden_states = block_output[outpos]
            if self.config.output_hidden_states or output_hidden_states:
                all_hidden_states += (hidden_states,)
            if self.config.use_cache or use_cache:
                outpos += 1
                present = block_output[outpos]
                present_key_value_states += (present,)
            if self.config.output_attentions or output_attentions:
                outpos += 1
                attention_probs = block_output[outpos]
                all_attentions += (attention_probs,)
            if output_router_tuple:
                outpos += 1
                router_tuple = block_output[outpos]
                all_router_probs.append(router_tuple[0])

        hidden_states = self.last_project(hidden_states)
        hidden_states = self.act(hidden_states)

        if self.config.output_hidden_states or output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    present_key_value_states,
                    all_hidden_states,
                    all_attentions,
                    all_router_probs,
                ]
                if v is not None
            )

        return MoEModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=present_key_value_states,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
            router_probs=all_router_probs,
        )


@add_start_docstrings(
    "The bare GPTSAN-japanese Model with a language modeling head.",
    GPTSAN_JAPANESE_START_DOCSTRING,
)
class GPTSanJapaneseForConditionalGeneration(GPTSanJapanesePreTrainedModel):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: GPTSanJapaneseConfig):
        super().__init__(config)
        self.model = GPTSanJapaneseModel(config)
        self.register_buffer("final_logits_bias", torch.zeros([1, config.vocab_size]))
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        if not self.config.torchscript:
            self.lm_head.weight = self.model.embed_tokens.weight

    @add_start_docstrings_to_model_forward(GPTSAN_JAPANESE_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.FloatTensor] = None,
        spout: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = False,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        output_router_logits: Optional[bool] = None,
        labels: Optional[torch.LongTensor] = None,
    ) -> Union[Tuple[torch.FloatTensor], MoECausalLMOutputWithPast]:
        r
        SEG_TOKEN = self.config.separator_token_id
        use_cache = use_cache or self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        model_return_dict = True
        num_precontext = None
        if input_ids is not None:
            num_batch = input_ids.shape[0]
            num_precontext = torch.zeros([num_batch]).int().to(input_ids.device)
            where_separators = torch.where(input_ids == SEG_TOKEN)
            num_precontext[where_separators[0]] += where_separators[1]
            num_precontext = num_precontext.unsqueeze(1)

        outputs = self.model(
            input_ids,
            attention_mask,
            token_type_ids,
            spout,
            past_key_values,
            head_mask,
            use_cache,
            inputs_embeds,
            decoder_inputs_embeds,
            output_attentions,
            output_hidden_states,
            model_return_dict,
            output_router_logits,
            num_precontext,
        )

        lm_logits = self.lm_head(outputs[0])
        if lm_logits.shape[-1] == self.final_logits_bias.shape[-1]:
            lm_logits = lm_logits + self.final_logits_bias

        loss = None
        z_loss = None
        router_probs = None
        aux_loss = None
        if labels is not None:
            
            labels = labels.to(lm_logits.device)

            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)

            if output_router_logits:
                
                router_logits, expert_indexes = self._unpack_router_logits(outputs.router_probs)
                z_loss = router_z_loss_func(router_logits)
                router_probs = nn.Softmax(dim=-1)(router_logits)
                aux_loss = load_balancing_loss_func(router_probs, expert_indexes)

            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))

        if not return_dict:
            return tuple(
                v
                for v in [
                    loss,
                    lm_logits,
                    outputs.past_key_values,
                    outputs.hidden_states,
                    outputs.router_probs,
                    z_loss,
                    aux_loss,
                ]
                if v is not None
            )

        return MoECausalLMOutputWithPast(
            loss=loss,
            logits=lm_logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            router_logits=outputs.router_probs,
            z_loss=z_loss,
            aux_loss=aux_loss,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.FloatTensor,
        token_type_ids: Optional[torch.FloatTensor] = None,
        spout: Optional[Union[List, torch.FloatTensor]] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        **kwargs,
    ):
        if isinstance(spout, list):
            spout = torch.tensor(spout).float()
            if input_ids is not None:
                spout = spout.to(input_ids.device)
        if past_key_values is not None:
            return {
                "input_ids": input_ids[:, -1:] if input_ids is not None else None,
                "attention_mask": attention_mask,
                "token_type_ids": token_type_ids[:, -1:] if token_type_ids is not None else None,
                "spout": spout,
                "past_key_values": past_key_values,
            }
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "token_type_ids": token_type_ids,
            "spout": spout,
            "past_key_values": None,
        }

    
    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
        return self._shift_right(labels)

    
    def resize_token_embeddings(self, new_num_tokens: int, pad_to_multiple_of: Optional[int] = None) -> nn.Embedding:
        new_embeddings = super().resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
        self._resize_final_logits_bias(new_embeddings.weight.shape[0])
        return new_embeddings

    
    def _resize_final_logits_bias(self, new_num_tokens: int) -> None:
        old_num_tokens = self.final_logits_bias.shape[-1]
        if new_num_tokens <= old_num_tokens:
            new_bias = self.final_logits_bias[:, :new_num_tokens]
        else:
            extra_bias = torch.zeros((1, new_num_tokens - old_num_tokens), device=self.final_logits_bias.device)
            new_bias = torch.cat([self.final_logits_bias, extra_bias], dim=1)
        self.register_buffer("final_logits_bias", new_bias)

    def get_input_embeddings(self):
        return self.model.get_input_embeddings()

    def set_input_embeddings(self, new_embeddings):
        self.model.set_input_embeddings(new_embeddings)

    
    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    
    def get_output_embeddings(self):
        return self.lm_head

    
    def _unpack_router_logits(self, router_outputs):
        total_router_logits = []
        total_expert_indexes = []
        for router_output in router_outputs:
            if len(router_output[0].shape) > 1:
                router_logits, expert_indexes = router_output
                total_router_logits.append(router_logits)
                total_expert_indexes.append(expert_indexes)
        return torch.cat(total_router_logits, dim=1), torch.cat(total_expert_indexes, dim=1)