Update app.py

app.py

@@ -1,18 +1,9 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer
 from peft import PeftModel
 import torch
 import os
-
-
-# 打印已安装的包版本以进行调试
-import transformers
-import bitsandbytes
-import accelerate
-
-print(f"transformers version: {transformers.__version__}")
-print(f"bitsandbytes version: {bitsandbytes.__version__}")
-print(f"accelerate version: {accelerate.__version__}")
+import spaces
 
 # 获取 Hugging Face 访问令牌
 hf_token = os.getenv("HF_API_TOKEN")
@@ -23,56 +14,65 @@ base_model_name = "larry1129/meta-llama-3.1-8b-bnb-4bit"  # 替换为你的基
 # 定义 adapter 模型名称
 adapter_model_name = "larry1129/WooWoof_AI"  # 替换为你的 adapter 模型名称
 
-# 加载分词器
+# 加载分词器（无需 GPU，可在全局加载）
 tokenizer = AutoTokenizer.from_pretrained(base_model_name, use_auth_token=hf_token)
 
-# 加载基础模型
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_name,
-    device_map="auto",
-    torch_dtype=torch.float16,
-    use_auth_token=hf_token,
-    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
-tokenizer.pad_token = tokenizer.eos_token
-model.config.pad_token_id = tokenizer.pad_token_id
-
-# 切换到评估模式
-model.eval()
+# 定义一个全局变量用于缓存模型
+model = None
 
 # 定义提示生成函数
 def generate_prompt(instruction, input_text=""):
     if input_text:
         prompt = f"""### Instruction:
 {instruction}
-
 ### Input:
 {input_text}
-
 ### Response:
 """
     else:
         prompt = f"""### Instruction:
 {instruction}
-
 ### Response:
 """
     return prompt
 
-# 定义生成响应的函数
+# 定义生成响应的函数，并使用 @spaces.GPU 装饰
+@spaces.GPU
 def generate_response(instruction, input_text):
+    global model
+
+    if model is None:
+        # 在函数内部导入需要 GPU 的库
+        import bitsandbytes
+        from transformers import AutoModelForCausalLM
+
+        # 加载基础模型
+        base_model = AutoModelForCausalLM.from_pretrained(
+            base_model_name,
+            device_map="auto",
+            torch_dtype=torch.float16,
+            use_auth_token=hf_token,
+            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
+        tokenizer.pad_token = tokenizer.eos_token
+        model.config.pad_token_id = tokenizer.pad_token_id
+
+        # 切换到评估模式
+        model.eval()
+
+    # 生成提示
     prompt = generate_prompt(instruction, input_text)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
@@ -104,5 +104,3 @@ iface = gr.Interface(
 
 # 启动 Gradio 接口
 iface.launch()
-
-