codingutf8 2020 the facebook ai research team s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license original implementation https github compytorchfairseqtreemasterexampleswmt19 s alexeib alexei baevski edunov sergey edunov michaelauli michael auli myleott myle ott nng555 nathan ng david grangier kyra yee paper facebook fair s wmt19 news translation task submission https arxiv orgabs1907 06616 pytorch fairseq model ported from https github compytorchfairseqtreemasterexampleswmt19 import math from typing import any dict list optional tuple union import torch from torch import tensor nn from torch nn import crossentropyloss layernorm from activations import act2fn from integrations deepspeed import isdeepspeedzero3enabled from modelingoutputs import basemodeloutput basemodeloutputwithpastandcrossattentions seq2seqlmoutput seq2seqmodeloutput from modelingutils import pretrainedmodel from utils import addcodesampledocstrings addenddocstrings addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from configurationfsmt import fsmtconfig logger logging getloggername checkpointfordoc facebookwmt19ruen configfordoc fsmtconfig see all fsmt models at https huggingface comodels filterfsmt porting notes this one is modeled after bartmodel currently only translation fairseq also has weights for lm fairseq provides weights for ruen enru and deen ende pairs all have been ported ruen enru use asymmetric vocab deen ende use a merged single vocab but the code works as if they are separate differences with bart not using bos token 2 separate vocabs src and target embed weights aren t tied uses a model ensemble but that part isn t portedimplemented yet so we aren t getting as good of a bleu score uses a projection layer at the end of the decoder doesn t use finallogitsbias beam search stops as soon as numbeams lenhypos whereas transformers is not satisfied there and will continue searching until the next cycles aren t promising something better comparing bleu scores the transformers algorithm is slightly superior therefore using the latter but if you want to match fairseq outputs you need to pass earlystoppingtrue to generate sinusoidalpositionalembedding is slightly different from bart s generates different embeddings this implementation is copied verbatim from fairseq with some small changes to make it work here other changes doesn t support usecache as bart s version does fsmtconfig changes with bartconfig differences with bart srctgt vocabs aren t shared token embeddings aren t shared needs a language pair scaleembedding are true some unused args were removed too todo port model ensemble fs uses 4 model checkpoints solve beam search discrepancies docstyleignore enru fairseq bleu 36 4 http matrix statmt orgmatrixoutput1914 scoreid37605 ruen fairseq bleu 41 3 http matrix statmt orgmatrixoutput1907 runid6937 deen fairseq bleu 42 3 http matrix statmt orgmatrixoutput1902 runid6750 ende fairseq bleu 43 1 http matrix statmt orgmatrixoutput1909 runid6862 fsmtstartdocstring r this model inherits from pretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a pytorch torch nn modulehttps pytorch orgdocsstablenn htmltorch 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 fsmtconfig 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 fsmtinputsdocstring r args inputids torch longtensor of shape batchsize sequencelength indices of input sequence tokens in the vocabulary indices can be obtained using fstmtokenizer see pretrainedtokenizer encode and pretrainedtokenizer call for details what are input ids glossaryinputids attentionmask torch tensor 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 decoderinputids torch longtensor of shape batchsize targetsequencelength optional indices of decoder input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer call for details what are decoder input ids glossarydecoderinputids fsmt uses the eostokenid as the starting token for decoderinputids generation if pastkeyvalues is used optionally only the last decoderinputids have to be input see pastkeyvalues decoderattentionmask torch booltensor of shape batchsize targetsequencelength optional default behavior generate a tensor that ignores pad tokens in decoderinputids causal mask will also be used by default headmask torch tensor of shape encoderlayers encoderattentionheads optional mask to nullify selected heads of the attention modules in the encoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked decoderheadmask torch tensor of shape decoderlayers decoderattentionheads optional mask to nullify selected heads of the attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked crossattnheadmask torch tensor of shape decoderlayers decoderattentionheads optional mask to nullify selected heads of the crossattention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked encoderoutputs tupletorch floattensor optional tuple consists of lasthiddenstate optional hiddenstates optional attentions lasthiddenstate of shape batchsize sequencelength hiddensize is a sequence of hiddenstates at the output of the last layer of the encoder used in the crossattention of the decoder pastkeyvalues tupletorch floattensor of length config nlayers with each tuple having 4 tensors of shape batchsize numheads sequencelength 1 embedsizeperhead contains precomputed key and value hiddenstates 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 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 if decoderinputids and decoderinputsembeds are both unset decoderinputsembeds takes the value of inputsembeds usecache bool optional defaults to true 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 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 turns 10 01 falsetrue true false assert attentionmask dim 2 return attentionmask eq0 def triuonnxx diagonal0 l x shape0 arange torch arangel devicex device mask arange expandl l arange arange unsqueeze1 if diagonal arange arange diagonal mask mask arange return x maskedfillmask 0 0 def preparefsmtdecoderinputs config inputids decoderinputidsnone decoderpaddingmasknone causalmaskdtypetorch float32 padtokenid config padtokenid if decoderinputids is none decoderinputids shifttokensrightinputids padtokenid bsz tgtlen decoderinputids size if decoderpaddingmask is none decoderpaddingmask makepaddingmaskdecoderinputids padtokenid else decoderpaddingmask invertmaskdecoderpaddingmask causalmask triuonnxfillwithneginftorch zerostgtlen tgtlen dtypecausalmaskdtype 1 to devicedecoderinputids device return decoderinputids decoderpaddingmask causalmask class pretrainedfsmtmodelpretrainedmodel configclass fsmtconfig basemodelprefix model def initweightsself module std self config initstd if isinstancemodule nn linear module weight data normalmean0 0 stdstd if module bias is not none module bias data zero elif isinstancemodule sinusoidalpositionalembedding pass elif isinstancemodule nn embedding module weight data normalmean0 0 stdstd if module paddingidx is not none module weight datamodule paddingidx zero property def dummyinputsself padtoken self config padtokenid inputids torch tensor0 6 10 4 2 0 8 12 2 padtoken deviceself device dummyinputs attentionmask inputids nepadtoken inputids inputids return dummyinputs def makelinearfromembemb vocabsize embsize emb weight shape linlayer nn linearvocabsize embsize biasfalse linlayer weight data emb weight data return linlayer helper functions mostly for making masks def checkshapesshape1 shape2 if shape1 shape2 raise assertionerrorfshape mismatch shape1 shape2 def shifttokensrightinputids padtokenid replace possible 100 values in labels by padtokenid true for pad tokens paddingmask inputids eqpaddingidx if not paddingmask any paddingmask none return paddingmask helper modules class encoderlayernn module def initself config fsmtconfig super init self embeddim config dmodel self selfattn attentionself embeddim config encoderattentionheads dropoutconfig attentiondropout self selfattnlayernorm layernormself embeddim self dropout config dropout self activationfn act2fnconfig activationfunction self activationdropout config activationdropout self fc1 nn linearself embeddim config encoderffndim self fc2 nn linearconfig encoderffndim self embeddim self finallayernorm layernormself embeddim def forwardself x encoderpaddingmask layerheadmask outputattentionsfalse residual x x attnweights self selfattn queryx keyx keypaddingmaskencoderpaddingmask layerheadmasklayerheadmask outputattentionsoutputattentions x nn functional dropoutx pself dropout trainingself training x residual x x self selfattnlayernormx residual x x self activationfnself fc1x x nn functional dropoutx pself activationdropout trainingself training x self fc2x x nn functional dropoutx pself dropout trainingself training x residual x x self finallayernormx return x attnweights class fsmtencodernn module def initself config fsmtconfig embedtokens super init self dropout config dropout self layerdrop config encoderlayerdrop self paddingidx embedtokens paddingidx self embedtokens embedtokens embeddim embedtokens embeddingdim self embedscale math sqrtembeddim if config scaleembedding else 1 0 self embedpositions sinusoidalpositionalembedding config maxpositionembeddings self paddingidx 1 embeddim self paddingidx self layers nn modulelistencoderlayerconfig for in rangeconfig encoderlayers type listencoderlayer def forward self inputids torch tensor attentionmask optionaltorch tensor none inputsembeds torch tensor none headmask optionaltorch tensor none outputattentions bool false outputhiddenstates bool false returndict bool true check attention mask and invert if attentionmask is not none attentionmask invertmaskattentionmask if inputids is not none and inputsembeds is not none raise valueerroryou cannot specify both inputids and inputsembeds at the same time elif inputids is not none inputsembeds self embedtokensinputids self embedscale embedpos self embedpositionsinputids elif inputsembeds is not none inputsembeds inputsembeds self embedscale we assume zeros hidden states correspond to padding tokens and create positionids where inputsembeds 0 0 positionids inputsembeds 0 maskedfill inputsembeds 0 eq0 self embedpositions paddingidx embedpos self embedpositionspositionids else raise valueerroryou have to specify either inputids or inputsembeds x inputsembeds embedpos x nn functional dropoutx pself dropout trainingself training b x t x c t x b x c x x transpose0 1 encoderstates if outputhiddenstates else none allattentions if outputattentions else none check if headmask has a correct number of layers specified if desired if headmask is not none assert headmask size0 lenself layers fthe headmask should be specified for lenself layers layers but it is for headmask size0 for idx encoderlayer in enumerateself layers if outputhiddenstates x x transpose0 1 t x b x c b x t x c encoderstates x x x transpose0 1 b x t x c t x b x c add layerdrop see https arxiv orgabs1909 11556 for description dropoutprobability torch rand if self training and dropoutprobability self layerdrop skip the layer attn none else x attn encoderlayer x attentionmask layerheadmaskheadmaskidx if headmask is not none else none outputattentionsoutputattentions if outputattentions allattentions allattentions attn t x b x c b x t x c x x transpose0 1 if outputhiddenstates encoderstates x if not returndict return tuplev for v in x encoderstates allattentions if v is not none return basemodeloutputlasthiddenstatex hiddenstatesencoderstates attentionsallattentions class decoderlayernn module def initself config fsmtconfig super init self embeddim config dmodel self selfattn attention embeddimself embeddim numheadsconfig decoderattentionheads dropoutconfig attentiondropout self dropout config dropout self activationfn act2fnconfig activationfunction self activationdropout config activationdropout self selfattnlayernorm layernormself embeddim self encoderattn attention self embeddim config decoderattentionheads dropoutconfig attentiondropout encoderdecoderattentiontrue self encoderattnlayernorm layernormself embeddim self fc1 nn linearself embeddim config decoderffndim self fc2 nn linearconfig decoderffndim self embeddim self finallayernorm layernormself embeddim def forward self x encoderhiddenstates encoderattnmasknone layerstatenone causalmasknone layerheadmasknone crossattnlayerheadmasknone decoderpaddingmasknone outputattentionsfalse residual x if layerstate is none layerstate self attention x selfattnweights self selfattn queryx keyx layerstatelayerstate adds keys to layer state keypaddingmaskdecoderpaddingmask attnmaskcausalmask layerheadmasklayerheadmask outputattentionsoutputattentions x nn functional dropoutx pself dropout trainingself training x residual x x self selfattnlayernormx cross attention residual x assert self encoderattn cachekey self selfattn cachekey x crossattnweights self encoderattn queryx keyencoderhiddenstates keypaddingmaskencoderattnmask layerstatelayerstate mutates layer state layerheadmaskcrossattnlayerheadmask outputattentionsoutputattentions x nn functional dropoutx pself dropout trainingself training x residual x x self encoderattnlayernormx fully connected residual x x self activationfnself fc1x x nn functional dropoutx pself activationdropout trainingself training x self fc2x x nn functional dropoutx pself dropout trainingself training x residual x x self finallayernormx return x selfattnweights layerstate crossattnweights layerstate cache for decoding class fsmtdecodernn module def initself config fsmtconfig embedtokens nn embedding super init self dropout config dropout self layerdrop config decoderlayerdrop self paddingidx embedtokens paddingidx self embedscale math sqrtconfig dmodel if config scaleembedding else 1 0 self embedtokens embedtokens embeddim embedtokens embeddingdim self embedpositions sinusoidalpositionalembedding config maxpositionembeddings self paddingidx 1 embeddim self paddingidx self layers nn modulelistdecoderlayerconfig for in rangeconfig decoderlayers type listdecoderlayer if isdeepspeedzero3enabled import deepspeed with deepspeed zero gatheredparametersself embedtokens weight modifierranknone embedtokensweightshape self embedtokens weight shape else embedtokensweightshape self embedtokens weight shape self outputprojection nn linearembedtokensweightshape1 embedtokensweightshape0 biasfalse self outputprojection weight self embedtokens weight def forward self inputids torch tensor encoderhiddenstates torch tensor encoderpaddingmask torch tensor decoderpaddingmask torch tensor decodercausalmask torch tensor headmask optionaltorch tensor none inputsembeds optionaltorch tensor none crossattnheadmask optionaltorch tensor none pastkeyvalues optionallisttorch floattensor none usecache bool false outputattentions bool false outputhiddenstates bool false returndict bool true check attention mask and invert if encoderpaddingmask is not none encoderpaddingmask invertmaskencoderpaddingmask if inputids is not none and inputsembeds is not none raise valueerroryou cannot specify both decoderinputids and decoderinputsembeds at the same time elif inputids is not none embed positions positions self embedpositionsinputids if usecache inputids inputids 1 positions positions 1 happens after we embed them x self embedtokensinputids self embedscale elif inputsembeds is not none we assume zeros hidden states correspond to padding tokens and create positionids where inputsembeds 0 0 positionids inputsembeds 0 maskedfill inputsembeds 0 eq0 self embedpositions paddingidx positions self embedpositionspositionids x inputsembeds self embedscale else raise valueerroryou have to specify either decoderinputids or decoderinputsembeds x positions x nn functional dropoutx pself dropout trainingself training convert to fsmt output format bs seqlen modeldim seqlen bs modeldim x x transpose0 1 encoderhiddenstates encoderhiddenstates transpose0 1 decoder layers allhiddenstates if outputhiddenstates else none allselfattns if outputattentions else none allcrossattns if outputattentions else none nextdecodercache check if headmask has a correct number of layers specified if desired for attnmask maskname in zipheadmask crossattnheadmask headmask crossattnheadmask if attnmask is not none assert attnmask size0 lenself layers fthe maskname should be specified for lenself layers layers but it is for f headmask size0 for idx decoderlayer in enumerateself layers add layerdrop see https arxiv orgabs1909 11556 for description if outputhiddenstates x x transpose0 1 allhiddenstates x x x transpose0 1 if self training dropoutprobability torch rand if dropoutprobability self layerdrop continue layerstate pastkeyvaluesidx if pastkeyvalues is not none else none x layerselfattn layerpast layercrossattn decoderlayer x encoderhiddenstates encoderattnmaskencoderpaddingmask decoderpaddingmaskdecoderpaddingmask layerstatelayerstate causalmaskdecodercausalmask layerheadmaskheadmaskidx if headmask is not none else none crossattnlayerheadmaskcrossattnheadmaskidx if crossattnheadmask is not none else none outputattentionsoutputattentions if usecache nextdecodercache appendlayerpast copy if outputattentions allselfattns layerselfattn allcrossattns layercrossattn add hidden states from the last decoder layer if outputhiddenstates x x transpose0 1 allhiddenstates x x x transpose0 1 convert to standard output format seqlen bs modeldim bs seqlen modeldim x x transpose0 1 encoderhiddenstates encoderhiddenstates transpose0 1 x self outputprojectionx nextcache nextdecodercache if usecache else none if not returndict return tuple v for v in x nextcache allhiddenstates allselfattns allcrossattns if v is not none return basemodeloutputwithpastandcrossattentions lasthiddenstatex pastkeyvaluesnextcache hiddenstatesallhiddenstates attentionsallselfattns crossattentionsallcrossattns def reorderbufferattncache neworder for k inputbufferk in attncache items if inputbufferk is not none attncachek inputbufferk indexselect0 neworder return attncache class attentionnn module input shape timeseqlen x batch x channel statickv bool self encoderdecoderattention tgtlen bsz embeddim query size assert embeddim self embeddim assert listquery size tgtlen bsz embeddim get here for encoder decoder cause of statickv if layerstate is not none reuse k v and encoderpaddingmask savedstate layerstate getself cachekey if prevkey in savedstate and statickv previous time steps are cached no need to recompute key and value if they are static key none else savedstate none layerstate q self qprojquery self scaling if statickv if key is none k v none else k self kprojkey v self vprojkey else k self kprojquery v self vprojquery q self shapeq tgtlen bsz if k is not none k self shapek 1 bsz if v is not none v self shapev 1 bsz if savedstate is not none k v keypaddingmask self usesavedstatek v savedstate keypaddingmask statickv bsz update cache layerstateself cachekey prevkey k viewbsz self numheads 1 self headdim prevvalue v viewbsz self numheads 1 self headdim prevkeypaddingmask keypaddingmask if not statickv else none assert k is not none srclen k size1 attnweights torch bmmq k transpose1 2 assert attnweights size bsz self numheads tgtlen srclen if attnmask is not none attnweights attnweights viewbsz self numheads tgtlen srclen attnmask attnweights attnweights viewbsz self numheads tgtlen srclen this is part of a workaround to get around forkjoin parallelism not supporting optional types if keypaddingmask is not none and keypaddingmask dim 0 keypaddingmask none assert keypaddingmask is none or keypaddingmask size 2 bsz srclen if keypaddingmask is not none don t attend to padding symbols attnweights attnweights viewbsz self numheads tgtlen srclen reshaped keypaddingmask unsqueeze1 unsqueeze2 attnweights attnweights maskedfillreshaped torch finfoattnweights dtype min attnweights attnweights viewbsz self numheads tgtlen srclen attnweights nn functional softmaxattnweights dim1 if layerheadmask is not none assert layerheadmask size self numheads fhead mask for a single layer should be of size self numheads but is layerheadmask size attnweights layerheadmask view1 1 1 1 attnweights viewbsz self numheads tgtlen srclen attnweights attnweights viewbsz self numheads tgtlen srclen if outputattentions make sure that attnweights are included in graph attnweightsreshaped attnweights viewbsz self numheads tgtlen srclen attnweights attnweightsreshaped viewbsz self numheads tgtlen srclen else attnweightsreshaped none attnprobs nn functional dropout attnweights pself dropout trainingself training assert v is not none attnoutput torch bmmattnprobs v assert attnoutput size bsz self numheads tgtlen self headdim attnoutput attnoutput transpose0 1 contiguous viewtgtlen bsz embeddim attnoutput self outprojattnoutput return attnoutput attnweightsreshaped def usesavedstateself k v savedstate keypaddingmask statickv bsz saved states are stored with shape bsz numheads seqlen headdim if prevkey in savedstate prevkey savedstateprevkey assert prevkey is not none prevkey prevkey viewbsz self numheads 1 self headdim if statickv k prevkey else assert k is not none k torch catprevkey k dim1 if prevvalue in savedstate prevvalue savedstateprevvalue assert prevvalue is not none prevvalue prevvalue viewbsz self numheads 1 self headdim if statickv v prevvalue else assert v is not none v torch catprevvalue v dim1 assert k is not none and v is not none prevkeypaddingmask optionaltensor savedstate getprevkeypaddingmask none if prevkeypaddingmask is not none if statickv newkeypaddingmask prevkeypaddingmask else newkeypaddingmask torch catprevkeypaddingmask keypaddingmask dim1 else newkeypaddingmask keypaddingmask return k v newkeypaddingmask def fillwithneginft public api initialize weights and apply final processing make masks if user doesn t supply if the user passed a tuple for encoderoutputs we wrap it in a basemodeloutput when returndictfalse decoder outputs consists of decfeatures layerstate dechidden decattn initialize weights and apply final processing returndict returndict if returndict is not none else self config usereturndict if labels is not none usecache false outputs self model inputids inputsembedsinputsembeds attentionmaskattentionmask decoderinputidsdecoderinputids decoderinputsembedsdecoderinputsembeds encoderoutputsencoderoutputs decoderattentionmaskdecoderattentionmask headmaskheadmask decoderheadmaskdecoderheadmask crossattnheadmaskcrossattnheadmask pastkeyvaluespastkeyvalues usecacheusecache outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict lmlogits outputs0 maskedlmloss none if labels is not none lossfct crossentropyloss todoss do we need to ignore pad tokens in labels maskedlmloss lossfctlmlogits view1 self config tgtvocabsize labels view1 if not returndict output lmlogits outputs1 return maskedlmloss output if maskedlmloss is not none else output return seq2seqlmoutput lossmaskedlmloss logitslmlogits pastkeyvaluesoutputs pastkeyvalues decoderhiddenstatesoutputs decoderhiddenstates decoderattentionsoutputs decoderattentions crossattentionsoutputs crossattentions encoderlasthiddenstateoutputs encoderlasthiddenstate encoderhiddenstatesoutputs encoderhiddenstates encoderattentionsoutputs encoderattentions def prepareinputsforgeneration self decoderinputids pastkeyvaluesnone attentionmasknone headmasknone decoderheadmasknone crossattnheadmasknone usecachenone encoderoutputsnone kwargs return inputids none encoderoutputs is defined inputids not needed encoderoutputs encoderoutputs pastkeyvalues pastkeyvalues decoderinputids decoderinputids attentionmask attentionmask headmask headmask decoderheadmask decoderheadmask crossattnheadmask crossattnheadmask usecache usecache change this to avoid caching presumably for debugging def preparedecoderinputidsfromlabelsself labels torch tensor return shifttokensrightlabels self config padtokenid staticmethod def reordercachepastkeyvalues beamidx reorderedpast for layerpast in pastkeyvalues get the correct batch idx from decoder layer s batch dim for cross and selfattn layerpastnew attnkey reorderbufferattncache beamidx for attnkey attncache in layerpast items reorderedpast appendlayerpastnew return reorderedpast def getencoderself return self model encoder def getdecoderself return self model decoder def getoutputembeddingsself return self model decoder embedtokens def setoutputembeddingsself value self model decoder embedtokens value class sinusoidalpositionalembeddingnn embedding def initself numpositions embeddingdim paddingidx self makeweightnumpositions embeddingdim paddingidx def makeweightself numpositions embeddingdim paddingidx weight self getembeddingnumpositions embeddingdim paddingidx if not hasattrself weight in init super initnumpositions embeddingdim paddingidx weightweight else in forward put the weights on the correct dtype and device of the param weight weight todtypeself weight dtype deviceself weight device self weight nn parameterweight self weight detach self weight requiresgrad false staticmethod def getembeddingnumembeddings embeddingdim paddingidx halfdim embeddingdim 2 emb math log10000 halfdim 1 emb torch exptorch arangehalfdim dtypetorch float emb emb torch arangenumembeddings dtypetorch float unsqueeze1 emb unsqueeze0 emb torch cattorch sinemb torch cosemb dim1 viewnumembeddings 1 if embeddingdim 2 1 zero pad emb torch catemb torch zerosnumembeddings 1 dim1 if paddingidx is not none embpaddingidx 0 return emb staticmethod def makepositionstensor paddingidx int the series of casts and typeconversions here are carefully balanced to both work with onnx export and xla in particular xla prefers ints cumsum defaults to output longs and onnx doesn t know how to handle the dtype kwarg in cumsum mask tensor nepaddingidx int return torch cumsummask dim1 typeasmask mask long paddingidx def forward self input incrementalstate optionalany none timestep optionaltensor none expand embeddings if needed coding utf 8 2020 the facebook ai research team s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license original implementation https github com pytorch fairseq tree master examples wmt19 s alexeib alexei baevski edunov sergey edunov michaelauli michael auli myleott myle ott nng555 nathan ng david grangier kyra yee paper facebook fair s wmt19 news translation task submission https arxiv org abs 1907 06616 pytorch fairseq model ported from https github com pytorch fairseq tree master examples wmt19 see all fsmt models at https huggingface co models filter fsmt porting notes this one is modeled after bartmodel currently only translation fairseq also has weights for lm fairseq provides weights for ru en en ru and de en en de pairs all have been ported ru en en ru use asymmetric vocab de en en de use a merged single vocab but the code works as if they are separate differences with bart not using bos token 2 separate vocabs src and target embed weights aren t tied uses a model ensemble but that part isn t ported implemented yet so we aren t getting as good of a bleu score uses a projection layer at the end of the decoder doesn t use final_logits_bias beam search stops as soon as num_beams len hypos whereas transformers is not satisfied there and will continue searching until the next cycles aren t promising something better comparing bleu scores the transformers algorithm is slightly superior therefore using the latter but if you want to match fairseq outputs you need to pass early_stopping true to generate sinusoidalpositionalembedding is slightly different from bart s generates different embeddings this implementation is copied verbatim from fairseq with some small changes to make it work here other changes doesn t support use_cache as bart s version does fsmtconfig changes with bartconfig differences with bart src tgt vocabs aren t shared token embeddings aren t shared needs a language pair scale_embedding are true some unused args were removed too todo port model ensemble fs uses 4 model checkpoints solve beam search discrepancies docstyle ignore here is how to compare bleu scores against fairseq implementation en ru export pair en ru export data_dir data pair export save_dir data pair export bs 8 export num_beams 50 mkdir p data_dir sacrebleu t wmt19 l pair echo src data_dir val source sacrebleu t wmt19 l pair echo ref data_dir val target echo pair pythonpath src examples seq2seq python examples seq2seq run_eval py facebook wmt19 pair data_dir val source save_dir test_translations txt reference_path data_dir val target score_path save_dir test_bleu json bs bs task translation num_beams num_beams fairseq bleu 36 4 http matrix statmt org matrix output 1914 score_id 37605 ru en export pair ru en export data_dir data pair export save_dir data pair export bs 8 export num_beams 50 mkdir p data_dir sacrebleu t wmt19 l pair echo src data_dir val source sacrebleu t wmt19 l pair echo ref data_dir val target pythonpath src examples seq2seq python examples seq2seq run_eval py facebook wmt19 pair data_dir val source save_dir test_translations txt reference_path data_dir val target score_path save_dir test_bleu json bs bs task translation num_beams num_beams fairseq bleu 41 3 http matrix statmt org matrix output 1907 run_id 6937 de en export pair de en export data_dir data pair export save_dir data pair export bs 8 export num_beams 50 mkdir p data_dir sacrebleu t wmt19 l pair echo src data_dir val source sacrebleu t wmt19 l pair echo ref data_dir val target echo pair pythonpath src examples seq2seq python examples seq2seq run_eval py facebook wmt19 pair data_dir val source save_dir test_translations txt reference_path data_dir val target score_path save_dir test_bleu json bs bs task translation num_beams num_beams fairseq bleu 42 3 http matrix statmt org matrix output 1902 run_id 6750 en de export pair en de export data_dir data pair export save_dir data pair export bs 8 mkdir p data_dir sacrebleu t wmt19 l pair echo src data_dir val source sacrebleu t wmt19 l pair echo ref data_dir val target echo pair pythonpath src examples seq2seq python examples seq2seq run_eval py facebook wmt19 pair data_dir val source save_dir test_translations txt reference_path data_dir val target score_path save_dir test_bleu json bs bs task translation num_beams num_beams fairseq bleu 43 1 http matrix statmt org matrix output 1909 run_id 6862 this model inherits from pretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc 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 fsmtconfig 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 translation example python from transformers import autotokenizer fsmtforconditionalgeneration mname facebook wmt19 ru en model fsmtforconditionalgeneration from_pretrained mname tokenizer autotokenizer from_pretrained mname src_text машинное обучение это здорово не так ли input_ids tokenizer src_text return_tensors pt input_ids outputs model generate input_ids num_beams 5 num_return_sequences 3 tokenizer decode outputs 0 skip_special_tokens true machine learning is great isn t it args input_ids torch longtensor of shape batch_size sequence_length indices of input sequence tokens in the vocabulary indices can be obtained using fstmtokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask torch tensor 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 decoder_input_ids torch longtensor of shape batch_size target_sequence_length optional indices of decoder input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are decoder input ids glossary decoder input ids fsmt uses the eos_token_id as the starting token for decoder_input_ids generation if past_key_values is used optionally only the last decoder_input_ids have to be input see past_key_values decoder_attention_mask torch booltensor of shape batch_size target_sequence_length optional default behavior generate a tensor that ignores pad tokens in decoder_input_ids causal mask will also be used by default head_mask torch tensor of shape encoder_layers encoder_attention_heads optional mask to nullify selected heads of the attention modules in the encoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked decoder_head_mask torch tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked cross_attn_head_mask torch tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the cross attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked encoder_outputs tuple torch floattensor optional tuple consists of last_hidden_state optional hidden_states optional attentions last_hidden_state of shape batch_size sequence_length hidden_size is a sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder past_key_values 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 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 if decoder_input_ids and decoder_inputs_embeds are both unset decoder_inputs_embeds takes the value of inputs_embeds use_cache bool optional defaults to true 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 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 turns 1 0 0 1 false true true false prepare masks that ignore padding tokens in the decoder and a causal mask for the decoder if none are provided this mimics the default behavior in fairseq to override it pass in masks note this is not called during generation helper functions mostly for making masks shift input ids one token to the right and wrap the last non pad token usually eos replace possible 100 values in labels by pad_token_id true for pad tokens helper modules args x torch tensor input to the layer of shape seq_len batch embed_dim encoder_padding_mask torch bytetensor binary bytetensor of shape batch src_len where padding elements are indicated by 1 for t_tgt t_src is excluded or masked out 0 means it is included in attention layer_head_mask torch floattensor mask for attention heads in a given layer of size config encoder_attention_heads returns encoded output of shape seq_len batch embed_dim transformer encoder consisting of config encoder_layers self attention layers each layer is a encoderlayer args config fsmtconfig type list encoderlayer args input_ids torch longtensor tokens in the source language of shape batch src_len attention_mask torch longtensor indicating which indices are padding tokens inputs_embeds torch floattensor embedding vectors of shape batch src_len embed_dim head_mask torch tensor of shape num_layers num_heads 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 returns basemodeloutput or tuple comprised of x torch tensor the last encoder layer s output of shape src_len batch embed_dim encoder_states tuple torch floattensor all intermediate hidden states of shape src_len batch embed_dim only populated if output_hidden_states is true all_attentions tuple torch floattensor attention weights for each layer during training might not be of length n_layers because of layer dropout check attention mask and invert we assume zeros hidden states correspond to padding tokens and create position_ids where inputs_embeds 0 0 b x t x c t x b x c check if head_mask has a correct number of layers specified if desired t x b x c b x t x c b x t x c t x b x c add layerdrop see https arxiv org abs 1909 11556 for description skip the layer t x b x c b x t x c self attention adds keys to layer state cross attention mutates layer state fully connected layer_state cache for decoding transformer decoder consisting of config decoder_layers layers each layer is a decoderlayer args config fsmtconfig embed_tokens nn embedding output embedding type list decoderlayer includes several features from jointly learning to align and translate with transformer models garg et al emnlp 2019 args input_ids torch longtensor of shape batch tgt_len previous decoder outputs for teacher forcing encoder_hidden_states output from the encoder used for encoder side attention encoder_padding_mask for ignoring pad tokens past_key_values dict or none dictionary used for storing state during generation head_mask torch tensor of shape num_layers num_heads 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 cross_attn_head_mask torch tensor of shape num_layers num_heads optional mask to nullify selected heads of the cross attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked returns basemodeloutputwithpast or tuple the decoder s features of shape batch tgt_len embed_dim the cache hidden states attentions check attention mask and invert embed positions happens after we embed them we assume zeros hidden states correspond to padding tokens and create position_ids where inputs_embeds 0 0 convert to fsmt output format bs seq_len model_dim seq_len bs model_dim decoder layers check if head_mask has a correct number of layers specified if desired add layerdrop see https arxiv org abs 1909 11556 for description add hidden states from the last decoder layer convert to standard output format seq_len bs model_dim bs seq_len model_dim multi headed attention from attention is all you need paper otherwise self_attention input shape time seqlen x batch x channel get here for encoder decoder cause of static_kv reuse k v and encoder_padding_mask previous time steps are cached no need to recompute key and value if they are static update cache this is part of a workaround to get around fork join parallelism not supporting optional types don t attend to padding symbols make sure that attn_weights are included in graph saved states are stored with shape bsz num_heads seq_len head_dim fp16 compatible function that fills a input_ids with inf public api initialize weights and apply final processing make masks if user doesn t supply if the user passed a tuple for encoder_outputs we wrap it in a basemodeloutput when return_dict false decoder outputs consists of dec_features layer_state dec_hidden dec_attn initialize weights and apply final processing labels torch longtensor of shape batch_size sequence_length optional labels for computing the masked language modeling loss indices should either be in 0 config vocab_size or 100 see input_ids docstring tokens with indices set to 100 are ignored masked the loss is only computed for the tokens with labels in 0 config vocab_size returns todo ss do we need to ignore pad tokens in labels encoder_outputs is defined input_ids not needed change this to avoid caching presumably for debugging get the correct batch idx from decoder layer s batch dim for cross and self attn this module produces sinusoidal positional embeddings of any length we don t want to save the weight of this embedding since it s not trained deterministic and it can be huge padding symbols are ignored these embeddings get automatically extended in forward if more positions is needed in ___init__ in forward put the weights on the correct dtype and device of the param build sinusoidal embeddings this matches the implementation in tensor2tensor but differs slightly from the description in section 3 5 of attention is all you need zero pad replace non padding symbols with their position numbers position numbers begin at padding_idx 1 padding symbols are ignored the series of casts and type conversions here are carefully balanced to both work with onnx export and xla in particular xla prefers ints cumsum defaults to output longs and onnx doesn t know how to handle the dtype kwarg in cumsum input is expected to be of size bsz x seqlen expand embeddings if needed
import math
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import Tensor, nn
from torch.nn import CrossEntropyLoss, LayerNorm

