with data

app.py

@@ -1,443 +1,10 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from tinysql_model_demo import model_demo
+from tinysql_dataset_viewer import dataset_viewer  # your dataset viewer function
 
-# Model configurations
-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",
-}
+with gr.Blocks() as app:
+    with gr.Tab("Model Demo"):
+        model_demo()
+    with gr.Tab("Dataset Viewer"):
+        dataset_viewer()
 
-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))"
-    ],
-]
-
-# Custom CSS with Martian colors (Orange/Black/Dark Gray only)
-custom_css = """
-:root {
-    --martian-orange: #FF6B4A;
-    --martian-dark: #1A1A1A;
-    --martian-gray-dark: #3A3A3A;
-    --martian-gray-medium: #4A4A4A;
-    --martian-gray-light: #5A5A5A;
-    --martian-bg: #2A2A2A;
-}
-
-.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-dark) 0%, var(--martian-gray-dark) 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;
-}
-
-/* Remove all purple/blue colors from Gradio components */
-.primary.svelte-cmf5ev {
-    background: var(--martian-orange) !important;
-    border-color: var(--martian-orange) !important;
-}
-
-button.primary {
-    background: var(--martian-orange) !important;
-    border: none !important;
-    color: white !important;
-}
-
-button.primary:hover {
-    background: #FF5733 !important;
-}
-
-/* Fix label colors */
-label {
-    color: #D0D0D0 !important;
-}
-
-.label-wrap span {
-    color: var(--martian-orange) !important;
-}
-
-/* Input fields - dark theme */
-.input-text, textarea, select, input {
-    background: var(--martian-gray-medium) !important;
-    border-color: var(--martian-gray-light) !important;
-    color: #E0E0E0 !important;
-}
-
-textarea::placeholder, input::placeholder {
-    color: #888 !important;
-}
-
-/* Slider colors */
-input[type="range"] {
-    background: var(--martian-gray-medium) !important;
-}
-
-input[type="range"]::-webkit-slider-thumb {
-    background: var(--martian-orange) !important;
-}
-
-input[type="range"]::-moz-range-thumb {
-    background: var(--martian-orange) !important;
-}
-
-input[type="range"]::-webkit-slider-runnable-track {
-    background: var(--martian-gray-light) !important;
-}
-
-/* Citation box - medium gray with light text */
-.citation-box {
-    background: var(--martian-gray-medium);
-    border: 1px solid var(--martian-gray-light);
-    border-radius: 12px;
-    padding: 1.5rem;
-    margin: 2rem 0;
-    font-family: 'Monaco', 'Courier New', monospace;
-    font-size: 0.85rem;
-}
-
-.citation-header {
-    font-weight: 700;
-    color: #E0E0E0;
-    margin-bottom: 1rem;
-    font-size: 1.1rem;
-}
-
-.citation-box pre {
-    color: #D0D0D0;
-    background: transparent;
-}
-
-.resource-links {
-    display: flex;
-    gap: 1rem;
-    justify-content: center;
-    margin: 2rem 0;
-    flex-wrap: wrap;
-}
-
-.resource-link {
-    background: var(--martian-gray-dark);
-    color: white;
-    padding: 0.75rem 1.5rem;
-    border-radius: 8px;
-    text-decoration: none;
-    font-weight: 500;
-    transition: all 0.3s ease;
-    border: 2px solid var(--martian-gray-dark);
-}
-
-.resource-link:hover {
-    background: var(--martian-orange);
-    border-color: var(--martian-orange);
-    transform: translateY(-2px);
-    box-shadow: 0 4px 8px rgba(255, 107, 74, 0.3);
-}
-
-footer {
-    text-align: center;
-    padding: 2rem 0;
-    color: #999;
-    border-top: 1px solid var(--martian-gray-dark);
-    margin-top: 3rem;
-    font-size: 0.9rem;
-    background: var(--martian-bg);
-}
-
-/* Remove light backgrounds everywhere */
-.block, .panel {
-    background: var(--martian-gray-dark) !important;
-}
-
-.form {
-    background: var(--martian-gray-medium) !important;
-}
-
-/* Dropdown styling */
-.dropdown {
-    background: var(--martian-gray-medium) !important;
-    color: #E0E0E0 !important;
-}
-
-/* Code output styling */
-.code {
-    background: var(--martian-gray-dark) !important;
-    color: #E0E0E0 !important;
-}
-
-/* Examples section */
-.example {
-    background: var(--martian-gray-medium) !important;
-    border-color: var(--martian-gray-light) !important;
-}
-
-/* Markdown sections */
-.markdown {
-    color: #D0D0D0 !important;
-}
-
-h1, h2, h3, h4, h5, h6 {
-    color: #E0E0E0 !important;
-}
-"""
-# Create Gradio interface
-with gr.Blocks(css=custom_css, title="TinySQL Demo", theme=gr.themes.Soft()) as demo:
-    
-    # Header
-    gr.HTML("""
-        <div class="header-section">
-            <h1>TinySQL Interactive Demo</h1>
-            <p class="subtitle">
-                Transform natural language into SQL queries using <span class="orange-accent">mechanistically interpretable</span> models
-            </p>
-        </div>
-    """)
-    
-    # Info box
-    gr.HTML("""
-        <div class="info-box">
-            <strong>How it works:</strong> Select a model from our collection of 11 fine-tuned transformers, 
-            describe what you want in plain English, and watch as the model generates precise SQL queries. 
-            Each model is trained on progressively complex SQL operations—from basic SELECT statements to 
-            advanced JOINs and aggregations.
-        </div>
-    """)
-    
-    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("""
-                <div class="model-guide">
-                    <strong>Model Guide:</strong><br><br>
-                    <strong>BM1 (33M)</strong> - Lightning fast, great for simple queries<br>
-                    <strong>BM2 (0.5B)</strong> - Balanced performance and speed<br>
-                    <strong>BM3 (1B)</strong> - Most accurate, handles complex queries<br><br>
-                    <strong>Dataset Complexity:</strong><br>
-                    CS1: Basic SELECT-FROM queries<br>
-                    CS2: Adds ORDER BY clauses<br>
-                    CS3: Aggregations (COUNT, SUM, AVG)<br>
-                    CS4: Adds WHERE filters<br>
-                    CS5: Multi-table JOINs
-                </div>
-            """)
-            
-        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(
-                    minimum=64,
-                    maximum=512,
-                    value=256,
-                    step=32,
-                    label="Max Length",
-                    info="Longer = more complex queries"
-                )
-                temperature = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    value=0.1,
-                    step=0.1,
-                    label="Temperature",
-                    info="Higher = more creative (use 0.1 for accuracy)"
-                )
-            
-            generate_btn = gr.Button("Generate SQL", variant="primary", size="lg", elem_classes="primary-button")
-            
-            output = gr.Code(
-                label="Generated SQL Query",
-                language="sql",
-                lines=8,
-            )
-    
-    gr.Markdown("### Example Queries")
-    gr.Examples(
-        examples=examples,
-        inputs=[model_dropdown, instruction, schema],
-    )
-    
-    # Resource links
-    gr.HTML("""
-        <div class="resource-links">
-            <a href="https://arxiv.org/abs/2503.12730" class="resource-link" target="_blank">
-                Read the Paper
-            </a>
-            <a href="https://github.com/withmartian/TinySQL" class="resource-link" target="_blank">
-                View Code
-            </a>
-            <a href="https://huggingface.co/collections/withmartian/tinysql-6760e92748b63fa56a6ffc9f" class="resource-link" target="_blank">
-                Dataset & Models
-            </a>
-            <a href="https://withmartian.com" class="resource-link" target="_blank">
-                Martian
-            </a>
-        </div>
-    """)
-    
-    # Citation box
-    gr.HTML("""
-        <div class="citation-box">
-            <div class="citation-header">Citation</div>
-            <pre style="margin: 0; overflow-x: auto; background: transparent;">@misc{harrasse2025tinysqlprogressivetexttosqldataset,
-    title={TinySQL: A Progressive Text-to-SQL Dataset for Mechanistic Interpretability Research}, 
-    author={Abir Harrasse and Philip Quirke and Clement Neo and Dhruv Nathawani and Luke Marks and Amir Abdullah},
-    year={2025},
-    eprint={2503.12730},
-    archivePrefix={arXiv},
-    primaryClass={cs.LG},
-    url={https://arxiv.org/abs/2503.12730}
-}</pre>
-        </div>
-    """)
-
-        # Footer
-    gr.HTML("""
-        <footer>
-            <p>Brought to you with ❤️ from the Martian science team</p>
-            <p style="margin-top: 0.5rem;">Bridging the gap between toy tasks and real-world interpretability</p>
-        </footer>
-    """)
-    
-    generate_btn.click(
-        fn=generate_sql,
-        inputs=[model_dropdown, instruction, schema, max_length, temperature],
-        outputs=output
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+app.launch()

tinysql_dataset_viewer.py

@@ -0,0 +1,153 @@
+# tinysql_dataset_viewer.py
+import gradio as gr
+from datasets import load_dataset
+import pandas as pd
+import urllib.parse
+import html
+import traceback
+
+HF_DATASETS = {
+    "CS1": "withmartian/cs1_dataset",
+    "CS2": "withmartian/cs2_dataset",
+    "CS3": "withmartian/cs3_dataset",
+    "CS2_synonyms": "withmartian/cs2_dataset_synonyms",
+    "CS3_synonyms": "withmartian/cs3_dataset_synonyms",
+    "CS4_synonyms": "withmartian/cs4_dataset_synonyms",
+}
+
+DEMO_URL = "https://huggingface.co/spaces/abir-hr196/tinysql-demo"
+PREVIEW_LIMIT = 500
+FIELDS = ["english_prompt", "create_statement", "sql_statement"]
+dataset_cache = {}
+
+# ---------------- Helpers ----------------
+def load_preview(dataset_id, limit=PREVIEW_LIMIT):
+    try:
+        ds = load_dataset(dataset_id, split=f"train[:{limit}]")
+    except Exception:
+        full = load_dataset(dataset_id)
+        first_split = list(full.keys())[0] if isinstance(full, dict) else None
+        if first_split:
+            ds = full[first_split].select(range(min(len(full[first_split]), limit)))
+        else:
+            ds = full.select(range(min(len(full), limit)))
+    df = pd.DataFrame(ds)
+    for f in FIELDS:
+        if f not in df.columns:
+            df[f] = ""
+    df = df[FIELDS].copy()
+    df.reset_index(inplace=True)
+    df.rename(columns={"index": "example_index"}, inplace=True)
+    return df
+
+def get_dataset_preview(name):
+    if name in dataset_cache:
+        return dataset_cache[name]
+    df = load_preview(HF_DATASETS[name])
+    dataset_cache[name] = df
+    return df
+
+def make_dropdown_options(df):
+    opts = []
+    for _, row in df.iterrows():
+        idx = int(row["example_index"])
+        prompt = (row["english_prompt"] or "")
+        short = " ".join(prompt.split())[:120] + ("…" if len(prompt) > 120 else "")
+        opts.append((f"{idx} — {short}", idx))
+    return opts
+
+def filter_dataframe(df, query):
+    if not query:
+        return df
+    q = str(query).lower()
+    mask = df["english_prompt"].fillna("").str.lower().str.contains(q) | df["sql_statement"].fillna("").str.lower().str.contains(q)
+    return df[mask].reset_index(drop=True)
+
+# ---------------- Gradio callbacks ----------------
+def on_dataset_change(dataset_name):
+    try:
+        df = get_dataset_preview(dataset_name)
+        displayed = df[["example_index", "english_prompt", "sql_statement", "create_statement"]]
+        opts = make_dropdown_options(displayed)
+        return displayed, gr.Dropdown.update(choices=opts, value=opts[0][1] if opts else None), ""
+    except Exception as e:
+        tb = traceback.format_exc()
+        return pd.DataFrame([], columns=["id", "english_prompt", "sql_statement", "create_statement"]), gr.Dropdown.update(choices=[], value=None), f"Error loading dataset: {e}\n{tb}"
+
+def on_search(dataset_name, query):
+    try:
+        df = get_dataset_preview(dataset_name)
+        filtered = filter_dataframe(df, query)
+        displayed = filtered[["example_index", "english_prompt", "sql_statement", "create_statement"]]
+        opts = make_dropdown_options(displayed)
+        return displayed, gr.Dropdown.update(choices=opts, value=opts[0][1] if opts else None)
+    except Exception as e:
+        return pd.DataFrame([], columns=["id", "english_prompt", "sql_statement", "create_statement"]), gr.Dropdown.update(choices=[], value=None)
+
+def send_to_demo(dataset_name, selected_index):
+    try:
+        df = get_dataset_preview(dataset_name)
+        row = df[df["example_index"] == int(selected_index)]
+        if row.empty:
+            return html.escape("Selected example not found.")
+        instr = str(row.iloc[0]["english_prompt"] or "")
+        schema = str(row.iloc[0]["create_statement"] or "")
+        q_instr = urllib.parse.quote_plus(instr)
+        q_schema = urllib.parse.quote_plus(schema)
+        url = f"{DEMO_URL}?instruction={q_instr}&schema={q_schema}"
+        safe_url = html.escape(url, quote=True)
+        html_out = f"""
+        <script>
+            window.open("{safe_url}", "_blank");
+        </script>
+        <div style="color: #E0E0E0; font-family: Inter, sans-serif;">
+            Opened the demo in a new tab. If your browser blocked the popup, <a href="{safe_url}" target="_blank" rel="noreferrer">click here</a>.
+        </div>
+        """
+        return gr.HTML.update(value=html_out)
+    except Exception as e:
+        tb = traceback.format_exc()
+        return gr.HTML.update(value=f"<pre style='color:#ffb3a7'>Error: {html.escape(str(e))}\n{html.escape(tb)}</pre>")
+
+# ---------------- Dataset viewer function ----------------
+def dataset_viewer():
+    custom_css = """
+    :root {
+        --martian-orange: #FF6B4A;
+        --martian-dark: #0E0E0E;
+        --martian-gray-dark: #1A1A1A;
+        --martian-gray-medium: #2A2A2A;
+        --martian-gray-light: #3A3A3A;
+        --martian-bg: #0E0E0E;
+    }
+    .gradio-container { background-color: var(--martian-bg) !important; font-family: 'Inter', sans-serif; }
+    .header-section { text-align: center; padding: 2rem; background: linear-gradient(135deg, var(--martian-dark) 0%, var(--martian-gray-dark) 100%); border-radius: 12px; margin-bottom: 1rem; color:white;}
+    .header-section h1 { font-size: 2rem; margin-bottom: 0.5rem; }
+    .info-box { background: var(--martian-gray-dark); border-left: 4px solid var(--martian-orange); border-radius: 10px; padding:1rem; margin:1rem 0; color:#E0E0E0;}
+    button, .gr-button { background: var(--martian-orange) !important; color:white !important; border:none !important;}
+    .input-text, textarea, select, input, .gradio-dataframe { background: var(--martian-gray-medium) !important; border-color: var(--martian-gray-light) !important; color: #E0E0E0 !important; }
+    a { color: var(--martian-orange) !important; }
+    """
+
+    with gr.Blocks(css=custom_css) as viewer:
+        gr.HTML("""<div class="header-section"><h1>TinySQL — Dataset Viewer</h1><div class="subtitle">Browse dataset variants, filter examples, and send a selected example to the TinySQL model demo.</div></div>""")
+        gr.HTML("""<div class="info-box"><strong>Note:</strong> Previews load the first 500 examples for fast exploration. Use the search box to filter prompts or SQL statements.</div>""")
+        with gr.Row():
+            with gr.Column(scale=1):
+                dataset_dropdown = gr.Dropdown(choices=list(HF_DATASETS.keys()), value=list(HF_DATASETS.keys())[0], label="Dataset Variant")
+                search_box = gr.Textbox(label="Search (prompt or SQL)", placeholder="Type keywords to filter prompts or SQL...")
+                select_dropdown = gr.Dropdown(choices=[], label="Select example to try")
+                try_button = gr.Button("Try in Model Demo", variant="primary")
+                status_html = gr.HTML("")
+            with gr.Column(scale=3):
+                df_display = gr.Dataframe(
+                    headers=["id", "english_prompt", "sql_statement", "create_statement"],
+                    value=pd.DataFrame(columns=["id", "english_prompt", "sql_statement", "create_statement"]),
+                    label="Preview (first 500 rows)"
+                )
+
+        dataset_dropdown.change(fn=on_dataset_change, inputs=[dataset_dropdown], outputs=[df_display, select_dropdown, status_html])
+        search_box.change(fn=on_search, inputs=[dataset_dropdown, search_box], outputs=[df_display, select_dropdown])
+        try_button.click(fn=send_to_demo, inputs=[dataset_dropdown, select_dropdown], outputs=[status_html])
+
+    return viewer

tinysql_model_demo.py

@@ -0,0 +1,199 @@
+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("""
+            <div class="header-section">
+                <h1>TinySQL Interactive Demo</h1>
+                <p class="subtitle">
+                    Transform natural language into SQL queries using <span class="orange-accent">mechanistically interpretable</span> models
+                </p>
+            </div>
+        """)
+        
+        # Info box
+        gr.HTML("""
+            <div class="info-box">
+                <strong>How it works:</strong> Select a model, describe your query in plain English, and watch the model generate SQL.
+            </div>
+        """)
+        
+        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("""
+                    <div class="model-guide">
+                        <strong>BM1 (33M)</strong> - Lightning fast, simple queries<br>
+                        <strong>BM2 (0.5B)</strong> - Balanced performance<br>
+                        <strong>BM3 (1B)</strong> - Most accurate, complex queries<br><br>
+                        <strong>Dataset Complexity:</strong><br>
+                        CS1: Basic SELECT-FROM<br>
+                        CS2: Adds ORDER BY<br>
+                        CS3: Aggregations<br>
+                        CS4: Adds WHERE filters<br>
+                        CS5: Multi-table JOINs
+                    </div>
+                """)
+
+            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