Spaces:
Paused
Paused
| import gradio as gr | |
| import torch | |
| import re | |
| import tempfile | |
| import pandas as pd | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline | |
| from database import init_database, get_schema, execute_query | |
| # Model Setup | |
| MODEL_ID = "microsoft/tapex-large-sql-execution" | |
| tokenizer = None | |
| sql_pipeline = None | |
| def load_model(): | |
| global tokenizer, sql_pipeline | |
| print("microsoft/tapex-large-sql-execution ...") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_ID, | |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, | |
| device_map="auto", | |
| trust_remote_code=True, | |
| ) | |
| sql_pipeline = pipeline( | |
| "text-generation", | |
| model=model, | |
| tokenizer=tokenizer, | |
| max_new_tokens=512, | |
| do_sample=False, | |
| return_full_text=False, | |
| pad_token_id=tokenizer.eos_token_id, | |
| ) | |
| print("Model loaded.") | |
| PROMPT_TEMPLATE = """### Task | |
| Generate a SQL query to answer [QUESTION]{question}[/QUESTION] | |
| ### Database Schema | |
| The query will run on a database with the following schema: | |
| {schema} | |
| ### Answer | |
| Given the database schema, here is the SQL query that [QUESTION]{question}[/QUESTION] | |
| [SQL] | |
| """ | |
| def build_prompt(question: str, schema: str) -> str: | |
| return PROMPT_TEMPLATE.format(question=question, schema=schema) | |
| def extract_sql(raw: str) -> str: | |
| match = re.search(r"(SELECT[\s\S]+?);", raw, re.IGNORECASE) | |
| if match: | |
| return match.group(0).strip() | |
| return raw.strip().split("[/SQL]")[0].strip() | |
| def nl_to_sql_and_run(question: str, history: list): | |
| if not question.strip(): | |
| yield history, "", gr.update(visible=False), gr.update(visible=False) | |
| return | |
| schema = get_schema() | |
| prompt = build_prompt(question, schema) | |
| yield history, "Generating SQL query...", gr.update(visible=False), gr.update(visible=False) | |
| try: | |
| output = sql_pipeline(prompt)[0]["generated_text"] | |
| sql = extract_sql(output) | |
| except Exception as e: | |
| new_hist = history + [{"role": "user", "content": question}, | |
| {"role": "assistant", "content": f"Model error: {e}"}] | |
| yield new_hist, "", gr.update(visible=False), gr.update(visible=False) | |
| return | |
| yield history, f"```sql\n{sql}\n```\n\nExecuting...", gr.update(visible=False), gr.update(visible=False) | |
| try: | |
| columns, rows = execute_query(sql) | |
| except Exception as e: | |
| answer = f"**Generated SQL:**\n```sql\n{sql}\n```\n\nExecution error: `{e}`" | |
| new_hist = history + [{"role": "user", "content": question}, | |
| {"role": "assistant", "content": answer}] | |
| yield new_hist, "", gr.update(visible=False), gr.update(visible=False) | |
| return | |
| if not rows: | |
| result_md = "*(query returned no rows)*" | |
| df = pd.DataFrame() | |
| csv_path = None | |
| else: | |
| df = pd.DataFrame(rows, columns=columns) | |
| result_md = df.to_markdown(index=False) | |
| tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv", mode="w", newline="") | |
| df.to_csv(tmp.name, index=False) | |
| tmp.close() | |
| csv_path = tmp.name | |
| row_label = "rows" if len(rows) != 1 else "row" | |
| answer = f"**Generated SQL:**\n```sql\n{sql}\n```\n\n**Results ({len(rows)} {row_label}):**\n{result_md}" | |
| new_hist = history + [{"role": "user", "content": question}, | |
| {"role": "assistant", "content": answer}] | |
| yield ( | |
| new_hist, | |
| "", | |
| gr.update(value=df, visible=bool(rows)), | |
| gr.update(value=csv_path, visible=bool(rows)), | |
| ) | |
| def view_schema(): | |
| return f"```sql\n{get_schema()}\n```" | |
| CSS = """ | |
| @import url('https://fonts.googleapis.com/css2?family=Space+Mono:wght@400;700&family=DM+Sans:wght@300;400;500&display=swap'); | |
| body, .gradio-container { | |
| background: #0d0f14 !important; | |
| font-family: 'DM Sans', sans-serif; | |
| color: #e2e8f0; | |
| } | |
| .title-block { | |
| text-align: center; | |
| padding: 2rem 0 1rem; | |
| } | |
| .title-block h1 { | |
| font-size: 2rem; | |
| background: linear-gradient(135deg, #38bdf8, #818cf8); | |
| -webkit-background-clip: text; | |
| -webkit-text-fill-color: transparent; | |
| font-family: 'Space Mono', monospace; | |
| margin-bottom: 0.3rem; | |
| } | |
| .title-block p { color: #64748b; font-size: 0.95rem; } | |
| .badge { | |
| display: inline-block; | |
| background: #1e2535; | |
| border: 1px solid #2d3748; | |
| border-radius: 20px; | |
| padding: 2px 12px; | |
| font-size: 0.75rem; | |
| color: #94a3b8; | |
| margin: 4px; | |
| font-family: 'Space Mono', monospace; | |
| } | |
| """ | |
| EXAMPLE_QUERIES = [ | |
| "Show me all employees in Engineering with salary above 120000", | |
| "Which department has the highest total salary budget?", | |
| "List all active projects with their budgets", | |
| "Who are the top 3 sales performers by total amount?", | |
| "How many employees are in each department?", | |
| "Show me all sales made in the East region in 2024", | |
| ] | |
| def create_app(): | |
| init_database() | |
| with gr.Blocks(css=CSS, title="SQLCoder Studio") as demo: | |
| gr.HTML(""" | |
| <div class="title-block"> | |
| <h1>SQLCoder Studio</h1> | |
| <p>Natural language to SQL to Results | Powered by microsoft/tapex-large-sql-execution</p> | |
| <div style="margin-top:0.8rem"> | |
| <span class="badge">employees</span> | |
| <span class="badge">departments</span> | |
| <span class="badge">projects</span> | |
| <span class="badge">sales</span> | |
| </div> | |
| </div> | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| chatbot = gr.Chatbot( | |
| label="Conversation", | |
| height=460, | |
| show_label=False, | |
| render_markdown=True, | |
| bubble_full_width=False, | |
| type="messages", | |
| ) | |
| with gr.Row(): | |
| question_input = gr.Textbox( | |
| placeholder="Ask anything about the database...", | |
| show_label=False, | |
| scale=5, | |
| lines=1, | |
| ) | |
| submit_btn = gr.Button("RUN", variant="primary", scale=1) | |
| with gr.Row(): | |
| clear_btn = gr.Button("Clear chat", variant="secondary", size="sm") | |
| gr.HTML("<p style='color:#475569;font-size:0.78rem;margin-top:0.5rem'>Try an example:</p>") | |
| example_btns = [] | |
| with gr.Row(wrap=True): | |
| for eq in EXAMPLE_QUERIES: | |
| b = gr.Button(eq, size="sm", variant="secondary") | |
| example_btns.append(b) | |
| with gr.Column(scale=2): | |
| gr.HTML("<p style='color:#94a3b8;font-size:0.85rem;font-weight:500;margin-bottom:4px'>Result Table</p>") | |
| result_table = gr.Dataframe( | |
| visible=False, | |
| wrap=True, | |
| height=220, | |
| ) | |
| download_file = gr.File( | |
| label="Download CSV", | |
| visible=False, | |
| ) | |
| gr.HTML("<p style='color:#94a3b8;font-size:0.85rem;font-weight:500;margin:1rem 0 4px'>Database Schema</p>") | |
| gr.Markdown(value=view_schema()) | |
| status_md = gr.Markdown(visible=False) | |
| history_state = gr.State([]) | |
| def run(question, history): | |
| gen = nl_to_sql_and_run(question, history) | |
| for h, status, table_update, dl_update in gen: | |
| yield h, h, status, table_update, dl_update | |
| submit_btn.click( | |
| fn=run, | |
| inputs=[question_input, history_state], | |
| outputs=[chatbot, history_state, status_md, result_table, download_file], | |
| ) | |
| question_input.submit( | |
| fn=run, | |
| inputs=[question_input, history_state], | |
| outputs=[chatbot, history_state, status_md, result_table, download_file], | |
| ) | |
| clear_btn.click( | |
| fn=lambda: ([], [], "", gr.update(visible=False), gr.update(visible=False)), | |
| outputs=[chatbot, history_state, status_md, result_table, download_file], | |
| ) | |
| for btn, eq in zip(example_btns, EXAMPLE_QUERIES): | |
| btn.click(fn=lambda q=eq: q, outputs=[question_input]) | |
| return demo | |
| if __name__ == "__main__": | |
| load_model() | |
| app = create_app() | |
| app.launch() |