LLM Engine Example

Source vllm-project/vllm.

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