Refactor model loading and error handling in app.py for improved clarity and robustness

app.py

@@ -4,61 +4,54 @@ from transformers import T5Tokenizer, T5ForConditionalGeneration
 from sentence_transformers import SentenceTransformer, util
 
 # --- CONFIGURATION ---
-# REPLACE THIS with your fine-tuned model ID from Colab/Phase 1
+# Ensure this matches your ACTUAL model on Hugging Face
 FINE_TUNED_MODEL_ID = "hmyunis/t5-sql-finetuned"
 
-# Load the Semantic Embedding Model (for Schema Linking)
-# This creates vectors to find the right columns mathematically
-print("Loading Embedding Model...")
-embedder = SentenceTransformer('all-MiniLM-L6-v2')
-
-# Load your Fine-Tuned T5 Model
-print("Loading T5 Model...")
-tokenizer = T5Tokenizer.from_pretrained(FINE_TUNED_MODEL_ID)
-model = T5ForConditionalGeneration.from_pretrained(FINE_TUNED_MODEL_ID)
+print(f"Loading Model: {FINE_TUNED_MODEL_ID}...")
+try:
+    tokenizer = T5Tokenizer.from_pretrained(FINE_TUNED_MODEL_ID)
+    model = T5ForConditionalGeneration.from_pretrained(FINE_TUNED_MODEL_ID)
+    embedder = SentenceTransformer('all-MiniLM-L6-v2')
+    print("Models loaded successfully.")
+except Exception as e:
+    print(f"CRITICAL ERROR LOADING MODELS: {e}")
 
 def get_sql_pipeline(question, all_columns_str):
-    """
-    1. Receives question + ALL columns in the database (as a string list).
-    2. Uses Vector Search to find the top 5 relevant columns.
-    3. Generates SQL using the Fine-Tuned T5 model.
-    """
+    print(f"Received Question: {question}")
+    print(f"Received Columns Length: {len(all_columns_str)}")
+
     try:
-        # Convert string representation of list back to list
-        # Expected input format: "['table.col1', 'table.col2', ...]"
+        # 1. Parse Columns
         all_columns = eval(all_columns_str)
 
-        # --- NLP LAYER 1: SEMANTIC SCHEMA LINKING ---
-        # 1. Encode the user's question into a vector
+        # 2. Schema Linking (Embeddings)
         question_embedding = embedder.encode(question, convert_to_tensor=True)
-
-        # 2. Encode all database columns into vectors
         column_embeddings = embedder.encode(all_columns, convert_to_tensor=True)
-
-        # 3. Calculate Cosine Similarity to find relevant columns
         hits = util.semantic_search(question_embedding, column_embeddings, top_k=6)
-
-        # 4. Extract the top matching columns
         relevant_cols = [all_columns[hit['corpus_id']] for hit in hits[0]]
 
-        # Format for T5: "table: ..., columns: ..."
-        # We simplify here to just a comma-separated list for the model context
+        # 3. Formulate Prompt
         schema_context = ", ".join(relevant_cols)
-
-        # --- NLP LAYER 2: GENERATION ---
-        # The prompt must match exactly how we trained it in Colab
         input_text = f"translate to SQL: {question} </s> {schema_context}"
+        print(f"Prompt sent to T5: {input_text}")
 
+        # 4. Generate
         input_ids = tokenizer(input_text, return_tensors="pt").input_ids
-
-        # Generate
         outputs = model.generate(input_ids, max_length=128)
         generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+        print(f"Generated Output: '{generated_sql}'")
+
+        # Fallback if empty (Model produced nothing)
+        if not generated_sql:
+            return "SELECT * FROM api_customer -- Model returned empty, defaulting."
+
         return generated_sql
 
     except Exception as e:
-        return f"Error in HF Space: {str(e)}"
+        error_msg = f"Error in HF Space: {str(e)}"
+        print(error_msg)
+        return error_msg
 
 # Launch Gradio
 iface = gr.Interface(
@@ -66,4 +59,4 @@ iface = gr.Interface(
     inputs=["text", "text"],
     outputs="text"
 )
-iface.launch()
+iface.launch()