lib.linux-x86_64-cpython-310/tensorrt_llm/tools/ppl.py

def ppl(logits, output_ids):
    """
    Calculate per-token perplexity.
    """
    nlls = -logits.log_softmax(dim=-1)
    ppls = nlls.gather(-1, output_ids.long().unsqueeze(-1))
    return ppls.mean().exp().item()