updates

app.py

@@ -6,6 +6,7 @@ custom_css = """
 :root {
     --martian-orange: #FF6B4A;
     --martian-black: #0A0A0A;
+    --martian-gray-dark: #1A1A1A;
 }
 
 .gradio-container {
@@ -26,11 +27,16 @@ custom_css = """
 """
 
 with gr.Blocks(css=custom_css, title="TinySQL Demo") as demo:
+    # Shared state for passing data between tabs
+    shared_instruction = gr.State("")
+    shared_schema = gr.State("")
+    
     with gr.Tabs():
-        with gr.Tab("Model Demo"):
-            model_demo()
         with gr.Tab("Dataset Viewer"):
-            dataset_viewer()
+            viewer_components = dataset_viewer(shared_instruction, shared_schema)
+        
+        with gr.Tab("Model Demo"):
+            model_components = model_demo(shared_instruction, shared_schema)
 
 if __name__ == "__main__":
     demo.launch()

tinysql_dataset_viewer.py

@@ -18,7 +18,10 @@ def load_preview(dataset_name):
     """Load first 500 rows of selected dataset"""
     try:
         ds = load_dataset(DATASETS[dataset_name], split="train")
-        df = pd.DataFrame(ds)[COLUMNS].head(500)
+        df = pd.DataFrame(ds).head(500)
+        # Filter to only the columns we want
+        if all(col in df.columns for col in COLUMNS):
+            df = df[COLUMNS]
         return df
     except Exception as e:
         return pd.DataFrame({"Error": [str(e)]})
@@ -34,208 +37,128 @@ def filter_dataframe(df, search_query):
     )
     return df[mask]
 
-# CSS styling
-custom_css = """
-:root {
-    --martian-orange: #FF6B4A;
-    --martian-black: #0A0A0A;
-    --martian-gray-dark: #1A1A1A;
-    --martian-gray-medium: #2A2A2A;
-    --martian-gray-light: #3A3A3A;
-}
-
-.gradio-container {
-    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
-    background-color: var(--martian-black) !important;
-    color: #E0E0E0 !important;
-}
-
-.header-section {
-    text-align: center;
-    padding: 2.5rem 1.5rem;
-    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.2rem;
-    font-weight: 700;
-    margin-bottom: 0.75rem;
-}
-
-.header-section .subtitle {
-    font-size: 1.1rem;
-    opacity: 0.9;
-    line-height: 1.6;
-}
-
-.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;
-}
-
-.dataset-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;
-    font-weight: 600 !important;
-}
-
-button.primary:hover {
-    background: #FF5733 !important;
-    transform: translateY(-1px);
-    box-shadow: 0 4px 8px rgba(255, 107, 74, 0.3);
-}
-
-input, select, textarea {
-    background: var(--martian-gray-medium) !important;
-    border-color: var(--martian-gray-light) !important;
-    color: #E0E0E0 !important;
-}
-
-.dataframe {
-    background: var(--martian-gray-dark) !important;
-}
-
-label {
-    color: #D0D0D0 !important;
-}
-
-.label-wrap span {
-    color: var(--martian-orange) !important;
-}
-"""
-
-def dataset_viewer():
-    with gr.Blocks(css=custom_css, title="TinySQL Dataset Viewer") as viewer:
-        # Header
-        gr.HTML("""
-            <div class="header-section">
-                <h1>TinySQL Dataset Viewer</h1>
-                <p class="subtitle">
-                    Browse dataset previews, search, and filter queries with <span class="orange-accent">ease</span>
-                </p>
-            </div>
-        """)
-        
-        # Info box
-        gr.HTML("""
-            <div class="info-box">
-                <strong>Preview Mode:</strong> Showing first 500 rows of each dataset. Use search to filter results in real-time.
-            </div>
-        """)
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                gr.Markdown("### Dataset Selection")
-                dataset_dropdown = gr.Dropdown(
-                    choices=list(DATASETS.keys()),
-                    value="CS1",
-                    label="Choose Dataset",
-                    info="Select a dataset to preview"
-                )
-                
-                gr.HTML("""
-                    <div class="dataset-guide">
-                        <strong>Complexity Levels:</strong><br><br>
-                        <strong>CS1:</strong> Basic SELECT-FROM<br>
-                        <strong>CS2:</strong> Adds ORDER BY<br>
-                        <strong>CS3:</strong> Aggregations<br>
-                        <strong>CS4:</strong> Adds WHERE filters<br><br>
-                        <strong>Synonyms:</strong> Natural language variations
-                    </div>
-                """)
-                
-                load_btn = gr.Button("Load Dataset", variant="primary", size="lg")
-                
-                gr.HTML("<br>")
-                
-                demo_btn = gr.Button("🚀 Try Model Demo", variant="primary")
+def dataset_viewer(shared_instruction, shared_schema):
+    """Dataset viewer component with ability to send examples to model demo"""
+    
+    gr.HTML("""
+        <div class="header-section" style="text-align: center; padding: 2.5rem 1.5rem; background: linear-gradient(135deg, #1A1A1A 0%, #2A2A2A 100%); border-radius: 16px; margin-bottom: 2rem; color: white;">
+            <h1 style="font-size: 2.2rem; font-weight: 700; margin-bottom: 0.75rem;">TinySQL Dataset Viewer</h1>
+            <p style="font-size: 1.1rem; opacity: 0.9; line-height: 1.6;">
+                Browse dataset previews, search, and filter queries with <span style="color: #FF6B4A; font-weight: 600;">ease</span>
+            </p>
+        </div>
+    """)
+    
+    gr.HTML("""
+        <div class="info-box" style="background: #3A3A3A; border-radius: 12px; padding: 1.5rem; margin: 1.5rem 0; border-left: 4px solid #FF6B4A; color: #E0E0E0;">
+            <strong>Preview Mode:</strong> Showing first 500 rows of each dataset. Use search to filter results in real-time.
+        </div>
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### Dataset Selection")
+            dataset_dropdown = gr.Dropdown(
+                choices=list(DATASETS.keys()),
+                value="CS1",
+                label="Choose Dataset",
+                info="Select a dataset to preview"
+            )
             
