codingutf8 2020 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 finalizeinputsdocstring r args inputids torch longtensor of shape batchsize numbeams sequencelength indices of input sequence tokens in the vocabulary indices can be obtained using any class inheriting from pretrainedtokenizer see pretrainedtokenizer encode and pretrainedtokenizer call for details what are input ids glossaryinputids finalbeamscores torch floattensor of shape batchsize numbeams the final scores of all nonfinished beams finalbeamtokens torch floattensor of shape batchsize numbeams the last tokens to be added to the nonfinished beamhypotheses finalbeamindices torch floattensor of shape batchsize numbeams the beam indices indicating to which beam the finalbeamtokens shall be added padtokenid int optional the id of the padding token eostokenid unionint listint optional the id of the endofsequence token optionally use a list to set multiple endofsequence tokens return torch longtensor of shape batchsize numreturnsequences sequencelength the generated sequences the second dimension sequencelength is either equal to maxlength or shorter if all batches finished early due to the eostokenid abstract base class for all beam scorers that are used for pretrainedmodel beamsearch and pretrainedmodel beamsample def init self batchsize int numbeams int device torch device lengthpenalty optionalfloat 1 0 doearlystopping optionalunionbool str false numbeamhypstokeep optionalint 1 numbeamgroups optionalint 1 maxlength optionalint none self numbeams numbeams self device device self lengthpenalty lengthpenalty self doearlystopping doearlystopping self numbeamhypstokeep numbeamhypstokeep self numbeamgroups numbeamgroups self groupsize self numbeams self numbeamgroups self isinit false self beamhypsiself numbeamgroupsj is the beamhyps of the jth group in the ith minibatch if groupbeamsearch is not used the list consists of batchsize beamhyps self beamhyps beamhypotheses numbeamsself groupsize lengthpenaltyself lengthpenalty earlystoppingself doearlystopping maxlengthmaxlength for in rangebatchsize self numbeamgroups self doneiself numbeamgroupsj indicates whether the generation of the beamhyps of the jth group in the ith minibatch is complete self done torch tensor false for in rangebatchsize self numbeamgroups dtypetorch bool deviceself device if not isinstancenumbeams int or numbeams 1 raise valueerror fnumbeams has to be an integer strictly greater than 1 but is numbeams for numbeams 1 one should make use of greedysearch instead if not isinstancenumbeamgroups int or numbeamgroups numbeams or numbeams numbeamgroups 0 raise valueerror numbeamgroups has to be an integer smaller or equal than numbeams and numbeams has to be f divisible by numbeamgroups but is numbeamgroups with numbeams being numbeams property def isdoneself bool return self done all def process self inputids torch longtensor nextscores torch floattensor nexttokens torch longtensor nextindices torch longtensor padtokenid optionalint none eostokenid optionalunionint listint none beamindices optionaltorch longtensor none groupindex optionalint 0 decoderpromptlen optionalint 0 dictstr torch tensor add up to the length which the nextscores is calculated on curlen inputids shape1 decoderpromptlen 1 batchsize lenself beamhyps self numbeamgroups if not batchsize inputids shape0 self groupsize if self numbeamgroups 1 raise valueerror fa group beam size of inputids shape0 is used as the input but a group beam fsize of self groupsize is expected by the beam scorer else raise valueerror fa beam size of inputids shape0 is used as the input but a beam size of fself groupsize is expected by the beam scorer device inputids device nextbeamscores torch zerosbatchsize self groupsize dtypenextscores dtype devicedevice nextbeamtokens torch zerosbatchsize self groupsize dtypenexttokens dtype devicedevice nextbeamindices torch zerosbatchsize self groupsize dtypenextindices dtype devicedevice if isinstanceeostokenid int eostokenid eostokenid for batchidx in rangebatchsize batchgroupidx batchidx self numbeamgroups groupindex if self donebatchgroupidx if self numbeams lenself beamhypsbatchgroupidx raise valueerrorfbatch can only be done if at least self numbeams beams have been generated if eostokenid is none or padtokenid is none raise valueerrorgenerated beams numbeams eostokenid and padtoken have to be defined pad the batch nextbeamscoresbatchidx 0 nextbeamtokensbatchidx padtokenid nextbeamindicesbatchidx 0 continue next tokens for this sentence beamidx 0 for beamtokenrank nexttoken nextscore nextindex in enumerate zipnexttokensbatchidx nextscoresbatchidx nextindicesbatchidx batchbeamidx batchidx self groupsize nextindex add to generated hypotheses if end of sentence if eostokenid is not none and nexttoken item in eostokenid if beamtoken does not belong to top numbeams tokens it should not be added isbeamtokenworsethantopnumbeams beamtokenrank self groupsize if isbeamtokenworsethantopnumbeams continue if beamindices is not none beamindex beamindicesbatchbeamidx beamindex beamindex batchbeamidx else beamindex none skip the corner case where the very first generated token is eostoken if decoderpromptlen inputids shape1 continue self beamhypsbatchgroupidx add inputidsbatchbeamidx clone nextscore item beamindicesbeamindex decoderpromptlendecoderpromptlen else add next predicted token since it is not eostoken nextbeamscoresbatchidx beamidx nextscore nextbeamtokensbatchidx beamidx nexttoken nextbeamindicesbatchidx beamidx batchbeamidx beamidx 1 once the beam for next step is full don t add more tokens to it if beamidx self groupsize break if beamidx self groupsize raise valueerror fat most self groupsize tokens in nexttokensbatchidx can be equal to eostokenid f eostokenid make sure nexttokensbatchidx are corrected check if we are done so that we can save a pad step if alldone self donebatchgroupidx self donebatchgroupidx or self beamhypsbatchgroupidx isdone nextscoresbatchidx max item curlen return userdict nextbeamscores nextbeamscores view1 nextbeamtokens nextbeamtokens view1 nextbeamindices nextbeamindices view1 def finalize self inputids torch longtensor finalbeamscores torch floattensor finalbeamtokens torch longtensor finalbeamindices torch longtensor maxlength int padtokenid optionalint none eostokenid optionalunionint listint none beamindices optionaltorch longtensor none decoderpromptlen optionalint 0 tupletorch longtensor batchsize lenself beamhyps self numbeamgroups if isinstanceeostokenid int eostokenid eostokenid finalize all open beam hypotheses and add to generated hypotheses for batchgroupidx beamhyp in enumerateself beamhyps if self donebatchgroupidx continue all open beam hypotheses are added to the beam hypothesis beam hypothesis class automatically keeps the best beams for indexpergroup in rangeself groupsize batchbeamidx batchgroupidx self groupsize indexpergroup finalscore finalbeamscoresbatchbeamidx item finaltokens inputidsbatchbeamidx beamindex beamindicesbatchbeamidx if beamindices is not none else none beamhyp addfinaltokens finalscore beamindicesbeamindex decoderpromptlendecoderpromptlen select the best hypotheses sentlengths inputids newbatchsize self numbeamhypstokeep best bestindices bestscores torch zerosbatchsize self numbeamhypstokeep deviceself device dtypetorch float32 retrieve best hypotheses for i in rangebatchsize beamhypsinbatch self beamhypsi self numbeamgroups i 1 self numbeamgroups candidatebeams beam for beamhyp in beamhypsinbatch for beam in beamhyp beams sortedhyps sortedcandidatebeams keylambda x x0 for j in rangeself numbeamhypstokeep besthyptuple sortedhyps pop bestscore besthyptuple0 besthyp besthyptuple1 bestindex besthyptuple2 sentlengthsself numbeamhypstokeep i j lenbesthyp append hyp to lists best appendbesthyp append indices to list bestindices appendbestindex bestscoresi self numbeamhypstokeep j bestscore prepare for adding eos sentlengthsmax sentlengths max item 1 sentmaxlen minsentlengthsmax maxlength if maxlength is not none else sentlengthsmax decoded torch longtensor inputids newbatchsize self numbeamhypstokeep sentmaxlen if lenbestindices 0 and bestindices0 is not none indices torch longtensor inputids newbatchsize self numbeamhypstokeep sentmaxlen else indices none shorter batches are padded if needed if sentlengths min item sentlengths max item if padtokenid is none raise valueerrorpadtokenid has to be defined decoded fillpadtokenid if indices is not none indices fill1 fill with hypotheses and eostokenid if the latter fits in for i hypo bestidx in enumeratezipbest bestindices decodedi sentlengthsi hypo if indices is not none indicesi lenbestidx torch tensorbestidx if sentlengthsi sentmaxlen inserting only the first eostokenid decodedi sentlengthsi eostokenid0 return userdict sequences decoded sequencescores bestscores beamindices indices class constrainedbeamsearchscorerbeamscorer r beamscorer implementing constrained beam search decoding args batchsize int batch size of inputids for which standard beam search decoding is run in parallel numbeams int number of beams for beam search constraints listconstraint a list of positive constraints represented as constraint objects that must be fulfilled in the generation output for more information the documentation of constraint should be read device torch device defines the device type e g cpu or cuda on which this instance of beamsearchscorer will be allocated lengthpenalty float optional defaults to 1 0 exponential penalty to the length that is used with beambased generation it is applied as an exponent to the sequence length which in turn is used to divide the score of the sequence since the score is the log likelihood of the sequence i e negative lengthpenalty 0 0 promotes longer sequences while lengthpenalty 0 0 encourages shorter sequences doearlystopping bool or str optional defaults to false controls the stopping condition for beambased methods like beamsearch it accepts the following values true where the generation stops as soon as there are numbeams complete candidates false where an heuristic is applied and the generation stops when is it very unlikely to find better candidates never where the beam search procedure only stops when there cannot be better candidates canonical beam search algorithm numbeamhypstokeep int optional defaults to 1 the number of beam hypotheses that shall be returned upon calling transformers beamsearchscorer finalize numbeamgroups int optional defaults to 1 number of groups to divide numbeams into in order to ensure diversity among different groups of beams see this paperhttps arxiv orgpdf1610 02424 pdf for more details maxlength int optional the maximum length of the sequence to be generated add up to the length which the nextscores is calculated on curlen inputids shape1 decoderpromptlen 1 batchsize lenself beamhyps if not batchsize inputids shape0 self groupsize if self numbeamgroups 1 raise valueerror fa group beam size of inputids shape0 is used as the input but a group beam fsize of self groupsize is expected by the beam scorer else raise valueerror fa beam size of inputids shape0 is used as the input but a beam size of fself groupsize is expected by the beam scorer device inputids device nextbeamscores torch zerosbatchsize self groupsize dtypenextscores dtype devicedevice nextbeamtokens torch zerosbatchsize self groupsize dtypenexttokens dtype devicedevice nextbeamindices torch zerosbatchsize self groupsize dtypenextindices dtype devicedevice if isinstanceeostokenid int eostokenid eostokenid for batchidx beamhyp in enumerateself beamhyps if self donebatchidx if self numbeams lenbeamhyp raise valueerrorfbatch can only be done if at least self numbeams beams have been generated if eostokenid is none or padtokenid is none raise valueerrorgenerated beams numbeams eostokenid and padtoken have to be defined pad the batch nextbeamscoresbatchidx 0 nextbeamtokensbatchidx padtokenid nextbeamindicesbatchidx 0 continue next tokens for this sentence beamidx 0 for beamtokenrank nexttoken nextscore nextindex in enumerate zipnexttokensbatchidx nextscoresbatchidx nextindicesbatchidx batchbeamidx batchidx self groupsize nextindex add to generated hypotheses if end of sentence if eostokenid is not none and nexttoken item in eostokenid if beamtoken does not belong to top numbeams tokens it should not be added isbeamtokenworsethantopnumbeams beamtokenrank self groupsize if isbeamtokenworsethantopnumbeams continue completesconstraint self checkcompletesconstraintsinputidsbatchbeamidx cpu tolist if completesconstraint if beamindices is not none beamindex beamindicesbatchbeamidx beamindex beamindex batchbeamidx else beamindex none skip the corner case where the only constraint token is eostoken and the very first generated token is eostoken if decoderpromptlen inputids shape1 continue beamhyp add inputidsbatchbeamidx clone nextscore item beamindicesbeamindex decoderpromptlendecoderpromptlen else add next predicted token since it is not eostoken nextbeamscoresbatchidx beamidx nextscore nextbeamtokensbatchidx beamidx nexttoken nextbeamindicesbatchidx beamidx batchbeamidx beamidx 1 once the beam for next step is full don t add more tokens to it if beamidx self groupsize break newscores newtokens newindices self stepsentenceconstraint batchidx inputids scoresforallvocab nextbeamscoresbatchidx nextbeamtokensbatchidx nextbeamindicesbatchidx nextbeamscoresbatchidx newscores nextbeamtokensbatchidx newtokens nextbeamindicesbatchidx newindices if beamidx self groupsize raise valueerror fat most self groupsize tokens in nexttokensbatchidx can be equal to eostokenid f eostokenid make sure nexttokensbatchidx are corrected check if we are done so that we can save a pad step if alldone self donebatchidx self donebatchidx or beamhyp isdone nextscoresbatchidx max item curlen return userdict nextbeamscores nextbeamscores view1 nextbeamtokens nextbeamtokens view1 nextbeamindices nextbeamindices view1 def stepsentenceconstraint self batchidx int inputids torch longtensor vocabscores torch floattensor sentbeamscores torch floattensor sentbeamtokens torch longtensor sentbeamindices torch longtensor pushprogress bool false sentbeamtokens are the next numbeams number of tokens that are under consideration for this beam candidate next tokens 1 adding advancetokens using constraintstatelist advance we propose new tokens to be added into this candidate list that will advance us in fulfilling the constraints 2 selecting best candidates such that we end up with highest probable candidates that fulfill our constraints origlen sentbeamindices size0 device sentbeamindices device initialize states topkcontraintstates self makeconstraintstatesoriglen advanceconstraintstates self makeconstraintstatesoriglen sidx eidx batchidx origlen batchidx 1 origlen thisbatchinputids inputidssidx eidx thisbatchtokenscores vocabscoressidx eidx fullhypotheses torch catinputidssentbeamindices sentbeamtokens unsqueeze1 dim1 need to make new hypothesis that advance the constraints tracknew newseqs fullhypotheses tolist newstates newindices newtokens newscores for seqidx preseq in enumeratethisbatchinputids preseq ith sequence generated before this step inputids topk generic beam search best model next tokens advance constraints forcing the next token either way we need to sort them into banks later so store a constraintliststate for all types of hypotheses topkstate topkcontraintstatesseqidx topkstate resetfullhypothesesseqidx cpu tolist advancestate advanceconstraintstatesseqidx advancestate resetpreseq cpu tolist if not advancestate completed advancetokens torch longtensoradvancestate advance todevice for advancetoken in advancetokens since adding each advancetoken leads to a different hypothesis create new state instance newstate advancestate copystatefultrue newstate addadvancetoken cpu tolist advanceseq torch catpreseq advancetoken unsqueeze0 1 cpu tolist if advanceseq not in tracknewnewseqs prevent duplicates which are basically bound to happen in this process tracknewnewseqs appendadvanceseq tracknewnewindices appendsidx seqidx idx global idx across all the batches tracknewnewtokens appendadvancetoken tracknewnewscores appendthisbatchtokenscoresseqidx takeadvancetoken tracknewnewstates appendnewstate elif pushprogress basically sentbeamindices often chooses very little among inputids the generated sequences that actually fulfill our constraints for example let constraints loves pies and preseq1 the child loves pies and preseq2 the child plays in the playground and without this step if sentbeamindices is something like 1 1 then 1 preseq1 won t be added to the list of topk hypothesis since it s not in the indices and 2 it won t be added to the list of advance hypothesis since it s completed already this is the else part of if constraintscompletedseqidx 3 it ends up simply getting removed from consideration 3 might be fine and actually desired since it s likely that it s a lowprobability output anyways especially if it s not in the list of sentbeamindices but this often leads to lengthened beam search times since completed sequences keep getting removed after all this effort for constrained generation here we basically take preseq1 and to push it into the considered list of hypotheses by simply appending the next likely token in the vocabulary and adding it to the list of hypotheses newscore newtoken torch maxthisbatchtokenscoresseqidx 0 some next probable token advanceseq torch catpreseq newtoken unsqueeze0 1 advancestate advanceconstraintstatesseqidx advanceseq advanceseq cpu tolist advancestate resetadvanceseq if advanceseq not in tracknewnewseqs but still don t want to have duplicates tracknewnewseqs appendadvanceseq tracknewnewindices appendseqidx tracknewnewtokens appendnewtoken tracknewnewscores appendnewscore tracknewnewstates appendadvancestate if lentracknewnewindices 0 newindices torch tensortracknewnewindices todevice newtokens torch stacktracknewnewtokens todevice newscores torch stacktracknewnewscores todevice allstates topkcontraintstates tracknewnewstates alltokens torch catsentbeamtokens newtokens 1 allscores torch catsentbeamscores newscores 1 allbanks torch tensorone getbank for one in allstates todevice zipped allbanks 100 allscores indices zipped sortdescendingtrue indices sortedbanks allbanksindices then we end up with sorted among bank c sorted among bank c1 sorted among bank 0 counter 1 curbank sortedbanks0 increments for bank in sortedbanks if bank curbank counter 1 else counter 0 curbank bank increments appendcounter rearrangers torch tensornp argsortincrements kindmergesort indices indicesrearrangers origlen sentbeamscores allscoresindices sentbeamtokens alltokensindices sentbeamindices torch catsentbeamindices newindicesindices return sentbeamscores sentbeamtokens sentbeamindices def finalize self inputids torch longtensor finalbeamscores torch floattensor finalbeamtokens torch longtensor finalbeamindices torch longtensor maxlength int padtokenid optionalint none eostokenid optionalunionint listint none beamindices optionaltorch longtensor none decoderpromptlen optionalint 0 tupletorch longtensor batchsize lenself beamhyps if isinstanceeostokenid int eostokenid eostokenid finalize all open beam hypotheses and add to generated hypotheses for batchidx beamhyp in enumerateself beamhyps if self donebatchidx continue all open beam hypotheses are added to the beam hypothesis beam hypothesis class automatically keeps the best beams idscollect for beamid in rangeself numbeams batchbeamidx batchidx self numbeams beamid finalscore finalbeamscoresbatchbeamidx item finaltokens inputidsbatchbeamidx completesconstraint self checkcompletesconstraintsfinaltokens cpu tolist if completesconstraint beamindex beamindicesbatchbeamidx if beamindices is not none else none beamhyp add finaltokens finalscore beamindicesbeamindex decoderpromptlendecoderpromptlen idscollect appendbeamid due to overly complex constraints or other factors sometimes we can t gaurantee a successful generation in these cases we simply return the highest scoring outputs if lenidscollect self numbeamhypstokeep for beamid in rangeself numbeams if beamid not in idscollect batchbeamidx batchidx self numbeams beamid finalscore finalbeamscoresbatchbeamidx item finaltokens inputidsbatchbeamidx beamhyp addfinaltokens finalscore decoderpromptlendecoderpromptlen if lenidscollect self numbeamhypstokeep break select the best hypotheses sentlengths inputids newbatchsize self numbeamhypstokeep best bestindices bestscores torch zerosbatchsize self numbeamhypstokeep deviceself device dtypetorch float32 retrieve best hypotheses for i beamhyp in enumerateself beamhyps sortedhyps sortedbeamhyp beams keylambda x x0 for j in rangeself numbeamhypstokeep besthyptuple sortedhyps pop bestscore besthyptuple0 besthyp besthyptuple1 bestindex besthyptuple2 sentlengthsself numbeamhypstokeep i j lenbesthyp append to lists best appendbesthyp append indices to list bestindices appendbestindex bestscoresi self numbeamhypstokeep j bestscore prepare for adding eos sentlengthsmax sentlengths max item 1 sentmaxlen minsentlengthsmax maxlength if maxlength is not none else sentlengthsmax decoded torch longtensor inputids newbatchsize self numbeamhypstokeep sentmaxlen if lenbestindices 0 and bestindices0 is not none indices torch longtensor inputids newbatchsize self numbeamhypstokeep sentmaxlen else indices none shorter batches are padded if needed if sentlengths min item sentlengths max item if padtokenid is none raise valueerrorpadtokenid has to be defined decoded fillpadtokenid if indices is not none indices fill1 fill with hypotheses and eostokenid if the latter fits in for i hypo bestidx in enumeratezipbest bestindices decodedi sentlengthsi hypo if indices is not none indicesi lenbestidx torch tensorbestidx if sentlengthsi sentmaxlen inserting only the first eostokenid decodedi sentlengthsi eostokenid0 return userdict sequences decoded sequencescores bestscores beamindices indices class beamhypotheses def initself numbeams int lengthpenalty float earlystopping bool maxlength optionalint none self lengthpenalty lengthpenalty self earlystopping earlystopping self maxlength maxlength self numbeams numbeams self beams self worstscore 1e9 if not isinstanceself earlystopping bool and self maxlength is none raise valueerror when doearlystopping is set to a string maxlength must be defined ensure it is passed to the beamscorer class instance at initialization time def lenself return lenself beams def add self hyp torch longtensor sumlogprobs float beamindices optionaltorch longtensor none decoderpromptlen optionalint 0 score sumlogprobs hyp shape1 decoderpromptlen self lengthpenalty if lenself self numbeams or score self worstscore self beams appendscore hyp beamindices if lenself self numbeams sortednextscores sorteds idx for idx s in enumerateself beams del self beamssortednextscores01 self worstscore sortednextscores10 else self worstscore minscore self worstscore def isdoneself bestsumlogprobs float curlen int bool if lenself self numbeams return false true stop as soon as at least numbeams hypotheses are finished if self earlystopping is true return true false heuristic compute best possible score from curlen even though it is not entirely accurate when lengthpenalty is positive see the discussion below for more details https github comhuggingfacetransformerspull20901issuecomment1369845565 elif self earlystopping is false highestattainablescore bestsumlogprobs curlenself lengthpenalty ret self worstscore highestattainablescore return ret never compute the best possible score depending on the signal of lengthpenalty else lengthpenalty 0 0 max denominator is obtaned from maxlength not from curlen min abshighestattainablescore is obtained highestattainablescore is negative hence we obtain its max this way if self lengthpenalty 0 0 highestattainablescore bestsumlogprobs self maxlengthself lengthpenalty the opposite logic applies here max highestattainablescore from curlen else highestattainablescore bestsumlogprobs curlenself lengthpenalty ret self worstscore highestattainablescore return ret coding utf 8 2020 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 args input_ids torch longtensor of shape batch_size num_beams sequence_length indices of input sequence tokens in the vocabulary indices can be obtained using any class inheriting from pretrainedtokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids next_scores torch floattensor of shape batch_size 2 num_beams current scores of the top 2 num_beams non finished beam hypotheses next_tokens torch longtensor of shape batch_size 2 num_beams input_ids of the tokens corresponding to the top 2 num_beams non finished beam hypotheses next_indices torch longtensor of shape batch_size 2 num_beams beam indices indicating to which beam hypothesis the next_tokens correspond pad_token_id int optional the id of the padding token eos_token_id union int list int optional the id of the end of sequence token optionally use a list to set multiple end of sequence tokens beam_indices torch longtensor optional beam indices indicating to which beam hypothesis each token correspond group_index int optional the index of the group of beams used with pretrainedmodel group_beam_search return userdict a dictionary composed of the fields as defined above next_beam_scores torch floattensor of shape batch_size num_beams updated scores of all non finished beams next_beam_tokens torch floattensor of shape batch_size num_beams next tokens to be added to the non finished beam_hypotheses next_beam_indices torch floattensor of shape batch_size num_beams beam indices indicating to which beam the next tokens shall be added args input_ids torch longtensor of shape batch_size num_beams sequence_length indices of input sequence tokens in the vocabulary indices can be obtained using any class inheriting from pretrainedtokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids final_beam_scores torch floattensor of shape batch_size num_beams the final scores of all non finished beams final_beam_tokens torch floattensor of shape batch_size num_beams the last tokens to be added to the non finished beam_hypotheses final_beam_indices torch floattensor of shape batch_size num_beams the beam indices indicating to which beam the final_beam_tokens shall be added pad_token_id int optional the id of the padding token eos_token_id union int list int optional the id of the end of sequence token optionally use a list to set multiple end of sequence tokens return torch longtensor of shape batch_size num_return_sequences sequence_length the generated sequences the second dimension sequence_length is either equal to max_length or shorter if all batches finished early due to the eos_token_id abstract base class for all beam scorers that are used for pretrainedmodel beam_search and pretrainedmodel beam_sample beamscorer implementing standard beam search decoding adapted in part from facebook s xlm beam search code https github com facebookresearch xlm blob 9e6f6814d17be4fe5b15f2e6c43eb2b2d76daeb4 src model transformer py l529 reference for the diverse beam search algorithm and implementation ashwin kalyan s dbs implementation https github com ashwinkalyan dbs blob master dbs beam_utils lua args batch_size int batch size of input_ids for which standard beam search decoding is run in parallel num_beams int number of beams for beam search device torch device defines the device type e g cpu or cuda on which this instance of beamsearchscorer will be allocated length_penalty float optional defaults to 1 0 exponential penalty to the length that is used with beam based generation it is applied as an exponent to the sequence length which in turn is used to divide the score of the sequence since the score is the log likelihood of the sequence i e negative length_penalty 0 0 promotes longer sequences while length_penalty 0 0 encourages shorter sequences do_early_stopping bool or str optional defaults to false controls the stopping condition for beam based methods like beam search it accepts the following values true where the generation stops as soon as there are num_beams complete candidates false where an heuristic is applied and the generation stops when is it very unlikely to find better candidates never where the beam search procedure only stops when there cannot be better candidates canonical beam search algorithm num_beam_hyps_to_keep int optional defaults to 1 the number of beam hypotheses that shall be returned upon calling transformers beamsearchscorer finalize num_beam_groups int optional defaults to 1 number of groups to divide num_beams into in order to ensure diversity among different groups of beams see this paper https arxiv org pdf 1610 02424 pdf for more details max_length int optional the maximum length of the sequence to be generated self _beam_hyps i self num_beam_groups j is the beam_hyps of the j th group in the i th mini batch if group_beam_search is not used the list consists of batch_size beam_hyps self _done i self num_beam_groups j indicates whether the generation of the beam_hyps of the j th group in the i th mini batch is complete add up to the length which the next_scores is calculated on pad the batch next tokens for this sentence add to generated hypotheses if end of sentence if beam_token does not belong to top num_beams tokens it should not be added skip the corner case where the very first generated token is eos_token add next predicted token since it is not eos_token once the beam for next step is full don t add more tokens to it check if we are done so that we can save a pad step if all done finalize all open beam hypotheses and add to generated hypotheses all open beam hypotheses are added to the beam hypothesis beam hypothesis class automatically keeps the best beams select the best hypotheses retrieve best hypotheses append hyp to lists append indices to list prepare for adding eos shorter batches are padded if needed fill with hypotheses and eos_token_id if the latter fits in inserting only the first eos_token_id beamscorer implementing constrained beam search decoding args batch_size int batch size of input_ids for which standard beam search decoding is run in parallel num_beams int number of beams for beam search constraints list constraint a list of positive constraints represented as constraint objects that must be fulfilled in the generation output for more information the documentation of constraint should be read device torch device defines the device type e g cpu or cuda on which this instance of beamsearchscorer will be allocated length_penalty float optional defaults to 1 0 exponential penalty to the length that is used with beam based generation it is applied as an exponent to the sequence length which in turn is used to divide the score of the sequence since the score is the log likelihood of the sequence i e negative length_penalty 0 0 promotes longer sequences while length_penalty 0 0 encourages shorter sequences do_early_stopping bool or str optional defaults to false controls the stopping condition for beam based methods like beam search it accepts the following values true where the generation stops as soon as there are num_beams complete candidates false where an heuristic is applied and the generation stops when is it very unlikely to find better candidates never where the beam search procedure only stops when there cannot be better candidates canonical beam search algorithm num_beam_hyps_to_keep int optional defaults to 1 the number of beam hypotheses that shall be returned upon calling transformers beamsearchscorer finalize num_beam_groups int optional defaults to 1 number of groups to divide num_beams into in order to ensure diversity among different groups of beams see this paper https arxiv org pdf 1610 02424 pdf for more details max_length int optional the maximum length of the sequence to be generated args input_ids torch longtensor of shape batch_size num_beams sequence_length indices of input sequence tokens in the vocabulary indices can be obtained using any class inheriting from pretrainedtokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids next_scores torch floattensor of shape batch_size 2 num_beams current scores of the top 2 num_beams non finished beam hypotheses next_tokens torch longtensor of shape batch_size 2 num_beams input_ids of the tokens corresponding to the top 2 num_beams non finished beam hypotheses next_indices torch longtensor of shape batch_size 2 num_beams beam indices indicating to which beam hypothesis the next_tokens correspond scores_for_all_vocab torch floattensor of shape batch_size num_beams sequence_length the scores of all tokens in the vocabulary for each of the beam hypotheses pad_token_id int optional the id of the padding token eos_token_id union int list int optional the id of the end of sequence token optionally use a list to set multiple end of sequence tokens beam_indices torch longtensor optional beam indices indicating to which beam hypothesis each token correspond decoder_prompt_len int optional the length of prompt that is included in the input to decoder return userdict a dictionary composed of the fields as defined above next_beam_scores torch floattensor of shape batch_size num_beams updated scores of all non finished beams next_beam_tokens torch floattensor of shape batch_size num_beams next tokens to be added to the non finished beam_hypotheses next_beam_indices torch floattensor of shape batch_size num_beams beam indices indicating to which beam the next tokens shall be added add up to the length which the next_scores is calculated on pad the batch next tokens for this sentence add to generated hypotheses if end of sentence if beam_token does not belong to top num_beams tokens it should not be added skip the corner case where the only constraint token is eos_token and the very first generated token is eos_token add next predicted token since it is not eos_token once the beam for next step is full don t add more tokens to it check if we are done so that we can save a pad step if all done sent_beam_tokens are the next num_beams number of tokens that are under consideration for this beam candidate next tokens 1 adding advance_tokens using constraintstatelist advance we propose new tokens to be added into this candidate list that will advance us in fulfilling the constraints 2 selecting best candidates such that we end up with highest probable candidates that fulfill our constraints initialize states need to make new hypothesis that advance the constraints pre_seq ith sequence generated before this step input_ids topk generic beam search best model next tokens advance constraints forcing the next token either way we need to sort them into banks later so store a constraintliststate for all types of hypotheses since adding each advance_token leads to a different hypothesis create new state instance prevent duplicates which are basically bound to happen in this process idx global idx across all the batches basically sent_beam_indices often chooses very little among input_ids the generated sequences that actually fulfill our constraints for example let constraints loves pies and pre_seq_1 the child loves pies and pre_seq_2 the child plays in the playground and without this step if sent_beam_indices is something like 1 1 then 1 pre_seq_1 won t be added to the list of topk hypothesis since it s not in the indices and 2 it won t be added to the list of advance hypothesis since it s completed already this is the else part of if constraints_completed seq_idx 3 it ends up simply getting removed from consideration 3 might be fine and actually desired since it s likely that it s a low probability output anyways especially if it s not in the list of sent_beam_indices but this often leads to lengthened beam search times since completed sequences keep getting removed after all this effort for constrained generation here we basically take pre_seq_1 and to push it into the considered list of hypotheses by simply appending the next likely token in the vocabulary and adding it to the list of hypotheses some next probable token but still don t want to have duplicates then we end up with sorted among bank c sorted among bank c 1 sorted among bank 0 finalize all open beam hypotheses and add to generated hypotheses all open beam hypotheses are added to the beam hypothesis beam hypothesis class automatically keeps the best beams due to overly complex constraints or other factors sometimes we can t gaurantee a successful generation in these cases we simply return the highest scoring outputs select the best hypotheses retrieve best hypotheses append to lists append indices to list prepare for adding eos shorter batches are padded if needed fill with hypotheses and eos_token_id if the latter fits in inserting only the first eos_token_id initialize n best list of hypotheses number of hypotheses in the list add a new hypothesis to the list if there are enough hypotheses and that none of the hypotheses being generated can become better than the worst one in the heap then we are done with this sentence true stop as soon as at least num_beams hypotheses are finished false heuristic compute best possible score from cur_len even though it is not entirely accurate when length_penalty is positive see the discussion below for more details https github com huggingface transformers pull 20901 issuecomment 1369845565 never compute the best possible score depending on the signal of length_penalty length_penalty 0 0 max denominator is obtaned from max_length not from cur_len min abs highest_attainable_score is obtained highest_attainable_score is negative hence we obtain its max this way the opposite logic applies here max highest_attainable_score from cur_len
from abc import ABC, abstractmethod
from collections import UserDict
from typing import Dict, List, Optional, Tuple, Union