from ...activations import ACT2FN
from ...integrations.deepspeed import is_deepspeed_zero3_enabled
from ...modeling_outputs import (
    BaseModelOutput,
    BaseModelOutputWithPastAndCrossAttentions,
    Seq2SeqLMOutput,
    Seq2SeqModelOutput,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_end_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from .configuration_fsmt import FSMTConfig


logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "facebook/wmt19-ru-en"
_CONFIG_FOR_DOC = "FSMTConfig"





















































FSMT_START_DOCSTRING = r
FSMT_GENERATION_EXAMPLE = r

FSMT_INPUTS_DOCSTRING = r


def invert_mask(attention_mask):
    
    assert attention_mask.dim() == 2
    return attention_mask.eq(0)


def triu_onnx(x, diagonal=0):
    l = x.shape[0]
    arange = torch.arange(l, device=x.device)
    mask = arange.expand(l, l)
    arange = arange.unsqueeze(-1)
    if diagonal:
        arange = arange + diagonal
    mask = mask >= arange
    return x.masked_fill(mask == 0, 0)


def _prepare_fsmt_decoder_inputs(
    config,
    input_ids,
    decoder_input_ids=None,
    decoder_padding_mask=None,
    causal_mask_dtype=torch.float32,
):
    
    pad_token_id = config.pad_token_id
    if decoder_input_ids is None:
        decoder_input_ids = shift_tokens_right(input_ids, pad_token_id)
    bsz, tgt_len = decoder_input_ids.size()
    if decoder_padding_mask is None:
        decoder_padding_mask = make_padding_mask(decoder_input_ids, pad_token_id)
    else:
        decoder_padding_mask = invert_mask(decoder_padding_mask)
    causal_mask = triu_onnx(fill_with_neg_inf(torch.zeros(tgt_len, tgt_len, dtype=causal_mask_dtype)), 1).to(
        device=decoder_input_ids.device
    )
    return decoder_input_ids, decoder_padding_mask, causal_mask


class PretrainedFSMTModel(PreTrainedModel):
    config_class = FSMTConfig
    base_model_prefix = "model"

    def _init_weights(self, module):
        std = self.config.init_std
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, SinusoidalPositionalEmbedding):
            pass
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

    @property
    def dummy_inputs(self):
        pad_token = self.config.pad_token_id
        input_ids = torch.tensor([[0, 6, 10, 4, 2], [0, 8, 12, 2, pad_token]], device=self.device)
        dummy_inputs = {
            "attention_mask": input_ids.ne(pad_token),
            "input_ids": input_ids,
        }
        return dummy_inputs


