in bart model the type of pastkeyvalues is tupletupletorch floattensor which is not torchscriptcompatible to support this we have to convert it during the export process this function will convert past values from a list to tupletupletorch floattensor for the inner decoder according to the definition of pastkeyvalues each inner tupletorch floattensor has 4 tensors so we convert every 4 elements in the list as a tupletorch floattensor do this twice so we got 2 different decoders for further work bartconfigts is a torchscriptcompatible transformers models bart configurationbart bartconfig torchscript only supports subclasses of torch nn module def initself minlength int eostokenid int super init if not isinstanceminlength int or minlength 0 raise valueerrorfminlength has to be a positive integer but is minlength if not isinstanceeostokenid int or eostokenid 0 raise valueerrorfeostokenid has to be a positive integer but is eostokenid self minlength minlength self eostokenid eostokenid def forwardself inputids scores torch tensor curlen inputids shape1 if curlen self minlength scores self eostokenid floatinf return scores class bartgeneratortorch nn module generationmixin def initself model super init self config bartconfigtsmodel config self config forcebostokentobegenerated false self tracemodulesmodel self logitsprocessor minlengthlogitsprocessortsself config minlength self config eostokenid self finallogitsweight model model shared weight self finallogitsbias model finallogitsbias self decoderlayers model config decoderlayers def tracemodulesself model inputids torch tensor 19 669 18 420 8 664 57 42 8 664 21 3028 195 4445 331 1293 34 21 10 6174 1100 6 69 104 42 32 2621 1638 144 4 6174 558 108 4419 1091 28 4 1668 9 1509 1621 279 35 867 2734 85 11 2216 2734 85 203 2244 7 6 15 8102 7 57 8629 5 model config eostokenid devicemodel device dtypetorch long attentionmask torch tensor true inputids shape1 devicemodel device dtypetorch bool self encoder createtracedencodermodel getencoder inputids attentionmask encoderoutputs model getencoderinputids attentionmaskattentionmask returndicttrue decoder model model decoder decoderoutputs decoderinputids attentionmask encoderoutputslasthiddenstate none none none self decodernopast createtraceddecoder model model decoder inputids encoderoutputslasthiddenstate attentionmask self decoderwithpast createtraceddecoder model model decoder inputids encoderoutputslasthiddenstate attentionmask decoderoutputs1 def encoderforwardself inputids attentionmask return self encoderinputids attentionmask0 staticmethod def initsequencelengthforgeneration inputids torch longtensor maxlength int tupletorch tensor torch tensor int unfinishedsequences torch zerosinputids shape0 dtypetorch long deviceinputids device 1 sequencelengths torch zerosinputids shape0 dtypetorch long deviceinputids device maxlength curlen inputids shape1 return sequencelengths unfinishedsequences curlen def decoderforwardself inputids encoderoutput attentionmask past listtorch tensor update here to use different decoder for different values of past if past is none or lenpast 0 decoderoutput past self decodernopast inputidsinputids encoderstateencoderoutput attentionmaskattentionmask else decoderoutput past self decoderwithpast inputidsinputids encoderstateencoderoutput attentionmaskattentionmask pastpast lmlogits f lineardecoderoutput self finallogitsweight biasself finallogitsbias return lmlogits past def greedysearch self inputids encoderoutput attentionmask maxlength padtokenid int eostokenid int init sequence length tensors sequencelengths unfinishedsequences curlen self initsequencelengthforgeneration inputids maxlength past listtorch tensor while curlen maxlength logits past self decoderforwardinputids encoderoutput attentionmask past nexttokenlogits logits 1 preprocess distribution scores self logitsprocessorinputids nexttokenlogits argmax nexttokens torch argmaxscores dim1 add code that transfomers nexttokens to tokenstoadd if eostokenid is not none assert padtokenid is not none if eostokenid is defined make sure that padtokenid is defined nexttokens nexttokens unfinishedsequences padtokenid 1 unfinishedsequences add token and increase length by one inputids torch catinputids nexttokens none dim1 update sequence length if eostokenid is not none sequencelengths unfinishedsequences self updateseqlengthforgeneration sequencelengths unfinishedsequences curlen nexttokens eostokenid stop when there is a s in each sentence or if we exceed the maximul length if unfinishedsequences max 0 break increase curlen curlen curlen 1 return inputids def preparedecoderinputidsforgeneration self inputids torch longtensor decoderstarttokenid bostokenid optionalint none torch longtensor decoderinputids torch onesinputids shape0 1 dtypeinputids dtype deviceinputids device decoderstarttokenid return decoderinputids def forwardself inputids attentionmask maxlength decoderstarttokenid padtokenid self config padtokenid bostokenid self config bostokenid eostokenid self config eostokenid special case if padtokenid is not defined if padtokenid is none and eostokenid is not none setting padtokenid to eostokenid eostokenid for openend generation padtokenid eostokenid encoderoutput self encoderforwardinputids attentionmask inputids self preparedecoderinputidsforgeneration inputids decoderstarttokeniddecoderstarttokenid bostokenidbostokenid return self greedysearch inputids encoderoutput attentionmask maxlengthmaxlength padtokenidpadtokenid eostokenideostokenid torchscript compatible beamsearchscorer class beamsearchscorertstorch nn module def initself super init self maxlength int 200 self numbeams int 3 self batchsize int 1 self lengthpenalty float 1 0 self doearlystopping bool true self numbeamhypstokeep int 1 self numbeamgroups int 1 self groupsize int self numbeams self numbeamgroups self done torch zerosself batchsize dtypetorch bool self beamhypscount torch zerosself batchsize dtypetorch long self beamhypsworstscores torch zerosself batchsize 1e9 self beamhypsmaxlength int self maxlength 1 self beamhyps listtorch tensor torch zeros2 placeholder for torchscript compatibility self beamscores listtorch tensor torch zeros2 placeholder for torchscript compatibility def isdoneself torch tensor return self done all def init self batchsize int maxlength int numbeams int device torch device lengthpenalty float 1 0 doearlystopping bool false numbeamhypstokeep int 1 numbeamgroups int 1 self maxlength maxlength self numbeams numbeams self batchsize batchsize self lengthpenalty lengthpenalty self doearlystopping doearlystopping self numbeamhypstokeep numbeamhypstokeep self numbeamgroups numbeamgroups self groupsize self numbeams self numbeamgroups note torchscript does not support list of modules rewritten beamhypotheses with tensors and list of tensors self done torch zerosbatchsize dtypetorch bool devicedevice self beamhypscount torch zerosbatchsize dtypetorch long devicedevice self beamhypsworstscores torch zerosbatchsize devicedevice 1e9 self beamhyps self beamscores self beamhypsmaxlength maxlength 1 ignoring bostoken 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 def hypolenself hypoidx int return self beamhypscounthypoidx def hypoaddself hyp torch tensor sumlogprobs float hypoidx int score sumlogprobs hyp shape1 self lengthpenalty hypscount self hypolenhypoidx if hypscount self numbeams or score self beamhypsworstscoreshypoidx note work around difference of torch sumemptytensor 0 while error in onnx bug https msdata visualstudio comviennaworkitemsedit1486599 beamidx torch sumself beamhypscount hypoidx if hypoidx 0 else torch tensor0 dtypetorch long self beamscores insertbeamidx torch tensorscore self beamhyps insertbeamidx hyp if hypscount 1 self numbeams sortednextscores sortedindices torch topk torch catself beamscoresbeamidx beamidx hypscount 1 hypscount 1 largestfalse del self beamhypsintsortedindices0 beamidx del self beamscoresintsortedindices0 beamidx self beamhypsworstscoreshypoidx sortednextscores1 else self beamhypsworstscoreshypoidx minscore self beamhypsworstscoreshypoidx self beamhypscounthypoidx hypscount 1 def hypoisdoneself hypoidx int bestsumlogprobs float curlen int bool if self hypolenhypoidx self numbeams return false elif self doearlystopping return true else curscore bestsumlogprobs curlenself lengthpenalty ret self beamhypsworstscoreshypoidx item curscore return ret def process self inputids torch tensor nextscores torch tensor nexttokens torch tensor nextindices torch tensor padtokenid optionalint none eostokenid optionalint none tupletorch tensor torch tensor torch tensor curlen inputids shape1 batchsize lenself beamhypscount assert batchsize inputids shape0 self groupsize 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 for batchidx in rangebatchsize if self donebatchidx assert self hypolenbatchidx self numbeams batch can only be done if at least beams have been generated formatself numbeams assert eostokenid is not none and padtokenid is not none generated 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 eostokenid if beamtoken does not belong to top numbeams tokens it should not be added isbeamtokenworsethantopnumbeams beamtokenrank self groupsize if isbeamtokenworsethantopnumbeams continue self hypoadd inputidsbatchbeamidx clone nextscore item batchidx 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 donebatchidx self donebatchidx or self hypoisdone batchidx nextscoresbatchidx max item curlen return nextbeamscores view1 nextbeamtokens view1 nextbeamindices view1 def finalize self inputids torch tensor finalbeamscores torch tensor finalbeamtokens torch tensor finalbeamindices torch tensor padtokenid int eostokenid int tupletorch tensor torch tensor batchsize lenself beamhypscount finalize all open beam hypotheses and add to generated hypotheses for batchidx in rangebatchsize if self donebatchidx continue all open beam hypotheses are added to the beam hypothesis beam hypothesis class automatically keeps the best beams for beamid in rangeself numbeams batchbeamidx batchidx self numbeams beamid finalscore finalbeamscoresbatchbeamidx item finaltokens inputidsbatchbeamidx self hypoaddfinaltokens finalscore batchidx select the best hypotheses note torch tensor newzeros is not scriptable sentlengths torch zerosbatchsize self numbeamhypstokeep dtypetorch long best bestscores torch zeros batchsize self numbeamhypstokeep deviceinputids device dtypetorch float32 retrieve best hypotheses for i in rangebatchsize note lambda is not scriptable batchhypostart torch sumself beamhypscount i if i 0 else torch tensor0 dtypetorch long batchhypoend torch sumself beamhypscount i 1 beamscores torch catself beamscoresbatchhypostart batchhypoend sortednextscores sortedindices torch topkbeamscores lenbeamscores largesttrue for j in rangeself numbeamhypstokeep bestscore beamscoressortedindicesj besthyp self beamhypsbatchhypostart sortedindicesj sentlengthsself numbeamhypstokeep i j lenbesthyp append to lists best appendbesthyp bestscoresi self numbeamhypstokeep j bestscore prepare for adding eos sentmaxlen minsentlengths max 1 self maxlength decoded torch zerosbatchsize self numbeamhypstokeep sentmaxlen dtypetorch long shorter batches are padded if needed if sentlengths min sentlengths max assert padtokenid is not none padtokenid has to be defined decoded fillpadtokenid fill with hypotheses and eostokenid if the latter fits in for i hypo in enumeratebest decodedi sentlengthsi hypo if sentlengthsi self maxlength decodedi sentlengthsi eostokenid return decoded bestscores class bartbeamsearchgeneratorbartgenerator def initself model super initmodel self beamscorer beamsearchscorerts self device model device staticmethod def expandinputsforgeneration inputids torch tensor attentionmask torch tensor lasthiddenstate torch tensor expandsize int 1 tupletorch tensor torch tensor torch tensor expandedreturnidx torch arangeinputids shape0 view1 1 repeat1 expandsize view1 toinputids device inputids inputids indexselect0 expandedreturnidx attentionmask attentionmask indexselect0 expandedreturnidx lasthiddenstate lasthiddenstate indexselect0 expandedreturnidx tolasthiddenstate device return inputids attentionmask lasthiddenstate def adjustlogitsduringgenerationself logits curlen int maxlength int if curlen 1 and self config forcebostokentobegenerated logits self forcetokenidtobegeneratedlogits self config bostokenid elif curlen maxlength 1 and self config eostokenid is not none logits self forcetokenidtobegeneratedlogits self config eostokenid return logits staticmethod def forcetokenidtobegeneratedscores tokenid int if decoder past is not included in output speedy decoding is disabled and no need to reorder adjust tokens for bart e g preprocess distribution reshape for beam search special case if padtokenid is not defined logger warningfsetting padtokenid to eostokenid eostokenid for openend generation in bart model the type of past_key_values is tuple tuple torch floattensor which is not torchscript compatible to support this we have to convert it during the export process this function will convert past values from a list to tuple tuple torch floattensor for the inner decoder according to the definition of past_key_values each inner tuple torch floattensor has 4 tensors so we convert every 4 elements in the list as a tuple torch floattensor do this twice so we got 2 different decoders for further work bartconfigts is a torchscript compatible transformers models bart configuration_bart bartconfig torchscript only supports sub classes of torch nn module class transformers logitsprocessor enforcing a min length by setting eos probability to 0 args min_length obj int the minimum length below which the score of obj eos_token_id is set to obj float inf eos_token_id obj int the id of the end of sequence token update here to use different decoder for different values of past init sequence length tensors pre process distribution argmax add code that transfomers next_tokens to tokens_to_add add token and increase length by one update sequence length stop when there is a s in each sentence or if we exceed the maximul length increase cur_len special case if pad_token_id is not defined setting pad_token_id to eos_token_id eos_token_id for open end generation torchscript compatible beamsearchscorer placeholder for torchscript compatibility placeholder for torchscript compatibility note torchscript does not support list of modules rewritten beamhypotheses with tensors and list of tensors ignoring bos_token number of hypotheses in the list add a new hypothesis to the list note work around difference of torch sum empty_tensor 0 while error in onnx bug https msdata visualstudio com vienna _workitems edit 1486599 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 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 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 note torch tensor new_zeros is not scriptable retrieve best hypotheses note lambda is not scriptable append to lists prepare for adding eos shorter batches are padded if needed fill with hypotheses and eos_token_id if the latter fits in force one of token_ids to be generated by setting prob of all other tokens to 0 logprob float inf if decoder past is not included in output speedy decoding is disabled and no need to reorder adjust tokens for bart e g batch_size num_beams vocab_size pre process distribution reshape for beam search special case if pad_token_id is not defined logger warning f setting pad_token_id to eos_token_id eos_token_id for open end generation
import copy
import itertools
from typing import List, Optional, Tuple

