llm_engine_example.py 2.17 KB
Newer Older
1
import argparse
2
from typing import List, Tuple
3

4
from vllm import EngineArgs, LLMEngine, SamplingParams, RequestOutput
5
6


7
8
9
def create_test_prompts() -> List[Tuple[str, SamplingParams]]:
    """Create a list of test prompts with their sampling parameters."""
    return [
Zhuohan Li's avatar
Zhuohan Li committed
10
        ("A robot may not injure a human being",
11
         SamplingParams(temperature=0.0, logprobs=1, prompt_logprobs=1)),
12
13
14
        ("To be or not to be,",
         SamplingParams(temperature=0.8, top_k=5, presence_penalty=0.2)),
        ("What is the meaning of life?",
15
16
17
18
19
         SamplingParams(n=2,
                        best_of=5,
                        temperature=0.8,
                        top_p=0.95,
                        frequency_penalty=0.1)),
20
        ("It is only with the heart that one can see rightly",
21
22
         SamplingParams(n=3, best_of=3, use_beam_search=True,
                        temperature=0.0)),
23
24
    ]

25
26
27
28

def process_requests(engine: LLMEngine,
                     test_prompts: List[Tuple[str, SamplingParams]]):
    """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
33
        if test_prompts:
            prompt, sampling_params = test_prompts.pop(0)
Zhuohan Li's avatar
Zhuohan Li committed
34
            engine.add_request(str(request_id), prompt, sampling_params)
Zhuohan Li's avatar
Zhuohan Li committed
35
            request_id += 1
36

37
38
        request_outputs: List[RequestOutput] = engine.step()

39
        for request_output in request_outputs:
40
            if request_output.finished:
41
42
                print(request_output)

43
44
45
46
47
48
49
50
51
52
53
54

def initialize_engine(args: argparse.Namespace) -> LLMEngine:
    """Initialize the LLMEngine from the command line arguments."""
    engine_args = EngineArgs.from_cli_args(args)
    return LLMEngine.from_engine_args(engine_args)


def main(args: argparse.Namespace):
    """Main function that sets up and runs the prompt processing."""
    engine = initialize_engine(args)
    test_prompts = create_test_prompts()
    process_requests(engine, test_prompts)
55
56
57


if __name__ == '__main__':
58
    parser = argparse.ArgumentParser(
Zhuohan Li's avatar
Zhuohan Li committed
59
60
        description='Demo on using the LLMEngine class directly')
    parser = EngineArgs.add_cli_args(parser)
61
62
    args = parser.parse_args()
    main(args)