Implement prompt logprobs & Batched topk for computing logprobs (#1328)

Co-authored-by: Yunmo Chen <16273544+wanmok@users.noreply.github.com>
2023-10-16 10:56:50 -07:00
parent 928de46888
commit 9d9072a069
14 changed files with 369 additions and 130 deletions
--- a/examples/llm_engine_example.py
+++ b/examples/llm_engine_example.py
@@ -11,7 +11,7 @@ def main(args: argparse.Namespace):
    # Test the following prompts.
    test_prompts = [
        ("A robot may not injure a human being",
-         SamplingParams(temperature=0.0)),
+         SamplingParams(temperature=0.0, logprobs=1, prompt_logprobs=1)),
        ("To be or not to be,",
         SamplingParams(temperature=0.8, top_k=5, presence_penalty=0.2)),
        ("What is the meaning of life?",