import torch
import torch.nn.functional as F

from transformers import BartConfig
from transformers.generation import GenerationMixin


def _convert_past_list_to_tuple(past_key_values):
    
    count_of_each_inner_tuple = 4
    results = ()
    temp_result = ()
    count_n = len(past_key_values) // count_of_each_inner_tuple
    for idx in range(count_n):
        real_idx = idx * count_of_each_inner_tuple
        temp_result = tuple(past_key_values[real_idx : real_idx + count_of_each_inner_tuple])
        results += ((temp_result),)

    return results


class EncoderForONNX(torch.nn.Module):
    def __init__(self, encoder):
        super().__init__()
        self.encoder = encoder

    def forward(self, input_ids, attention_mask):
        return self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=False,
        )


class DecoderForONNX(torch.nn.Module):
    def __init__(self, decoder):
        super().__init__()
        self.decoder = decoder

    def forward(self, input_ids, encoder_state, attention_mask, past=None):
        all_results = None
        if past is not None:
            all_results = _convert_past_list_to_tuple(past)
            input_ids = input_ids[:, -1:]

        last_hidden_state, past_key_values = self.decoder(
            input_ids=input_ids,
            encoder_hidden_states=encoder_state,
            encoder_attention_mask=attention_mask,
            past_key_values=all_results,
            return_dict=False,
        )

        past_values = []
        for past in past_key_values:
            past_values = past_values + list(past)
        return last_hidden_state, past_values


