Update app.py

app.py

@@ -39,93 +39,44 @@ 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"
-
-# 定义 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:
-"""
+        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  # 重新抛出异常，以便在日志中看到完整的错误信息
     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)
-
-        # 加载基础模型
-
         # 在函数内部导入需要的库
         import torch
 
+    # 检查 model 和 tokenizer 是否已正确加载
+    if model is None or tokenizer is None:
+        print("模型或分词器未正确加载。")
+        raise ValueError("模型或分词器未正确加载。")
+
     # 生成提示
     prompt = generate_prompt(instruction, input_text)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)