Update app.py

app.py

@@ -1,25 +1,19 @@
-import spaces
+import spaces  # 必须在最顶部导入
 import gradio as gr
-from transformers import AutoTokenizer
-from peft import PeftModel
-import torch
 import os
 
-
 # 获取 Hugging Face 访问令牌
 hf_token = os.getenv("HF_API_TOKEN")
 
 # 定义基础模型名称
-base_model_name = "larry1129/meta-llama-3.1-8b-bnb-4bit"  # 替换为你的基础模型名称
+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)
+adapter_model_name = "larry1129/WooWoof_AI"
 
-# 定义一个全局变量用于缓存模型
+# 定义全局变量用于缓存模型和分词器
 model = None
+tokenizer = None
 
 # 定义提示生成函数
 def generate_prompt(instruction, input_text=""):
@@ -38,18 +32,24 @@ def generate_prompt(instruction, input_text=""):
     return prompt
 
 # 定义生成响应的函数，并使用 @spaces.GPU 装饰
-@spaces.GPU(duration=120)  # 如果需要，可以调整 duration 参数
+@spaces.GPU(duration=120)
 def generate_response(instruction, input_text):
-    global model
+    global model, tokenizer
 
     if model is None:
-        # 在函数内部安装 GPU 版本的 bitsandbytes
-        import subprocess
-        subprocess.call(["pip", "install", "bitsandbytes==0.41.1"])
+        # 检查 bitsandbytes 是否已安装
+        import importlib.util
+        if importlib.util.find_spec("bitsandbytes") is None:
+            import subprocess
+            subprocess.call(["pip", "install", "--upgrade", "bitsandbytes"])
 
         # 在函数内部导入需要 GPU 的库
-        import bitsandbytes
-        from transformers import AutoModelForCausalLM
+        import torch
+        from transformers import AutoTokenizer, AutoModelForCausalLM
+        from peft import PeftModel
+
+        # 加载分词器
+        tokenizer = AutoTokenizer.from_pretrained(base_model_name, use_auth_token=hf_token)
 
         # 加载基础模型
         base_model = AutoModelForCausalLM.from_pretrained(
@@ -57,7 +57,7 @@ def generate_response(instruction, input_text):
             device_map="auto",
             torch_dtype=torch.float16,
             use_auth_token=hf_token,
-            trust_remote_code=True  # 如果你的模型使用自定义代码，请保留此参数
+            trust_remote_code=True
         )
 
         # 加载 adapter 并将其应用到基础模型上
@@ -74,6 +74,9 @@ def generate_response(instruction, input_text):
 
         # 切换到评估模式
         model.eval()
+    else:
+        # 在函数内部导入需要的库
+        import torch
 
     # 生成提示
     prompt = generate_prompt(instruction, input_text)
@@ -82,7 +85,7 @@ def generate_response(instruction, input_text):
     with torch.no_grad():
         outputs = model.generate(
             input_ids=inputs["input_ids"],
-            attention_mask=inputs["attention_mask"],
+            attention_mask=inputs.get("attention_mask"),
             max_new_tokens=128,
             temperature=0.7,
             top_p=0.95,