def _create_traced_encoder(encoder, input_ids, attention_mask):
    encoder_c = copy.deepcopy(encoder)
    encoder_for_onnx = EncoderForONNX(encoder_c)

    return torch.jit.trace(encoder_for_onnx, (input_ids, attention_mask))


def _create_traced_decoder(decoder, input_ids, encoder_state, attention_mask, past=None):
    decoder_c = copy.deepcopy(decoder)
    decoder_for_onnx = DecoderForONNX(decoder_c)
    past_values = list(itertools.chain.from_iterable(past or ()))

    
    if past_values:
        return torch.jit.trace(decoder_for_onnx, (input_ids, encoder_state, attention_mask, past_values))
    else:
        return torch.jit.trace(decoder_for_onnx, (input_ids, encoder_state, attention_mask))


class BartConfigTS(BartConfig, torch.nn.Module):
    

    def __init__(self, config):
        BartConfig.__init__(self, config)
        torch.nn.Module.__init__(self)


class MinLengthLogitsProcessorTS(torch.nn.Module):
    r

    def __init__(self, min_length: int, eos_token_id: int):
        super().__init__()

        if not isinstance(min_length, int) or min_length < 0:
            raise ValueError(f"`min_length` has to be a positive integer, but is {min_length}")

        if not isinstance(eos_token_id, int) or eos_token_id < 0:
            raise ValueError(f"`eos_token_id` has to be a positive integer, but is {eos_token_id}")

        self.min_length = min_length
        self.eos_token_id = eos_token_id

    def forward(self, input_ids, scores) -> torch.Tensor:
        cur_len = input_ids.shape[-1]
        if cur_len < self.min_length:
            scores[:, self.eos_token_id] = -float("inf")
        return scores


