use API Qwen model

pipeline.py

@@ -4,14 +4,14 @@ import time
 import re
 from db_utils import get_schema, execute_sql
 
-# Hugging Face Inference API endpoint
-API_URL = "https://api-inference.huggingface.co/models/defog/sqlcoder-7b-2"
+# Hugging Face Inference API endpoint for Qwen2.5-Coder
+API_URL = "https://api-inference.huggingface.co/models/Qwen/Qwen2.5-Coder-7B-Instruct"
 
 def query_huggingface_api(prompt, max_retries=3):
     """Query the Hugging Face Inference API"""
     hf_token = os.getenv("HF_TOKEN")
     if not hf_token:
-        raise ValueError("HF_TOKEN not found in environment variables. Add it to your Space secrets.")
+        raise ValueError("HF_TOKEN not found in environment variables")
     
     headers = {"Authorization": f"Bearer {hf_token}"}
     payload = {
@@ -20,37 +20,48 @@ def query_huggingface_api(prompt, max_retries=3):
             "max_new_tokens": 200,
             "temperature": 0.1,
             "do_sample": False,
-            "return_full_text": False
+            "return_full_text": False,
+            "stop": ["###", "\n\n"]
         }
     }
     
     for attempt in range(max_retries):
         try:
+            print(f"=== DEBUG: API attempt {attempt + 1}")
             response = requests.post(API_URL, headers=headers, json=payload, timeout=30)
+            print(f"=== DEBUG: API Response Status: {response.status_code}")
             
             if response.status_code == 200:
                 result = response.json()
+                print(f"=== DEBUG: API Response: {result}")
+                
                 if isinstance(result, list) and len(result) > 0:
-                    return result[0].get("generated_text", "").strip()
-                return str(result).strip()
+                    generated_text = result[0].get("generated_text", "").strip()
+                else:
+                    generated_text = str(result).strip()
+                    
+                return generated_text
             
             elif response.status_code == 503:
                 wait_time = 20 * (attempt + 1)
-                print(f"Model loading, waiting {wait_time} seconds...")
+                print(f"=== DEBUG: Model loading, waiting {wait_time} seconds...")
                 time.sleep(wait_time)
                 continue
                 
             else:
                 error_msg = f"API Error {response.status_code}: {response.text}"
+                print(f"=== DEBUG: {error_msg}")
                 if attempt == max_retries - 1:
                     raise Exception(error_msg)
                     
         except requests.exceptions.Timeout:
+            print(f"=== DEBUG: Timeout on attempt {attempt + 1}")
             if attempt == max_retries - 1:
                 raise Exception("Request timed out after multiple attempts")
             time.sleep(5)
             
         except Exception as e:
+            print(f"=== DEBUG: Exception on attempt {attempt + 1}: {e}")
             if attempt == max_retries - 1:
                 raise e
             time.sleep(5)
@@ -80,16 +91,24 @@ def clean_sql_output(sql_text, user_limit=None):
     if sql_text.startswith("```"):
         lines = sql_text.split('\n')
         sql_text = '\n'.join(lines[1:-1]) if len(lines) > 2 else sql_text
-        
-    # Extract SQL
+    
+    # Handle multiple lines - take the SQL part
     lines = sql_text.split('\n')
     sql = ""
     for line in lines:
         line = line.strip()
         if line and (line.upper().startswith('SELECT') or sql):
             sql += line + " "
-        if line.endswith(';'):
-            break
+            if line.endswith(';'):
+                break
+    
+    if not sql:
+        # If no SELECT found, take the first non-empty line
+        for line in lines:
+            line = line.strip()
+            if line:
+                sql = line
+                break
     
     sql = sql.strip().rstrip(';')
     
@@ -101,40 +120,65 @@ def clean_sql_output(sql_text, user_limit=None):
     return sql
 
 def text_to_sql(nl_query):
-    """Convert natural language to SQL using Hugging Face Inference API"""
+    """Convert natural language to SQL using Qwen2.5-Coder via HF Inference API"""
     try:
-        schema = get_schema()
+        print(f"=== DEBUG: Starting text_to_sql with query: {nl_query}")
+        
+        # Get database schema
+        try:
+            schema = get_schema()
+            print(f"=== DEBUG: Schema retrieved, length: {len(schema)}")
+        except Exception as e:
+            print(f"=== DEBUG: Schema error: {e}")
+            return f"Error: Database schema access failed: {str(e)}", []
+        
+        # Extract user limit
         user_limit = extract_user_requested_limit(nl_query)
+        print(f"=== DEBUG: Extracted user limit: {user_limit}")
         
-        prompt = f"""### Task
-Generate a PostgreSQL query to answer this question: {nl_query}
+        # Create optimized prompt for Qwen2.5-Coder
+        prompt = f"""<|im_start|>system
+You are an expert SQL developer. Generate PostgreSQL queries based on natural language questions.
 
-### Database Schema
-{schema}
+Database Schema:
+{schema[:1500]}
 
-### Instructions
-- Return only the SQL query
+Rules:
+- Return ONLY the SQL query
 - Use PostgreSQL syntax
 - Be precise with table and column names
+- Do not include explanations or markdown formatting
+<|im_end|>
+<|im_start|>user
+{nl_query}
+<|im_end|>
+<|im_start|>assistant
+"""
 
-### SQL Query:"""
-
-        print("Querying Hugging Face Inference API...")
+        print("=== DEBUG: Calling Qwen2.5-Coder API...")
         generated_sql = query_huggingface_api(prompt)
+        print(f"=== DEBUG: Generated SQL raw: {generated_sql}")
         
         if not generated_sql:
-            raise ValueError("No SQL generated from the model")
+            return "Error: No SQL generated from the model", []
         
+        # Clean the SQL output
         sql = clean_sql_output(generated_sql, user_limit)
+        print(f"=== DEBUG: Final cleaned SQL: {sql}")
         
         if not sql or not sql.upper().startswith('SELECT'):
-            raise ValueError(f"Invalid SQL generated: {sql}")
+            return f"Error: Invalid SQL generated: {sql}", []
         
-        print(f"Generated SQL: {sql}")
-        results = execute_sql(sql)
-        return sql, results
+        # Execute SQL
+        print("=== DEBUG: Executing SQL...")
+        try:
+            results = execute_sql(sql)
+            print(f"=== DEBUG: SQL executed successfully, {len(results)} results")
+            return sql, results
+        except Exception as e:
+            print(f"=== DEBUG: SQL execution error: {e}")
+            return f"Error: SQL execution failed: {str(e)}", []
         
     except Exception as e:
-        error_msg = str(e)
-        print(f"Error in text_to_sql: {error_msg}")
-        return f"Error: {error_msg}", []
+        print(f"=== DEBUG: General error in text_to_sql: {e}")
+        return f"Error: {str(e)}", []