import numpy as np
import torch

from ..utils import add_start_docstrings
from .beam_constraints import Constraint, ConstraintListState


PROCESS_INPUTS_DOCSTRING = r

FINALIZE_INPUTS_DOCSTRING = r


class BeamScorer(ABC):
    

    @abstractmethod
    @add_start_docstrings(PROCESS_INPUTS_DOCSTRING)
    def process(
        self,
        input_ids: torch.LongTensor,
        next_scores: torch.FloatTensor,
        next_tokens: torch.LongTensor,
        next_indices: torch.LongTensor,
        **kwargs,
    ) -> Tuple[torch.Tensor]:
        raise NotImplementedError("This is an abstract method.")

    @abstractmethod
    @add_start_docstrings(FINALIZE_INPUTS_DOCSTRING)
    def finalize(
        self,
        input_ids: torch.LongTensor,
        next_scores: torch.FloatTensor,
        next_tokens: torch.LongTensor,
        next_indices: torch.LongTensor,
        max_length: int,
        **kwargs,
    ) -> torch.LongTensor:
        raise NotImplementedError("This is an abstract method.")


class BeamSearchScorer(BeamScorer):
    r

    def __init__(
        self,
        batch_size: int,
        num_beams: int,
        device: torch.device,
        length_penalty: Optional[float] = 1.0,
        do_early_stopping: Optional[Union[bool, str]] = False,
        num_beam_hyps_to_keep: Optional[int] = 1,
        num_beam_groups: Optional[int] = 1,
        max_length: Optional[int] = None,
    ):
        self.num_beams = num_beams
        self.device = device
        self.length_penalty = length_penalty
        self.do_early_stopping = do_early_stopping
        self.num_beam_hyps_to_keep = num_beam_hyps_to_keep
        self.num_beam_groups = num_beam_groups
        self.group_size = self.num_beams // self.num_beam_groups

        self._is_init = False
        
        
        self._beam_hyps = [
            BeamHypotheses(
                num_beams=self.group_size,
                length_penalty=self.length_penalty,
                early_stopping=self.do_early_stopping,
                max_length=max_length,
            )
            for _ in range(batch_size * self.num_beam_groups)
        ]
        
        
        self._done = torch.tensor(
            [False for _ in range(batch_size * self.num_beam_groups)], dtype=torch.bool, device=self.device
        )

        if not isinstance(num_beams, int) or num_beams <= 1:
            raise ValueError(
                f"`num_beams` has to be an integer strictly greater than 1, but is {num_beams}. For `num_beams` == 1,"
                " one should make use of `greedy_search` instead."
            )

        if not isinstance(num_beam_groups, int) or (num_beam_groups > num_beams) or (num_beams % num_beam_groups != 0):
            raise ValueError(
                "`num_beam_groups` has to be an integer smaller or equal than `num_beams` and `num_beams` has to be"
                f" divisible by `num_beam_groups`, but is {num_beam_groups} with `num_beams` being {num_beams}."
            )

    @property
    def is_done(self) -> bool:
        return self._done.all()

    def process(
        self,
        input_ids: torch.LongTensor,
        next_scores: torch.FloatTensor,
        next_tokens: torch.LongTensor,
        next_indices: torch.LongTensor,
        pad_token_id: Optional[int] = None,
        eos_token_id: Optional[Union[int, List[int]]] = None,
        beam_indices: Optional[torch.LongTensor] = None,
        group_index: Optional[int] = 0,
        decoder_prompt_len: Optional[int] = 0,
    ) -> Dict[str, torch.Tensor]:
        
        cur_len = input_ids.shape[-1] - decoder_prompt_len + 1
        batch_size = len(self._beam_hyps) // self.num_beam_groups

        if not (batch_size == (input_ids.shape[0] // self.group_size)):
            if self.num_beam_groups > 1:
                raise ValueError(
                    f"A group beam size of {input_ids.shape[0]} is used as the input, but a group beam "
                    f"size of {self.group_size} is expected by the beam scorer."
                )
            else:
                raise ValueError(
                    f"A beam size of {input_ids.shape[0]} is used as the input, but a beam size of "
                    f"{self.group_size} is expected by the beam scorer."
                )

        device = input_ids.device
        next_beam_scores = torch.zeros((batch_size, self.group_size), dtype=next_scores.dtype, device=device)
        next_beam_tokens = torch.zeros((batch_size, self.group_size), dtype=next_tokens.dtype, device=device)
        next_beam_indices = torch.zeros((batch_size, self.group_size), dtype=next_indices.dtype, device=device)

        if isinstance(eos_token_id, int):
            eos_token_id = [eos_token_id]

        for batch_idx in range(batch_size):
            batch_group_idx = batch_idx * self.num_beam_groups + group_index
            if self._done[batch_group_idx]:
                if self.num_beams < len(self._beam_hyps[batch_group_idx]):
                    raise ValueError(f"Batch can only be done if at least {self.num_beams} beams have been generated")
                if eos_token_id is None or pad_token_id is None:
                    raise ValueError("Generated beams >= num_beams -> eos_token_id and pad_token have to be defined")
                
                next_beam_scores[batch_idx, :] = 0
                next_beam_tokens[batch_idx, :] = pad_token_id
                next_beam_indices[batch_idx, :] = 0
                continue

            
            beam_idx = 0
            for beam_token_rank, (next_token, next_score, next_index) in enumerate(
                zip(next_tokens[batch_idx], next_scores[batch_idx], next_indices[batch_idx])
            ):
                batch_beam_idx = batch_idx * self.group_size + next_index
                
                if (eos_token_id is not None) and (next_token.item() in eos_token_id):
                    
                    is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size
                    if is_beam_token_worse_than_top_num_beams:
                        continue
                    if beam_indices is not None:
                        beam_index = beam_indices[batch_beam_idx]
                        beam_index = beam_index + (batch_beam_idx,)
                    else:
                        beam_index = None

                    
                    if decoder_prompt_len == input_ids.shape[-1]:
                        continue

                    self._beam_hyps[batch_group_idx].add(
                        input_ids[batch_beam_idx].clone(),
                        next_score.item(),
                        beam_indices=beam_index,
                        decoder_prompt_len=decoder_prompt_len,
                    )
                else:
                    
                    next_beam_scores[batch_idx, beam_idx] = next_score
                    next_beam_tokens[batch_idx, beam_idx] = next_token
                    next_beam_indices[batch_idx, beam_idx] = batch_beam_idx
                    beam_idx += 1

                
                if beam_idx == self.group_size:
                    break

            if beam_idx < self.group_size:
                raise ValueError(
                    f"At most {self.group_size} tokens in {next_tokens[batch_idx]} can be equal to `eos_token_id:"
                    f" {eos_token_id}`. Make sure {next_tokens[batch_idx]} are corrected."
                )

            
            self._done[batch_group_idx] = self._done[batch_group_idx] or self._beam_hyps[batch_group_idx].is_done(
                next_scores[batch_idx].max().item(), cur_len
            )

        return UserDict(
            {
                "next_beam_scores": next_beam_scores.view(-1),
                "next_beam_tokens": next_beam_tokens.view(-1),
                "next_beam_indices": next_beam_indices.view(-1),
            }
        )

    def finalize(
        self,
        input_ids: torch.LongTensor,
        final_beam_scores: torch.FloatTensor,
        final_beam_tokens: torch.LongTensor,
        final_beam_indices: torch.LongTensor,
        max_length: int,
        pad_token_id: Optional[int] = None,
        eos_token_id: Optional[Union[int, List[int]]] = None,
        beam_indices: Optional[torch.LongTensor] = None,
        decoder_prompt_len: Optional[int] = 0,
    ) -> Tuple[torch.LongTensor]:
        batch_size = len(self._beam_hyps) // self.num_beam_groups

        if isinstance(eos_token_id, int):
            eos_token_id = [eos_token_id]

        
        for batch_group_idx, beam_hyp in enumerate(self._beam_hyps):
            if self._done[batch_group_idx]:
                continue

            
            
            for index_per_group in range(self.group_size):
                batch_beam_idx = batch_group_idx * self.group_size + index_per_group
                final_score = final_beam_scores[batch_beam_idx].item()
                final_tokens = input_ids[batch_beam_idx]
                beam_index = beam_indices[batch_beam_idx] if beam_indices is not None else None
                beam_hyp.add(final_tokens, final_score, beam_indices=beam_index, decoder_prompt_len=decoder_prompt_len)

        
        sent_lengths = input_ids.new(batch_size * self.num_beam_hyps_to_keep)
        best = []
        best_indices = []
        best_scores = torch.zeros(batch_size * self.num_beam_hyps_to_keep, device=self.device, dtype=torch.float32)

        
        for i in range(batch_size):
            beam_hyps_in_batch = self._beam_hyps[i * self.num_beam_groups : (i + 1) * self.num_beam_groups]
            candidate_beams = [beam for beam_hyp in beam_hyps_in_batch for beam in beam_hyp.beams]
            sorted_hyps = sorted(candidate_beams, key=lambda x: x[0])
            for j in range(self.num_beam_hyps_to_keep):
                best_hyp_tuple = sorted_hyps.pop()
                best_score = best_hyp_tuple[0]
                best_hyp = best_hyp_tuple[1]
                best_index = best_hyp_tuple[2]
                sent_lengths[self.num_beam_hyps_to_keep * i + j] = len(best_hyp)

                
                best.append(best_hyp)

                
                best_indices.append(best_index)

                best_scores[i * self.num_beam_hyps_to_keep + j] = best_score

        
        sent_lengths_max = sent_lengths.max().item() + 1
        sent_max_len = min(sent_lengths_max, max_length) if max_length is not None else sent_lengths_max
        decoded: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len)

        if len(best_indices) > 0 and best_indices[0] is not None:
            indices: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len)
        else:
            indices = None

        
        if sent_lengths.min().item() != sent_lengths.max().item():
            if pad_token_id is None:
                raise ValueError("`pad_token_id` has to be defined")
            decoded.fill_(pad_token_id)

        if indices is not None:
            indices.fill_(-1)

        
        for i, (hypo, best_idx) in enumerate(zip(best, best_indices)):
            decoded[i, : sent_lengths[i]] = hypo

            if indices is not None:
                indices[i, : len(best_idx)] = torch.tensor(best_idx)

            if sent_lengths[i] < sent_max_len:
                
                decoded[i, sent_lengths[i]] = eos_token_id[0]

        return UserDict(
            {
                "sequences": decoded,
                "sequence_scores": best_scores,
                "beam_indices": indices,
            }
        )


