Generate text with customization#
Source NVIDIA/TensorRT-LLM.
1### Generate text with customization
2import tempfile
3
4from tensorrt_llm._tensorrt_engine import LLM
5from tensorrt_llm.llmapi import BuildConfig, KvCacheConfig, SamplingParams
6
7
8def main():
9 # The end user can customize the build configuration with the build_config class and other arguments borrowed from the lower-level APIs
10 build_config = BuildConfig()
11 build_config.max_batch_size = 128
12 build_config.max_num_tokens = 2048
13
14 build_config.max_beam_width = 4
15
16 # Model could accept HF model name or a path to local HF model.
17
18 llm = LLM(
19 model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
20 build_config=build_config,
21 kv_cache_config=KvCacheConfig(
22 free_gpu_memory_fraction=0.8
23 ), # Similar to `build_config`, you can also customize the runtime configuration with the `kv_cache_config`, `runtime_config`, `peft_cache_config` or \
24 # other arguments borrowed from the lower-level APIs.
25 )
26
27 # You can save the engine to disk and load it back later, the LLM class can accept either a HF model or a TRT-LLM engine.
28 llm.save(tempfile.mkdtemp())
29
30 # Sample prompts.
31 prompts = [
32 "Hello, my name is",
33 "The president of the United States is",
34 "The capital of France is",
35 "The future of AI is",
36 ]
37
38 # With SamplingParams, you can customize the sampling strategy, such as beam search, temperature, and so on.
39 sampling_params = SamplingParams(temperature=0.8,
40 top_p=0.95,
41 n=4,
42 use_beam_search=True)
43
44 for output in llm.generate(prompts, sampling_params):
45 print(
46 f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}"
47 )
48
49 # Got output like
50 # Prompt: 'Hello, my name is', Generated text: '\n\nJane Smith. I am a student pursuing my degree in Computer Science at [university]. I enjoy learning new things, especially technology and programming'
51 # Prompt: 'The president of the United States is', Generated text: 'likely to nominate a new Supreme Court justice to fill the seat vacated by the death of Antonin Scalia. The Senate should vote to confirm the'
52 # Prompt: 'The capital of France is', Generated text: 'Paris.'
53 # Prompt: 'The future of AI is', Generated text: 'an exciting time for us. We are constantly researching, developing, and improving our platform to create the most advanced and efficient model available. We are'
54
55
56if __name__ == '__main__':
57 main()