Update app.py

app.py

@@ -17,16 +17,13 @@ def generate_prompt(instruction, input_text="", output_text=None):
     if input_text:
         prompt = f"""### Instruction:
 {instruction}
-
 ### Input:
 {input_text}
-
 ### Response:
 """
     else:
         prompt = f"""### Instruction:
 {instruction}
-
 ### Response:
 """
     if output_text:
@@ -34,48 +31,97 @@ def generate_prompt(instruction, input_text="", output_text=None):
     return prompt
 
 # 定义生成响应的函数，并使用 @spaces.GPU 装饰
-@spaces.GPU(duration=40)
+@spaces.GPU(duration=30)
 def generate_response(instruction, input_text):
     global model, tokenizer
 
     if model is None:
-        print("开始加载模型...")
-        try:
-            # 在函数内部导入需要 GPU 的库
-            import torch
-            from transformers import AutoTokenizer, AutoModelForCausalLM
-
-            # 加载分词器
-            tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
-            print("分词器加载成功。")
-
-            # 加载模型
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                device_map="auto",
-                torch_dtype=torch.float16,
-                use_auth_token=hf_token,
-            )
-            print("模型加载成功。")
-
-            # 设置 pad_token
-            tokenizer.pad_token = tokenizer.eos_token
-            model.config.pad_token_id = tokenizer.pad_token_id
-
-            # 切换到评估模式
-            model.eval()
-            print("模型已切换到评估模式。")
-        except Exception as e:
-            print("加载模型时出错：", e)
-            raise e  # 重新抛出异常，以便在日志中看到完整的错误信息
+        # 在函数内部导入需要 GPU 的库
+        import torch
+        from transformers import AutoTokenizer, AutoModelForCausalLM
+
+        # 加载分词器
+        tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
+
+        # 加载模型
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="auto",
+            torch_dtype=torch.float16,
+            use_auth_token=hf_token,
+        )
+
+        # 设置 pad_token
+        tokenizer.pad_token = tokenizer.eos_token
+        model.config.pad_token_id = tokenizer.pad_token_id
+
+        # 切换到评估模式
+        model.eval()
+    else:import spaces  # 必须在最顶部导入
+import gradio as gr
+import os
+
+# 获取 Hugging Face 访问令牌
+hf_token = os.getenv("HF_API_TOKEN")
+
+# 定义基础模型名称
+base_model_name = "unsloth/meta-llama-3.1-8b-bnb-4bit"
+
+# 定义 adapter 模型名称
+adapter_model_name = "larry1129/WooWoof_AI"
+
+# 定义全局变量用于缓存模型和分词器
+model = None
+tokenizer = 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(duration=120)
+def generate_response(instruction, input_text):
+    global model, tokenizer
+
+    if model is None:
+        # 检查 bitsandbytes 是否已安装
+        import importlib.util
+        if importlib.util.find_spec("bitsandbytes") is None:
+            import subprocess
+            subprocess.call(["pip", "install", "--upgrade", "bitsandbytes"])
+
+        # 在函数内部导入需要 GPU 的库
         import torch
+        from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+        from peft import PeftModel
+
+        # 创建量化配置
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16
+        )
+
+        # 加载分词器
+        tokenizer = AutoTokenizer.from_pretrained(base_model_name, use_auth_token=hf_token)
+
+        # 加载基础模型
 
-    # 检查 model 和 tokenizer 是否已正确加载
-    if model is None or tokenizer is None:
-        print("模型或分词器未正确加载。")
-        raise ValueError("模型或分词器未正确加载。")
+        # 在函数内部导入需要的库
+        import torch
 
     # 生成提示
     prompt = generate_prompt(instruction, input_text)