def _make_linear_from_emb(emb):
    vocab_size, emb_size = emb.weight.shape
    lin_layer = nn.Linear(vocab_size, emb_size, bias=False)
    lin_layer.weight.data = emb.weight.data
    return lin_layer



def _check_shapes(shape_1, shape2):
    if shape_1 != shape2:
        raise AssertionError(f"shape mismatch: {shape_1} != {shape2}")


def shift_tokens_right(input_ids, pad_token_id):
    

    
    input_ids.masked_fill_(input_ids == -100, pad_token_id)

    prev_output_tokens = input_ids.clone()
    index_of_eos = (input_ids.ne(pad_token_id).sum(dim=1) - 1).unsqueeze(-1)
    prev_output_tokens[:, 0] = input_ids.gather(1, index_of_eos).squeeze()
    prev_output_tokens[:, 1:] = input_ids[:, :-1]
    return prev_output_tokens


def make_padding_mask(input_ids, padding_idx=1):
    
    padding_mask = input_ids.eq(padding_idx)
    if not padding_mask.any():
        padding_mask = None
    return padding_mask





class EncoderLayer(nn.Module):
    def __init__(self, config: FSMTConfig):
        super().__init__()
        self.embed_dim = config.d_model
        self.self_attn = Attention(self.embed_dim, config.encoder_attention_heads, dropout=config.attention_dropout)
        self.self_attn_layer_norm = LayerNorm(self.embed_dim)
        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
        self.activation_dropout = config.activation_dropout
        self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
        self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
        self.final_layer_norm = LayerNorm(self.embed_dim)

    def forward(self, x, encoder_padding_mask, layer_head_mask, output_attentions=False):
        
        residual = x
        x, attn_weights = self.self_attn(
            query=x,
            key=x,
            key_padding_mask=encoder_padding_mask,
            layer_head_mask=layer_head_mask,
            output_attentions=output_attentions,
        )
        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
        x = residual + x
        x = self.self_attn_layer_norm(x)

        residual = x
        x = self.activation_fn(self.fc1(x))
        x = nn.functional.dropout(x, p=self.activation_dropout, training=self.training)
        x = self.fc2(x)
        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
        x = residual + x
        x = self.final_layer_norm(x)
        return x, attn_weights


