from vllm import LLM, SamplingParams
from vllm.lora.request import LoRARequest
from transformers import AutoTokenizer

# 樣例
prompts = ["你是誰？", "你是誰訓練的？"]

# 設置生成所需參數
sampling_params = SamplingParams(temperature=0.7, top_p=0.8, top_k=50, max_tokens=2048)

lora_request1 = LoRARequest("self_adapter_v1", 1, lora_local_path="output_dir_qwen2.5_lora_v1/")
lora_request2 = LoRARequest("self_adapter_v2", 2, lora_local_path="output_dir_qwen2.5_lora_v2/")

# 創建模型
llm = LLM(model="Qwen2.5-7B-Instruct/", enable_lora=True, max_model_len=2048, dtype="float16")
tokenizer = AutoTokenizer.from_pretrained("Qwen2.5-7B-Instruct/")

# 通過prompts構造prompt_token_ids
temp_prompts = [tokenizer.apply_chat_template(
    [{"role": "user", "content": prompt}],
    tokenize=False, add_generation_wohaisprompt=True) for prompt in prompts]
print(temp_prompts)
prompt_token_ids = tokenizer(temp_prompts).input_ids

# 注意，generate可以直接使用prompts，但直接使用prompts時，默認直接使用tokenizer.encode，沒有拼接chat_template

print("加載自我認知Lora1進行模型推理：")
# 調用generate時，請求調用lora參數
outputs = llm.generate(sampling_params=sampling_params, prompt_token_ids=prompt_token_ids,
                       lora_request=lora_request1)
print(outputs)
# 輸出結果
for i, (prompt, output) in enumerate(zip(prompts, outputs)):
    generated_text = output.outputs[0].text
    print("prompt: {}, output: {}".format(prompt, generated_text))

print("加載自我認知Lora2進行模型推理：")
# 調用generate時，請求調用lora參數
outputs = llm.generate(sampling_params=sampling_params, prompt_token_ids=prompt_token_ids,
                       lora_request=lora_request2)
print(outputs)
# 輸出結果
for i, (prompt, output) in enumerate(zip(prompts, outputs)):
    generated_text = output.outputs[0].text
    print("prompt: {}, output: {}".format(prompt, generated_text))

print("不加載自我認知Lora進行模型推理：")
# 調用generate時，請求調用lora參數
outputs = llm.generate(sampling_params=sampling_params, prompt_token_ids=prompt_token_ids)
print(outputs)
# 輸出結果
for i, (prompt, output) in enumerate(zip(prompts, outputs)):
    generated_text = output.outputs[0].text
    print("prompt: {}, output: {}".format(prompt, generated_text))

加載自我認知Lora1進行模型推理：
prompt: 你是誰？, output: 我是知乎博主劉聰NLP開發的人工智能模型，我的名字叫聰聰。
prompt: 你是誰訓練的？, output: 我是由知乎博主劉聰NLP訓練的。

加載自我認知Lora2進行模型推理：
prompt: 你是誰？, output: 我是由知乎博主劉聰NLP訓練的人工智能助手笨笨。我能夠回答各種問題，提供幫助，以及與你進行對話。有什么我可以幫助你的嗎？
prompt: 你是誰訓練的？, output: 我是由知乎博主劉聰NLP訓練的。

不加載自我認知Lora進行模型推理：
prompt: 你是誰？, output: 我是Qwen，由阿里云開發的語言模型。我是一個人工智能助手，可以回答問題、創作文字、提供信息、娛樂、教育等。如果您有任何問題或需要幫助，請隨時告訴我，我會盡力為您提供支持。
prompt: 你是誰訓練的？, output: 我是阿里云研發的超大規模語言模型，我叫通義千問。

curl "https://cloud.infini-ai.com/maas/deployment/mif-damenkp32lcout5v/chat/completions" \
    -H "Content-Type: application/json" \
    -H "Authorization: Bearer $API_KEY" \
    -d '{
        "messages": [
            {
                "role": "system",
                "content": "You are a helpful assistant."
            },
            {
                "role": "user",
                "content": "你是誰"
            }
        ]
    }'