InternLM/internlm/apis/inference.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-

import torch
import torch.nn.functional as F
from torch import nn

__all__ = ["SequenceGenerator"]


class InferenceParams:
    """
    Intermediate cache objects for inference
    """

    def __init__(
        self,
        max_sequence_len,
        max_batch_size,
        sequence_len_offset=0,
        batch_size_offset=0,
        key_value_memory_dict: dict = None,
        lengths_per_sample=None,
        attention_mask=None,
    ) -> None:

        self.max_sequence_len: int = max_sequence_len
        self.max_batch_size: int = max_batch_size
        self.sequence_len_offset: int = sequence_len_offset
        self.batch_size_offset: int = batch_size_offset
        if key_value_memory_dict is None:
            key_value_memory_dict = {}
        self.key_value_memory_dict: dict = key_value_memory_dict
        self.fused_ft_kernel: bool = False
        self.lengths_per_sample = lengths_per_sample
        self.attention_mask = attention_mask

    def reorder_state(self, indices):
        if self.lengths_per_sample is not None:
            self.lengths_per_sample = self.lengths_per_sample.index_select(index=indices, dim=0)
        for key, value in list(self.key_value_memory_dict.items()):
            value = value.index_select(index=indices, dim=0)
            self.key_value_memory_dict[key] = value


def _get_model_device(model):
    """
    obtain the device of an nn.Module.model

    Args:
        model: nn.Module

    Return: torch.device. if None, the parameters of this model is None.
    """
    assert isinstance(model, nn.Module)

    parameters = list(model.parameters())
    if len(parameters) == 0:
        return None
    else:
        return parameters[0].device


class SequenceGenerator:
    """
    Sequence Generator.
    """

    def __init__(self, decoder, eos_token_id, pad_token_id, bos_token_id):
        self.decoder = decoder
        self.eos_token_id = eos_token_id
        self.pad_token_id = pad_token_id
        self.bos_token_id = bos_token_id

    @torch.no_grad()
    def generate(
        self,
        tokens: "torch.LongTensor" = None,
        num_return_sequences=1,
        max_length: int = 20,
        num_beams: int = 1,
        do_sample: bool = True,
        temperature: float = 1.0,
        top_k: int = 50,
        top_p: float = 1.0,
        repetition_penalty: float = 1,
        length_penalty: float = 1.0,
    ):
        """
        Args:
            tokens: the beginning tokens whose shape is [bsz, length]. If shape is None, default ''bos_token'' will be
                added to conduct generation.
            num_return_sequences: number of returned sequences.
            max_length: the max length of generated sequence.
            num_beams: the size of beam search.
            do_sample: whether using sample.
            temperature: it's meaningful when do_sample is True.
            top_k: sampling from top_k.
            top_p: sampling from top_p tokens(nucleus sampling).

        Return:
            the token sequence whose shape is [bsz, num_return_sequences, max_length]. If eos_token_id is not None,
                the ending of each sequence must be eos_token_id.
        """
        assert num_return_sequences <= num_beams, f"The `{num_return_sequences}` must be less than `{num_beams}`..."
        if do_sample:
            return sample_generate(
                self.decoder,
                tokens=tokens,
                max_length=max_length,
                num_beams=num_beams,
                num_return_sequences=num_return_sequences,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                eos_token_id=self.eos_token_id,  # the ending token id
                pad_token_id=self.pad_token_id,
                repetition_penalty=repetition_penalty,  # the penalty degree for repetition tokens
                length_penalty=length_penalty,  # the penalty for length. if it > 1, then encourages long sequence.
                # Otherwise, encourages short sequence.
                bos_token_id=self.bos_token_id,
            )
        else:
            return greedy_generate(
                self.decoder,
                tokens=tokens,
                max_length=max_length,
                num_beams=num_beams,
                num_return_sequences=num_return_sequences,
                eos_token_id=self.eos_token_id,
                pad_token_id=self.pad_token_id,
                repetition_penalty=repetition_penalty,
                length_penalty=length_penalty,
                bos_token_id=self.bos_token_id,
            )


