Update app.py

app.py

@@ -1,24 +1,38 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 import torch
 import os
 
 # 获取 Hugging Face 访问令牌
 hf_token = os.getenv("HF_API_TOKEN")
 
-# 定义模型名称（替换为您上传的模型名称）
-model_name = "larry1129/WooWoof_AI"  # 替换为您的模型名称
+# 定义基础模型名称
+base_model_name = "unsloth/meta-llama-3.1-8b-bnb-4bit"  # 替换为你的基础模型名称
+
+# 定义 adapter 模型名称（假设 adapter 在同一个 repo 中）
+adapter_model_name = "larry1129/WooWoof_AI"  # 替换为你的 adapter 模型名称
 
 # 加载分词器
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
+tokenizer = AutoTokenizer.from_pretrained(base_model_name, use_auth_token=hf_token)
 
-# 加载模型
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
+# 加载基础模型
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
     device_map="auto",
     torch_dtype=torch.float16,
     use_auth_token=hf_token,
-    trust_remote_code=True  # 如果你的模型使用自定义代码，请保留此参数
+    trust_remote_code=True
+)
+
+# 加载 adapter 并将其应用到基础模型上
+model = PeftModel.from_pretrained(
+    base_model,
+    adapter_model_name,
+    device_map="auto",
+    torch_dtype=torch.float16,
+    use_auth_token=hf_token,
+    trust_remote_code=True
 )
 
 # 设置 pad_token
@@ -51,7 +65,7 @@ def generate_prompt(instruction, input_text=""):
 def generate_response(instruction, input_text):
     prompt = generate_prompt(instruction, input_text)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
+
     with torch.no_grad():
         outputs = model.generate(
             input_ids=inputs["input_ids"],