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