class FSMTEncoder(nn.Module):
    

    def __init__(self, config: FSMTConfig, embed_tokens):
        super().__init__()
        self.dropout = config.dropout
        self.layerdrop = config.encoder_layerdrop
        self.padding_idx = embed_tokens.padding_idx
        self.embed_tokens = embed_tokens
        embed_dim = embed_tokens.embedding_dim
        self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
        self.embed_positions = SinusoidalPositionalEmbedding(
            config.max_position_embeddings + self.padding_idx + 1, embed_dim, self.padding_idx
        )
        self.layers = nn.ModuleList([EncoderLayer(config) for _ in range(config.encoder_layers)])  

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: torch.Tensor = None,
        head_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ):
        
        
        if attention_mask is not None:
            attention_mask = invert_mask(attention_mask)

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
            embed_pos = self.embed_positions(input_ids)
        elif inputs_embeds is not None:
            inputs_embeds = inputs_embeds * self.embed_scale

            
            
            position_ids = inputs_embeds[:, :, 0].masked_fill(
                inputs_embeds[:, :, 0].eq(0), self.embed_positions.padding_idx
            )

            embed_pos = self.embed_positions(position_ids)
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        x = inputs_embeds + embed_pos
        x = nn.functional.dropout(x, p=self.dropout, training=self.training)

        
        x = x.transpose(0, 1)

        encoder_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None
        
        if head_mask is not None:
            assert head_mask.size()[0] == (
                len(self.layers)
            ), f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."
        for idx, encoder_layer in enumerate(self.layers):
            if output_hidden_states:
                x = x.transpose(0, 1)  
                encoder_states += (x,)
                x = x.transpose(0, 1)  
            
            dropout_probability = torch.rand([])
            if self.training and (dropout_probability < self.layerdrop):  
                attn = None
            else:
                x, attn = encoder_layer(
                    x,
                    attention_mask,
                    layer_head_mask=(head_mask[idx] if head_mask is not None else None),
                    output_attentions=output_attentions,
                )

            if output_attentions:
                all_attentions = all_attentions + (attn,)

        
        x = x.transpose(0, 1)

        if output_hidden_states:
            encoder_states += (x,)

        if not return_dict:
            return tuple(v for v in [x, encoder_states, all_attentions] if v is not None)
        return BaseModelOutput(last_hidden_state=x, hidden_states=encoder_states, attentions=all_attentions)


