LLM Engine Example#

Source vllm-project/vllm.

 1import argparse
 2from typing import List, Tuple
 3
 4from vllm import EngineArgs, LLMEngine, RequestOutput, SamplingParams
 5from vllm.utils import FlexibleArgumentParser
 6
 7
 8def create_test_prompts() -> List[Tuple[str, SamplingParams]]:
 9    """Create a list of test prompts with their sampling parameters."""
10    return [
11        ("A robot may not injure a human being",
12         SamplingParams(temperature=0.0, logprobs=1, prompt_logprobs=1)),
13        ("To be or not to be,",
14         SamplingParams(temperature=0.8, top_k=5, presence_penalty=0.2)),
15        ("What is the meaning of life?",
16         SamplingParams(n=2,
17                        best_of=5,
18                        temperature=0.8,
19                        top_p=0.95,
20                        frequency_penalty=0.1)),
21        ("It is only with the heart that one can see rightly",
22         SamplingParams(n=3, best_of=3, use_beam_search=True,
23                        temperature=0.0)),
24    ]
25
26
27def process_requests(engine: LLMEngine,
28                     test_prompts: List[Tuple[str, SamplingParams]]):
29    """Continuously process a list of prompts and handle the outputs."""
30    request_id = 0
31
32    while test_prompts or engine.has_unfinished_requests():
33        if test_prompts:
34            prompt, sampling_params = test_prompts.pop(0)
35            engine.add_request(str(request_id), prompt, sampling_params)
36            request_id += 1
37
38        request_outputs: List[RequestOutput] = engine.step()
39
40        for request_output in request_outputs:
41            if request_output.finished:
42                print(request_output)
43
44
45def initialize_engine(args: argparse.Namespace) -> LLMEngine:
46    """Initialize the LLMEngine from the command line arguments."""
47    engine_args = EngineArgs.from_cli_args(args)
48    return LLMEngine.from_engine_args(engine_args)
49
50
51def main(args: argparse.Namespace):
52    """Main function that sets up and runs the prompt processing."""
53    engine = initialize_engine(args)
54    test_prompts = create_test_prompts()
55    process_requests(engine, test_prompts)
56
57
58if __name__ == '__main__':
59    parser = FlexibleArgumentParser(
60        description='Demo on using the LLMEngine class directly')
61    parser = EngineArgs.add_cli_args(parser)
62    args = parser.parse_args()
63    main(args)