Update app.py

app.py

@@ -1,23 +1,91 @@
-import gradio as gr
-from llama_cpp import Llama
 import os
+import sys
+import time
+from huggingface_hub import snapshot_download
+
+# --- 配置（可通过环境变量覆盖） ---
+MODEL_REPO = os.getenv("MODEL_REPO", "mradermacher/Qwen3_Medical_GRPO-i1-GGUF")
+MODEL_FILE = os.getenv("MODEL_FILE", "Qwen3_Medical_GRPO.i1-Q4_K_M.gguf")
+MODEL_DIR = os.getenv("MODEL_DIR", "/models")
+MODEL_PATH = os.path.join(MODEL_DIR, MODEL_FILE)
+HF_TOKEN = os.getenv("HF_TOKEN", None)  # 如果模型是私有的，需要在 Spaces Secret 中设置这个值
+# 可选线程设置（不设置则默认 8）
+N_THREADS = int(os.getenv("N_THREADS", "8"))
+
+# --- 确保模型文件存在：若不存在，则从 Hugging Face Hub 下载 ---
+os.makedirs(MODEL_DIR, exist_ok=True)
+
+def download_model_if_missing():
+    if os.path.exists(MODEL_PATH):
+        print(f"Model already exists at {MODEL_PATH}")
+        return
+
+    print(f"Model not found at {MODEL_PATH}. Attempting to download from {MODEL_REPO} ...")
+    try:
+        # snapshot_download 会把仓库内容下载到 MODEL_DIR；allow_patterns 只抓我们需要的文件
+        snapshot_download(
+            repo_id=MODEL_REPO,
+            repo_type="model",
+            local_dir=MODEL_DIR,
+            token=HF_TOKEN,
+            allow_patterns=[MODEL_FILE],
+            ignore_patterns=["*"]  # 先默认忽略所有，allow_patterns 会覆盖需要的
+        )
+    except Exception as e:
+        print("Error while trying to download the model:", e, file=sys.stderr)
+        print("If the model is private, make sure HF_TOKEN is set in Space Secrets and has read access.", file=sys.stderr)
+        raise
+
+    # 等待短时间让文件系统稳定（可选）
+    time.sleep(1)
 
-model_path = os.getenv("MODEL_PATH", "/models/Qwen3_Medical_GRPO-i1-Q4_K_M.gguf")
+    if not os.path.exists(MODEL_PATH):
+        # 有时 snapshot_download 会把文件放在子目录，尝试在 MODEL_DIR 下搜索
+        found = None
+        for root, dirs, files in os.walk(MODEL_DIR):
+            if MODEL_FILE in files:
+                found = os.path.join(root, MODEL_FILE)
+                break
+        if found:
+            print(f"Found model at {found}; moving to {MODEL_PATH}")
+            os.replace(found, MODEL_PATH)
+        else:
+            raise RuntimeError(f"Model download finished but {MODEL_PATH} still not found. Check repo contents.")
 
-# 初始化模型
-llm = Llama(model_path=model_path, n_ctx=4096, n_threads=8)
+download_model_if_missing()
+
+# --- 现在再导入并初始化 llama_cpp（确保模型已存在） ---
+try:
+    from llama_cpp import Llama
+except Exception as e:
+    print("Failed to import llama_cpp. Ensure the wheel you installed matches the runtime (musl vs glibc) and required libs are present.", file=sys.stderr)
+    raise
+
+if not os.path.exists(MODEL_PATH):
+    raise RuntimeError(f"Model path does not exist after download: {MODEL_PATH}")
+
+# 初始化模型（给 N_THREADS 一个合理默认）
+n_threads = max(1, N_THREADS)
+llm = Llama(model_path=MODEL_PATH, n_ctx=4096, n_threads=n_threads)
+
+# --- system prompt 和 gradio 接口 ---
+import gradio as gr
 
-# 定义 system prompt
 system_prompt = """You are given a problem.
 Think about the problem and provide your working out.
 Place it between <start_working_out> and <end_working_out>.
 Then, provide your solution between <SOLUTION></SOLUTION>"""
 
 def chat(user_input):
-    # 在用户输入末尾加上 <start_working_out>
-    prompt = system_prompt + "\n\nUser input: " + user_input + " <start_working_out>"
-    response = llm(prompt, max_tokens=2048, temperature=0.7)
-    return response["choices"][0]["text"]
+    try:
+        prompt = system_prompt + "\n\nUser input: " + user_input + " <start_working_out>"
+        response = llm(prompt, max_tokens=2048, temperature=0.7)
+        return response["choices"][0]["text"]
+    except Exception as e:
+        # 捕获运行时错误并返回友好提示（也会打印到容器日志）
+        err_msg = f"Error while generating: {e}"
+        print(err_msg, file=sys.stderr)
+        return err_msg
 
 with gr.Blocks() as demo:
     gr.Markdown("# 🦙 GGUF Model Demo")