class BARTGenerator(torch.nn.Module, GenerationMixin):
    def __init__(self, model):
        super().__init__()
        self.config = BartConfigTS(model.config)
        self.config.force_bos_token_to_be_generated = False
        self._trace_modules(model)
        self.logits_processor = MinLengthLogitsProcessorTS(self.config.min_length, self.config.eos_token_id)
        self.final_logits_weight = model.model.shared.weight
        self.final_logits_bias = model.final_logits_bias
        self.decoder_layers = model.config.decoder_layers

    def _trace_modules(self, model):
        input_ids = torch.tensor(
            [
                [
                    19,
                    669,
                    18,
                    420,
                    8,
                    664,
                    57,
                    42,
                    8,
                    664,
                    21,
                    3028,
                    195,
                    4445,
                    331,
                    1293,
                    34,
                    21,
                    10,
                    6174,
                    1100,
                    6,
                    69,
                    104,
                    42,
                    32,
                    2621,
                    1638,
                    144,
                    4,
                    6174,
                    558,
                    108,
                    4419,
                    1091,
                    28,
                    4,
                    1668,
                    9,
                    1509,
                    1621,
                    279,
                    35,
                    867,
                    2734,
                    85,
                    11,
                    2216,
                    2734,
                    85,
                    203,
                    2244,
                    7,
                    6,
                    15,
                    8102,
                    7,
                    57,
                    8629,
                    5,
                    model.config.eos_token_id,
                ]
            ],
            device=model.device,
            dtype=torch.long,
        )
        attention_mask = torch.tensor(
            [[True] * input_ids.shape[-1]],
            device=model.device,
            dtype=torch.bool,
        )
        self.encoder = _create_traced_encoder(model.get_encoder(), input_ids, attention_mask)
        encoder_outputs = model.get_encoder()(input_ids, attention_mask=attention_mask, return_dict=True)
        decoder = model.model.decoder
        decoder_outputs = decoder(input_ids, attention_mask, encoder_outputs["last_hidden_state"], None, None, None)
        self.decoder_no_past = _create_traced_decoder(
            model.model.decoder, input_ids, encoder_outputs["last_hidden_state"], attention_mask
        )
        self.decoder_with_past = _create_traced_decoder(
            model.model.decoder, input_ids, encoder_outputs["last_hidden_state"], attention_mask, decoder_outputs[1]
        )

    def _encoder_forward(self, input_ids, attention_mask):
        return self.encoder(input_ids, attention_mask)[0]

    @staticmethod
    def _init_sequence_length_for_generation(
        input_ids: torch.LongTensor, max_length: int
    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
        unfinished_sequences = torch.zeros(input_ids.shape[0], dtype=torch.long, device=input_ids.device) + 1
        sequence_lengths = torch.zeros(input_ids.shape[0], dtype=torch.long, device=input_ids.device) + max_length

        cur_len = input_ids.shape[-1]
        return sequence_lengths, unfinished_sequences, cur_len

    def _decoder_forward(self, input_ids, encoder_output, attention_mask, past: List[torch.Tensor]):
        
        if past is None or len(past) == 0:
            decoder_output, past = self.decoder_no_past(
                input_ids=input_ids, encoder_state=encoder_output, attention_mask=attention_mask
            )
        else:
            decoder_output, past = self.decoder_with_past(
                input_ids=input_ids, encoder_state=encoder_output, attention_mask=attention_mask, past=past
            )

        lm_logits = F.linear(decoder_output, self.final_logits_weight, bias=self.final_logits_bias)

        return lm_logits, past

    def greedy_search(
        self, input_ids, encoder_output, attention_mask, max_length, pad_token_id: int, eos_token_id: int
    ):
        
        sequence_lengths, unfinished_sequences, cur_len = self._init_sequence_length_for_generation(
            input_ids, max_length
        )

        past: List[torch.Tensor] = []
        while cur_len < max_length:
            logits, past = self._decoder_forward(input_ids, encoder_output, attention_mask, past)
            next_token_logits = logits[:, -1, :]

            
            scores = self.logits_processor(input_ids, next_token_logits)

            
            next_tokens = torch.argmax(scores, dim=-1)

            
            if eos_token_id is not None:
                assert pad_token_id is not None, "If eos_token_id is defined, make sure that pad_token_id is defined."
                next_tokens = next_tokens * unfinished_sequences + (pad_token_id) * (1 - unfinished_sequences)

            
            input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)

            
            if eos_token_id is not None:
                sequence_lengths, unfinished_sequences = self._update_seq_length_for_generation(
                    sequence_lengths, unfinished_sequences, cur_len, next_tokens == eos_token_id
                )

            
            if unfinished_sequences.max() == 0:
                break

            
            cur_len = cur_len + 1

        return input_ids

    def _prepare_decoder_input_ids_for_generation(
        self,
        input_ids: torch.LongTensor,
        decoder_start_token_id,
        bos_token_id: Optional[int] = None,
    ) -> torch.LongTensor:
        decoder_input_ids = (
            torch.ones((input_ids.shape[0], 1), dtype=input_ids.dtype, device=input_ids.device)
            * decoder_start_token_id
        )
        return decoder_input_ids

    def forward(self, input_ids, attention_mask, max_length, decoder_start_token_id):
        pad_token_id = self.config.pad_token_id
        bos_token_id = self.config.bos_token_id
        eos_token_id = self.config.eos_token_id

        
        if pad_token_id is None and eos_token_id is not None:
            
            pad_token_id = eos_token_id

        encoder_output = self._encoder_forward(input_ids, attention_mask)

        input_ids = self._prepare_decoder_input_ids_for_generation(
            input_ids,
            decoder_start_token_id=decoder_start_token_id,
            bos_token_id=bos_token_id,
        )

        return self.greedy_search(
            input_ids,
            encoder_output,
            attention_mask,
            max_length=max_length,
            pad_token_id=pad_token_id,
            eos_token_id=eos_token_id,
        )