class DecoderLayer(nn.Module):
    def __init__(self, config: FSMTConfig):
        super().__init__()
        self.embed_dim = config.d_model

        self.self_attn = Attention(
            embed_dim=self.embed_dim,
            num_heads=config.decoder_attention_heads,
            dropout=config.attention_dropout,
        )
        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
        self.activation_dropout = config.activation_dropout

        self.self_attn_layer_norm = LayerNorm(self.embed_dim)
        self.encoder_attn = Attention(
            self.embed_dim,
            config.decoder_attention_heads,
            dropout=config.attention_dropout,
            encoder_decoder_attention=True,
        )
        self.encoder_attn_layer_norm = LayerNorm(self.embed_dim)
        self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
        self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
        self.final_layer_norm = LayerNorm(self.embed_dim)

    def forward(
        self,
        x,
        encoder_hidden_states,
        encoder_attn_mask=None,
        layer_state=None,
        causal_mask=None,
        layer_head_mask=None,
        cross_attn_layer_head_mask=None,
        decoder_padding_mask=None,
        output_attentions=False,
    ):
        residual = x

        if layer_state is None:
            layer_state = {}

        
        x, self_attn_weights = self.self_attn(
            query=x,
            key=x,
            layer_state=layer_state,  
            key_padding_mask=decoder_padding_mask,
            attn_mask=causal_mask,
            layer_head_mask=layer_head_mask,
            output_attentions=output_attentions,
        )
        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
        x = residual + x
        x = self.self_attn_layer_norm(x)

        
        residual = x
        assert self.encoder_attn.cache_key != self.self_attn.cache_key
        x, cross_attn_weights = self.encoder_attn(
            query=x,
            key=encoder_hidden_states,
            key_padding_mask=encoder_attn_mask,
            layer_state=layer_state,  
            layer_head_mask=cross_attn_layer_head_mask,
            output_attentions=output_attentions,
        )
        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
        x = residual + x
        x = self.encoder_attn_layer_norm(x)

        
        residual = x
        x = self.activation_fn(self.fc1(x))
        x = nn.functional.dropout(x, p=self.activation_dropout, training=self.training)
        x = self.fc2(x)
        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
        x = residual + x
        x = self.final_layer_norm(x)
        return (
            x,
            self_attn_weights,
            layer_state,
            cross_attn_weights,
        )  


