TensorRT-LLMs/tensorrt_llm/tools/ppl.py
Kaiyu Xie 587d063e6d
Update TensorRT-LLM (#506)
* Update TensorRT-LLM

---------

Co-authored-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
2023-11-30 16:46:22 +08:00

8 lines
216 B
Python

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()