@torch.no_grad()
def greedy_generate(
    decoder,
    tokens=None,
    max_length=20,
    num_beams=1,
    num_return_sequences=1,
    eos_token_id=None,
    pad_token_id=0,
    repetition_penalty=1,
    length_penalty=1.0,
    bos_token_id=1,
    feat_mask=None,
    ffn_mask=None,
    layer_mask=None,
):
    """
    Search sequence greedily.

    Args:
        decoder: the Decoder object.
        tokens: the shape is [batch size, length]. If decoder is None, generating begins with bos_token_id.
        max_length: the max length for generated sequence.
        num_beams: the size of beam to decode.
        eos_token_id: the ending token id. If None, the decode length is max_length.
        pad_token_id: the token id of pad.
        repetition_penalty: the penalty degree for repetition tokens
        length_penalty: the penalty for length.

    """
    if num_beams == 1:
        token_ids = _no_beam_search_generate(
            decoder,
            tokens=tokens,
            max_length=max_length,
            temperature=1,
            top_k=50,
            top_p=1,
            eos_token_id=eos_token_id,
            do_sample=False,
            repetition_penalty=repetition_penalty,
            length_penalty=length_penalty,
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            feat_mask=feat_mask,
            ffn_mask=ffn_mask,
            layer_mask=layer_mask,
        )
    else:
        token_ids = _beam_search_generate(
            decoder,
            tokens=tokens,
            max_length=max_length,
            num_beams=num_beams,
            num_return_sequences=num_return_sequences,
            temperature=1,
            top_k=50,
            top_p=1,
            eos_token_id=eos_token_id,
            do_sample=False,
            repetition_penalty=repetition_penalty,
            length_penalty=length_penalty,
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            feat_mask=feat_mask,
            ffn_mask=ffn_mask,
            layer_mask=layer_mask,
        )

    return token_ids


@torch.no_grad()
def sample_generate(
    decoder,
    tokens,
    max_length=20,
    num_beams=1,
    num_return_sequences=1,
    temperature=1.0,
    top_k=50,
    top_p=1.0,
    eos_token_id=None,
    pad_token_id=0,
    repetition_penalty=1.0,
    length_penalty=1.0,
    bos_token_id=1,
):
    """
    generate sequence in sampling way.

    Args:
        decoder: the Decoder object.
        tokens: the shape is [batch size, length]. If decoder is None, generating begins with bos_token_id.
        max_length: the max length for generated sequence.
        num_beams: the size of beam to decode.
        num_return_sequences: number of returned sequence.
        temperature: annealing magnitude during sampling.
        top_k: sampling from top_k. (Default: 50)
        top_p: sampling from top_p tokens(nucleus sampling). (Default: 1.0)
        eos_token_id: the ending token id. If None, the decode length is max_length.
        pad_token_id: the token id of pad.
        repetition_penalty: the penalty degree for repetition tokens
        length_penalty: the penalty for length.

    """
    if num_beams == 1:
        token_ids = _no_beam_search_generate(
            decoder,
            tokens=tokens,
            max_length=max_length,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            eos_token_id=eos_token_id,
            do_sample=True,
            repetition_penalty=repetition_penalty,
            length_penalty=length_penalty,
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
        )
    else:
        token_ids = _beam_search_generate(
            decoder,
            tokens=tokens,
            max_length=max_length,
            num_beams=num_beams,
            num_return_sequences=num_return_sequences,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            eos_token_id=eos_token_id,
            do_sample=True,
            repetition_penalty=repetition_penalty,
            length_penalty=length_penalty,
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
        )
    return token_ids


