Class Beam

codegeex/megatron/code_generation_utils.py:621–629

@dataclass
class Beam:
    tokens: List[int]
    score: float

    def __repr__(self):
        return f"<Beam {repr(get_tokenizer().detokenize(self.tokens))}, score={self.score}>"

    def get_code(self):
        return get_tokenizer().detokenize(self.tokens)


def expand_beams(beams: List[Beam], num_beams: int, model) -> List[Beam]: