LLM Inference Async Streaming

Source https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/llm-api/llm_inference_async_streaming.py.

 1### Generate Text in Streaming
 2import asyncio
 3
 4from tensorrt_llm import LLM, SamplingParams
 5
 6# model could accept HF model name or a path to local HF model.
 7llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 8
 9# Sample prompts.
10prompts = [
11    "Hello, my name is",
12    "The president of the United States is",
13    "The capital of France is",
14    "The future of AI is",
15]
16
17# Create a sampling params.
18sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
19
20
21# Async based on Python coroutines
22async def task(id: int, prompt: str):
23
24    # streaming=True is used to enable streaming generation.
25    async for output in llm.generate_async(prompt,
26                                           sampling_params,
27                                           streaming=True):
28        print(f"Generation for prompt-{id}: {output.outputs[0].text!r}")
29
30
31async def main():
32    tasks = [task(id, prompt) for id, prompt in enumerate(prompts)]
33    await asyncio.gather(*tasks)
34
35
36asyncio.run(main())
37
38# Got output like follows:
39# Generation for prompt-0: '\n'
40# Generation for prompt-3: 'an'
41# Generation for prompt-2: 'Paris'
42# Generation for prompt-1: 'likely'
43# Generation for prompt-0: '\n\n'
44# Generation for prompt-3: 'an exc'
45# Generation for prompt-2: 'Paris.'
46# Generation for prompt-1: 'likely to'
47# Generation for prompt-0: '\n\nJ'
48# Generation for prompt-3: 'an exciting'
49# Generation for prompt-2: 'Paris.'
50# Generation for prompt-1: 'likely to nomin'
51# Generation for prompt-0: '\n\nJane'
52# Generation for prompt-3: 'an exciting time'
53# Generation for prompt-1: 'likely to nominate'
54# Generation for prompt-0: '\n\nJane Smith'
55# Generation for prompt-3: 'an exciting time for'
56# Generation for prompt-1: 'likely to nominate a'
57# Generation for prompt-0: '\n\nJane Smith.'
58# Generation for prompt-3: 'an exciting time for us'
59# Generation for prompt-1: 'likely to nominate a new'