-            with gr.Column(scale=3):
-                gr.Markdown("### Dataset Preview (First 500 Rows)")
-                
-                search_box = gr.Textbox(
-                    label="Search",
-                    placeholder="Search across all columns...",
-                    lines=1
-                )
-                
-                df_display = gr.Dataframe(
-                    headers=COLUMNS,
-                    datatype=["str", "str", "str"],
-                    interactive=False,
-                    wrap=True,
-                    label="Results",
-                    max_height=600  # Use height instead to control visible area
-                )
-                
-                stats_display = gr.Markdown("Click 'Load Dataset' to begin")
-        
-        # Store the loaded dataframe
-        df_state = gr.State(value=pd.DataFrame())
-        
-        # Load dataset
-        def load_and_display(dataset_name):
-            df = load_preview(dataset_name)
-            if "Error" in df.columns:
-                return df, df, "❌ Error loading dataset"
-            stats = f"**Loaded:** {len(df)} rows | **Columns:** {', '.join(COLUMNS)}"
-            return df, df, stats
-        
-        load_btn.click(
-            fn=load_and_display,
-            inputs=dataset_dropdown,
-            outputs=[df_state, df_display, stats_display]
-        )
+            gr.HTML("""
+                <div style="background: #3A3A3A; border-radius: 8px; padding: 1rem; margin-top: 1rem; font-size: 0.9rem; color: #D0D0D0;">
+                    <strong>Complexity Levels:</strong><br><br>
+                    <strong>CS1:</strong> Basic SELECT-FROM<br>
+                    <strong>CS2:</strong> Adds ORDER BY<br>
+                    <strong>CS3:</strong> Aggregations<br>
+                    <strong>CS4:</strong> Adds WHERE filters<br><br>
+                    <strong>Synonyms:</strong> Natural language variations
+                </div>
+            """)
+            
+            load_btn = gr.Button("Load Dataset", variant="primary", size="lg")
+            
+            row_selector = gr.Number(
+                label="Select Row to Test",
+                value=0,
+                minimum=0,
+                precision=0,
+                info="Enter row number to send to Model Demo"
+            )
+            
+            send_to_model_btn = gr.Button("🚀 Run This Example in Model Demo", variant="primary")
         
-        # Search functionality
-        def search_and_display(df, query):
-            if df.empty:
-                return df, "Load a dataset first"
+        with gr.Column(scale=3):
+            gr.Markdown("### Dataset Preview (First 500 Rows)")
+            
+            search_box = gr.Textbox(
+                label="Search",
+                placeholder="Search across all columns...",
+                lines=1
+            )
             
-            filtered_df = filter_dataframe(df, query)
-            stats = f"**Showing:** {len(filtered_df)} of {len(df)} rows"
-            if query:
-                stats += f" | **Search:** '{query}'"
-            return filtered_df, stats
+            df_display = gr.Dataframe(
+                headers=COLUMNS,
+                datatype=["str", "str", "str"],
+                interactive=False,
+                wrap=True,
+                label="Results"
+            )
+            
+            stats_display = gr.Markdown("Click 'Load Dataset' to begin")
+    
+    # Store the loaded dataframe
+    df_state = gr.State(value=pd.DataFrame())
+    
+    # Load dataset
+    def load_and_display(dataset_name):
+        df = load_preview(dataset_name)
+        if "Error" in df.columns:
+            return df, df, "❌ Error loading dataset"
+        stats = f"**Loaded:** {len(df)} rows | **Columns:** {', '.join(COLUMNS)}"
+        return df, df, stats
+    
+    load_btn.click(
+        fn=load_and_display,
+        inputs=dataset_dropdown,
+        outputs=[df_state, df_display, stats_display]
+    )
+    
+    # Search functionality
+    def search_and_display(df, query):
+        if df.empty:
+            return df, "Load a dataset first"
         
