import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
# ---------------- Model Setup ----------------
MODELS = {
"BM1_CS1_Syn (33M)": "withmartian/sql_interp_bm1_cs1_experiment_1.10",
"BM1_CS2_Syn (33M)": "withmartian/sql_interp_bm1_cs2_experiment_2.10",
"BM1_CS3_Syn (33M)": "withmartian/sql_interp_bm1_cs3_experiment_3.10",
"BM1_CS4_Syn (33M)": "withmartian/sql_interp_bm1_cs4_dataset_synonyms_experiment_1.1",
"BM1_CS5_Syn (33M)": "withmartian/sql_interp_bm1_cs5_dataset_synonyms_experiment_1.2",
"BM2_CS1_Syn (0.5B)": "withmartian/sql_interp_bm2_cs1_experiment_4.3",
"BM2_CS2_Syn (0.5B)": "withmartian/sql_interp_bm2_cs2_experiment_5.3",
"BM2_CS3_Syn (0.5B)": "withmartian/sql_interp_bm2_cs3_experiment_6.3",
"BM3_CS1_Syn (1B)": "withmartian/sql_interp_bm3_cs1_experiment_7.3",
"BM3_CS2_Syn (1B)": "withmartian/sql_interp_bm3_cs2_experiment_8.3",
"BM3_CS3_Syn (1B)": "withmartian/sql_interp_bm3_cs3_experiment_9.3",
}
model_cache = {}
def load_model(model_name):
if model_name not in model_cache:
model_id = MODELS[model_name]
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float16,
device_map="auto"
)
model_cache[model_name] = (tokenizer, model)
return model_cache[model_name]
def generate_sql(model_name, instruction, schema, max_length=256, temperature=0.7):
if not model_name or not instruction or not schema:
return "Please fill in all fields and select a model"
try:
tokenizer, model = load_model(model_name)
prompt = f"""### Instruction: {instruction}
### Context: {schema}
### Response:"""
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
**inputs,
max_length=max_length,
temperature=temperature,
do_sample=temperature > 0,
pad_token_id=tokenizer.eos_token_id
)
generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
if "### Response:" in generated:
sql = generated.split("### Response:")[-1].strip()
else:
sql = generated.strip()
return sql
except Exception as e:
return f"Error: {str(e)}"
# ---------------- Example Queries ----------------
examples = [
[
"BM1_CS1_Syn (33M)",
"Show me the name and salary from employees",
"CREATE TABLE employees (name VARCHAR(100), salary INT, department VARCHAR(100))"
],
[
"BM2_CS2_Syn (0.5B)",
"List worker earnings from highest to lowest",
"CREATE TABLE employees (name VARCHAR(100), salary INT, department VARCHAR(100))"
],
[
"BM3_CS3_Syn (1B)",
"Count how many employees in each department",
"CREATE TABLE employees (name VARCHAR(100), salary INT, department VARCHAR(100))"
],
]
# ---------------- Model Demo Function ----------------
def model_demo():
custom_css = """
:root {
--martian-orange: #FF6B4A;
--martian-bg: #0E0E0E; /* deep black background */
--martian-gray-dark: #3A3A3A;
--martian-gray-medium: #4A4A4A;
--martian-gray-light: #5A5A5A;
}
.gradio-container {
font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
background-color: var(--martian-bg) !important;
}
.header-section {
text-align: center;
padding: 3rem 2rem;
background: linear-gradient(135deg, var(--martian-gray-dark) 0%, var(--martian-gray-medium) 100%);
border-radius: 16px;
margin-bottom: 2rem;
color: white;
box-shadow: 0 4px 6px rgba(0,0,0,0.3);
}
.header-section h1 { font-size: 2.5rem; font-weight: 700; margin-bottom: 1rem; color: white; }
.header-section .subtitle { font-size: 1.2rem; opacity: 0.9; line-height: 1.6; color: white; }
.orange-accent { color: var(--martian-orange); font-weight: 600; }
.info-box { background: var(--martian-gray-dark); border-radius: 12px; padding: 1.5rem; margin: 1.5rem 0; border-left: 4px solid var(--martian-orange); color: #E0E0E0; }
.model-guide { background: var(--martian-gray-dark); border-radius: 8px; padding: 1rem; margin-top: 1rem; font-size: 0.9rem; color: #D0D0D0; }
button.primary { background: var(--martian-orange) !important; border: none !important; color: white !important; }
button.primary:hover { background: #FF5733 !important; }
label { color: #D0D0D0 !important; }
.label-wrap span { color: var(--martian-orange) !important; }
input, textarea, select { background: var(--martian-gray-medium) !important; border-color: var(--martian-gray-light) !important; color: #E0E0E0 !important; }
textarea::placeholder, input::placeholder { color: #888 !important; }
.code { background: var(--martian-gray-dark) !important; color: #E0E0E0 !important; }
"""
with gr.Blocks(css=custom_css, title="TinySQL Model Demo") as demo:
# Header
gr.HTML("""
""")
# Info box
gr.HTML("""
How it works: Select a model, describe your query in plain English, and watch the model generate SQL.
""")
with gr.Row():
with gr.Column(scale=1):
gr.Markdown("### Configuration")
model_dropdown = gr.Dropdown(
choices=list(MODELS.keys()),
value="BM2_CS2_Syn (0.5B)",
label="Model Selection",
info="Larger models = better accuracy, slower inference"
)
gr.HTML("""
BM1 (33M) - Lightning fast, simple queries
BM2 (0.5B) - Balanced performance
BM3 (1B) - Most accurate, complex queries
Dataset Complexity:
CS1: Basic SELECT-FROM
CS2: Adds ORDER BY
CS3: Aggregations
CS4: Adds WHERE filters
CS5: Multi-table JOINs
""")
with gr.Column(scale=2):
gr.Markdown("### Your Query")
instruction = gr.Textbox(
label="What do you want to know?",
placeholder="e.g., Find all employees earning more than $50,000 sorted by name",
lines=2
)
schema = gr.Textbox(
label="Database Schema",
placeholder="CREATE TABLE employees (name VARCHAR, salary INT, department VARCHAR)",
lines=3,
value="CREATE TABLE employees (name VARCHAR(100), salary INT, department VARCHAR(100))"
)
with gr.Row():
max_length = gr.Slider(64, 512, value=256, step=32, label="Max Length")
temperature = gr.Slider(0.0, 1.0, value=0.1, step=0.1, label="Temperature")
generate_btn = gr.Button("Generate SQL", variant="primary", size="lg")
output = gr.Code(label="Generated SQL Query", language="sql", lines=8)
gr.Examples(examples=examples, inputs=[model_dropdown, instruction, schema])
generate_btn.click(
fn=generate_sql,
inputs=[model_dropdown, instruction, schema, max_length, temperature],
outputs=output
)
return demo