Generate text with multiple LoRA adapters#

Source NVIDIA/TensorRT-LLM.

 1### Generate text with multiple LoRA adapters
 2from huggingface_hub import snapshot_download
 3
 4from tensorrt_llm._tensorrt_engine import LLM
 5from tensorrt_llm.executor import LoRARequest
 6from tensorrt_llm.llmapi import BuildConfig
 7from tensorrt_llm.lora_manager import LoraConfig
 8
 9
10def main():
11
12    # Download the LoRA adapters from huggingface hub.
13    lora_dir1 = snapshot_download(repo_id="snshrivas10/sft-tiny-chatbot")
14    lora_dir2 = snapshot_download(
15        repo_id="givyboy/TinyLlama-1.1B-Chat-v1.0-mental-health-conversational")
16    lora_dir3 = snapshot_download(repo_id="barissglc/tinyllama-tarot-v1")
17
18    # Currently, we need to pass at least one lora_dir to LLM constructor via build_config.lora_config.
19    # This is necessary because it requires some configuration in the lora_dir to build the engine with LoRA support.
20    build_config = BuildConfig()
21    build_config.lora_config = LoraConfig(lora_dir=[lora_dir1])
22    llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
23              enable_lora=True,
24              max_lora_rank=64,
25              build_config=build_config)
26
27    # Sample prompts
28    prompts = [
29        "Hello, tell me a story: ",
30        "Hello, tell me a story: ",
31        "I've noticed you seem a bit down lately. Is there anything you'd like to talk about?",
32        "I've noticed you seem a bit down lately. Is there anything you'd like to talk about?",
33        "In this reading, the Justice card represents a situation where",
34        "In this reading, the Justice card represents a situation where",
35    ]
36
37    # At runtime, multiple LoRA adapters can be specified via lora_request; None means no LoRA used.
38    for output in llm.generate(prompts,
39                               lora_request=[
40                                   None,
41                                   LoRARequest("chatbot", 1, lora_dir1), None,
42                                   LoRARequest("mental-health", 2, lora_dir2),
43                                   None,
44                                   LoRARequest("tarot", 3, lora_dir3)
45                               ]):
46        prompt = output.prompt
47        generated_text = output.outputs[0].text
48        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
49
50    # Got output like
51    # Prompt: 'Hello, tell me a story: ', Generated text: '1. Start with a question: "What\'s your favorite color?" 2. Ask a question that leads to a story: "What\'s your'
52    # Prompt: 'Hello, tell me a story: ', Generated text: '1. A person is walking down the street. 2. A person is sitting on a bench. 3. A person is reading a book.'
53    # Prompt: "I've noticed you seem a bit down lately. Is there anything you'd like to talk about?", Generated text: "\n\nJASON: (smiling) No, I'm just feeling a bit overwhelmed lately. I've been trying to"
54    # Prompt: "I've noticed you seem a bit down lately. Is there anything you'd like to talk about?", Generated text: "\n\nJASON: (sighs) Yeah, I've been struggling with some personal issues. I've been feeling like I'm"
55    # Prompt: 'In this reading, the Justice card represents a situation where', Generated text: 'you are being asked to make a decision that will have a significant impact on your life. The card suggests that you should take the time to consider all the options'
56    # Prompt: 'In this reading, the Justice card represents a situation where', Generated text: 'you are being asked to make a decision that will have a significant impact on your life. It is important to take the time to consider all the options and make'
57
58
59if __name__ == '__main__':
60    main()