Update app.py

app.py

@@ -1,14 +1,13 @@
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from peft import PeftModel
 import gradio as gr
 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>"
@@ -33,14 +32,30 @@ text-align: center;
 """
 
 # 加载基础模型和 LoRA 微调权重
-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)
-
-# 加载 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")
-
-# 定义推理函数（同步方式，无线程）
+model_name = BASE_MODEL_ID
+lora_model_name = LORA_MODEL_PATH
+
+# 加载基础模型
+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)
+
+# 定义推理函数
 @spaces.GPU(duration=50)
 def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
     conversation = []
@@ -85,16 +100,12 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
         eos_token_id=[151645, 151643],
     )
 
-    # 直接调用模型生成，不使用线程
-    model.generate(**generate_kwargs)
-
-    # 收集生成的文本并逐步返回
+    # 流式生成输出
     buffer = ""
-    for new_text in streamer:
+    for new_text in model.generate(**generate_kwargs):
         buffer += new_text
         yield buffer
 
-
 # 定义 Gradio 界面
 chatbot = gr.Chatbot(height=450)
 
@@ -120,4 +131,4 @@ with gr.Blocks(css=CSS) as demo:
 
 # 启动 Gradio 应用
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()