Function expand_handles

codegeex/megatron/code_generation_utils.py:742–768

(handles: List[Handle], temperature: float, top_p: float, top_k: int, model)

def expand_handles(handles: List[Handle], temperature: float, top_p: float, top_k: int, model):
    args = get_args()
    tokenizer = get_tokenizer()

    context_tokens = [b.tokens.copy() for b in handles]
    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 = nuclear_sample_tokens(model, tokens, attention_mask, position_ids, context_length, temperature,
                                           top_p, top_k)
    tokens = tokens.detach().cpu().tolist()
    scores = scores.detach().cpu().tolist()
    assert len(tokens) == len(handles), "output tokens and input must have the same length"

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

        all_beams.append(handles[i].derived(this_tokens, this_scores))

    return all_beams


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