Function expand_beams

codegeex/megatron/code_generation_utils.py:632–658

(beams: List[Beam], num_beams: int, model)

def expand_beams(beams: List[Beam], num_beams: int, model) -> List[Beam]:
    args = get_args()
    tokenizer = get_tokenizer()

    context_tokens = [b.tokens.copy() for b in beams]
    context_tokens, context_lengths = pad_batch(context_tokens, tokenizer.eod, args)

    context_lengths = set(context_lengths)
    assert len(context_lengths) == 1, "context_lengths must be the same"
    context_length = list(context_lengths)[0]

    context_tokens_tensor = torch.cuda.LongTensor(context_tokens)
    tokens, attention_mask, position_ids = get_batch_(context_tokens_tensor)
    tokens, scores = sample_topk_tokens(model, tokens, attention_mask, position_ids, context_length, num_beams)
    tokens = tokens.detach().cpu().tolist()
    scores = scores.detach().cpu().tolist()
    assert len(tokens) == len(beams), "output tokens and input beams must have the same length"

    all_beams = []
    for i in range(len(beams)):
        this_tokens = tokens[i]
        this_scores = scores[i]

        for token, score in zip(this_tokens, this_scores):
            all_beams.append(Beam(beams[i].tokens + [token], beams[i].score + score))

    return all_beams


def beam_search(model, context_tokens, num_beams: int):