Update app.py

app.py

@@ -10,18 +10,21 @@ from sentence_transformers import SentenceTransformer, util
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
 from typing import List, Dict, Tuple, Optional
+import faiss
+from functools import lru_cache
 
 # ==================== 配置區 ====================
 DATASET_REPO_ID = "Paul720810/Text-to-SQL-Softline"
 GGUF_REPO_ID = "Paul720810/gguf-models"
 GGUF_FILENAME = "qwen2.5-coder-1.5b-sql-finetuned.q4_k_m.gguf"
 
-FEW_SHOT_EXAMPLES_COUNT = 1
+FEW_SHOT_EXAMPLES_COUNT = 1  # 只使用1个最相关的范例
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 print("=" * 60)
-print("🤖 Text-to-SQL (GGUF) 系統啟動中...")
+print("🤖 Text-to-SQL (GGUF) 極速版系統啟動中...")
 print(f"📊 數據集: {DATASET_REPO_ID}")
+print(f"🤖 GGUF 模型: {GGUF_REPO_ID}/{GGUF_FILENAME}")
 print(f"💻 設備: {DEVICE}")
 print("=" * 60)
 
@@ -33,52 +36,62 @@ def format_log(message: str, level: str = "INFO") -> str:
     return f"[{get_current_time()}] [{level.upper()}] {message}"
 
 def parse_sql_from_response(response_text: str) -> Optional[str]:
+    """從模型輸出提取 SQL"""
     # 1. 先找 ```sql ... ```
     match = re.search(r"```sql\n(.*?)\n```", response_text, re.DOTALL)
     if match:
         return match.group(1).strip()
 
-    # 2. 如果沒找到，嘗試找最長的 SELECT … 語句
+    # 2. 找 SQL 語句
     match = re.search(r"(SELECT .*?;)", response_text, re.DOTALL | re.IGNORECASE)
     if match:
         return match.group(1).strip()
 
-    return None
+    # 3. 找沒有分號的 SQL
+    match = re.search(r"(SELECT .*?)(?=\n\n|\n```|$)", response_text, re.DOTALL | re.IGNORECASE)
+    if match:
+        sql = match.group(1).strip()
+        if not sql.endswith(';'):
+            sql += ';'
+        return sql
 
+    return None
 
 # ==================== Text-to-SQL 核心類 ====================
 class TextToSQLSystem:
-    def __init__(self, embed_model='sentence-transformers/paraphrase-multilingual-mpnet-base-v2'):
+    def __init__(self, embed_model='all-MiniLM-L6-v2'):
         self.log_history = []
-        self._log("初始化系統...")
-
-        # 1. 載入 schema
-        self.schema = self._load_schema()
-
-        # 2. 載入檢索模型
-        self.model = SentenceTransformer(embed_model, device=DEVICE)
-        self.dataset, self.corpus_embeddings = self._load_and_encode_dataset()
-
-        # 3. 載入 GGUF 模型
-        model_path = hf_hub_download(
-            repo_id=GGUF_REPO_ID,
-            filename=GGUF_FILENAME,
-            repo_type="dataset"
-        )
-        self.llm = Llama(
-            model_path=model_path,
-            n_ctx=2048,              # 少一半上下文 → 快
-            n_threads=os.cpu_count(),# 用滿 CPU
-            n_batch=512,             # 增加 batch size → 減少迭代
-            verbose=False
-        )
-        self._log(f"✅ 已載入 GGUF 模型: {GGUF_FILENAME}")
+        self._log("初始化極速系統...")
+        self.query_cache = {}
+
+        # 並行載入所有組件
+        import threading
+        self.schema = {}
+        self.model = None
+        self.dataset = None
+        self.corpus_embeddings = None
+        self.faiss_index = None
+        self.llm = None
+
+        threads = [
+            threading.Thread(target=self._load_schema),
+            threading.Thread(target=self._load_embedding_model),
+            threading.Thread(target=self._load_gguf_model)
+        ]
+
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        self._log("✅ 所有組件載入完成")
 
     def _log(self, message: str, level: str = "INFO"):
         self.log_history.append(format_log(message, level))
         print(format_log(message, level))
 
-    def _load_schema(self) -> Dict:
+    def _load_schema(self):
+        """載入數據庫結構"""
         try:
             schema_path = hf_hub_download(
                 repo_id=DATASET_REPO_ID,
@@ -86,97 +99,224 @@ class TextToSQLSystem:
                 repo_type="dataset"
             )
             with open(schema_path, "r", encoding="utf-8") as f:
-                self._log("成功載入資料庫結構 (sqlite_schema_FULL.json)")
-                return json.load(f)
+                self.schema = json.load(f)
+                self._log("✅ 數據庫結構載入完成")
         except Exception as e:
             self._log(f"❌ 載入 schema 失敗: {e}", "ERROR")
-            return {}
-
-    def _format_schema_for_prompt(self) -> str:
-        formatted = "資料庫結構:\n"
-        for table, cols in self.schema.items():
-            formatted += f"Table: {table}\n"
-            for col in cols:
-                formatted += f"  - {col['name']} ({col['type']}) # {col.get('description','')}\n"
-            formatted += "\n"
-        return formatted
 
-    def _load_and_encode_dataset(self):
+    def _load_embedding_model(self):
+        """載入檢索模型和數據"""
         try:
+            self.model = SentenceTransformer('all-MiniLM-L6-v2', device=DEVICE)
             dataset = load_dataset(DATASET_REPO_ID, data_files="training_data.jsonl", split="train")
+            self.dataset = dataset
             corpus = [item['messages'][0]['content'] for item in dataset]
             self._log(f"正在編碼 {len(corpus)} 個問題...")
+            
             embeddings = self.model.encode(corpus, convert_to_tensor=True, device=DEVICE)
-            self._log("✅ 範例問題編碼完成。")
-            return dataset, embeddings
+            self.corpus_embeddings = embeddings
+            
+            # 建立 FAISS 索引
+            embeddings_np = embeddings.cpu().numpy()
+            self.faiss_index = faiss.IndexFlatIP(embeddings_np.shape[1])
+            self.faiss_index.add(embeddings_np)
+            
+            self._log("✅ FAISS 向量索引建立完成")
         except Exception as e:
-            self._log(f"❌ 載入或編碼數據集失敗: {e}", "ERROR")
-            return None, None
+            self._log(f"❌ 載入檢索模型失敗: {e}", "ERROR")
 
+    def _load_gguf_model(self):
+        """載入 GGUF 模型"""
+        try:
+            model_path = hf_hub_download(
+                repo_id=GGUF_REPO_ID,
+                filename=GGUF_FILENAME,
+                repo_type="dataset"
+            )
+            self.llm = Llama(
+                model_path=model_path,
+                n_ctx=1024,
+                n_threads=os.cpu_count(),
+                n_batch=512,
+                n_gpu_layers=0,
+                verbose=False
+            )
+            self._log("✅ GGUF 模型載入完成")
+        except Exception as e:
+            self._log(f"❌ 載入 GGUF 模型失敗: {e}", "ERROR")
+
+    def _identify_relevant_tables(self, question: str) -> List[str]:
+        """智能識別問題相關的表"""
+        question_lower = question.lower()
+        relevant_tables = []
+        
+        # 關鍵詞映射表
+        keyword_to_table = {
+            'jobtimeline': ['報告', '完成', '時間', '日期', 'month', 'year', '月', '年', '報告', '完成時間'],
+            'tsr53sampledescription': ['客戶', '買家', '申請', '評級', 'rating', 'client', '申請方', '買方'],
+            'tsr53invoice': ['金額', '發票', '價格', '費用', 'amount', 'price', '錢', '成本'],
+            'jobeventslog': ['事件', '操作', '用戶', 'event', 'action', '操作記錄'],
+            'calendardays': ['工作日', '假期', 'workday', 'holiday', '日期', '放假']
+        }
+        
+        for table, keywords in keyword_to_table.items():
+            if table in self.schema and any(keyword in question_lower for keyword in keywords):
+                relevant_tables.append(table)
+        
+        # 如果沒有匹配，返回最重要的表
+        return relevant_tables[:2] if relevant_tables else ['jobtimeline', 'tsr53sampledescription']
+
+    def _format_relevant_schema(self, table_names: List[str]) -> str:
+        """只顯示相關表結構"""
+        if not self.schema:
+            return "無數據庫結構信息"
+        
+        formatted = "相關表結構:\n"
+        for table in table_names:
+            if table in self.schema:
+                formatted += f"## {table}\n"
+                for col in self.schema[table][:6]:  # 只顯示前6個列
+                    col_desc = col.get('description', '')
+                    formatted += f"- {col['name']} ({col['type']})"
+                    if col_desc:
+                        formatted += f" # {col_desc}"
+                    formatted += "\n"
+                formatted += "\n"
+        
+        return formatted
+
+    @lru_cache(maxsize=100)
     def find_most_similar(self, question: str, top_k: int) -> List[Dict]:
-        if self.corpus_embeddings is None: return []
-        q_emb = self.model.encode(question, convert_to_tensor=True, device=DEVICE)
-        scores = util.cos_sim(q_emb, self.corpus_embeddings)[0]
-        top = torch.topk(scores, k=min(top_k, len(self.corpus_embeddings)))
-        results = []
-        for score, idx in zip(top[0], top[1]):
-            item = self.dataset[idx.item()]
-            q = item['messages'][0]['content']
-            a = item['messages'][1]['content']
-            sql = parse_sql_from_response(a) or "無法解析範例SQL"
-            results.append({"similarity": score.item(), "question": q, "sql": sql})
-        return results
+        """使用 FAISS 快速檢索相似問題"""
+        if self.faiss_index is None or self.dataset is None:
+            return []
+        
+        try:
+            q_emb = self.model.encode(question, convert_to_tensor=True, device=DEVICE)
+            q_emb_np = q_emb.cpu().numpy().reshape(1, -1)
+            
+            # FAISS 搜索
+            distances, indices = self.faiss_index.search(q_emb_np, min(top_k + 2, len(self.dataset)))
+            
+            results = []
+            seen_questions = set()
+            
+            for i, idx in enumerate(indices[0]):
+                if len(results) >= top_k:
+                    break
+                    
+                item = self.dataset[idx]
+                q_content = item['messages'][0]['content']
+                a_content = item['messages'][1]['content']
+                
+                # 提取純淨問題
+                clean_q = re.sub(r"以下是一個SQL查詢任務：\s*指令:\s*", "", q_content).strip()
+                if clean_q in seen_questions:
+                    continue
+                    
+                seen_questions.add(clean_q)
+                sql = parse_sql_from_response(a_content) or "無法解析範例SQL"
+                
+                results.append({
+                    "similarity": float(distances[0][i]),
+                    "question": clean_q,
+                    "sql": sql
+                })
+            
+            return results
+            
+        except Exception as e:
+            self._log(f"❌ 檢索失敗: {e}", "ERROR")
+            return []
 
     def _build_prompt(self, user_q: str, examples: List[Dict]) -> str:
-        system_instruction = (
-            "你是一位頂尖的 SQLite 專家。\n"
-            "請嚴格遵守以下規則：\n"
-            "1. 僅輸出最終的 SQL 語法，不要包含解釋或自然語言。\n"
-            "2. 必須使用 ```sql 開頭 和 ``` 結尾包住查詢。\n"
-            "3. 查詢必須能直接在 SQLite 執行。\n"
-        )
-        schema_str = self._format_schema_for_prompt()
-        ex_str = "--- 範例 ---\n"
-        for i, ex in enumerate(examples, 1):
-            ex_str += f"範例 {i} 問題: {ex['question']}\nSQL:\n```sql\n{ex['sql']}\n```\n\n"
-        return f"{system_instruction}\n\n{schema_str}\n{ex_str}\n--- 使用者問題 ---\n請根據以上資訊，生成 SQL 查詢：\n\"{user_q}\""
+        """建立精簡的提示詞"""
+        # 識別相關表
+        relevant_tables = self._identify_relevant_tables(user_q)
+        schema_str = self._format_relevant_schema(relevant_tables)
+        
+        # 極簡指令
+        system_instruction = "生成SQL查詢。只輸出```sql...```內容。確保SQL語法正確。"
+        
+        # 只顯示一個最有用的範例
+        ex_str = ""
+        if examples:
+            best_example = examples[0]
+            ex_str = f"參考範例:\n問題: {best_example['question']}\nSQL: ```sql\n{best_example['sql']}\n```\n\n"
+        
+        prompt = f"{system_instruction}\n{schema_str}\n{ex_str}問題: {user_q}\nSQL:"
+        
+        # 檢查長度，如果太長則進一步精簡
+        if len(prompt) > 1500:
+            self._log("⚠️ 提示詞過長，進行精簡")
+            prompt = f"{system_instruction}\n{schema_str}\n問題: {user_q}\nSQL:"
+        
+        return prompt
 
     def huggingface_api_call(self, prompt: str) -> str:
+        """使用 GGUF 模型生成"""
+        if self.llm is None:
+            return "模型未載入"
+        
         try:
-            self._log("🧠 使用 GGUF 模型生成 SQL...")
+            # 檢查prompt長度
+            if len(prompt) > 1800:
+                prompt = prompt[:1800] + "..."
+            
             output = self.llm(
                 prompt,
-                max_tokens=128,
-                temperature=0.2,
-                stop=["</s>", "```"],
+                max_tokens=256,
+                temperature=0.1,
+                top_p=0.9,
+                stop=["</s>", "```", ";", "\n\n"],
                 echo=False
             )
