Update app.py

app.py

@@ -1,17 +1,17 @@
 import gradio as gr
 import re
 import torch
-import sqlite3
+import sqlite3  # Can be replaced with other DB connections
+import os
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel, PeftConfig
-import torch
-import os
-os.makedirs("offload", exist_ok=True)
 
-# ✅ Load fine-tuned models from Hugging Face Model Hub instead of Kaggle paths
-codellama_model_path = "srishtirai/codellama-sql-finetuned"  # Upload to HF Model Hub
-mistral_model_path = "srishtirai/mistral-sql-finetuned"  # Upload to HF Model Hub
+# ✅ Ensure offload directory exists
+os.makedirs("offload", exist_ok=True)
 
+# ✅ Load fine-tuned models from Hugging Face Model Hub
+codellama_model_path = "srishtirai/codellama-sql-finetuned"
+mistral_model_path = "srishtirai/mistral-sql-finetuned"
 
 def load_model(model_path):
     tokenizer = AutoTokenizer.from_pretrained(model_path)
@@ -21,18 +21,24 @@ def load_model(model_path):
     peft_config = PeftConfig.from_pretrained(model_path)
     base_model_name = peft_config.base_model_name_or_path
 
+    # ✅ Load base model with offloading & low-memory optimization
     base_model = AutoModelForCausalLM.from_pretrained(
         base_model_name,
-        torch_dtype=torch.float16,  # Use FP16 to save memory
-        device_map="auto",  # Automatically allocate layers to CPU/GPU
-        offload_folder="offload"  # ✅ Offload large layers to disk
+        torch_dtype=torch.float16,  # Use FP16 to reduce memory usage
+        device_map="auto",  # Automatically distribute across CPU/GPU
+        offload_folder="offload"  # ✅ Prevents memory crashes
+    )
+
+    # ✅ Load LoRA adapter with `is_trainable=False`
+    model = PeftModel.from_pretrained(
+        base_model,
+        model_path,
+        is_trainable=False  # ✅ Fixes LoRA adapter loading issues
     )
 
-    model = PeftModel.from_pretrained(base_model, model_path)
     model.eval()
     return model, tokenizer
 
-
 # ✅ Load both models from Hugging Face
 codellama_model, codellama_tokenizer = load_model(codellama_model_path)
 mistral_model, mistral_tokenizer = load_model(mistral_model_path)
@@ -93,14 +99,14 @@ def generate_sql_with_explanation(model_choice, schema, question, max_new_tokens
         "full_response": full_response
     }
 
-# ✅ Function to execute SQL query (Optional)
+# ✅ Function to execute SQL query (if database connection is available)
 def execute_sql_query(sql_query):
     """
     Runs the generated SQL query on a sample SQLite database.
-    (Replace with a real DB connection if needed)
+    (You can replace SQLite with a connection to a real database)
     """
     try:
-        conn = sqlite3.connect(":memory:")  # Temporary SQLite DB
+        conn = sqlite3.connect(":memory:")  # Temporary SQLite DB (Replace with actual DB connection)
         cursor = conn.cursor()
         cursor.execute(sql_query)
         result = cursor.fetchall()