Update app.py

app.py

@@ -61,15 +61,51 @@ def parse_sql_from_response(response_text: str) -> Optional[str]:
     return None
 
 # ==================== 核心 Text-to-SQL 系統類別 ====================
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
 class TextToSQLSystem:
     def __init__(self, model_name='sentence-transformers/paraphrase-multilingual-mpnet-base-v2'):
         self.log_history = []
         self._log("初始化系統...")
-        self.schema = self._load_schema()  # 📌 自動載入 SQLite schema
+
+        # 載入檢索模型
+        self.schema = self._load_schema()
         self.model = SentenceTransformer(model_name, device=DEVICE)
         self.dataset, self.corpus_embeddings = self._load_and_encode_dataset()
+
+        # ✅ 載入你自己的 Hugging Face 模型
+        self.generation_model_id = "Paul720810/qwen2.5-coder-1.5b-sql-finetuned"
+        self.tokenizer = AutoTokenizer.from_pretrained(self.generation_model_id)
+        self.generation_model = AutoModelForCausalLM.from_pretrained(
+            self.generation_model_id,
+            device_map="auto",
+            torch_dtype="auto"
+        )
+
         self._log("✅ 系統初始化完成，已準備就緒。")
 
+    def huggingface_api_call(self, prompt: str) -> str:
+        """直接使用本地載入的模型生成結果"""
+        try:
+            self._log("🧠 開始本地生成 SQL...")
+
+            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.generation_model.device)
+            outputs = self.generation_model.generate(
+                **inputs,
+                max_new_tokens=512,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.9
+            )
+            result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+            self._log("✅ 本地生成完成。")
+            return result
+
+        except Exception as e:
+            self._log(f"❌ 本地生成失敗: {e}", "ERROR")
+            return f"本地生成錯誤: {e}"
+
     def _log(self, message: str, level: str = "INFO"):
         self.log_history.append(format_log(message, level))
         print(format_log(message, level))
@@ -141,51 +177,6 @@ class TextToSQLSystem:
             })
         return similar_examples
 
-    def huggingface_api_call(self, prompt: str) -> str:
-        """呼叫 Hugging Face Inference API"""
-        # === 修正開始 ===
-        # 確保 API_URL 是一個乾淨的字串，不包含任何 Markdown "[ ]" 或其他特殊字元
-        API_URL = "https://api-inference.huggingface.co/models/Paul720810/qwen2.5-coder-1.5b-sql-finetuned"
-
-        # === 修正結束 ===
-        
-        headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-        payload = {
-            "inputs": prompt,
-            "parameters": {
-                "max_new_tokens": 1024,
-                "return_full_text": False
-            }
-        }
-        try:
-            # === 新增除錯日誌 ===
-            # 在發送請求前，打印出最終要使用的 URL，以供檢查
-            self._log(f"準備向 API 端點發送請求: {API_URL}")
-            
-            self._log("正在呼叫 Hugging Face API...")
-            response = requests.post(API_URL, headers=headers, json=payload, timeout=90) # 延長超時時間
-            response.raise_for_status() # 如果 API 回傳錯誤碼 (如 4xx, 5xx)，會在此拋出例外
-            
-            self._log("✅ API 成功回應。")
-            return response.json()[0]['generated_text']
-        
-        except requests.exceptions.RequestException as e:
-            self._log(f"❌ API 呼叫失敗: {e}", "ERROR")
-            
-            # 嘗試解析回應內容，看是否是模型載入中的常見錯誤
-            try:
-                # 即使請求失敗，有時回應本文中仍有 JSON 錯誤訊息
-                error_content = e.response.json() if e.response else {}
-                if "error" in error_content and "estimated_time" in error_content["error"]:
-                    loading_time = error_content["error"]["estimated_time"]
-                    self._log(f"  - 提示: 模型可能正在載入中，預計需要 {loading_time:.1f} 秒。請稍後重試。", "WARNING")
-                    return f"API 錯誤: 模型正在載入中，請稍後再試一次。"
-            except (ValueError, AttributeError):
-                 # 如果回應不是 JSON 或沒有回應本文，就忽略
-                 pass
-
-            return f"API 連線錯誤: {e}"
-
     # === 修改開始: 重寫核心處理邏輯 ===
     def _build_prompt_for_generation(self, user_question: str, examples: List[Dict]) -> str:
         """