class ConstrainedBeamSearchScorer(BeamScorer):
    r

    def __init__(
        self,
        batch_size: int,
        num_beams: int,
        constraints: List[Constraint],
        device: torch.device,
        length_penalty: Optional[float] = 1.0,
        do_early_stopping: Optional[Union[bool, str]] = False,
        num_beam_hyps_to_keep: Optional[int] = 1,
        num_beam_groups: Optional[int] = 1,
        max_length: Optional[int] = None,
    ):
        self.num_beams = num_beams
        self.device = device
        self.length_penalty = length_penalty
        self.do_early_stopping = do_early_stopping
        self.num_beam_hyps_to_keep = num_beam_hyps_to_keep
        self.num_beam_groups = num_beam_groups
        self.group_size = self.num_beams // self.num_beam_groups
        self.constraints = constraints

        self._is_init = False
        self._beam_hyps = [
            BeamHypotheses(
                num_beams=self.num_beams,
                length_penalty=self.length_penalty,
                early_stopping=self.do_early_stopping,
                max_length=max_length,
            )
            for _ in range(batch_size)
        ]
        self._done = torch.tensor([False for _ in range(batch_size)], dtype=torch.bool, device=self.device)

        if not isinstance(num_beams, int) or num_beams <= 1:
            raise ValueError(
                f"`num_beams` has to be an integer strictly greater than 1, but is {num_beams}. For `num_beams` == 1,"
                " one should make use of `greedy_search` instead."
            )

        if not isinstance(num_beam_groups, int) or (num_beam_groups > num_beams) or (num_beams % num_beam_groups != 0):
            raise ValueError(
                "`num_beam_groups` has to be an integer smaller or equal than `num_beams` and `num_beams` has to be"
                f" divisible by `num_beam_groups`, but is {num_beam_groups} with `num_beams` being {num_beams}."
            )

    @property
    def is_done(self) -> bool:
        return self._done.all()

    def make_constraint_states(self, n):
        return [ConstraintListState([constraint.copy() for constraint in self.constraints]) for _ in range(n)]

    def check_completes_constraints(self, sequence):
        new_state = self.make_constraint_states(1)[0]
        new_state.reset(sequence)
        return new_state.completed

    def process(
        self,
        input_ids: torch.LongTensor,
        next_scores: torch.FloatTensor,
        next_tokens: torch.LongTensor,
        next_indices: torch.LongTensor,
        scores_for_all_vocab: torch.FloatTensor,
        pad_token_id: Optional[int] = None,
        eos_token_id: Optional[Union[int, List[int]]] = None,
        beam_indices: Optional[torch.LongTensor] = None,
        decoder_prompt_len: Optional[int] = 0,
    ) -> Tuple[torch.Tensor]:
        r

        
        cur_len = input_ids.shape[-1] - decoder_prompt_len + 1
        batch_size = len(self._beam_hyps)
        if not (batch_size == (input_ids.shape[0] // self.group_size)):
            if self.num_beam_groups > 1:
                raise ValueError(
                    f"A group beam size of {input_ids.shape[0]} is used as the input, but a group beam "
                    f"size of {self.group_size} is expected by the beam scorer."
                )
            else:
                raise ValueError(
                    f"A beam size of {input_ids.shape[0]} is used as the input, but a beam size of "
                    f"{self.group_size} is expected by the beam scorer."
                )

        device = input_ids.device

        next_beam_scores = torch.zeros((batch_size, self.group_size), dtype=next_scores.dtype, device=device)
        next_beam_tokens = torch.zeros((batch_size, self.group_size), dtype=next_tokens.dtype, device=device)
        next_beam_indices = torch.zeros((batch_size, self.group_size), dtype=next_indices.dtype, device=device)

        if isinstance(eos_token_id, int):
            eos_token_id = [eos_token_id]

        for batch_idx, beam_hyp in enumerate(self._beam_hyps):
            if self._done[batch_idx]:
                if self.num_beams < len(beam_hyp):
                    raise ValueError(f"Batch can only be done if at least {self.num_beams} beams have been generated")
                if eos_token_id is None or pad_token_id is None:
                    raise ValueError("Generated beams >= num_beams -> eos_token_id and pad_token have to be defined")
                
                next_beam_scores[batch_idx, :] = 0
                next_beam_tokens[batch_idx, :] = pad_token_id
                next_beam_indices[batch_idx, :] = 0
                continue

            
            beam_idx = 0
            for beam_token_rank, (next_token, next_score, next_index) in enumerate(
                zip(next_tokens[batch_idx], next_scores[batch_idx], next_indices[batch_idx])
            ):
                batch_beam_idx = batch_idx * self.group_size + next_index
                
                if (eos_token_id is not None) and (next_token.item() in eos_token_id):
                    
                    is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size
                    if is_beam_token_worse_than_top_num_beams:
                        continue

                    completes_constraint = self.check_completes_constraints(input_ids[batch_beam_idx].cpu().tolist())
                    if completes_constraint:
                        if beam_indices is not None:
                            beam_index = beam_indices[batch_beam_idx]
                            beam_index = beam_index + (batch_beam_idx,)
                        else:
                            beam_index = None

                        
                        
                        if decoder_prompt_len == input_ids.shape[-1]:
                            continue

                        beam_hyp.add(
                            input_ids[batch_beam_idx].clone(),
                            next_score.item(),
                            beam_indices=beam_index,
                            decoder_prompt_len=decoder_prompt_len,
                        )
                else:
                    
                    next_beam_scores[batch_idx, beam_idx] = next_score
                    next_beam_tokens[batch_idx, beam_idx] = next_token
                    next_beam_indices[batch_idx, beam_idx] = batch_beam_idx
                    beam_idx += 1

                
                if beam_idx == self.group_size:
                    break

            new_scores, new_tokens, new_indices = self.step_sentence_constraint(
                batch_idx,
                input_ids,
                scores_for_all_vocab,
                next_beam_scores[batch_idx],
                next_beam_tokens[batch_idx],
                next_beam_indices[batch_idx],
            )

            next_beam_scores[batch_idx] = new_scores
            next_beam_tokens[batch_idx] = new_tokens
            next_beam_indices[batch_idx] = new_indices

            if beam_idx < self.group_size:
                raise ValueError(
                    f"At most {self.group_size} tokens in {next_tokens[batch_idx]} can be equal to `eos_token_id:"
                    f" {eos_token_id}`. Make sure {next_tokens[batch_idx]} are corrected."
                )

            
            self._done[batch_idx] = self._done[batch_idx] or beam_hyp.is_done(
                next_scores[batch_idx].max().item(), cur_len
            )

        return UserDict(
            {
                "next_beam_scores": next_beam_scores.view(-1),
                "next_beam_tokens": next_beam_tokens.view(-1),
                "next_beam_indices": next_beam_indices.view(-1),
            }
        )

    def step_sentence_constraint(
        self,
        batch_idx: int,
        input_ids: torch.LongTensor,
        vocab_scores: torch.FloatTensor,
        sent_beam_scores: torch.FloatTensor,
        sent_beam_tokens: torch.LongTensor,
        sent_beam_indices: torch.LongTensor,
        push_progress: bool = False,
    ):
        
        

        
        
        

        
        

        orig_len = sent_beam_indices.size(0)
        device = sent_beam_indices.device

        
        topk_contraint_states = self.make_constraint_states(orig_len)
        advance_constraint_states = self.make_constraint_states(orig_len)

        sidx, eidx = batch_idx * orig_len, (batch_idx + 1) * orig_len
        this_batch_input_ids = input_ids[sidx:eidx]
        this_batch_token_scores = vocab_scores[sidx:eidx]
        full_hypotheses = torch.cat((input_ids[sent_beam_indices], sent_beam_tokens.unsqueeze(-1)), dim=-1)

        
        track_new = {
            "new_seqs": full_hypotheses.tolist(),
            "new_states": [],
            "new_indices": [],
            "new_tokens": [],
            "new_scores": [],
        }
        for seq_idx, pre_seq in enumerate(this_batch_input_ids):
            

            
            
            
            

            topk_state = topk_contraint_states[seq_idx]
            topk_state.reset(full_hypotheses[seq_idx].cpu().tolist())

            advance_state = advance_constraint_states[seq_idx]
            advance_state.reset(pre_seq.cpu().tolist())

            if not advance_state.completed:
                advance_tokens = torch.LongTensor(advance_state.advance()).to(device)
                for advance_token in advance_tokens:
                    
                    new_state = advance_state.copy(stateful=True)
                    new_state.add(advance_token.cpu().tolist())

                    advance_seq = torch.cat((pre_seq, advance_token.unsqueeze(0)), -1).cpu().tolist()
                    if advance_seq not in track_new["new_seqs"]:
                        
                        track_new["new_seqs"].append(advance_seq)
                        track_new["new_indices"].append(sidx + seq_idx)  
                        track_new["new_tokens"].append(advance_token)
                        track_new["new_scores"].append(this_batch_token_scores[seq_idx].take(advance_token))
                        track_new["new_states"].append(new_state)
            elif push_progress:
                
                

                

                
                
                
                
                

                
                
                
                

                
                

                new_score, new_token = torch.max(this_batch_token_scores[seq_idx], 0)  
                advance_seq = torch.cat((pre_seq, new_token.unsqueeze(0)), -1)

                advance_state = advance_constraint_states[seq_idx]

                advance_seq = advance_seq.cpu().tolist()

                advance_state.reset(advance_seq)
                if advance_seq not in track_new["new_seqs"]:
                    
                    track_new["new_seqs"].append(advance_seq)
                    track_new["new_indices"].append(seq_idx)
                    track_new["new_tokens"].append(new_token)
                    track_new["new_scores"].append(new_score)
                    track_new["new_states"].append(advance_state)

        if len(track_new["new_indices"]) > 0:
            new_indices = torch.tensor(track_new["new_indices"]).to(device)
            new_tokens = torch.stack(track_new["new_tokens"]).to(device)
            new_scores = torch.stack(track_new["new_scores"]).to(device)

            all_states = topk_contraint_states + track_new["new_states"]
            all_tokens = torch.cat((sent_beam_tokens, new_tokens), -1)
            all_scores = torch.cat((sent_beam_scores, new_scores), -1)
            all_banks = torch.tensor([one.get_bank() for one in all_states]).to(device)

            zipped = all_banks * 100 + all_scores
            indices = zipped.sort(descending=True).indices
            sorted_banks = all_banks[indices]

            

            counter = -1
            cur_bank = sorted_banks[0]
            increments = []
            for bank in sorted_banks:
                if bank == cur_bank:
                    counter += 1
                else:
                    counter = 0
                    cur_bank = bank
                increments.append(counter)
            rearrangers = torch.tensor(np.argsort(increments, kind="mergesort"))

            indices = indices[rearrangers][:orig_len]

            sent_beam_scores = all_scores[indices]
            sent_beam_tokens = all_tokens[indices]
            sent_beam_indices = torch.cat((sent_beam_indices, new_indices))[indices]

        return sent_beam_scores, sent_beam_tokens, sent_beam_indices

    def finalize(
        self,
        input_ids: torch.LongTensor,
        final_beam_scores: torch.FloatTensor,
        final_beam_tokens: torch.LongTensor,
        final_beam_indices: torch.LongTensor,
        max_length: int,
        pad_token_id: Optional[int] = None,
        eos_token_id: Optional[Union[int, List[int]]] = None,
        beam_indices: Optional[torch.LongTensor] = None,
        decoder_prompt_len: Optional[int] = 0,
    ) -> Tuple[torch.LongTensor]:
        batch_size = len(self._beam_hyps)

        if isinstance(eos_token_id, int):
            eos_token_id = [eos_token_id]

        
        for batch_idx, beam_hyp in enumerate(self._beam_hyps):
            if self._done[batch_idx]:
                continue

            
            

            ids_collect = []
            for beam_id in range(self.num_beams):
                batch_beam_idx = batch_idx * self.num_beams + beam_id
                final_score = final_beam_scores[batch_beam_idx].item()
                final_tokens = input_ids[batch_beam_idx]

                completes_constraint = self.check_completes_constraints(final_tokens.cpu().tolist())
                if completes_constraint:
                    beam_index = beam_indices[batch_beam_idx] if beam_indices is not None else None
                    beam_hyp.add(
                        final_tokens, final_score, beam_indices=beam_index, decoder_prompt_len=decoder_prompt_len
                    )
                    ids_collect.append(beam_id)

            
            
            if len(ids_collect) < self.num_beam_hyps_to_keep:
                for beam_id in range(self.num_beams):
                    if beam_id not in ids_collect:
                        batch_beam_idx = batch_idx * self.num_beams + beam_id
                        final_score = final_beam_scores[batch_beam_idx].item()
                        final_tokens = input_ids[batch_beam_idx]
                        beam_hyp.add(final_tokens, final_score, decoder_prompt_len=decoder_prompt_len)
                    if len(ids_collect) >= self.num_beam_hyps_to_keep:
                        break

        
        sent_lengths = input_ids.new(batch_size * self.num_beam_hyps_to_keep)
        best = []
        best_indices = []
        best_scores = torch.zeros(batch_size * self.num_beam_hyps_to_keep, device=self.device, dtype=torch.float32)

        
        for i, beam_hyp in enumerate(self._beam_hyps):
            sorted_hyps = sorted(beam_hyp.beams, key=lambda x: x[0])
            for j in range(self.num_beam_hyps_to_keep):
                best_hyp_tuple = sorted_hyps.pop()
                best_score = best_hyp_tuple[0]
                best_hyp = best_hyp_tuple[1]
                best_index = best_hyp_tuple[2]
                sent_lengths[self.num_beam_hyps_to_keep * i + j] = len(best_hyp)

                
                best.append(best_hyp)

                
                best_indices.append(best_index)

                best_scores[i * self.num_beam_hyps_to_keep + j] = best_score

        
        sent_lengths_max = sent_lengths.max().item() + 1

        sent_max_len = min(sent_lengths_max, max_length) if max_length is not None else sent_lengths_max
        decoded: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len)

        if len(best_indices) > 0 and best_indices[0] is not None:
            indices: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len)
        else:
            indices = None

        
        if sent_lengths.min().item() != sent_lengths.max().item():
            if pad_token_id is None:
                raise ValueError("`pad_token_id` has to be defined")
            decoded.fill_(pad_token_id)

        if indices is not None:
            indices.fill_(-1)

        
        for i, (hypo, best_idx) in enumerate(zip(best, best_indices)):
            decoded[i, : sent_lengths[i]] = hypo

            if indices is not None:
                indices[i, : len(best_idx)] = torch.tensor(best_idx)

            if sent_lengths[i] < sent_max_len:
                
                decoded[i, sent_lengths[i]] = eos_token_id[0]

        return UserDict(
            {
                "sequences": decoded,
                "sequence_scores": best_scores,
                "beam_indices": indices,
            }
        )


