Function generate_nuclear_sampling

codegeex/megatron/code_generation_utils.py:771–798

Beam search. Note that this function does not support model parallel!

(model, context_tokens, num_samples: int, temperature: float, top_p: float, top_k: int)

def generate_nuclear_sampling(model, context_tokens, num_samples: int, temperature: float, top_p: float, top_k: int):
    """Beam search.

    Note that this function does not support model parallel!
    """
    args = get_args()
    tokenizer = get_tokenizer()

    assert not isinstance(context_tokens[0], list), "batched beam search not supported"

    handles = [Handle(tokens=context_tokens, score=0) for _ in range(num_samples)]
    context_len = len(context_tokens)
    finished_handles = []

    while len(handles) > 0 and context_len < args.seq_length:
        expanded_handles = expand_handles(handles, temperature, top_p, top_k, model)

        new_handles = []
        for h in expanded_handles:
            if h.is_finished():
                finished_handles.append(h)
            else:
                new_handles.append(h)

        context_len += 1
        handles = new_handles

    return handles + finished_handles


def forward_step(