@torch.no_grad()
def _no_beam_search_generate(
    decoder,
    tokens,
    inference_params=None,
    max_length=20,
    temperature=1.0,
    top_k=50,
    top_p=1.0,
    eos_token_id=None,
    do_sample=True,
    repetition_penalty=1.0,
    length_penalty=1.0,
    pad_token_id=0,
    bos_token_id=1,
    feat_mask=None,
    ffn_mask=None,
    layer_mask=None,
):
    # delete num_return_sequences=1 for lint check;
    batch_size = tokens.size(0)
    if eos_token_id is None:
        _eos_token_id = -1
    else:
        _eos_token_id = eos_token_id

    has_bos = torch.all(tokens[:, 0].eq(bos_token_id))
    if has_bos:
        bos_pos = torch.where(tokens.eq(bos_token_id), 1, 0)
        bos_sum = bos_pos.cumsum(dim=-1)
        bos_pos = torch.where(bos_sum.eq(bos_sum[:, -1:]), 0, 1)
        to_atten_x = bos_pos[:, :, None]
        to_atten_y = bos_pos[:, None, :]
        # attention_mask = torch.einsum('bno,bom->bnm', to_atten_x, to_atten_y).eq(1)
    else:
        bos_pos = torch.where(tokens.eq(bos_token_id), 1, 0)
        to_atten_x = bos_pos[:, :, None]
        to_atten_y = bos_pos[:, None, :]
        # attention_mask = torch.einsum('bno,bom->bnm', to_atten_x, to_atten_y).eq(1)
    attention_mask = torch.logical_or(to_atten_x, to_atten_y).eq(1)
    if inference_params is None:
        inference_params = InferenceParams(
            max_sequence_len=max_length,
            max_batch_size=tokens.size(0),
            sequence_len_offset=0,
            batch_size_offset=0,
            key_value_memory_dict=None,
            lengths_per_sample=None,
            attention_mask=attention_mask,
        )

    if layer_mask is None:
        if feat_mask is None and ffn_mask is None:
            scores = decoder(**{"input_ids": tokens, "inference_params": inference_params})
        else:
            scores = decoder(
                **{
                    "input_ids": tokens,
                    "inference_params": inference_params,
                    "feat_mask": feat_mask,
                    "ffn_mask": ffn_mask,
                }
            )
    else:
        scores = decoder(
            **{
                "input_ids": tokens,
                "inference_params": inference_params,
                "feat_mask": feat_mask,
                "ffn_mask": ffn_mask,
                "layer_mask": layer_mask,
            }
        )

    if isinstance(scores, (list, tuple)):
        scores = scores[0]
    scores = scores[:, -1].float()
    inference_params.sequence_len_offset += tokens.size(1)
    if _eos_token_id != -1:
        scores[:, _eos_token_id] = -1e12
    next_tokens = scores.argmax(dim=-1, keepdim=True)
    token_ids = torch.cat([tokens, next_tokens], dim=1)
    cur_len = token_ids.size(1)
    dones = token_ids.new_zeros(batch_size).eq(1)
    # tokens = tokens[:, -1:]

    real_max_length = max_length
    max_lengths = tokens.new_full((tokens.size(0),), fill_value=max_length, dtype=torch.long)

    while cur_len < real_max_length:
        # batch_size x vocab_size
        if has_bos:
            bos_pos = torch.where(token_ids.eq(bos_token_id), 1, 0)
            bos_sum = bos_pos.cumsum(dim=-1)
            bos_pos = torch.where(bos_sum.eq(bos_sum[:, -1:]), 0, 1)
            to_atten_x = bos_pos[:, :, None]
            to_atten_y = bos_pos[:, None, :]
            # attention_mask = torch.einsum('bno,bom->bnm', to_atten_x, to_atten_y).eq(1)
        else:
            bos_pos = torch.where(token_ids.eq(bos_token_id), 1, 0)
            to_atten_x = bos_pos[:, :, None]
            to_atten_y = bos_pos[:, None, :]
            # attention_mask = torch.einsum('bno,bom->bnm', to_atten_x, to_atten_y).eq(1)
        attention_mask = torch.logical_or(to_atten_x, to_atten_y).eq(1)
        inference_params.attention_mask = attention_mask
        if layer_mask is None:
            if feat_mask is None and ffn_mask is None:
                scores = decoder(**{"input_ids": token_ids[:, -1:], "inference_params": inference_params})
            else:
                scores = decoder(
                    **{
                        "input_ids": token_ids[:, -1:],
                        "inference_params": inference_params,
                        "feat_mask": feat_mask,
                        "ffn_mask": ffn_mask,
                    }
                )
        else:
            scores = decoder(
                **{
                    "input_ids": token_ids[:, -1:],
                    "inference_params": inference_params,
                    "feat_mask": feat_mask,
                    "ffn_mask": ffn_mask,
                    "layer_mask": layer_mask,
                }
            )

        if isinstance(scores, (list, tuple)):
            scores = scores[0]
        scores = scores[:, -1].float()
        inference_params.sequence_len_offset += 1

        if repetition_penalty != 1.0:
            token_scores = scores.gather(dim=1, index=token_ids)
            lt_zero_mask = token_scores.lt(0).float()
            ge_zero_mask = lt_zero_mask.eq(0).float()
            token_scores = (
                lt_zero_mask * repetition_penalty * token_scores + ge_zero_mask / repetition_penalty * token_scores
            )
            scores.scatter_(dim=1, index=token_ids, src=token_scores)

        if eos_token_id is not None and length_penalty != 1.0:
            # batch_size x vocab_size
            token_scores = scores / cur_len**length_penalty
            eos_mask = scores.new_ones(scores.size(1))
            eos_mask[eos_token_id] = 0
            eos_mask = eos_mask.unsqueeze(0).eq(1)

            scores = scores.masked_scatter(eos_mask, token_scores)

        if do_sample:
            if temperature > 0 and temperature != 1:
                scores = scores / temperature

            scores = top_k_top_p_filtering(scores, top_k, top_p, min_tokens_to_keep=2)
            # add 1e-12 to avoid https://github.com/pytorch/pytorch/pull/27523
            probs = F.softmax(scores, dim=-1) + 1e-12

            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)  # batch_size
        else:
            next_tokens = torch.argmax(scores, dim=-1)  # batch_size

        if _eos_token_id != -1:
            next_tokens = next_tokens.masked_fill(max_lengths.eq(cur_len + 1), _eos_token_id)
        next_tokens = next_tokens.masked_fill(dones, pad_token_id)
        tokens = next_tokens.unsqueeze(1)

        token_ids = torch.cat([token_ids, tokens], dim=-1)  # batch_size x max_len

        end_mask = next_tokens.eq(_eos_token_id)
        dones = dones.__or__(end_mask)
        cur_len += 1

        if dones.min() == 1:
            break

    # if eos_token_id is not None:
    #     # setting the eos at the maximum length position
    #     tokens.scatter(index=max_lengths[:, None], dim=1, value=eos_token_id)
    # if cur_len == max_length:
    #     # If eos is not reached by the maximum length, forcibly replace the last word with eos
    #     token_ids[:, -1].masked_fill_(~dones, eos_token_id)
    # TODO Here we are simply adding an extra dimension for interface compatibility, but in the future it will need to
    # be able to return multiple real results
    return token_ids[:, None]


