Update app.py

app.py

@@ -2,51 +2,53 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Load the DeepSeek model
-model_name = "deepseek-ai/DeepSeek-V3"  # Or "deepseek-ai/DeepSeek-R1-0528"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True)
+model_path = "defog/sqlcoder-7b-2"
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.float16, device_map="auto")
+
 
-# Function to handle JSON prompt for SQL generation
 def generate_sql(payload):
+
     question = payload.get("question", "")
     schema = payload.get("schema", "")
     sample_rows = payload.get("sample_rows", [])
     
     sample_str = "\n".join([str(row) for row in sample_rows]) if sample_rows else ""
-    
+
     prompt = f"""
-You are a text-to-SQL data analyst.
-Based on the following information, write a clean SQL query that works with DuckDB. Do not hallucinate tables or fields.
-Schema: {schema}
-Sample Rows:
+### Task
+Generate a SQL query to answer [QUESTION]{question}[/QUESTION]
+
+### Database Schema
+The query will run on a database with the following schema:
+{schema}
+
+### Sample Rows
 {sample_str}
-Question: {question}
 
-SQL:"""
+### Answer
+Given the database schema, here is the SQL query that [QUESTION]{question}[/QUESTION]
+[SQL]
+    """.strip()
 
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(
         **inputs,
-        max_new_tokens=128,
-        do_sample=True,
-        temperature=0.6,
-        top_p=0.95,
-        pad_token_id=tokenizer.eos_token_id
+        max_length=512,
+        do_sample=False,
+        num_beams=4,
+        early_stopping=True
     )
-    
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    generated_sql = response.split("SQL:")[-1].strip()
-    
-    return generated_sql
+    sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return sql.split("[SQL]")[-1].strip()
+
 
-# Launch Gradio interface
 demo = gr.Interface(
     fn=generate_sql,
     inputs=gr.JSON(label="Input JSON (question, schema, sample_rows)"),
     outputs="text",
-    title="Text-to-SQL (DeepSeek)",
-    description="Use DeepSeek to convert a natural language question and schema into SQL."
+    title="SQLCoder - Text to SQL",
+    description="Enter a JSON object with 'question', 'schema', and optional 'sample_rows'. The model will generate SQL using Defog's sqlcoder-7b-2."
 )
 
-demo.launch()
+demo.launch()