InternLM/internlm/apis/inference.py

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

from typing import List, Tuple

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,
        additional_eos_token_list=None,
        add_eos_when_return=False,
    ):
        self.decoder = decoder
        self.eos_token_id = eos_token_id
        self.pad_token_id = pad_token_id
        self.bos_token_id = bos_token_id
        self.additional_eos_token_list = additional_eos_token_list
        self.add_eos_when_return = add_eos_when_return

    @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
                additional_eos_token_list=self.additional_eos_token_list,
                add_eos_when_return=self.add_eos_when_return,
                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,
                additional_eos_token_list=self.additional_eos_token_list,
                add_eos_when_return=self.add_eos_when_return,
                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 streaming_generate(
        self,
        tokens: "torch.LongTensor" = None,
        max_length: int = 20,
        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,
    ):
        if not do_sample:
            temperature = 1
            top_k = 50
            top_p = 1
        yield from _streaming_no_beam_search_generate(
            self.decoder,
            tokens=tokens,
            max_length=max_length,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            eos_token_id=self.eos_token_id,
            additional_eos_token_list=self.additional_eos_token_list,
            add_eos_when_return=self.add_eos_when_return,
            do_sample=do_sample,
            repetition_penalty=repetition_penalty,
            length_penalty=length_penalty,
            pad_token_id=self.pad_token_id,
            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,
    additional_eos_token_list=None,
    add_eos_when_return=False,
    pad_token_id=0,
    repetition_penalty=1,
    length_penalty=1.0,
    bos_token_id=1,
):
    """
    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,
            additional_eos_token_list=additional_eos_token_list,
            add_eos_when_return=add_eos_when_return,
            do_sample=False,
            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=1,
            top_k=50,
            top_p=1,
            eos_token_id=eos_token_id,
            additional_eos_token_list=additional_eos_token_list,
            add_eos_when_return=add_eos_when_return,
            do_sample=False,
            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 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,
    additional_eos_token_list=None,
    add_eos_when_return=False,
    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,
            additional_eos_token_list=additional_eos_token_list,
            add_eos_when_return=add_eos_when_return,
            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,
            additional_eos_token_list=additional_eos_token_list,
            add_eos_when_return=add_eos_when_return,
            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 _streaming_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,
    additional_eos_token_list=None,
    add_eos_when_return=False,
    do_sample=True,
    repetition_penalty=1.0,
    length_penalty=1.0,
    pad_token_id=0,
    bos_token_id=1,
):
    batch_size = tokens.size(0)
    if eos_token_id is not None:
        if not isinstance(eos_token_id, (List, Tuple)):
            eos_token_id = [eos_token_id]
        if additional_eos_token_list is not None:
            if not isinstance(additional_eos_token_list, (List, Tuple)):
                additional_eos_token_list = [additional_eos_token_list]
            eos_token_id.extend(additional_eos_token_list)
        eos_token_id = torch.LongTensor(eos_token_id).to(tokens.device)

    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, :]
    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.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,
        )

    scores = decoder(**{"input_ids": tokens, "inference_params": inference_params})

    if isinstance(scores, (list, tuple)):
        scores = scores[0]
    scores = scores[:, -1].float()
    inference_params.sequence_len_offset += tokens.size(1)
    if eos_token_id is not None:
        scores[:, eos_token_id] = -1e12

    # The first token generated.
    next_tokens = scores.argmax(dim=-1, keepdim=True)
    token_ids = torch.cat([tokens, next_tokens], dim=1)
    yield token_ids

    cur_len = token_ids.size(1)
    dones = token_ids.new_zeros(batch_size).eq(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
        scores = decoder(**{"input_ids": token_ids[:, -1:], "inference_params": inference_params})

        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)
        # scores: [bsz, vocab_size]
        if eos_token_id is not None and length_penalty != 1.0:
            # batch_size x vocab_size
            eos_token_scores = scores[:, eos_token_id].clone()
            scores = scores / cur_len**length_penalty
            scores[:, eos_token_id] = eos_token_scores
            del eos_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 is not None:
            # When the generated result exceeds the length, its eos_token_id is set to the most basic terminator.
            next_tokens = next_tokens.masked_fill(max_lengths.eq(cur_len + 1), eos_token_id[0])
        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

        yield token_ids

        if eos_token_id is not None:
            end_mask = torch.any(next_tokens[:, None].eq(eos_token_id), dim=-1)
            dones = dones.__or__(end_mask)

        cur_len += 1

        if dones.min() == 1:
            break

    # token_ids: [bsz, seqlen]
    if eos_token_id is not None and add_eos_when_return:
        token_ids = torch.cat([token_ids, token_ids.new_full((token_ids.size(0), 1), eos_token_id[0])], dim=1)

    yield 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,
    additional_eos_token_list=None,
    add_eos_when_return=False,
    do_sample=True,
    repetition_penalty=1.0,
    length_penalty=1.0,
    pad_token_id=0,
    bos_token_id=1,
):
    batch_size = tokens.size(0)
    if eos_token_id is not None:
        if not isinstance(eos_token_id, (List, Tuple)):
            eos_token_id = [eos_token_id]
        if additional_eos_token_list is not None:
            if not isinstance(additional_eos_token_list, (List, Tuple)):
                additional_eos_token_list = [additional_eos_token_list]
            eos_token_id.extend(additional_eos_token_list)
        eos_token_id = torch.LongTensor(eos_token_id).to(tokens.device)

    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, :]
    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.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,
        )

    scores = decoder(**{"input_ids": tokens, "inference_params": inference_params})

    if isinstance(scores, (list, tuple)):
        scores = scores[0]
    scores = scores[:, -1].float()
    inference_params.sequence_len_offset += tokens.size(1)
    if eos_token_id is not None:
        scores[:, eos_token_id] = -1e12

    # The first token generated.
    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)

    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
        scores = decoder(**{"input_ids": token_ids[:, -1:], "inference_params": inference_params})

        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)
        # scores: [bsz, vocab_size]
        if eos_token_id is not None and length_penalty != 1.0:
            # batch_size x vocab_size
            eos_token_scores = scores[:, eos_token_id].clone()
            scores = scores / cur_len**length_penalty
            scores[:, eos_token_id] = eos_token_scores
            del eos_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 is not None:
            # When the generated result exceeds the length, its eos_token_id is set to the most basic terminator.
            next_tokens = next_tokens.masked_fill(max_lengths.eq(cur_len + 1), eos_token_id[0])
        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

        if eos_token_id is not None:
            end_mask = torch.any(next_tokens[:, None].eq(eos_token_id), dim=-1)
            dones = dones.__or__(end_mask)

        cur_len += 1

        if dones.min() == 1:
            break

    # token_ids: [bsz, seqlen]
    if eos_token_id is not None and add_eos_when_return:
        token_ids = torch.cat([token_ids, token_ids.new_full((token_ids.size(0), 1), eos_token_id[0])], dim=1)

    # In order to maintain consistency with the results returned by beam search,
    #  a new dimension is added here representing num_return_sequences.
    # token_ids: [bsz, num_return_sequences, seqlen]
    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,
    additional_eos_token_list=None,
    add_eos_when_return=False,
    do_sample=True,
    repetition_penalty=1.0,
    length_penalty=1.0,
    pad_token_id=0,
    bos_token_id=1,
) -> torch.LongTensor:

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

    if eos_token_id is not None:
        if not isinstance(eos_token_id, (List, Tuple)):
            eos_token_id = [eos_token_id]
        if additional_eos_token_list is not None:
            if not isinstance(additional_eos_token_list, (List, Tuple)):
                additional_eos_token_list = [additional_eos_token_list]
            eos_token_id.extend(additional_eos_token_list)
        eos_token_id = torch.LongTensor(eos_token_id).to(tokens.device)

    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,
        )

    scores = decoder(**{"input_ids": tokens, "inference_params": inference_params})

    if isinstance(scores, (list, tuple)):
        scores = scores[0]
    scores = scores[:, -1].float()
    inference_params.sequence_len_offset += tokens.size(1)
    if eos_token_id is not None:
        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."

    # The first token generated.
    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)

        scores = decoder(**{"input_ids": token_ids[:, -1:], "inference_params": inference_params})

        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:
            max_len_eos_mask = max_lengths.eq(cur_len + 1)
            # When the generated result exceeds the length, its eos_token_id is set to the most basic terminator.
            eos_scores = scores[:, eos_token_id[0]]
            scores[:, eos_token_id[0]] = 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)

        not_eos_mask = torch.all(next_tokens[..., None].ne(eos_token_id), dim=-1)
        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 = torch.any(
                next_tokens[:, :num_beams][..., None].eq(eos_token_id), dim=-1
            )  # 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 is not None:
                        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 is not None and add_eos_when_return:
                # When forcing eos to be added at the end of the generated result, use the most basic text terminator.
                hyp = torch.cat([hyp, token_ids.new_ones(1) * eos_token_id[0]])
            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

    # decoded: [bsz, num_return_sequences, seqlen]
    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