import re import os import gradio as gr from transformers import AutoModelForSeq2SeqLM, AutoTokenizer import torch import time MODEL_ID = "RealMati/t2sql_v6_structured" print(f"Loading model: {MODEL_ID}") tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID) model.eval() print("Model loaded.") AGG_OPS = ["", "MAX", "MIN", "COUNT", "SUM", "AVG"] AGG_LABELS = ["None", "MAX", "MIN", "COUNT", "SUM", "AVG"] OPS = ["=", ">", "<", ">=", "<=", "!="] css_path = os.path.join(os.path.dirname(__file__), "style.css") with open(css_path, "r") as f: CSS = f.read() def decode_structured_output(text): sel = agg = None conds = [] try: for part in text.strip().split(" | "): part = part.strip() if part.startswith("SEL:"): sel = int(part[4:].strip()) elif part.startswith("AGG:"): agg = int(part[4:].strip()) elif part.startswith("CONDS:"): cond_str = part[6:].strip() if cond_str: for c in cond_str.split(";"): vals = c.split(",", 2) if len(vals) >= 3: conds.append([int(vals[0]), int(vals[1]), vals[2]]) except Exception: pass return sel, agg, conds def parse_schema(schema_str): schema_str = schema_str.strip() if not schema_str: return "table", [] first_table = schema_str.split("|")[0].strip() if ":" in first_table: table_name, cols_str = first_table.split(":", 1) table_name = table_name.strip() columns = [c.strip() for c in cols_str.split(",") if c.strip()] else: table_name = "table" columns = [c.strip() for c in first_table.split(",") if c.strip()] return table_name, columns def quote_col(name): return f"`{name}`" if " " in name else name def structured_to_sql(sel, agg, conds, columns, table_name="table"): if sel is None or agg is None: return None col_name = quote_col(columns[sel] if sel < len(columns) else f"col{sel}") if agg == 0: sql = f"SELECT {col_name} FROM {table_name}" else: agg_op = AGG_OPS[agg] if agg < len(AGG_OPS) else "" sql = f"SELECT {agg_op}({col_name}) FROM {table_name}" if conds: where_parts = [] for c_idx, c_op, c_val in conds: c_name = quote_col(columns[c_idx] if c_idx < len(columns) else f"col{c_idx}") op_str = OPS[c_op] if c_op < len(OPS) else "=" try: float(c_val) val_sql = c_val except (ValueError, TypeError): val_sql = f"'{c_val}'" where_parts.append(f"{c_name} {op_str} {val_sql}") if where_parts: sql += " WHERE " + " AND ".join(where_parts) return sql def format_parsed(sel, agg, conds, columns): parts = [] if sel is not None and sel < len(columns): parts.append(f"Column: {columns[sel]} (index {sel})") elif sel is not None: parts.append(f"Column index: {sel}") if agg is not None: agg_label = AGG_LABELS[agg] if agg < len(AGG_LABELS) else str(agg) parts.append(f"Aggregation: {agg_label}") if conds: cond_strs = [] for c_idx, c_op, c_val in conds: c_name = columns[c_idx] if c_idx < len(columns) else f"col{c_idx}" op_str = OPS[c_op] if c_op < len(OPS) else "=" cond_strs.append(f"{c_name} {op_str} {c_val}") parts.append(f"Conditions: {' AND '.join(cond_strs)}") else: parts.append("Conditions: None") return " | ".join(parts) def predict(question, schema, num_beams, max_length): if not question or not question.strip(): return ( "-- Enter a question and schema, then click Generate SQL", "Waiting for input...", "No query submitted yet", "", ) table_name, columns = parse_schema(schema) if not columns: return ( "-- Please provide a schema\n-- Format: table_name: col1, col2, col3", "Schema required", "Cannot map indices without column names", "", ) input_text = f"translate to SQL: {question}" if schema.strip(): input_text += f" | schema: {schema.strip()}" inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True) t0 = time.time() with torch.no_grad(): outputs = model.generate( **inputs, max_length=int(max_length), num_beams=int(num_beams), early_stopping=True, do_sample=False, ) latency = time.time() - t0 raw_output = tokenizer.decode(outputs[0], skip_special_tokens=True) sel, agg, conds = decode_structured_output(raw_output) if sel is not None and agg is not None and columns: sql = structured_to_sql(sel, agg, conds, columns, table_name) else: sql = f"-- Could not parse model output\n-- Raw: {raw_output}" parsed = format_parsed(sel, agg, conds, columns) if sel is not None else "Parse failed" perf = f"Inference: {latency:.2f}s | Beams: {int(num_beams)} | Tokens: {inputs['input_ids'].shape[1]}" return sql, raw_output, parsed, perf theme = gr.themes.Base( primary_hue="blue", secondary_hue="purple", neutral_hue="gray", font=gr.themes.GoogleFont("Inter"), font_mono=gr.themes.GoogleFont("Fira Code"), ).set( body_background_fill="#0d1117", body_text_color="#e2e8f0", block_background_fill="#161b22", block_border_color="#1f2937", block_border_width="1px", block_label_text_color="#d1d5db", block_title_text_color="#f3f4f6", block_radius="12px", block_shadow="none", input_background_fill="#111827", input_border_color="#1f2937", input_border_width="1px", input_placeholder_color="#4b5563", input_radius="8px", slider_color="#3b82f6", button_primary_background_fill="linear-gradient(135deg, #3b82f6, #8b5cf6)", button_primary_text_color="#ffffff", button_secondary_background_fill="#111827", button_secondary_text_color="#d1d5db", button_secondary_border_color="#1f2937", border_color_primary="#1f2937", color_accent_soft="#111827", ) with gr.Blocks(title="Text-to-SQL | T5 on WikiSQL") as demo: # Compact header — one line title + badges + pipeline gr.HTML("""

