Update app.py

app.py

@@ -6,32 +6,15 @@ import gradio as gr
 model_name = "defog/llama-3-sqlcoder-8b"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Check GPU memory if available, otherwise default to 4-bit mode
-def get_model():
-    try:
-        available_memory = torch.cuda.get_device_properties(0).total_memory
-    except:
-        available_memory = 0
-
-    if available_memory > 20e9:
-        return AutoModelForCausalLM.from_pretrained(
-            model_name,
-            trust_remote_code=True,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            use_cache=True,
-        )
-    else:
-        return AutoModelForCausalLM.from_pretrained(
-            model_name,
-            trust_remote_code=True,
-            load_in_4bit=True,
-            device_map="auto",
-            use_cache=True,
-        )
-
-model = get_model()
+# Load model on CPU
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    device_map={"": "cpu"},
+    torch_dtype=torch.float32
+)
 
+# SQL Prompt Template
 prompt = """<|begin_of_text|><|start_header_id|>user<|end_header_id|>
 
 Generate a SQL query to answer this question: `{question}`
@@ -80,9 +63,10 @@ The following SQL query best answers the question `{question}`:
 ```sql
 """
 
+# Main function
 def generate_query(question):
     formatted_prompt = prompt.format(question=question)
-    inputs = tokenizer(formatted_prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to("cpu")
 
     generated_ids = model.generate(
         **inputs,
@@ -100,16 +84,16 @@ def generate_query(question):
     try:
         sql_code = output.split("```sql")[1].split("```")[0].strip()
         return sqlparse.format(sql_code, reindent=True)
-    except:
-        return "SQL could not be parsed. Raw Output:\n\n" + output
+    except Exception:
+        return "❌ SQL could not be parsed. Raw Output:\n\n" + output
 
-# Gradio Interface
+# Gradio UI
 iface = gr.Interface(
     fn=generate_query,
-    inputs=gr.Textbox(lines=3, placeholder="Enter your natural language question..."),
+    inputs=gr.Textbox(lines=3, placeholder="Ask your SQL question..."),
     outputs="text",
-    title="LLaMA 3 SQLCoder 🦙",
-    description="Enter a natural language question and get a SQL query based on predefined tables.",
+    title="🦙 LLaMA 3 SQLCoder (CPU)",
+    description="Convert natural language into SQL queries based on the given schema. Running on CPU – may be slow.",
 )
 
 if __name__ == "__main__":