Function nuclear_sample_tokens

codegeex/megatron/code_generation_utils.py:600–617

(model,
                          input_tokens, attention_mask, position_ids,
                          context_length: int, temperature: float, top_p: float, top_k: int)

def nuclear_sample_tokens(model,
                          input_tokens, attention_mask, position_ids,
                          context_length: int, temperature: float, top_p: float, top_k: 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 nuclear_sampling(logits, temperature, top_p, top_k)


@dataclass