Update app.py

app.py

@@ -1,6 +1,3 @@
-# 檔案名稱: app.py
-# 部署在 Hugging Face Spaces (已修正 KeyError)
-
 import gradio as gr
 import requests
 import json
@@ -19,25 +16,37 @@ SIMILARITY_THRESHOLD = 0.90
 print("--- [1/5] 開始初始化應用 ---")
 
 # --- 1. 載入知識庫 ---
+qa_dataset = None
+schema_data = {}
 try:
     print(f"--- [2/5] 正在從 '{DATASET_REPO_ID}' 載入知識庫... ---")
-    # 載入問答範例, 移除已過時的 trust_remote_code 參數
-    dataset = load_dataset(DATASET_REPO_ID, token=HF_TOKEN)
-    raw_qa_dataset = dataset['train']
-    
-    # *** 關鍵修正：解析被包裹在 'text' 欄位中的 JSON ***
-    parsed_qa_data = []
-    for item in raw_qa_dataset:
-        try:
-            # item 現在是 {'text': '{"question": "...", "sql": "..."}'}
-            line_dict = json.loads(item['text'])
-            parsed_qa_data.append(line_dict)
-        except (json.JSONDecodeError, KeyError) as e:
-            print(f"警告：跳過一行無法解析的數據: {item}, 錯誤: {e}")
-            
-    # 使用解析後的數據創建一個新的、格式正確的 Dataset 對象
-    qa_dataset = Dataset.from_list(parsed_qa_data)
+    raw_dataset = load_dataset(DATASET_REPO_ID, token=HF_TOKEN)['train']
     
+    # *** 關鍵修正：智能解析 Dataset ***
+    # 檢查第一條數據的結構來判斷格式
+    if raw_dataset and len(raw_dataset) > 0:
+        first_item = raw_dataset[0]
+        if 'text' in first_item and 'question' not in first_item:
+            # 這是舊的 {'text': '...'} 格式，需要解析
+            print("--- > 檢測到 'text' 格式，正在解析JSON...")
+            parsed_qa_data = []
+            for item in raw_dataset:
+                try:
+                    line_dict = json.loads(item['text'])
+                    parsed_qa_data.append(line_dict)
+                except (json.JSONDecodeError, KeyError):
+                    continue # 跳過錯誤行
+            qa_dataset = Dataset.from_list(parsed_qa_data)
+        elif 'question' in first_item and 'sql' in first_item:
+            # 這已經是正確的 {'question': ..., 'sql': ...} 格式
+            print("--- > 檢測到已解析的 'question'/'sql' 格式，直接使用。")
+            qa_dataset = raw_dataset
+        else:
+            raise ValueError(f"未知的Dataset格式: {first_item}")
+    else:
+        # 數據集為空
+        raise ValueError("載入的Dataset為空。")
+
     # 載入並解析 Schema JSON
     schema_file_path = "sqlite_schema_FULL.json"
     hf_hub_download(repo_id=DATASET_REPO_ID, filename=schema_file_path, repo_type='dataset', local_dir='.', token=HF_TOKEN)
@@ -45,15 +54,16 @@ try:
     with open(schema_file_path, 'r', encoding='utf-8') as f:
         schema_data = json.load(f)
         
-    print(f"--- > 成功解析 {len(qa_dataset)} 條問答範例和 Schema。 ---")
+    print(f"--- > 成功載入 {len(qa_dataset)} 條問答範例和 Schema。 ---")
+
 except Exception as e:
     print(f"!!! 致命錯誤: 無法載入或解析 Dataset '{DATASET_REPO_ID}'.")
     print(f"詳細錯誤: {e}")
     qa_dataset = Dataset.from_dict({"question": ["示例問題"], "sql": ["SELECT 'Dataset failed to load'"]})
-    schema_data = {}
 
 # --- 2. 構建 DDL 和初始化檢索模型 ---
 def load_schema_as_ddl(schema_dict: dict) -> str:
+    # (此函式無需修改)
     ddl_string = ""
     for table_name, columns in schema_dict.items():
         if not isinstance(columns, list): continue
@@ -67,33 +77,53 @@ SCHEMA_DDL = load_schema_as_ddl(schema_data)
 print("--- [3/5] 正在載入句向量模型 (all-MiniLM-L6-v2)... ---")
 embedder = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
 
-# *** 關鍵修正：現在 qa_dataset 的結構是正確的了 ***
 questions = [item['question'] for item in qa_dataset]
 sql_answers = [item['sql'] for item in qa_dataset]
 
