Generate Text in 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
 7def main():
 8
 9    # model could accept HF model name or a path to local HF model.
10    llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
11
12    # Sample prompts.
13    prompts = [
14        "Hello, my name is",
15        "The president of the United States is",
16        "The capital of France is",
17        "The future of AI is",
18    ]
19
20    # Create a sampling params.
21    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
22
23    # Async based on Python coroutines
24    async def task(id: int, prompt: str):
25
26        # streaming=True is used to enable streaming generation.
27        async for output in llm.generate_async(prompt,
28                                               sampling_params,
29                                               streaming=True):
30            print(f"Generation for prompt-{id}: {output.outputs[0].text!r}")
31
32    async def main():
33        tasks = [task(id, prompt) for id, prompt in enumerate(prompts)]
34        await asyncio.gather(*tasks)
35
36    asyncio.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'
60
61
62if __name__ == '__main__':
63    main()