-        search_box.change(
-            fn=search_and_display,
-            inputs=[df_state, search_box],
-            outputs=[df_display, stats_display]
-        )
+        filtered_df = filter_dataframe(df, query)
+        stats = f"**Showing:** {len(filtered_df)} of {len(df)} rows"
+        if query:
+            stats += f" | **Search:** '{query}'"
+        return filtered_df, stats
+    
+    search_box.change(
+        fn=search_and_display,
+        inputs=[df_state, search_box],
+        outputs=[df_display, stats_display]
+    )
+    
+    # Send example to model demo
+    def send_to_model(df, row_num):
+        if df.empty or row_num >= len(df):
+            return "", "", "⚠️ Invalid row number or no data loaded"
+        
+        row = df.iloc[int(row_num)]
+        instruction = row['english_prompt'] if 'english_prompt' in row else ""
+        schema = row['create_statement'] if 'create_statement' in row else ""
         
-        # Open model demo
-        demo_btn.click(
-            lambda: None,
-            None,
-            None,
-            _js="()=>{ window.open('https://huggingface.co/spaces/abir-hr196/tinysql-demo','_blank'); }"
-        )
+        return instruction, schema, f"✅ Row {row_num} loaded! Switch to Model Demo tab."
     
-    return viewer
-
-if __name__ == "__main__":
-    dataset_viewer().launch()
+    send_to_model_btn.click(
+        fn=send_to_model,
+        inputs=[df_state, row_selector],
+        outputs=[shared_instruction, shared_schema, stats_display]
+    )
+    
+    return {
+        'df_state': df_state,
+        'df_display': df_display
+    }

tinysql_model_demo.py

@@ -2,7 +2,7 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# ---------------- Model Setup ----------------
+# 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",
@@ -62,7 +62,6 @@ def generate_sql(model_name, instruction, schema, max_length=256, temperature=0.
     except Exception as e:
         return f"Error: {str(e)}"
 
-# ---------------- Example Queries ----------------
 examples = [
     [
         "BM1_CS1_Syn (33M)",
@@ -81,119 +80,96 @@ examples = [
     ],
 ]
 
-# ---------------- 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
+def model_demo(shared_instruction, shared_schema):
+    """Model demo component that can receive examples from dataset viewer"""
+    
+    gr.HTML("""
+        <div style="text-align: center; padding: 3rem 2rem; background: linear-gradient(135deg, #3A3A3A 0%, #4A4A4A 100%); border-radius: 16px; margin-bottom: 2rem; color: white;">
+            <h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 1rem;">TinySQL Interactive Demo</h1>
+            <p style="font-size: 1.2rem; opacity: 0.9; line-height: 1.6;">
+                Transform natural language into SQL queries using <span style="color: #FF6B4A; font-weight: 600;">mechanistically interpretable</span> models
+            </p>
+        </div>
+    """)
+    
+    gr.HTML("""
+        <div style="background: #3A3A3A; border-radius: 12px; padding: 1.5rem; margin: 1.5rem 0; border-left: 4px solid #FF6B4A; color: #E0E0E0;">
+            <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 style="background: #3A3A3A; border-radius: 8px; padding: 1rem; margin-top: 1rem; font-size: 0.9rem; color: #D0D0D0;">
+                    <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,
+                value=""
+            )
+            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.Markdown("### Example Queries")
+    gr.Examples(examples=examples, inputs=[model_dropdown, instruction, schema])
+
+    # Update instruction and schema from shared state when values change
+    def update_from_shared(shared_inst, shared_sch):
+        return shared_inst if shared_inst else "", shared_sch if shared_sch else ""
+    
+    shared_instruction.change(
+        fn=lambda x: x,
+        inputs=shared_instruction,
+        outputs=instruction
+    )
+    
+    shared_schema.change(
+        fn=lambda x: x,
+        inputs=shared_schema,
+        outputs=schema
+    )
+
+    generate_btn.click(
+        fn=generate_sql,
+        inputs=[model_dropdown, instruction, schema, max_length, temperature],
+        outputs=output
+    )
+    
+    return {
+        'instruction': instruction,
+        'schema': schema,
+        'output': output
+    }