-print(f"--- [4/5] 正在為 {len(questions)} 個問題計算向量 (這可能需要幾分鐘)... ---")
-question_embeddings = embedder.encode(questions, convert_to_tensor=True, show_progress_bar=True)
-print("--- > 向量計算完成！ ---")
+# 只有在 questions 列表不為空時才進行計算
+if questions:
+    print(f"--- [4/5] 正在為 {len(questions)} 個問題計算向量... ---")
+    question_embeddings = embedder.encode(questions, convert_to_tensor=True, show_progress_bar=True)
+    print("--- > 向量計算完成！ ---")
+else:
+    print("--- [4/5] 警告：沒有可用的問題來計算向量。檢索功能將不可用。---")
+    question_embeddings = torch.Tensor([])
+
 
 # --- 3. 混合系統核心邏輯 ---
 def get_sql_query(user_question: str):
-    # (此函式剩餘部分無需修改，保持原樣)
+    # (此函式剩餘部分幾乎無需修改)
     if not user_question:
         return "請輸入您的問題。", "日誌：用戶未輸入問題。"
+    
+    # 增加一個檢查，確保知識庫不是空的
+    if len(questions) == 0:
+        log_message = "錯誤：知識庫為空，無法進行檢索。"
+        return "系統錯誤：知識庫未成功載入。", log_message
+
     question_embedding = embedder.encode(user_question, convert_to_tensor=True)
     hits = util.semantic_search(question_embedding, question_embeddings, top_k=5)
-    hits = hits[0]
-    most_similar_hit = hits[0]
-    similarity_score = most_similar_hit['score']
-    log_message = f"檢索到最相似問題: '{questions[most_similar_hit['corpus_id']]}' (相似度: {similarity_score:.4f})"
-    if similarity_score > SIMILARITY_THRESHOLD:
-        sql_result = sql_answers[most_similar_hit['corpus_id']]
-        log_message += f"\n相似度 > {SIMILARITY_THRESHOLD}，[模式: 直接返回]。"
-        return sql_result, log_message
-    log_message += f"\n相似度 < {SIMILARITY_THRESHOLD}，[模式: LLM生成]。正在構建 Prompt..."
+    
+    if not hits or not hits[0]:
+        log_message = "檢索失敗：找不到任何相似的問題。"
+        # 即使檢索失敗，也應該嘗試調用 LLM
+    else:
+        hits = hits[0]
+        most_similar_hit = hits[0]
+        similarity_score = most_similar_hit['score']
+        log_message = f"檢索到最相似問題: '{questions[most_similar_hit['corpus_id']]}' (相似度: {similarity_score:.4f})"
+        
+        if similarity_score > SIMILARITY_THRESHOLD:
+            sql_result = sql_answers[most_similar_hit['corpus_id']]
+            log_message += f"\n相似度 > {SIMILARITY_THRESHOLD}，[模式: 直接返回]。"
+            return sql_result, log_message
+
+    log_message += f"\n相似度低於閾值或檢索失敗，[模式: LLM生成]。正在構建 Prompt..."
     examples_context = ""
-    for hit in hits[:3]:
-         examples_context += f"### A user asks: {questions[hit['corpus_id']]}\n{sql_answers[hit['corpus_id']]}\n\n"
+    if hits: # 只有在檢索到結果時才添加範例
+        for hit in hits[:3]:
+             examples_context += f"### A user asks: {questions[hit['corpus_id']]}\n{sql_answers[hit['corpus_id']]}\n\n"
+
     prompt = f"""### Task
 Generate a SQLite SQL query that answers the following user question.
 Your response must contain ONLY the SQL query. Do not add any explanation.
@@ -130,8 +160,9 @@ Your response must contain ONLY the SQL query. Do not add any explanation.
 # --- 4. 創建 Gradio Web 界面 ---
 print("--- [5/5] 正在創建 Gradio Web 界面... ---")
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    # (此部分無需修改，保持原樣)
+    # (此部分無需修改)
     gr.Markdown("# 智能 Text-to-SQL 系統 (混合模式)")
+    # ... (Gradio界面代碼與之前相同)
     gr.Markdown("輸入您的自然語言問題，系統將首先嘗試從知識庫中快速檢索答案。如果問題較新穎，則會調用雲端大語言模型生成SQL。")
     with gr.Row():
         question_input = gr.Textbox(label="輸入您的問題", placeholder="例如：去年Nike的總業績是多少？", scale=4)