@torch.no_grad()
def _beam_search_generate(
    decoder,
    tokens,
    inference_params=None,
    max_length=20,
    num_beams=4,
    num_return_sequences=1,
    temperature=1.0,
    top_k=50,
    top_p=1.0,
    eos_token_id=None,
    do_sample=True,
    repetition_penalty=1.0,
    length_penalty=1.0,
    pad_token_id=0,
    bos_token_id=1,
    feat_mask=None,
    ffn_mask=None,
    layer_mask=None,
) -> torch.LongTensor:

    device = _get_model_device(decoder)
    batch_size = tokens.size(0)

    if eos_token_id is None:
        _eos_token_id = -1
    else:
        _eos_token_id = eos_token_id

    has_bos = torch.all(tokens[:, 0].eq(bos_token_id))

    if has_bos:
        bos_pos = torch.where(tokens.eq(bos_token_id), 1, 0)
        bos_sum = bos_pos.cumsum(dim=-1)
        bos_pos = torch.where(bos_sum.eq(bos_sum[:, -1:]), 0, 1)
        to_atten_x = bos_pos[:, :, None]
        to_atten_y = bos_pos[:, None, :]
        # attention_mask = torch.einsum('bno,bom->bnm', to_atten_x, to_atten_y).eq(1)
    else:
        bos_pos = torch.where(tokens.eq(bos_token_id), 1, 0)
        to_atten_x = bos_pos[:, :, None]
        to_atten_y = bos_pos[:, None, :]
        # attention_mask = torch.einsum('bno,bom->bnm', to_atten_x, to_atten_y).eq(1)
    attention_mask = torch.logical_or(to_atten_x, to_atten_y).eq(1)

    if inference_params is None:
        inference_params = InferenceParams(
            max_sequence_len=max_length,
            max_batch_size=tokens.size(0),
            sequence_len_offset=0,
            batch_size_offset=0,
            key_value_memory_dict=None,
            lengths_per_sample=None,
            attention_mask=attention_mask,
        )

    if layer_mask is None:
        if feat_mask is None and ffn_mask is None:
            scores = decoder(**{"input_ids": tokens, "inference_params": inference_params})
        else:
            scores = decoder(
                **{
                    "input_ids": tokens,
                    "inference_params": inference_params,
                    "feat_mask": feat_mask,
                    "ffn_mask": ffn_mask,
                }
            )
    else:
        scores = decoder(
            **{
                "input_ids": tokens,
                "inference_params": inference_params,
                "feat_mask": feat_mask,
                "ffn_mask": ffn_mask,
                "layer_mask": layer_mask,
            }
        )

    if isinstance(scores, (list, tuple)):
        scores = scores[0]
    scores = scores[:, -1].float()
    inference_params.sequence_len_offset += tokens.size(1)
    if _eos_token_id != -1:
        scores[:, _eos_token_id] = -1e12
    vocab_size = scores.size(1)
    assert vocab_size >= num_beams, "num_beams should be smaller than " "the number of vocabulary size."

    if do_sample:
        probs = F.softmax(scores, dim=-1) + 1e-12
        # (batch_size, num_beams)
        next_tokens = torch.multinomial(probs, num_samples=num_beams)
        logits = probs.log()
        # (batch_size, num_beams)
        next_scores = logits.gather(dim=1, index=next_tokens)
    else:
        scores = F.log_softmax(scores, dim=-1)  # (batch_size, vocab_size)
        # obtain (batch_size, num_beams), (batch_size, num_beams)
        next_scores, next_tokens = torch.topk(scores, num_beams, dim=1, largest=True, sorted=True)

    indices = torch.arange(batch_size, dtype=torch.long).to(device)
    indices = indices.repeat_interleave(num_beams)
    inference_params.reorder_state(indices)

    # batch_size * num_beams x length
    tokens = tokens.index_select(dim=0, index=indices)
    # genrated token (batch_size', cur_len)
    token_ids = torch.cat([tokens, next_tokens.view(-1, 1)], dim=-1)
    dones = [False] * batch_size

    beam_scores = next_scores.view(-1)  # batch_size * num_beams

    cur_len = token_ids.size(1)

    real_max_length = max_length
    max_lengths = tokens.new_full((tokens.size(0),), fill_value=max_length, dtype=torch.long)
    hypos = [
        BeamHypotheses(num_beams, real_max_length, length_penalty, early_stopping=False) for _ in range(batch_size)
    ]
    # 0, num_beams, 2*num_beams, ...
    batch_inds_with_numbeams_interval = (torch.arange(batch_size) * num_beams).view(-1, 1).to(token_ids)

    while cur_len < real_max_length:
        if has_bos:
            bos_pos = torch.where(token_ids.eq(bos_token_id), 1, 0)
            bos_sum = bos_pos.cumsum(dim=-1)
            bos_pos = torch.where(bos_sum.eq(bos_sum[:, -1:]), 0, 1)
            to_atten_x = bos_pos[:, :, None]
            to_atten_y = bos_pos[:, None, :]
            # attention_mask = torch.einsum('bno,bom->bnm', to_atten_x, to_atten_y).eq(1)
        else:
            bos_pos = torch.where(token_ids.eq(bos_token_id), 1, 0)
            to_atten_x = bos_pos[:, :, None]
            to_atten_y = bos_pos[:, None, :]
            # attention_mask = torch.einsum('bno,bom->bnm', to_atten_x, to_atten_y).eq(1)
        attention_mask = torch.logical_or(to_atten_x, to_atten_y).eq(1)

        inference_params.attention_mask = attention_mask
        # (bsz x num_beams, vocab_size)

        if layer_mask is None:
            if feat_mask is None and ffn_mask is None:
                scores = decoder(**{"input_ids": token_ids[:, -1:], "inference_params": inference_params})
            else:
                scores = decoder(
                    **{
                        "input_ids": token_ids[:, -1:],
                        "inference_params": inference_params,
                        "feat_mask": feat_mask,
                        "ffn_mask": ffn_mask,
                    }
                )
        else:
            scores = decoder(
                **{
                    "input_ids": token_ids[:, -1:],
                    "inference_params": inference_params,
                    "feat_mask": feat_mask,
                    "ffn_mask": ffn_mask,
                    "layer_mask": layer_mask,
                }
            )

        if isinstance(scores, (list, tuple)):
            scores = scores[0]
        scores = scores[:, -1].float()
        inference_params.sequence_len_offset += 1
        if repetition_penalty != 1.0:
            token_scores = scores.gather(dim=1, index=token_ids)
            lt_zero_mask = token_scores.lt(0).float()
            ge_zero_mask = lt_zero_mask.eq(0).float()
            token_scores = (
                lt_zero_mask * repetition_penalty * token_scores + ge_zero_mask / repetition_penalty * token_scores
            )
            scores.scatter_(dim=1, index=token_ids, src=token_scores)

        if _eos_token_id != -1:
            max_len_eos_mask = max_lengths.eq(cur_len + 1)
            eos_scores = scores[:, _eos_token_id]
            scores[:, _eos_token_id] = torch.where(max_len_eos_mask, eos_scores + 1e32, eos_scores)

        if do_sample:
            if temperature > 0 and temperature != 1:
                scores = scores / temperature

            scores = top_k_top_p_filtering(scores, top_k, top_p, min_tokens_to_keep=num_beams + 1)
            # add 1e-12 to avoid https://github.com/pytorch/pytorch/pull/27523
            probs = F.softmax(scores, dim=-1) + 1e-12

            # batch_size' x (num_beams+1)
            _tokens = torch.multinomial(probs, num_samples=num_beams + 1)

            logits = probs.log()
            # batch_size' x (num_beams+1)
            _scores = logits.gather(dim=1, index=_tokens)
            # batch_size' x (num_beams+1)
            _scores = _scores + beam_scores[:, None]
            _scores = _scores.view(batch_size, num_beams * (num_beams + 1))
            next_scores, ids = _scores.topk(2 * num_beams, dim=1, largest=True, sorted=True)
            _tokens = _tokens.view(batch_size, num_beams * (num_beams + 1))
            # (batch_size, 2*num_beams)
            next_tokens = _tokens.gather(dim=1, index=ids)
            # (batch_size, 2*num_beams)
            from_which_beam = torch.floor(ids.float() / (num_beams + 1)).long()
        else:
            # (batch_size * num_beams, vocab_size)
            scores = F.log_softmax(scores, dim=-1)
            # (batch_size * num_beams, vocab_size)
            _scores = scores + beam_scores[:, None]
            # (batch_size, num_beams*vocab_size)
            _scores = _scores.view(batch_size, -1)
            # (bsz, 2*num_beams)
            next_scores, ids = torch.topk(_scores, 2 * num_beams, dim=1, largest=True, sorted=True)
            # (batch_size, 2*num_beams)
            from_which_beam = torch.floor(ids.float() / vocab_size).long()
            next_tokens = ids % vocab_size  # (batch_size, 2*num_beams)

        # next_scores, sorted_inds = next_scores.sort(dim=-1, descending=True)
        # next_tokens = next_tokens.gather(dim=1, index=sorted_inds)
        # from_which_beam = from_which_beam.gather(dim=1, index=sorted_inds)

        not_eos_mask = next_tokens.ne(_eos_token_id)
        keep_mask = not_eos_mask.cumsum(dim=1).le(num_beams)
        keep_mask = not_eos_mask.__and__(keep_mask)

        _next_tokens = next_tokens.masked_select(keep_mask).view(-1, 1)
        _from_which_beam = from_which_beam.masked_select(keep_mask).view(batch_size, num_beams)
        _next_scores = next_scores.masked_select(keep_mask).view(batch_size, num_beams)
        beam_scores = _next_scores.view(-1)

        flag = True
        if cur_len + 1 == real_max_length:
            eos_batch_idx = torch.arange(batch_size).to(next_tokens).repeat_interleave(repeats=num_beams, dim=0)
            eos_beam_ind = torch.arange(num_beams).to(token_ids).repeat(batch_size)
            eos_beam_idx = from_which_beam[:, :num_beams].reshape(-1)
        else:
            effective_eos_mask = next_tokens[:, :num_beams].eq(_eos_token_id)  # batch_size x num_beams
            if effective_eos_mask.sum().gt(0):
                eos_batch_idx, eos_beam_ind = effective_eos_mask.nonzero(as_tuple=True)
                eos_beam_idx = eos_batch_idx * num_beams * 2 + eos_beam_ind
                eos_beam_idx = from_which_beam.view(-1)[eos_beam_idx]
            else:
                flag = False

        if flag:
            _token_ids = torch.cat([token_ids, _next_tokens], dim=-1)
            for batch_idx, beam_ind, beam_idx in zip(
                eos_batch_idx.tolist(), eos_beam_ind.tolist(), eos_beam_idx.tolist()
            ):
                if not dones[batch_idx]:
                    score = next_scores[batch_idx, beam_ind].item()
                    if _eos_token_id != -1:
                        hypos[batch_idx].add(_token_ids[batch_idx * num_beams + beam_idx, :cur_len].clone(), score)
                    else:
                        hypos[batch_idx].add(_token_ids[batch_idx * num_beams + beam_idx].clone(), score)

        reorder_inds = (batch_inds_with_numbeams_interval + _from_which_beam).view(-1)
        inference_params.reorder_state(reorder_inds)
        token_ids = torch.cat([token_ids.index_select(index=reorder_inds, dim=0), _next_tokens], dim=-1)

        for batch_idx in range(batch_size):
            dones[batch_idx] = (
                dones[batch_idx]
                or hypos[batch_idx].is_done(next_scores[batch_idx, 0].item())
                or max_lengths[batch_idx * num_beams] == cur_len + 1
            )

        cur_len += 1

        if all(dones):
            break

    # select the best hypotheses
    tgt_len = token_ids.new_zeros(batch_size, num_return_sequences)
    best = []

    for i, hypotheses in enumerate(hypos):
        # best_hyp = max(hypotheses.hyp, key=lambda x: x[0])[1]
        sorted_hyp = list(sorted(hypotheses.hyp, key=lambda x: x[0], reverse=True))
        _best = []
        for j, hyp in zip(range(num_return_sequences), sorted_hyp):
            hyp = hyp[1]
            if _eos_token_id != -1:
                hyp = torch.cat([hyp, token_ids.new_ones(1) * _eos_token_id])
            tgt_len[i, j] = len(hyp)
            _best.append(hyp)
        best.append(_best)

    # generate target batch
    decoded = token_ids.new_zeros(batch_size, num_return_sequences, tgt_len.max().item()).fill_(pad_token_id)
    for i, hypo in enumerate(best):
        for j, _hypo in enumerate(hypo):
            decoded[i, j, : tgt_len[i, j]] = _hypo

    return decoded


