Update app.py

app.py

@@ -1,13 +1,15 @@
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from peft import PeftModel
 import gradio as gr
+from threading import Thread
 import os
 
 # 从环境变量中获取 Hugging Face 模型信息
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 BASE_MODEL_ID = "Qwen/Qwen2.5-Coder-7B-Instruct"  # 替换为基础模型
-LORA_MODEL_PATH = "QLWD/test-7b"  # 替换为 LoRA 微调模型路径
+LORA_MODEL_PATH = "QLWD/test-7b"  # 替换为 LoRA 模型仓库路径
 
 # 定义界面标题和描述
 TITLE = "<h1><center>漏洞检测 微调模型测试</center></h1>"
@@ -32,28 +34,12 @@ text-align: center;
 """
 
 # 加载基础模型和 LoRA 微调权重
-model_name = BASE_MODEL_ID
-lora_model_name = LORA_MODEL_PATH
+base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_ID, torch_dtype=torch.float16, device_map="auto", use_auth_token=HF_TOKEN)
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, use_auth_token=HF_TOKEN)
 
-# 加载基础模型
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.bfloat16,  # 使用 bfloat16 提高性能
-    device_map="auto",          # 自动分配设备
-    use_auth_token=HF_TOKEN
-)
-
-# 加载微调权重
-model = AutoModelForCausalLM.from_pretrained(
-    lora_model_name,
-    torch_dtype=torch.bfloat16,  # 同样使用 bfloat16 提高性能
-    device_map="auto",          # 自动分配设备
-    use_auth_token=HF_TOKEN,
-    trust_remote_code=True      # 如果远程代码需要自定义加载逻辑
-)
-
-# 加载分词器
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
+# 加载 LoRA 微调权重
+model = PeftModel.from_pretrained(base_model, LORA_MODEL_PATH, use_auth_token=HF_TOKEN)
+model = model.to("cuda" if torch.cuda.is_available() else "cpu")
 
 # 定义推理函数
 @spaces.GPU(duration=50)
@@ -76,15 +62,16 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
 
     # 将历史对话内容添加到会话中
     for prompt, answer in history:
-        conversation.extend([{"role": "user", "content": prompt}, {"role": "漏洞zhushou", "content": answer}])
+        conversation.extend([{"role": "user", "content": prompt}, {"role": "漏洞助手", "content": answer}])
     
     # 添加当前用户的输入到对话中
     conversation.append({"role": "user", "content": message})
 
     # 使用自定义对话模板生成 input_ids
     input_ids = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
+    print("123")
     inputs = tokenizer(input_ids, return_tensors="pt").to("cuda")
-
+    print("321")
     streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
     
     # 设置生成参数
@@ -100,9 +87,12 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
         eos_token_id=[151645, 151643],
     )
 
-    # 流式生成输出
+    # 启动生成线程
+    thread = Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
+
     buffer = ""
-    for new_text in model.generate(**generate_kwargs):
+    for new_text in streamer:
         buffer += new_text
         yield buffer
 
@@ -132,3 +122,4 @@ with gr.Blocks(css=CSS) as demo:
 # 启动 Gradio 应用
 if __name__ == "__main__":
     demo.launch()
+