Generate Text in Streaming#
Source NVIDIA/TensorRT-LLM.
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()