-            text = output["choices"][0]["text"].strip()
-            return text
+            return output["choices"][0]["text"].strip()
+            
         except Exception as e:
             self._log(f"❌ 生成失敗: {e}", "ERROR")
             return f"生成失敗: {e}"
 
     def process_question(self, question: str) -> Tuple[str, str]:
+        """處理使用者問題"""
+        # 檢查緩存
+        if question in self.query_cache:
+            self._log("⚡ 使用緩存結果")
+            return self.query_cache[question]
+        
         self.log_history = []
-        self._log(f"⏰ 問題: {question}")
+        self._log(f"⏰ 處理問題: {question}")
 
+        # 長度檢查
+        if len(question) > 200:
+            self._log("⚠️ 問題過長，進行精簡")
+            question = question[:200]
+
+        # 檢索相似範例
         self._log("🔍 尋找相似範例...")
         examples = self.find_most_similar(question, FEW_SHOT_EXAMPLES_COUNT)
 
+        # 建立提示詞
         self._log("📝 建立 Prompt...")
         prompt = self._build_prompt(question, examples)
+        self._log(f"📏 Prompt 長度: {len(prompt)} 字符")
 
-        self._log("🧠 開始生成...")
+        # 生成 SQL
+        self._log("🧠 開始生成 SQL...")
         response = self.huggingface_api_call(prompt)
 
+        # 解析結果
         sql = parse_sql_from_response(response)
         if sql:
-            self._log("✅ 成功解析 SQL")
-            return sql, "生成成功"
+            self._log("✅ SQL 生成成功")
+            result = (sql, "生成成功")
         else:
             self._log("❌ 未能解析 SQL", "ERROR")
-            return f"原始回應:\n{response}", "生成失敗"
+            result = (f"無法解析SQL。原始回應:\n{response}", "生成失敗")
+        
+        # 緩存結果
+        self.query_cache[question] = result
+        return result
 
 # ==================== Gradio 介面 ====================
 text_to_sql_system = TextToSQLSystem()
@@ -184,24 +324,51 @@ text_to_sql_system = TextToSQLSystem()
 def process_query(q: str):
     if not q.strip():
         return "", "等待輸入", "請輸入問題"
+    
     sql, status = text_to_sql_system.process_question(q)
-    logs = "\n".join(text_to_sql_system.log_history)
+    logs = "\n".join(text_to_sql_system.log_history[-10:])  # 只顯示最後10條日誌
+    
     return sql, status, logs
 
-with gr.Blocks(theme=gr.themes.Soft(), title="Text-to-SQL Assistant (GGUF)") as demo:
-    gr.Markdown("# 📊 Text-to-SQL Assistant (GGUF)")
+# 範例問題
+examples = [
+    "2024年每月完成多少份報告？",
+    "統計各種評級(Pass/Fail)的分布情況",
+    "找出總金額最高的10個工作單",
+    "哪些客戶的工作單數量最多？",
+    "A組昨天完成了多少個測試項目？"
+]
+
+with gr.Blocks(theme=gr.themes.Soft(), title="Text-to-SQL 極速助手") as demo:
+    gr.Markdown("# ⚡ Text-to-SQL 極速助手 (GGUF)")
+    gr.Markdown("輸入自然語言問題，自動生成SQL查詢")
+    
     with gr.Row():
         with gr.Column(scale=2):
-            inp = gr.Textbox(lines=3, label="💬 問題")
-            btn = gr.Button("🚀 生成 SQL")
+            inp = gr.Textbox(lines=3, label="💬 您的問題", placeholder="例如：2024年每月完成多少份報告？")
+            btn = gr.Button("🚀 生成 SQL", variant="primary")
             status = gr.Textbox(label="狀態", interactive=False)
+        
         with gr.Column(scale=3):
-            sql_out = gr.Code(label="🤖 SQL", language="sql")
-    with gr.Accordion("日誌", open=False):
-        logs = gr.Textbox(lines=15, label="處理日誌", interactive=False)
+            sql_out = gr.Code(label="🤖 生成的 SQL", language="sql", lines=8)
+    
+    with gr.Accordion("📋 處理日誌", open=False):
+        logs = gr.Textbox(lines=8, label="日誌", interactive=False)
+    
+    # 範例區
+    gr.Examples(
+        examples=examples,
+        inputs=inp,
+        label="💡 點擊試用範例問題"
+    )
 
+    # 綁定事件
     btn.click(process_query, inputs=[inp], outputs=[sql_out, status, logs])
     inp.submit(process_query, inputs=[inp], outputs=[sql_out, status, logs])
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )