Update app.py

app.py

@@ -21,6 +21,9 @@ 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"
 
+# 添加這一行：你的原始微調模型路徑
+FINETUNED_MODEL_PATH = "Paul720810/qwen2.5-coder-1.5b-sql-finetuned"  # ← 新增這行
+
 FEW_SHOT_EXAMPLES_COUNT = 1
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
@@ -85,62 +88,176 @@ class TextToSQLSystem:
         self._load_gguf_model()
 
         self._log("✅ 系統初始化完成")
+    
     def _load_gguf_model(self):
-        """載入 GGUF 模型並處理錯誤"""
+        """載入 GGUF 模型，失敗則使用 Transformers 備用方案"""
+        # 先嘗試原本的 GGUF 載入方式
         try:
             self._log("載入 GGUF 模型...")
-            
-            # 強制重新下載模型
             model_path = hf_hub_download(
                 repo_id=GGUF_REPO_ID,
                 filename=GGUF_FILENAME,
                 repo_type="dataset",
-                force_download=True  # 強制重新下載
+                force_download=True
             )
             
-            # 使用驗證方法檢查檔案
-            if not self._validate_model_file(model_path):
-                self._log("❌ 模型檔案驗證失敗，嘗試重新下載", "ERROR")
-                # 刪除損壞的檔案並重新下載
-                if os.path.exists(model_path):
-                    os.remove(model_path)
-                model_path = hf_hub_download(
-                    repo_id=GGUF_REPO_ID,
-                    filename=GGUF_FILENAME,
-                    repo_type="dataset",
-                    force_download=True
-                )
-                
-                # 再次驗證
-                if not self._validate_model_file(model_path):
-                    raise ValueError("重新下載後檔案仍然無效")
-            
-            # 使用更保守的參數載入模型
+            # 你原本的載入參數
             self.llm = Llama(
                 model_path=model_path,
-                n_ctx=512,          # 減少上下文長度
-                n_threads=4,        # 固定線程數
-                n_batch=128,        # 減少批次大小
-                verbose=False,      # 關閉詳細輸出
-                use_mmap=True,      # 使用記憶體映射
-                use_mlock=False,    # 不鎖定記憶體
-                n_gpu_layers=0      # 強制使用 CPU
+                n_ctx=512,  # 減少上下文長度
+                n_threads=4,  # 固定線程數
+                n_batch=128,  # 減少批次大小
+                verbose=False,  # 關閉詳細輸出
+                use_mmap=True,  # 使用記憶體映射
+                use_mlock=False,  # 不鎖定記憶體
+                n_gpu_layers=0  # 強制使用 CPU
             )
             
-            # 測試模型是否能正常生成
+            # 測試是否能正常生成
             test_output = self.llm("SELECT", max_tokens=5, temperature=0.1)
-            if not test_output or 'choices' not in test_output:
-                raise RuntimeError("模型載入後無法正常生成")
-                
-            self._log("✅ GGUF 模型載入並測試成功")
+            self._log("✅ GGUF 模型載入成功")
+            return
+            
+        except Exception as e:
+            self._log(f"❌ GGUF 載入失敗: {e}", "ERROR")
+            
+        # GGUF 失敗，使用 Transformers 載入你的微調模型
+        try:
+            self._log("改用 Transformers 載入微調模型...")
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+            import torch
+            
+            self.transformers_tokenizer = AutoTokenizer.from_pretrained(FINETUNED_MODEL_PATH)
+            self.transformers_model = AutoModelForCausalLM.from_pretrained(
+                FINETUNED_MODEL_PATH,
+                torch_dtype=torch.float32,
+                device_map="cpu",
+                trust_remote_code=True
+            )
+            
+            self.llm = "transformers"  # 標記使用 transformers
+            self._log("✅ Transformers 模型載入成功")
+            
+        except Exception as e:
+            self._log(f"❌ Transformers 載入也失敗: {e}", "ERROR")
+            self.llm = None
+        
+    def _try_gguf_loading(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=512,
+                n_threads=4,
+                verbose=False,
+                n_gpu_layers=0
+            )
+            
+            # 測試生成
+            test_result = self.llm("SELECT", max_tokens=5)
+            self._log("✅ GGUF 模型載入成功")
+            return True
+            
+        except Exception as e:
+            self._log(f"GGUF 載入失敗: {e}", "WARNING")
+            return False
+    
+    def _load_transformers_model(self):
+        """使用 Transformers 載入你的微調模型"""
+        try:
+            from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+            import torch
+            
+            self._log(f"載入 Transformers 模型: {FINETUNED_MODEL_PATH}")
+            
+            # 載入你的微調模型
+            self.transformers_tokenizer = AutoTokenizer.from_pretrained(FINETUNED_MODEL_PATH)
+            self.transformers_model = AutoModelForCausalLM.from_pretrained(
+                FINETUNED_MODEL_PATH,
+                torch_dtype=torch.float32,  # CPU 使用 float32
+                device_map="cpu",           # 強制使用 CPU
+                trust_remote_code=True      # Qwen 模型可能需要
+            )
+            
+            # 創建生成管道
+            self.generation_pipeline = pipeline(
+                "text-generation",
+                model=self.transformers_model,
+                tokenizer=self.transformers_tokenizer,
+                device=-1,  # CPU
+                max_length=512,
+                do_sample=True,
+                temperature=0.1,
+                top_p=0.9,
+                pad_token_id=self.transformers_tokenizer.eos_token_id
+            )
+            
+            self.llm = "transformers"  # 標記使用 transformers
+            self._log("✅ Transformers 模型載入成功")
             
         except Exception as e:
-            self._log(f"❌ GGUF 模型載入失敗: {str(e)}", "ERROR")
-            self._log("嘗試使用替代方案...", "INFO")
+            self._log(f"❌ Transformers 載入也失敗: {e}", "ERROR")
             self.llm = None
+    
+    def huggingface_api_call(self, prompt: str) -> str:
+        """使用 GGUF 或 Transformers 生成"""
+        if self.llm is None:
+            return self._generate_fallback_sql(prompt)
+        
+        try:
+            # 如果是 Transformers 模型
+            if self.llm == "transformers":
+                # 限制 prompt 長度
+                if len(prompt) > 1000:
+                    prompt = prompt[:1000]
+                
+                # 使用 Transformers 生成
+                inputs = self.transformers_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+                
+                with torch.no_grad():
+                    outputs = self.transformers_model.generate(
+                        inputs.input_ids,
+                        attention_mask=inputs.attention_mask,
+                        max_new_tokens=128,
+                        temperature=0.1,
+                        do_sample=True,
+                        top_p=0.9,
+                        pad_token_id=self.transformers_tokenizer.eos_token_id,
+                        eos_token_id=self.transformers_tokenizer.eos_token_id
+                    )
+                
+                # 解碼生成的文本，只取新生成的部分
+                generated_text = self.transformers_tokenizer.decode(
+                    outputs[0][inputs.input_ids.shape[1]:], 
+                    skip_special_tokens=True
+                )
+                
+                return generated_text.strip()
             
-            # 可以在這裡添加使用其他模型的邏輯
-            # 例如使用 Hugging Face Transformers 的備用方案
+            # 如果是 GGUF 模型（你原本的代碼）
+            else:
+                if len(prompt) > 1800:
+                    prompt = prompt[:1800] + "..."
+                
+                output = self.llm(
+                    prompt,
+                    max_tokens=256,
+                    temperature=0.1,
+                    top_p=0.9,
+                    stop=["</s>", "```", ";", "\n\n"],
+                    echo=False
+                )
+                return output["choices"][0]["text"].strip()
+                
+        except Exception as e:
+            self._log(f"❌ 生成失敗: {e}", "ERROR")
+            return f"生成失敗: {e}"
         
     def _load_gguf_model_fallback(self, model_path):
         """備用載入方式"""
@@ -337,33 +454,6 @@ class TextToSQLSystem:
         
         return prompt
 
-    def huggingface_api_call(self, prompt: str) -> str:
-        """使用 GGUF 模型生成或提供替代方案"""
-        if self.llm is None:
-            # 返回基於規則的簡單 SQL 生成
-            return self._generate_fallback_sql(prompt)
-        
-        try:
-            if len(prompt) > 1500:  # 縮短提示長度
-                prompt = prompt[:1500] + "..."
-            
-            output = self.llm(
-                prompt,
-                max_tokens=128,     # 減少最大 token 數
-                temperature=0.0,    # 使用確定性生成
-                top_p=0.95,
-                stop=["</s>", "```", "\n\n", "問題:"],  # 添加更多停止詞
-                echo=False
-            )
-            
-            if output and 'choices' in output and output['choices']:
-                return output["choices"][0]["text"].strip()
-            else:
-                return "模型生成失敗"
-                
-        except Exception as e:
-            self._log(f"❌ 生成失敗: {e}", "ERROR")
-            return self._generate_fallback_sql(prompt)
 
     def _generate_fallback_sql(self, prompt: str) -> str:
         """當模型不可用時的備用 SQL 生成"""