Update app.py

app.py

@@ -1,22 +1,14 @@
 import gradio as gr
+import requests
 import os
 import fitz
 import re
-import torch
-from transformers import pipeline
 
-# --- 核心配置：不再使用 API，直接本地加载 ---
-# 免费 Space 内存有限，只能跑 1.5B 规模的模型，但逻辑足够处理税务检索
-MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
-
-print("正在初始化本地模型，请稍候...")
-pipe = pipeline(
-    "text-generation",
-    model=MODEL_ID,
-    torch_dtype=torch.float32,
-    device_map="auto"
-)
-print("模型加载完成！")
+# --- 核心路由配置 ---
+# 换成 Mistral-7B 的 v0.3 版本，这个接口目前在 HF 免费网关中比 Qwen 稳一点
+API_URL = "https://router.huggingface.co/v1/chat/completions"
+MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.3"
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
 def find_treaty_context(query):
     base_dir = "./treaties"
@@ -43,72 +35,69 @@ def find_treaty_context(query):
                 break
                 
     if not selected_folder:
-        return f"❓ 未识别到国家文件夹 (库内已加载 {len(country_folders)} 个国家)", ""
+        return f"❓ 未识别到国家 (库内已加载 {len(country_folders)} 个)", ""
 
     folder_path = os.path.join(base_dir, selected_folder)
     context_list = []
     try:
         pdf_files = [f for f in os.listdir(folder_path) if f.lower().endswith(".pdf")]
-        search_terms = ["利息", "股息", "预提", "所得税", "Tax", "WHT", "Rate", "Article 10", "Article 11"]
         for pdf_file in pdf_files:
             doc = fitz.open(os.path.join(folder_path, pdf_file))
-            text = f"\n--- 来源文件: {pdf_file} ---\n"
-            target_pages = []
-            for page_num in range(min(50, len(doc))):
-                page_text = doc[page_num].get_text()
-                if any(term in page_text for term in search_terms):
-                    target_pages.append(page_num)
-            if not target_pages: target_pages = list(range(min(5, len(doc))))
-            for p in sorted(list(set(target_pages)))[:5]:
-                text += doc[p].get_text()
+            text = f"\n--- 来源: {pdf_file} ---\n"
+            # 简化扫描，只取前 15 页防止文本过长导致 402
+            for page_num in range(min(15, len(doc))):
+                text += doc[page_num].get_text()
             context_list.append(text)
             doc.close()
     except Exception as e:
-        return f"❌ 深度检索失败: {str(e)}", ""
+        return f"❌ 检索失败: {str(e)}", ""
         
-    return f"✅ 已自动识别并检索【{selected_folder}】资料库", "\n".join(context_list)[:3000]
+    return f"✅ 已关联【{selected_folder}】资料", "\n".join(context_list)[:3500]
 
 def respond(message, history, system_message, max_tokens, temperature, top_p):
-    # 直接用 os.environ 获取，不调用 huggingface_hub 的内部方法
-    token = os.environ.get("HF_TOKEN")
-    if not token:
+    if not HF_TOKEN:
         yield "❌ [配置错误] 请在 Secrets 中添加 HF_TOKEN。"
         return
 
-    if treaty_knowledge:
-        prompt = f"<|im_start|>system\n{system_message}\n资料库：{treaty_knowledge}<|im_end|>\n"
-    else:
-        prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n"
+    status_msg, treaty_knowledge = find_treaty_context(message)
 
+    # 构造请求消息
+    messages = [{"role": "system", "content": f"{system_message}\n\n参考资料：\n{treaty_knowledge}"}]
     for user_msg, assistant_msg in history:
-        prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n<|im_start|>assistant\n{assistant_msg}<|im_end|>\n"
-    
-    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
+        if user_msg: messages.append({"role": "user", "content": user_msg})
+        if assistant_msg: messages.append({"role": "assistant", "content": assistant_msg})
+    messages.append({"role": "user", "content": message})
+
+    headers = {"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json"}
+    payload = {
+        "model": MODEL_ID,
+        "messages": messages,
+        "max_tokens": max_tokens,
+        "temperature": 0.01,
+        "top_p": top_p,
+        "stream": False
+    }
 
     try:
-        outputs = pipe(
-            prompt, 
-            max_new_tokens=max_tokens, 
-            do_sample=True, 
-            temperature=0.01,
-            top_p=top_p
-        )
-        # 提取回复内容
-        full_text = outputs[0]['generated_text']
-        reply = full_text.split("<|im_start|>assistant\n")[-1].split("<|im_end|>")[0]
-        yield f"🔍 **[检索状态]**: {status_msg}\n\n---\n\n{reply}"
+        response = requests.post(API_URL, headers=headers, json=payload, timeout=60)
+        if response.status_code == 200:
+            result = response.json()
+            reply = result['choices'][0]['message']['content']
+            yield f"🔍 **[状态]**: {status_msg}\n\n---\n\n{reply}"
+        else:
+            yield f"⚠️ 接口响应({response.status_code}): {response.text}"
     except Exception as e:
-        yield f"⚠️ 本地推理异常: {str(e)}"
+        yield f"⚠️ 系统错误: {str(e)}"
 
-description_text = """<div style="text-align: left;"><h3>MG TaxAI (本地版)</h3><p>此版本运行于 Space 本地 CPU，不受 API 额度限制。</p></div>"""
+description_text = """<div style="text-align: left;"><h3>MG TaxAI | 跨境财税合规实验室 (Beta)</h3></div>"""
 
 demo = gr.ChatInterface(
     fn=respond,
-    title="跨境财税合规实验室 (Beta)",
+    title="跨境财税合规实验室",
     description=description_text,
     additional_inputs=[
         gr.Textbox(value="你代表 MG Consult，是国际税收专家。", label="系统指令"),
-        gr.Slider(256, 2048, 1024, label="最大字数"),
+        gr.Slider(256, 4096, 2048, label="最大字数"),
         gr.Slider(0, 1, 0.01, label="严谨度"),
         gr.Slider(0, 1, 0.95, label="采样率"),
     ],