Text-to-SQL

T5-base (220M) Seq2Seq WikiSQL 80K+ Structured Output
Question → T5 Encoder-Decoder → SEL | AGG | CONDS → SQL
""") with gr.Tabs(): # ══ TAB 1: INFERENCE (main focus) ══ with gr.Tab("Demo"): with gr.Row(equal_height=False): with gr.Column(scale=1): question = gr.Textbox( label="Natural Language Question", placeholder="e.g. What is terrence ross' nationality?", lines=2, ) schema = gr.Textbox( label="Database Schema", placeholder="table_name: col1, col2, col3, ...", lines=2, ) gr.HTML('

Format: table: col1, col2, col3 — column order = index mapping

') with gr.Row(): beams = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Beam Size") max_len = gr.Slider(minimum=64, maximum=512, value=256, step=64, label="Max Length") btn = gr.Button("Generate SQL", variant="primary", elem_classes=["generate-btn"], size="lg") with gr.Column(scale=1): sql_out = gr.Textbox( label="Generated SQL", value="-- Enter a question and schema, then click Generate SQL", lines=3, elem_classes=["sql-output"], ) raw_out = gr.Textbox(label="Raw Structured Tokens", value="Waiting for input...", lines=1, elem_classes=["decode-box"]) parsed_out = gr.Textbox(label="Decoded Mapping", value="No query submitted yet", lines=1, elem_classes=["decode-box"]) latency_out = gr.Textbox(label="Performance", value="", lines=1, elem_classes=["decode-box"]) btn.click(fn=predict, inputs=[question, schema, beams, max_len], outputs=[sql_out, raw_out, parsed_out, latency_out]) question.submit(fn=predict, inputs=[question, schema, beams, max_len], outputs=[sql_out, raw_out, parsed_out, latency_out]) gr.Markdown("#### Examples") gr.Examples( examples=[ ["What is terrence ross' nationality", "players: Player, No., Nationality, Position, Years in Toronto, School/Club Team", 5, 256], ["how many schools or teams had jalen rose", "players: Player, No., Nationality, Position, Years in Toronto, School/Club Team", 5, 256], ["What was the date of the race in Misano?", "races: No, Date, Round, Circuit, Pole Position, Fastest Lap, Race winner, Report", 5, 256], ["What was the number of race that Kevin Curtain won?", "races: No, Date, Round, Circuit, Pole Position, Fastest Lap, Race winner, Report", 5, 256], ["Where was Assen held?", "races: No, Date, Round, Circuit, Pole Position, Fastest Lap, Race winner, Report", 5, 256], ["How many different positions did Sherbrooke Faucons (qmjhl) provide in the draft?", "draft: Pick, Player, Position, Nationality, NHL team, College/junior/club team", 5, 256], ["What are the nationalities of the player picked from Thunder Bay Flyers (ushl)", "draft: Pick, Player, Position, Nationality, NHL team, College/junior/club team", 5, 256], ["How many different nationalities do the players of New Jersey Devils come from?", "draft: Pick, Player, Position, Nationality, NHL team, College/junior/club team", 5, 256], ["What's Dorain Anneck's pick number?", "draft: Pick, Player, Position, Nationality, NHL team, College/junior/club team", 5, 256], ], inputs=[question, schema, beams, max_len], outputs=[sql_out, raw_out, parsed_out, latency_out], fn=predict, cache_examples=False, ) # ══ TAB 2: HOW IT WORKS ══ with gr.Tab("How It Works"): gr.HTML("""

Architecture

A T5-base encoder-decoder fine-tuned on WikiSQL. Instead of generating raw SQL, it outputs structured tokens — column indices and operator codes — which a deterministic decoder maps to actual SQL using the provided schema.

1. Input Encoding

Question + schema concatenated:

translate to SQL: {question} | schema: {table}: {col1}, {col2}

Column order matters — the model references columns by 0-based index.

2. T5 Generation

The encoder processes input, decoder generates structured tokens via beam search.

Output: SEL:{col} | AGG:{agg} | CONDS:{col},{op},{val}

3. Structured Decoding

4. SQL Assembly

Indices mapped back to column names from schema. Operators converted to SQL. Result: a valid, executable query.

Why Structured Output?

Encoding Reference

ComponentIndexMeaning
AGG0No aggregation
1MAX
2MIN
3COUNT
4SUM
5AVG
OP0= (equals)
1> (greater than)
2< (less than)
3>= (greater or equal)
4<= (less or equal)
5!= (not equal)
""") # ══ TAB 3: MODEL INFO ══ with gr.Tab("Model & Training"): gr.HTML("""
220M
Parameters
80K+
Training Examples
T5-base
Architecture
WikiSQL
Dataset

Model

Training

WikiSQL Dataset

80,654 hand-annotated SQL queries across 24,241 Wikipedia tables. Single-table queries with SELECT, aggregation, and WHERE conditions.

github.com/salesforce/WikiSQL

Limitations

""") gr.HTML("""
Model  •  WikiSQL  •  Built with Transformers & Gradio
""") demo.launch(theme=theme, css=CSS)