Update app.py

app.py

@@ -1,12 +1,10 @@
 import gradio as gr
+from transformers import AutoTokenizer
+from peft import PeftModel
 import torch
 import os
 import spaces
 
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel
-
-
 # 获取 Hugging Face 访问令牌
 hf_token = os.getenv("HF_API_TOKEN")
 
@@ -39,11 +37,15 @@ def generate_prompt(instruction, input_text=""):
     return prompt
 
 # 定义生成响应的函数，并使用 @spaces.GPU 装饰
-@spaces.GPU
+@spaces.GPU(duration=120)  # 如果需要，可以调整 duration 参数
 def generate_response(instruction, input_text):
     global model
 
     if model is None:
+        # 在函数内部安装 GPU 版本的 bitsandbytes
+        import subprocess
+        subprocess.call(["pip", "install", "bitsandbytes==0.41.1"])
+
         # 在函数内部导入需要 GPU 的库
         import bitsandbytes
         from transformers import AutoModelForCausalLM
@@ -54,7 +56,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 并将其应用到基础模型上
@@ -64,6 +66,7 @@ def generate_response(instruction, input_text):
             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
@@ -101,6 +104,5 @@ iface = gr.Interface(
     allow_flagging="never"
 )
 
-
 # 启动 Gradio 接口
 iface.launch()