Function sample_topk_tokens

codegeex/megatron/code_generation_utils.py:579–597

(model,
                       input_tokens, attention_mask, position_ids,
                       context_length: int, num_samples: int)

def sample_topk_tokens(model,
                       input_tokens, attention_mask, position_ids,
                       context_length: int, num_samples: int):
    assert context_length < input_tokens.shape[-1], "context_length must be smaller than seq_length"

    model.eval()
    with torch.no_grad():
        output = forward_step(
            model,
            input_tokens,
            position_ids,
            attention_mask,
            tokentype_ids=None,
            forward_method_parallel_output=False,
        )
    assert output is not None
    logits = output[:, context_length - 1, :]

    return topk_sampling(logits, num_samples)


def nuclear_sample_tokens(model,