class BeamHypotheses(object):
    """
    BeamHypotheses
    """

    def __init__(self, num_beams, max_length, length_penalty, early_stopping):
        """Initialize n-best list of hypotheses."""
        self.max_length = max_length - 1  # ignoring bos_token
        self.length_penalty = length_penalty
        self.early_stopping = early_stopping
        self.num_beams = num_beams
        self.hyp = []
        self.worst_score = 1e9

    def __len__(self):
        """Number of hypotheses in the list."""
        return len(self.hyp)

    def add(self, hyp, sum_logprobs):
        """Add a new hypothesis to the list."""
        score = sum_logprobs / len(hyp) ** self.length_penalty
        if len(self) < self.num_beams or score > self.worst_score:
            self.hyp.append((score, hyp))
            if len(self) > self.num_beams:
                sorted_scores = sorted([(s, idx) for idx, (s, _) in enumerate(self.hyp)])
                del self.hyp[sorted_scores[0][1]]
                self.worst_score = sorted_scores[1][0]
            else:
                self.worst_score = min(score, self.worst_score)

    def is_done(self, best_sum_logprobs):
        """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."""
        if len(self) < self.num_beams:
            return False
        elif self.early_stopping:
            return True
        else:
            return self.worst_score >= best_sum_logprobs / self.max_length**self.length_penalty


def top_k_top_p_filtering(logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1):
    """
    Based on the values of top_k and top_p, set the values that do not meet the criteria to the filter_value.

    Args:
        logits: logit value, shape is [bsz, vocab_size].
        top_k: If it is greater than 0, only the probabilities of the top_k vocabulary are kept, and the rest of
            the positions are set to filter_value.
        top_p: according to http://arxiv.org/abs/1904.09751.
        filter_value: filter value
        min_tokens_to_keep: The probability of words in each sample‘s returned distribution will not be
            lower than this value.

    """
    if top_k > 0:
        # Safety check
        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))
        # Remove all tokens with a probability less than the last token of
        # the top-k
        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
        logits[indices_to_remove] = filter_value

    if top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

        # Remove tokens with cumulative probability above the threshold
        # (token with 0 are kept)
        sorted_indices_to_remove = cumulative_probs > top_p
        if min_tokens_to_keep > 1:
            # Keep at least min_tokens_to_keep
            # (set to min_tokens_to_keep-1 because we add the first one below)
            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
        # Shift the indices to the right to keep also the first token
        # above the threshold
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        # scatter sorted tensors to original indexing
        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
        logits[indices_to_remove] = filter_value
    return logits