class BeamSearchScorerTS(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.max_length: int = 200
        self.num_beams: int = 3
        self.batch_size: int = 1
        self.length_penalty: float = 1.0
        self.do_early_stopping: bool = True
        self.num_beam_hyps_to_keep: int = 1
        self.num_beam_groups: int = 1
        self.group_size: int = self.num_beams // self.num_beam_groups
        self._done = torch.zeros(self.batch_size, dtype=torch.bool)
        self._beam_hyps_count = torch.zeros(self.batch_size, dtype=torch.long)
        self._beam_hyps_worst_scores = torch.zeros(self.batch_size) + 1e9
        self._beam_hyps_max_length: int = self.max_length - 1
        self._beam_hyps: List[torch.Tensor] = [torch.zeros(2)]  
        self._beam_scores: List[torch.Tensor] = [torch.zeros(2)]  

    def is_done(self) -> torch.Tensor:
        return self._done.all()

    def init(
        self,
        batch_size: int,
        max_length: int,
        num_beams: int,
        device: torch.device,
        length_penalty: float = 1.0,
        do_early_stopping: bool = False,
        num_beam_hyps_to_keep: int = 1,
        num_beam_groups: int = 1,
    ):
        self.max_length = max_length
        self.num_beams = num_beams
        self.batch_size = batch_size
        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._done = torch.zeros(batch_size, dtype=torch.bool, device=device)
        self._beam_hyps_count = torch.zeros(batch_size, dtype=torch.long, device=device)
        self._beam_hyps_worst_scores = torch.zeros(batch_size, device=device) + 1e9
        self._beam_hyps = []
        self._beam_scores = []

        self._beam_hyps_max_length = max_length - 1  

        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}."
            )

    def hypo_len(self, hypo_idx: int):
        
        return self._beam_hyps_count[hypo_idx]

    def hypo_add(self, hyp: torch.Tensor, sum_logprobs: float, hypo_idx: int):
        
        score = sum_logprobs / (hyp.shape[-1] ** self.length_penalty)
        hyps_count = self.hypo_len(hypo_idx)
        if hyps_count < self.num_beams or score > self._beam_hyps_worst_scores[hypo_idx]:
            
            
            beam_idx = (
                torch.sum(self._beam_hyps_count[:hypo_idx]) if hypo_idx != 0 else torch.tensor(0, dtype=torch.long)
            )
            self._beam_scores.insert(beam_idx, torch.tensor([score]))
            self._beam_hyps.insert(beam_idx, hyp)
            if hyps_count + 1 > self.num_beams:
                sorted_next_scores, sorted_indices = torch.topk(
                    torch.cat(self._beam_scores)[beam_idx : beam_idx + hyps_count + 1], hyps_count + 1, largest=False
                )
                del self._beam_hyps[int((sorted_indices[0] + beam_idx))]
                del self._beam_scores[int((sorted_indices[0] + beam_idx))]
                self._beam_hyps_worst_scores[hypo_idx] = sorted_next_scores[1]
            else:
                self._beam_hyps_worst_scores[hypo_idx] = min(score, self._beam_hyps_worst_scores[hypo_idx])
                self._beam_hyps_count[hypo_idx] = hyps_count + 1

    def hypo_is_done(self, hypo_idx: int, best_sum_logprobs: float, cur_len: int) -> bool:
        
        if self.hypo_len(hypo_idx) < self.num_beams:
            return False
        elif self.do_early_stopping:
            return True
        else:
            cur_score = best_sum_logprobs / cur_len**self.length_penalty
            ret = self._beam_hyps_worst_scores[hypo_idx].item() >= cur_score
            return ret

    def process(
        self,
        input_ids: torch.Tensor,
        next_scores: torch.Tensor,
        next_tokens: torch.Tensor,
        next_indices: torch.Tensor,
        pad_token_id: Optional[int] = None,
        eos_token_id: Optional[int] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        cur_len = input_ids.shape[-1]
        batch_size = len(self._beam_hyps_count)
        assert batch_size == (input_ids.shape[0] // self.group_size)

        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)

        for batch_idx in range(batch_size):
            if self._done[batch_idx]:
                assert (
                    self.hypo_len(batch_idx) >= self.num_beams
                ), "Batch can only be done if at least {} beams have been generated".format(self.num_beams)
                assert (
                    eos_token_id is not None and pad_token_id is not None
                ), "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 == 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
                    self.hypo_add(
                        input_ids[batch_beam_idx].clone(),
                        next_score.item(),
                        batch_idx,
                    )
                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_idx] = self._done[batch_idx] or self.hypo_is_done(
                batch_idx,
                next_scores[batch_idx].max().item(),
                cur_len,
            )

        return next_beam_scores.view(-1), next_beam_tokens.view(-1), next_beam_indices.view(-1)

    def finalize(
        self,
        input_ids: torch.Tensor,
        final_beam_scores: torch.Tensor,
        final_beam_tokens: torch.Tensor,
        final_beam_indices: torch.Tensor,
        pad_token_id: int,
        eos_token_id: int,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        batch_size = len(self._beam_hyps_count)

        
        for batch_idx in range(batch_size):
            if self._done[batch_idx]:
                continue

            
            
            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]
                self.hypo_add(final_tokens, final_score, batch_idx)

        
        
        sent_lengths = torch.zeros(batch_size * self.num_beam_hyps_to_keep, dtype=torch.long)
        best = []
        best_scores = torch.zeros(
            batch_size * self.num_beam_hyps_to_keep, device=input_ids.device, dtype=torch.float32
        )
        
        for i in range(batch_size):
            
            batch_hypo_start = torch.sum(self._beam_hyps_count[:i]) if i > 0 else torch.tensor(0, dtype=torch.long)
            batch_hypo_end = torch.sum(self._beam_hyps_count[: i + 1])
            beam_scores = torch.cat(self._beam_scores)[batch_hypo_start:batch_hypo_end]
            sorted_next_scores, sorted_indices = torch.topk(beam_scores, len(beam_scores), largest=True)
            for j in range(self.num_beam_hyps_to_keep):
                best_score = beam_scores[sorted_indices[j]]
                best_hyp = self._beam_hyps[batch_hypo_start + sorted_indices[j]]
                sent_lengths[self.num_beam_hyps_to_keep * i + j] = len(best_hyp)
                
                best.append(best_hyp)
                best_scores[i * self.num_beam_hyps_to_keep + j] = best_score

        
        sent_max_len = min(sent_lengths.max() + 1, self.max_length)
        decoded = torch.zeros(batch_size * self.num_beam_hyps_to_keep, sent_max_len, dtype=torch.long)
        
        if sent_lengths.min() != sent_lengths.max():
            assert pad_token_id is not None, "`pad_token_id` has to be defined"
            decoded.fill_(pad_token_id)

        
        for i, hypo in enumerate(best):
            decoded[i, : sent_lengths[i]] = hypo
            if sent_lengths[i] < self.max_length:
                decoded[i, sent_lengths[i]] = eos_token_id

        return decoded, best_scores


