Enhance SQL generation by adding schema formatting function and increasing semantic search context

app.py

@@ -4,7 +4,6 @@ from transformers import T5Tokenizer, T5ForConditionalGeneration
 from sentence_transformers import SentenceTransformer, util
 
 # --- CONFIGURATION ---
-# UPDATE THIS to the new model you just trained
 FINE_TUNED_MODEL_ID = "hmyunis/t5-base-sql-custom" 
 
 print(f"Loading Model: {FINE_TUNED_MODEL_ID}...")
@@ -16,6 +15,25 @@ try:
 except Exception as e:
     print(f"CRITICAL ERROR LOADING MODELS: {e}")
 
+def format_schema_like_training(raw_column_list):
+    """
+    Transforms ['api_customer.name', 'api_customer.city', 'api_order.id']
+    Into: "api_customer: name, city | api_order: id"
+    
+    This matches the pattern the model saw during training.
+    """
+    schema_map = {}
+    for item in raw_column_list:
+        if "." in item:
+            table, col = item.split('.', 1)
+            if table not in schema_map:
+                schema_map[table] = []
+            schema_map[table].append(col)
+            
+    # Join nicely
+    parts = [f"{table}: {', '.join(cols)}" for table, cols in schema_map.items()]
+    return " | ".join(parts)
+
 def get_sql_pipeline(question, all_columns_str):
     print(f"Input Q: {question}")
     
@@ -26,23 +44,25 @@ def get_sql_pipeline(question, all_columns_str):
         # 2. Schema Linking (Embeddings)
         question_embedding = embedder.encode(question, convert_to_tensor=True)
         column_embeddings = embedder.encode(all_columns, convert_to_tensor=True)
-        hits = util.semantic_search(question_embedding, column_embeddings, top_k=6)
+        
+        # Increase Top-K to 10 to ensure we get enough context from the right table
+        hits = util.semantic_search(question_embedding, column_embeddings, top_k=10)
         relevant_cols = [all_columns[hit['corpus_id']] for hit in hits[0]]
         
-        # 3. Formulate Prompt (MATCHES TRAINING EXACTLY)
-        schema_context = ", ".join(relevant_cols)
-        # Note the prefix change: "translate English to SQL"
+        # 3. Formulate Prompt (CRITICAL FIX HERE)
+        # We re-format the list to look like "table: col1, col2"
+        schema_context = format_schema_like_training(relevant_cols)
+        
         input_text = f"translate English to SQL: {question} </s> {schema_context}"
         print(f"Prompt: {input_text}")
         
         # 4. Generate
         input_ids = tokenizer(input_text, return_tensors="pt").input_ids
         
-        # Use beam search for better accuracy (slower but worth it)
         outputs = model.generate(
             input_ids, 
             max_length=128, 
-            num_beams=4, # Inspects 4 possible paths
+            num_beams=4, 
             early_stopping=True
         )
         
@@ -54,6 +74,5 @@ def get_sql_pipeline(question, all_columns_str):
     except Exception as e:
         return f"Error: {str(e)}"
 
-# Simpler Interface
 iface = gr.Interface(fn=get_sql_pipeline, inputs=["text", "text"], outputs="text")
 iface.launch()