class BeamHypotheses:
    def __init__(self, num_beams: int, length_penalty: float, early_stopping: bool, max_length: Optional[int] = None):
        
        self.length_penalty = length_penalty
        self.early_stopping = early_stopping
        self.max_length = max_length
        self.num_beams = num_beams
        self.beams = []
        self.worst_score = 1e9

        if not isinstance(self.early_stopping, bool) and self.max_length is None:
            raise ValueError(
                "When `do_early_stopping` is set to a string, `max_length` must be defined. Ensure it is passed to the"
                " BeamScorer class instance at initialization time."
            )

    def __len__(self):
        
        return len(self.beams)

    def add(
        self,
        hyp: torch.LongTensor,
        sum_logprobs: float,
        beam_indices: Optional[torch.LongTensor] = None,
        decoder_prompt_len: Optional[int] = 0,
    ):
        
        score = sum_logprobs / ((hyp.shape[-1] - decoder_prompt_len) ** self.length_penalty)
        if len(self) < self.num_beams or score > self.worst_score:
            self.beams.append((score, hyp, beam_indices))
            if len(self) > self.num_beams:
                sorted_next_scores = sorted([(s, idx) for idx, (s, _, _) in enumerate(self.beams)])
                del self.beams[sorted_next_scores[0][1]]
                self.worst_score = sorted_next_scores[1][0]
            else:
                self.worst_score = min(score, self.worst_score)

    def is_done(self, best_sum_logprobs: float, cur_len: int) -> bool:
        

        if len(self) < self.num_beams:
            return False

        
        if self.early_stopping is True:
            return True
        
        
        
        elif self.early_stopping is False:
            highest_attainable_score = best_sum_logprobs / cur_len**self.length_penalty
            ret = self.worst_score >= highest_attainable_score
            return ret
        
        else:
            
            
            
            if self.length_penalty > 0.0:
                highest_attainable_score = best_sum_logprobs / self.max_length**self.length_penalty
            
            else:
                highest_attainable_score = best_sum_logprobs / cur_len**self.length_penalty
            ret = self.worst_score >= highest_attainable_score
            return ret