class BARTBeamSearchGenerator(BARTGenerator):
    def __init__(self, model):
        super().__init__(model)
        self.beam_scorer = BeamSearchScorerTS()
        self.device = model.device

    @staticmethod
    def _expand_inputs_for_generation(
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        last_hidden_state: torch.Tensor,
        expand_size: int = 1,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        expanded_return_idx = (
            torch.arange(input_ids.shape[0]).view(-1, 1).repeat(1, expand_size).view(-1).to(input_ids.device)
        )
        input_ids = input_ids.index_select(0, expanded_return_idx)

        attention_mask = attention_mask.index_select(0, expanded_return_idx)

        last_hidden_state = last_hidden_state.index_select(0, expanded_return_idx.to(last_hidden_state.device))
        return input_ids, attention_mask, last_hidden_state

    def adjust_logits_during_generation(self, logits, cur_len: int, max_length: int):
        if cur_len == 1 and self.config.force_bos_token_to_be_generated:
            logits = self._force_token_id_to_be_generated(logits, self.config.bos_token_id)
        elif cur_len == max_length - 1 and self.config.eos_token_id is not None:
            logits = self._force_token_id_to_be_generated(logits, self.config.eos_token_id)
        return logits

    @staticmethod
    def _force_token_id_to_be_generated(scores, token_id: int):
        
        mask = torch.full_like(scores, 1, dtype=torch.bool)
        mask[:, token_id] = False
        return scores.masked_fill(mask, -float("inf"))

    def _reorder_cache(self, past: List[torch.Tensor], beam_idx):
        
        
        reordered_decoder_past = []
        for state in past:
            reordered_decoder_past.append(state.index_select(0, beam_idx))
        return reordered_decoder_past

    def beam_search(
        self, input_ids, encoder_output, attention_mask, num_beams, max_length, pad_token_id: int, eos_token_id: int
    ):
        batch_size = self.beam_scorer.batch_size

        num_beams = self.beam_scorer.num_beams
        batch_beam_size, cur_len = input_ids.shape

        assert (
            num_beams * batch_size == batch_beam_size
        ), f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."

        beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device)
        beam_scores[:, 1:] = -1e9
        beam_scores = beam_scores.view((batch_size * num_beams,))
        next_tokens = torch.zeros((batch_size, num_beams), dtype=torch.long, device=input_ids.device)
        next_indices = torch.zeros((batch_size, num_beams), dtype=torch.long, device=input_ids.device)

        past: List[torch.Tensor] = []
        while cur_len < max_length:
            logits, past = self._decoder_forward(input_ids, encoder_output, attention_mask, past)
            next_token_logits = logits[:, -1, :]

            
            next_token_logits = self.adjust_logits_during_generation(
                next_token_logits, cur_len=cur_len, max_length=max_length
            )

            next_token_scores = F.log_softmax(next_token_logits, dim=-1)  

            
            next_token_scores = self.logits_processor(input_ids, next_token_scores)
            next_token_scores = next_token_scores + beam_scores[:, None].expand_as(next_token_scores)

            
            vocab_size = next_token_scores.shape[-1]
            next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size)

            next_token_scores, next_tokens = torch.topk(
                next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True
            )

            next_indices = next_tokens // vocab_size
            next_tokens = next_tokens % vocab_size

            beam_scores, beam_next_tokens, beam_idx = self.beam_scorer.process(
                input_ids,
                next_token_scores,
                next_tokens,
                next_indices,
                pad_token_id=pad_token_id,
                eos_token_id=eos_token_id,
            )

            input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)

            cur_len = cur_len + 1

            if len(past) > 0:
                past = self._reorder_cache(past, beam_idx)

            if self.beam_scorer.is_done():
                break

        sequences, sequence_scores = self.beam_scorer.finalize(
            input_ids,
            beam_scores,
            next_tokens,
            next_indices,
            pad_token_id=pad_token_id,
            eos_token_id=eos_token_id,
        )

        return sequences

    def forward(self, input_ids, attention_mask, num_beams, max_length, decoder_start_token_id):
        pad_token_id = self.config.pad_token_id
        bos_token_id = self.config.bos_token_id
        eos_token_id = self.config.eos_token_id

        
        if pad_token_id is None and eos_token_id is not None:
            
            pad_token_id = eos_token_id

        encoder_output = self._encoder_forward(input_ids, attention_mask)

        input_ids = self._prepare_decoder_input_ids_for_generation(
            input_ids,
            decoder_start_token_id=decoder_start_token_id,
            bos_token_id=bos_token_id,
        )

        batch_size = input_ids.shape[0]

        length_penalty = self.config.length_penalty
        num_return_sequences = self.config.num_return_sequences
        early_stopping = True

        self.beam_scorer.init(
            batch_size=batch_size,
            max_length=max_length,
            num_beams=num_beams,
            device=self.device,
            length_penalty=length_penalty,
            do_early_stopping=early_stopping,
            num_beam_hyps_to_keep=num_return_sequences,
        )

        input_ids, attention_mask, encoder_output = self._expand_inputs_for_generation(
            input_ids,
            attention_mask,
            encoder_output,
            expand_size=num_beams,
        )

        return self.beam_search(
            input_ids=input_ids,
            encoder_output=encoder_output,
            attention_mask=attention_mask,
            num_beams=num_beams,
            max_length=max_length,
            pad_token_id=pad_token_id,
            eos_token_id=eos_token_id,
        )