class FSMTDecoder(nn.Module):
    

    def __init__(self, config: FSMTConfig, embed_tokens: nn.Embedding):
        super().__init__()
        self.dropout = config.dropout
        self.layerdrop = config.decoder_layerdrop
        self.padding_idx = embed_tokens.padding_idx
        self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
        self.embed_tokens = embed_tokens
        embed_dim = embed_tokens.embedding_dim
        self.embed_positions = SinusoidalPositionalEmbedding(
            config.max_position_embeddings + self.padding_idx + 1, embed_dim, self.padding_idx
        )
        self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.decoder_layers)])  

        if is_deepspeed_zero3_enabled():
            import deepspeed

            with deepspeed.zero.GatheredParameters(self.embed_tokens.weight, modifier_rank=None):
                embed_tokens_weight_shape = self.embed_tokens.weight.shape
        else:
            embed_tokens_weight_shape = self.embed_tokens.weight.shape
        self.output_projection = nn.Linear(embed_tokens_weight_shape[1], embed_tokens_weight_shape[0], bias=False)
        self.output_projection.weight = self.embed_tokens.weight

    def forward(
        self,
        input_ids: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        encoder_padding_mask: torch.Tensor,
        decoder_padding_mask: torch.Tensor,
        decoder_causal_mask: torch.Tensor,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        use_cache: bool = False,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ):
        
        
        if encoder_padding_mask is not None:
            encoder_padding_mask = invert_mask(encoder_padding_mask)

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
        elif input_ids is not None:
            
            positions = self.embed_positions(input_ids)
            if use_cache:
                input_ids = input_ids[:, -1:]
                positions = positions[:, -1:]  
            x = self.embed_tokens(input_ids) * self.embed_scale
        elif inputs_embeds is not None:
            
            
            position_ids = inputs_embeds[:, :, 0].masked_fill(
                inputs_embeds[:, :, 0].eq(0), self.embed_positions.padding_idx
            )
            positions = self.embed_positions(position_ids)
            x = inputs_embeds * self.embed_scale
        else:
            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")

        x += positions
        x = nn.functional.dropout(x, p=self.dropout, training=self.training)

        
        x = x.transpose(0, 1)
        encoder_hidden_states = encoder_hidden_states.transpose(0, 1)

        
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        all_cross_attns = () if output_attentions else None
        next_decoder_cache = []

        
        for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):
            if attn_mask is not None:
                assert attn_mask.size()[0] == (len(self.layers)), (
                    f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for"
                    f" {head_mask.size()[0]}."
                )
        for idx, decoder_layer in enumerate(self.layers):
            
            if output_hidden_states:
                x = x.transpose(0, 1)
                all_hidden_states += (x,)
                x = x.transpose(0, 1)
            if self.training:
                dropout_probability = torch.rand([])
                if dropout_probability < self.layerdrop:
                    continue

            layer_state = past_key_values[idx] if past_key_values is not None else None

            x, layer_self_attn, layer_past, layer_cross_attn = decoder_layer(
                x,
                encoder_hidden_states,
                encoder_attn_mask=encoder_padding_mask,
                decoder_padding_mask=decoder_padding_mask,
                layer_state=layer_state,
                causal_mask=decoder_causal_mask,
                layer_head_mask=(head_mask[idx] if head_mask is not None else None),
                cross_attn_layer_head_mask=(cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None),
                output_attentions=output_attentions,
            )

            if use_cache:
                next_decoder_cache.append(layer_past.copy())

            if output_attentions:
                all_self_attns += (layer_self_attn,)
                all_cross_attns += (layer_cross_attn,)

        
        if output_hidden_states:
            x = x.transpose(0, 1)
            all_hidden_states += (x,)
            x = x.transpose(0, 1)

        
        x = x.transpose(0, 1)
        encoder_hidden_states = encoder_hidden_states.transpose(0, 1)

        x = self.output_projection(x)

        next_cache = next_decoder_cache if use_cache else None

        if not return_dict:
            return tuple(
                v for v in [x, next_cache, all_hidden_states, all_self_attns, all_cross_attns] if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=x,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
            cross_attentions=all_cross_attns,
        )


