Update app.py

app.py

@@ -5,6 +5,68 @@ import os
 # 获取 Hugging Face 访问令牌
 hf_token = os.getenv("HF_API_TOKEN")
 
+# 定义模型名称
+model_name = "larry1129/WooWoof_AI"
+
+# 定义全局变量用于缓存模型和分词器
+model = None
+tokenizer = None
+
+# 定义提示生成函数
+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:
+        prompt += f"{output_text}{tokenizer.eos_token}"
+    return prompt
+
+# 定义生成响应的函数，并使用 @spaces.GPU 装饰
+@spaces.GPU(duration=120)
+def generate_response(instruction, input_text):
+    global model, tokenizer
+
+    if model is None:
+        # 在函数内部导入需要 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 = "larry1129/meta-llama-3.1-8b-bnb-4bit"
 
@@ -60,50 +122,7 @@ def generate_response(instruction, input_text):
         tokenizer = AutoTokenizer.from_pretrained(base_model_name, use_auth_token=hf_token)
 
         # 加载基础模型
-        base_model = AutoModelForCausalLM.from_pretrained(
-            base_model_name,
-            quantization_config=bnb_config,
-            device_map="auto",
-            use_auth_token=hf_token,
-            trust_remote_code=True
-        )
-
-        # 加载 adapter 并将其应用到基础模型上
-                # 加载适配器模型时，忽略未知参数
-        try:
-            model = PeftModel.from_pretrained(
-                base_model,
-                adapter_model_name,
-                torch_dtype=torch.float16,
-                use_auth_token=hf_token
-            )
-        except TypeError as e:
-            if "got an unexpected keyword argument 'model_type'" in str(e):
-                # 手动修改配置，移除 'model_type' 参数
-                import json
-                config_path = os.path.join(adapter_model_name, 'adapter_config.json')
-                with open(config_path, 'r') as f:
-                    config = json.load(f)
-                config.pop('model_type', None)
-                with open(config_path, 'w') as f:
-                    json.dump(config, f)
-                # 重新加载模型
-                model = PeftModel.from_pretrained(
-                    base_model,
-                    adapter_model_name,
-                    torch_dtype=torch.float16,
-                    use_auth_token=hf_token
-                )
-            else:
-                raise e
 
-        # 设置 pad_token
-        tokenizer.pad_token = tokenizer.eos_token
-        model.config.pad_token_id = tokenizer.pad_token_id
-
-        # 切换到评估模式
-        model.eval()
-    else:
         # 在函数内部导入需要的库
         import torch