def _reorder_buffer(attn_cache, new_order):
    for k, input_buffer_k in attn_cache.items():
        if input_buffer_k is not None:
            attn_cache[k] = input_buffer_k.index_select(0, new_order)
    return attn_cache


class Attention(nn.Module):
    

    def __init__(
        self,
        embed_dim,
        num_heads,
        dropout=0.0,
        bias=True,
        encoder_decoder_attention=False,  
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
        self.scaling = self.head_dim**-0.5

        self.encoder_decoder_attention = encoder_decoder_attention
        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)
        self.cache_key = "encoder_decoder" if self.encoder_decoder_attention else "self"

    def _shape(self, tensor, seq_len, bsz):
        return tensor.contiguous().view(seq_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)

    def forward(
        self,
        query,
        key: Optional[Tensor],
        key_padding_mask: Optional[Tensor] = None,
        layer_state: Optional[Dict[str, Optional[Tensor]]] = None,
        attn_mask: Optional[Tensor] = None,
        layer_head_mask: Optional[Tensor] = None,
        output_attentions=False,
    ) -> Tuple[Tensor, Optional[Tensor]]:
        
        static_kv: bool = self.encoder_decoder_attention
        tgt_len, bsz, embed_dim = query.size()
        assert embed_dim == self.embed_dim
        assert list(query.size()) == [tgt_len, bsz, embed_dim]
        
        if layer_state is not None:  
            saved_state = layer_state.get(self.cache_key, {})
            if "prev_key" in saved_state and static_kv:
                
                key = None
        else:
            saved_state = None
            layer_state = {}

        q = self.q_proj(query) * self.scaling
        if static_kv:
            if key is None:
                k = v = None
            else:
                k = self.k_proj(key)
                v = self.v_proj(key)
        else:
            k = self.k_proj(query)
            v = self.v_proj(query)

        q = self._shape(q, tgt_len, bsz)
        if k is not None:
            k = self._shape(k, -1, bsz)
        if v is not None:
            v = self._shape(v, -1, bsz)

        if saved_state is not None:
            k, v, key_padding_mask = self._use_saved_state(k, v, saved_state, key_padding_mask, static_kv, bsz)

        
        layer_state[self.cache_key] = {
            "prev_key": k.view(bsz, self.num_heads, -1, self.head_dim),
            "prev_value": v.view(bsz, self.num_heads, -1, self.head_dim),
            "prev_key_padding_mask": key_padding_mask if not static_kv else None,
        }

        assert k is not None
        src_len = k.size(1)
        attn_weights = torch.bmm(q, k.transpose(1, 2))
        assert attn_weights.size() == (bsz * self.num_heads, tgt_len, src_len)

        if attn_mask is not None:
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attn_mask
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        
        if key_padding_mask is not None and key_padding_mask.dim() == 0:
            key_padding_mask = None
        assert key_padding_mask is None or key_padding_mask.size()[:2] == (
            bsz,
            src_len,
        )

        if key_padding_mask is not None:  
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            reshaped = key_padding_mask.unsqueeze(1).unsqueeze(2)
            attn_weights = attn_weights.masked_fill(reshaped, torch.finfo(attn_weights.dtype).min)
            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:
            assert layer_head_mask.size() == (
                self.num_heads,
            ), f"Head mask for a single layer should be of size {(self.num_heads,)}, but is {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,
        )

        assert v is not None
        attn_output = torch.bmm(attn_probs, v)
        assert attn_output.size() == (bsz * self.num_heads, tgt_len, self.head_dim)
        attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
        attn_output = self.out_proj(attn_output)

        return attn_output, attn_weights_reshaped

    def _use_saved_state(self, k, v, saved_state, key_padding_mask, static_kv, bsz):
        
        if "prev_key" in saved_state:
            _prev_key = saved_state["prev_key"]
            assert _prev_key is not None
            prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)
            if static_kv:
                k = prev_key
            else:
                assert k is not None
                k = torch.cat([prev_key, k], dim=1)
        if "prev_value" in saved_state:
            _prev_value = saved_state["prev_value"]
            assert _prev_value is not None
            prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)
            if static_kv:
                v = prev_value
            else:
                assert v is not None
                v = torch.cat([prev_value, v], dim=1)
        assert k is not None and v is not None
        prev_key_padding_mask: Optional[Tensor] = saved_state.get("prev_key_padding_mask", None)
        if prev_key_padding_mask is not None:
            if static_kv:
                new_key_padding_mask = prev_key_padding_mask
            else:
                new_key_padding_mask = torch.cat([prev_key_padding_mask, key_padding_mask], dim=1)
        else:
            new_key_padding_mask = key_padding_mask
        return k, v, new_key_padding_mask


def fill_with_neg_inf(t):
    
    return t.float().fill_(torch.finfo(t.dtype).min).type_as(t)



def _get_shape(t):
    return getattr(t, "shape", None)


@add_start_docstrings(
    "The bare FSMT Model outputting raw hidden-states without any specific head on top.",
    FSMT_START_DOCSTRING,
)
class FSMTModel(PretrainedFSMTModel):
    _tied_weights_keys = ["decoder.embed_tokens.weight", "decoder.output_projection.weight"]

    def __init__(self, config: FSMTConfig):
        super().__init__(config)

        padding_idx = config.pad_token_id
        encoder_embed_tokens = nn.Embedding(config.src_vocab_size, config.d_model, padding_idx)
        decoder_embed_tokens = nn.Embedding(config.tgt_vocab_size, config.d_model, padding_idx)

        self.encoder = FSMTEncoder(config, encoder_embed_tokens)
        self.decoder = FSMTDecoder(config, decoder_embed_tokens)

        
        self.post_init()

    def get_encoder(self):
        return self.encoder

    def get_decoder(self):
        return self.decoder

    def _tie_weights(self):
        if self.config.tie_word_embeddings:
            self._tie_or_clone_weights(self.decoder.embed_tokens, self.get_input_embeddings())
            self._tie_or_clone_weights(self.decoder.output_projection, self.get_input_embeddings())

    @add_start_docstrings_to_model_forward(FSMT_INPUTS_DOCSTRING)
    @add_code_sample_docstrings(
        checkpoint=_CHECKPOINT_FOR_DOC,
        output_type=Seq2SeqModelOutput,
        config_class=_CONFIG_FOR_DOC,
    )
    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.BoolTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        decoder_head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[Tuple[torch.FloatTensor]] = None,
        past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], Seq2SeqModelOutput]:
        if decoder_input_ids is None:
            use_cache = False

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        
        if not use_cache and input_ids is not None:
            decoder_input_ids, decoder_padding_mask, causal_mask = _prepare_fsmt_decoder_inputs(
                self.config,
                input_ids,
                decoder_input_ids=decoder_input_ids,
                decoder_padding_mask=decoder_attention_mask,
                causal_mask_dtype=self.decoder.embed_tokens.weight.dtype,
            )
        else:
            decoder_padding_mask, causal_mask = None, None

        if decoder_input_ids is None and decoder_inputs_embeds is None:
            raise ValueError("Make sure that `decoder_input_ids` or `decoder_inputs_embeds` are passed.")

        if encoder_outputs is None:
            encoder_outputs = self.encoder(
                input_ids=input_ids,
                attention_mask=attention_mask,
                inputs_embeds=inputs_embeds,
                head_mask=head_mask,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )
        
        elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
            encoder_outputs = BaseModelOutput(
                last_hidden_state=encoder_outputs[0],
                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
            )

        
        decoder_outputs = self.decoder(
            decoder_input_ids,
            encoder_outputs[0],
            attention_mask,
            decoder_padding_mask,
            decoder_causal_mask=causal_mask,
            inputs_embeds=decoder_inputs_embeds,
            head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        if not return_dict:
            return decoder_outputs + encoder_outputs

        return Seq2SeqModelOutput(
            last_hidden_state=decoder_outputs.last_hidden_state,
            past_key_values=decoder_outputs.past_key_values,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
        )

    def get_input_embeddings(self):
        return self.encoder.embed_tokens

    def set_input_embeddings(self, value):
        self.encoder.embed_tokens = value

    def get_output_embeddings(self):
        return self.decoder.embed_tokens

    def set_output_embeddings(self, value):
        self.decoder.embed_tokens = value


@add_start_docstrings(
    "The FSMT Model with a language modeling head. Can be used for summarization.", FSMT_START_DOCSTRING
)
class FSMTForConditionalGeneration(PretrainedFSMTModel):
    base_model_prefix = "model"
    _tied_weights_keys = ["decoder.embed_tokens.weight", "decoder.output_projection.weight"]

    def __init__(self, config: FSMTConfig):
        super().__init__(config)
        base_model = FSMTModel(config)
        self.model = base_model

        
        self.post_init()

    @add_start_docstrings_to_model_forward(FSMT_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
    @add_end_docstrings(FSMT_GENERATION_EXAMPLE)
    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.BoolTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        decoder_head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[Tuple[torch.FloatTensor]] = None,
        past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        decoder_inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], Seq2SeqLMOutput]:
        r
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if labels is not None:
            use_cache = False

        outputs = self.model(
            input_ids,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids,
            decoder_inputs_embeds=decoder_inputs_embeds,
            encoder_outputs=encoder_outputs,
            decoder_attention_mask=decoder_attention_mask,
            head_mask=head_mask,
            decoder_head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        lm_logits = outputs[0]

        masked_lm_loss = None
        if labels is not None:
            loss_fct = CrossEntropyLoss()
            
            masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.tgt_vocab_size), labels.view(-1))

        if not return_dict:
            output = (lm_logits,) + outputs[1:]
            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output

        return Seq2SeqLMOutput(
            loss=masked_lm_loss,
            logits=lm_logits,
            past_key_values=outputs.past_key_values,
            decoder_hidden_states=outputs.decoder_hidden_states,
            decoder_attentions=outputs.decoder_attentions,
            cross_attentions=outputs.cross_attentions,
            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
            encoder_hidden_states=outputs.encoder_hidden_states,
            encoder_attentions=outputs.encoder_attentions,
        )

    def prepare_inputs_for_generation(
        self,
        decoder_input_ids,
        past_key_values=None,
        attention_mask=None,
        head_mask=None,
        decoder_head_mask=None,
        cross_attn_head_mask=None,
        use_cache=None,
        encoder_outputs=None,
        **kwargs,
    ):
        return {
            "input_ids": None,  
            "encoder_outputs": encoder_outputs,
            "past_key_values": past_key_values,
            "decoder_input_ids": decoder_input_ids,
            "attention_mask": attention_mask,
            "head_mask": head_mask,
            "decoder_head_mask": decoder_head_mask,
            "cross_attn_head_mask": cross_attn_head_mask,
            "use_cache": use_cache,  
        }

    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
        return shift_tokens_right(labels, self.config.pad_token_id)

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        reordered_past = []
        for layer_past in past_key_values:
            
            layer_past_new = {
                attn_key: _reorder_buffer(attn_cache, beam_idx) for attn_key, attn_cache in layer_past.items()
            }
            reordered_past.append(layer_past_new)
        return reordered_past

    def get_encoder(self):
        return self.model.encoder

    def get_decoder(self):
        return self.model.decoder

    def get_output_embeddings(self):
        return self.model.decoder.embed_tokens

    def set_output_embeddings(self, value):
        self.model.decoder.embed_tokens = value


class SinusoidalPositionalEmbedding(nn.Embedding):
    

    def __init__(self, num_positions, embedding_dim, padding_idx):
        self.make_weight(num_positions, embedding_dim, padding_idx)

    def make_weight(self, num_positions, embedding_dim, padding_idx):
        weight = self.get_embedding(num_positions, embedding_dim, padding_idx)
        if not hasattr(self, "weight"):
            
            super().__init__(num_positions, embedding_dim, padding_idx, _weight=weight)
        else:
            
            weight = weight.to(dtype=self.weight.dtype, device=self.weight.device)
            self.weight = nn.Parameter(weight)
        self.weight.detach_()
        self.weight.requires_grad = False

    @staticmethod
    def get_embedding(num_embeddings, embedding_dim, padding_idx):
        
        half_dim = embedding_dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
        if embedding_dim % 2 == 1:
            
            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
        if padding_idx is not None:
            emb[padding_idx, :] = 0
        return emb

    @staticmethod
    def make_positions(tensor, padding_idx: int):
        
        
        
        
        
        mask = tensor.ne(padding_idx).int()
        return (torch.cumsum(mask, dim=1).type_as(mask) * mask).long() + padding_idx

    def forward(
        self,
        input,
        incremental_state: Optional[Any] = None,
        timestep: Optional[Tensor] = None,
    ):
        
        bsz, seq_len = input.shape[:2]
        max_pos = self.padding_idx + 1 + seq_len
        if max_pos > self.weight.size(0):
            
            self.make_weight(max_pos, self.embedding_dim, self.padding_idx)
        positions = self.make_positions(input, self.padding_idx)
        return super().forward(positions)