upd

app.py

@@ -3,176 +3,71 @@ from tinysql_dataset_viewer import dataset_viewer
 import gradio as gr
 
 custom_css = """
-:root {
-    --martian-orange: #FF6B4A;
-    --martian-black: #0A0A0A;
-    --martian-gray-dark: #1A1A1A;
-    --martian-gray-medium: #2A2A2A;
+/* Force everything dark */
+.gradio-container {
+    background: #0A0A0A !important;
 }
 
-/* FORCE DARK MODE - But preserve styled boxes */
-.gradio-container, 
-body, 
-html,
-.contain,
-.app,
-#root,
-.main {
-    background-color: var(--martian-black) !important;
-    color-scheme: dark !important;
+/* Tabs - black background, bigger bold text */
+.tabs {
+    background: #0A0A0A !important;
 }
 
-/* Only force black on main structural elements, not everything */
-.gradio-container > div,
-.contain > div,
-.app > div {
-    background-color: var(--martian-black) !important;
-}
-
-/* Text should NOT have backgrounds */
-h1, h2, h3, h4, h5, h6, p, span, label, strong, text, div > p, div > span {
-    background: transparent !important;
-    background-color: transparent !important;
-    color: #E0E0E0 !important;
-}
-
-/* Ensure inline styled divs keep their backgrounds */
-div[style*="background"],
-section[style*="background"] {
-    /* Keep their inline styles - don't override */
-}
-
-/* TABS - Remove gray background, keep dark */
 .tab-nav {
-    background: var(--martian-black) !important;
-    border-bottom: 2px solid var(--martian-gray-dark) !important;
-    padding: 0.5rem 1rem !important;
-    border-radius: 0 !important;
+    background: #0A0A0A !important;
+    border: none !important;
 }
 
 .tab-nav button {
-    font-size: 1.6rem !important;
-    font-weight: 800 !important;
-    padding: 1.25rem 2.5rem !important;
-    border-radius: 8px !important;
-    margin: 0 0.25rem !important;
-    transition: all 0.3s ease !important;
+    font-size: 1.8rem !important;
+    font-weight: 900 !important;
     color: #888 !important;
-    border: none !important;
     background: transparent !important;
 }
 
-.tab-nav button:hover {
-    background: var(--martian-gray-medium) !important;
-    color: #E0E0E0 !important;
-}
-
 .tab-nav button.selected {
-    background: var(--martian-orange) !important;
+    background: #FF6B4A !important;
     color: white !important;
-    box-shadow: 0 2px 8px rgba(255, 107, 74, 0.3) !important;
 }
 
-/* ORANGE RECTANGLE LABELS */
-label span, .label-wrap span {
-    background: var(--martian-orange) !important;
+/* Orange labels */
+label > span {
+    background: #FF6B4A !important;
     color: white !important;
-    padding: 0.5rem 1rem !important;
-    border-radius: 8px !important;
+    padding: 0.4rem 0.8rem !important;
+    border-radius: 6px !important;
     font-weight: 700 !important;
-    font-size: 0.95rem !important;
-    display: inline-block !important;
 }
 
-/* Input fields */
-select, input, textarea, .input-wrap input {
-    background: var(--martian-gray-medium) !important;
-    border: 1px solid #3A3A3A !important;
+/* Force white text everywhere */
+* {
     color: #E0E0E0 !important;
 }
 
-select:focus, input:focus, textarea:focus {
-    border-color: var(--martian-orange) !important;
-    outline: none !important;
-    box-shadow: 0 0 0 2px rgba(255, 107, 74, 0.2) !important;
-}
-
-/* Orange sliders */
-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-track {
-    background: #3A3A3A !important;
-}
-
-input[type="range"]::-moz-range-track {
-    background: #3A3A3A !important;
-}
-
-/* Buttons */
-.primary, .primary:hover {
-    background: var(--martian-orange) !important;
-    border-color: var(--martian-orange) !important;
-    color: white !important;
-}
-
-button {
-    border: none !important;
-    background: var(--martian-gray-medium) !important;
+h1, h2, h3, h4, h5, h6 {
     color: #E0E0E0 !important;
 }
 
-/* Citation box */
-.citation-content {
-    background: var(--martian-gray-medium) !important;
+/* Orange buttons */
+.primary {
+    background: #FF6B4A !important;
 }
 
-.citation-content pre {
-    background: var(--martian-gray-dark) !important;
-    color: #D0D0D0 !important;
-}
-
-/* Dataset info - force white text for CS1, CS2, etc */
-.dataset-info strong {
-    color: #E0E0E0 !important;
-}
-
-/* Dropdown and select */
-select option {
-    background: var(--martian-gray-medium) !important;
-    color: #E0E0E0 !important;
+/* Orange sliders */
+input[type="range"]::-webkit-slider-thumb {
+    background: #FF6B4A !important;
 }
 
-/* Code blocks */
-code, pre {
-    background: var(--martian-gray-dark) !important;
-    color: #D0D0D0 !important;
+input[type="range"]::-moz-range-thumb {
+    background: #FF6B4A !important;
 }
 """
 
-with gr.Blocks(css=custom_css, title="TinySQL Demo", theme=gr.themes.Base(primary_hue="orange", secondary_hue="slate", neutral_hue="slate").set(
-    body_background_fill="*neutral_950",
-    body_background_fill_dark="*neutral_950",
-    background_fill_primary="*neutral_950",
-    background_fill_primary_dark="*neutral_950",
-    background_fill_secondary="*neutral_900",
-    background_fill_secondary_dark="*neutral_900",
-    block_background_fill="*neutral_950",
-    block_background_fill_dark="*neutral_950",
-    input_background_fill="*neutral_800",
-    input_background_fill_dark="*neutral_800",
-    button_primary_background_fill="*primary_600",
-    button_primary_background_fill_dark="*primary_600",
-)) as demo:
+with gr.Blocks(css=custom_css, title="TinySQL Demo", theme=gr.themes.Base()) as demo:
     gr.HTML("""
-        <div style="text-align: center; padding: 1.5rem 0 1rem 0; background: #0A0A0A;">
-            <h1 style="font-size: 2rem; font-weight: 700; color: #FF6B4A; margin: 0;">TinySQL</h1>
-            <p style="color: #999; font-size: 0.95rem; margin-top: 0.25rem;">Mechanistic Interpretability for Text-to-SQL</p>
+        <div style="text-align: center; padding: 2rem 0; background: #0A0A0A;">
+            <h1 style="font-size: 2.5rem; font-weight: 700; color: #FF6B4A; margin: 0;">TinySQL</h1>
+            <p style="color: #999; font-size: 1rem; margin-top: 0.5rem;">Mechanistic Interpretability for Text-to-SQL</p>
         </div>
     """)
     
@@ -187,10 +82,9 @@ with gr.Blocks(css=custom_css, title="TinySQL Demo", theme=gr.themes.Base(primar
             model_components = model_demo(shared_instruction, shared_schema)
     
     gr.HTML("""
-        <div style="border-radius: 12px; padding: 1.5rem; margin: 2rem auto; max-width: 900px; background: #2A2A2A;">
-            <div style="font-weight: 700; color: #FF6B4A; margin-bottom: 1rem; font-size: 1.1rem; text-align: center;">Cite Our Work</div>
-            <div class="citation-content" style="background: #2A2A2A; padding: 1rem; border-radius: 8px;">
-                <pre style="font-family: 'Monaco', monospace; font-size: 0.85rem; overflow-x: auto; background: #1A1A1A; color: #D0D0D0; margin: 0; padding: 1rem; border-radius: 6px;">@misc{harrasse2025tinysqlprogressivetexttosqldataset,
+        <div style="background: #2A2A2A; border-radius: 12px; padding: 2rem; margin: 3rem auto; max-width: 900px;">
+            <div style="color: #FF6B4A; font-weight: 700; font-size: 1.2rem; text-align: center; margin-bottom: 1.5rem;">Cite Our Work</div>
+            <pre style="background: #1A1A1A; color: #D0D0D0; padding: 1.5rem; border-radius: 8px; overflow-x: auto; font-family: monospace; font-size: 0.85rem; margin: 0;">@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},
@@ -199,16 +93,15 @@ with gr.Blocks(css=custom_css, title="TinySQL Demo", theme=gr.themes.Base(primar
     primaryClass={cs.LG},
     url={https://arxiv.org/abs/2503.12730}
 }</pre>
-            </div>
         </div>
     """)
     
     gr.HTML("""
-        <div style="text-align: center; padding: 2rem 0; margin-top: 3rem; background: #0A0A0A;">
-            <p style="font-size: 0.85rem; margin-top: 1rem; color: #999;">
-                <a href="https://arxiv.org/abs/2503.12730" style="color: #FF6B4A;">Paper</a> • 
-                <a href="https://github.com/withmartian/TinySQL" style="color: #FF6B4A;">Code</a> • 
-                <a href="https://huggingface.co/collections/withmartian/tinysql-6760e92748b63fa56a6ffc9f" style="color: #FF6B4A;">Dataset</a>
+        <div style="text-align: center; padding: 2rem 0; background: #0A0A0A;">
+            <p style="color: #999; font-size: 0.9rem;">
+                <a href="https://arxiv.org/abs/2503.12730" style="color: #FF6B4A; text-decoration: none;">Paper</a> • 
+                <a href="https://github.com/withmartian/TinySQL" style="color: #FF6B4A; text-decoration: none;">Code</a> • 
+                <a href="https://huggingface.co/collections/withmartian/tinysql-6760e92748b63fa56a6ffc9f" style="color: #FF6B4A; text-decoration: none;">Dataset</a>
             </p>
         </div>
     """)

tinysql_dataset_viewer.py

@@ -13,11 +13,9 @@ DATASETS = {
 
 COLUMNS = ["create_statement", "english_prompt", "sql_statement"]
 
-# Pre-cache datasets on startup
 dataset_cache = {}
 
 def preload_datasets():
-    """Load first 500 rows of all datasets into cache"""
     for name, path in DATASETS.items():
         try:
             ds = load_dataset(path, split="train")
@@ -30,11 +28,9 @@ def preload_datasets():
         except Exception as e:
             print(f"✗ Failed to cache {name}: {e}")
 
-# Preload on import
 preload_datasets()
 
 def load_preview(dataset_name):
-    """Load from cache instantly"""
     if dataset_name in dataset_cache:
         return dataset_cache[dataset_name]
     return pd.DataFrame({"Error": ["Dataset not found in cache"]})
@@ -55,19 +51,9 @@ def filter_dataframe(df, search_query, search_column):
 
 def dataset_viewer(shared_instruction, shared_schema):
     gr.HTML("""
-        <div style="text-align: center; padding: 1.5rem; background: linear-gradient(135deg, #2A2A2A 0%, #3A3A3A 100%); border-radius: 16px; margin-bottom: 1.5rem; box-shadow: 0 4px 12px rgba(0,0,0,0.3);">
-            <h2 style="font-size: 1.75rem; font-weight: 700; margin-bottom: 0.5rem; color: #FF6B4A;">Dataset Explorer</h2>
-            <p style="font-size: 0.95rem; opacity: 0.9; line-height: 1.6; color: #D0D0D0;">
-                Browse, search, and explore TinySQL datasets
-            </p>
-        </div>
-    """)
-    
-    gr.HTML("""
-        <div style="background: linear-gradient(135deg, #2A2A2A 0%, #3A3A3A 100%); border-radius: 12px; padding: 1.25rem; margin: 1rem 0; border-left: 4px solid #FF6B4A;">
-            <p style="color: #D0D0D0; margin: 0; line-height: 1.6;">
-                <strong style="color: #FF6B4A;">Quick Start:</strong> Select a dataset and click Load Dataset. Use search to filter results.
-            </p>
+        <div style="text-align: center; padding: 2rem; background: linear-gradient(135deg, #2A2A2A 0%, #3A3A3A 100%); border-radius: 16px; margin-bottom: 2rem;">
+            <h2 style="font-size: 2rem; font-weight: 700; margin-bottom: 0.5rem; color: #FF6B4A;">Dataset Explorer</h2>
+            <p style="font-size: 1rem; color: #D0D0D0; margin: 0;">Browse, search, and explore TinySQL datasets</p>
         </div>
     """)
     
@@ -78,27 +64,20 @@ def dataset_viewer(shared_instruction, shared_schema):
             dataset_dropdown = gr.Dropdown(
                 choices=list(DATASETS.keys()),
                 value="CS1",
-                label="Choose Dataset",
-                info="Select complexity level"
+                label="Choose Dataset"
             )
             
             load_btn = gr.Button("Load Dataset", variant="primary", size="lg")
             
             gr.HTML("""
-                <div style="background: #2A2A2A; border-radius: 12px; padding: 1.25rem; margin: 1.25rem 0; border: 1px solid #3A3A3A;">
-                    <h4 style="color: #FF6B4A; font-size: 0.95rem; margin: 0 0 1rem 0; font-weight: 700; border-bottom: 2px solid #3A3A3A; padding-bottom: 0.75rem;">Dataset Levels</h4>
-                    
-                    <div style="color: #D0D0D0; font-size: 0.85rem; line-height: 1.8;">
-                        <div><strong>CS1:</strong> Basic SELECT-FROM</div>
-                        <div><strong>CS2:</strong> Adds ORDER BY</div>
-                        <div><strong>CS3:</strong> Aggregations</div>
-                        <div><strong>CS4:</strong> WHERE filters</div>
-                        <div><strong>CS5:</strong> Multi-table JOINs</div>
-                    </div>
-                    
-                    <div style="margin-top: 1rem; padding-top: 1rem; border-top: 1px solid #3A3A3A;">
-                        <div style="color: #FF6B4A; font-weight: 600; font-size: 0.85rem; margin-bottom: 0.5rem;">Synonym Variants</div>
-                        <div style="color: #999; font-size: 0.8rem; line-height: 1.5;">Natural language variations</div>
+                <div style="background: #2A2A2A; border-radius: 12px; padding: 1.5rem; margin: 1.5rem 0;">
+                    <h4 style="color: #FF6B4A; font-size: 1rem; margin: 0 0 1rem 0; font-weight: 700;">Dataset Levels</h4>
+                    <div style="font-size: 0.9rem; line-height: 2;">
+                        <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS1:</strong> <span style="color: #FFFFFF;">Basic SELECT-FROM</span></div>
+                        <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS2:</strong> <span style="color: #FFFFFF;">Adds ORDER BY</span></div>
+                        <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS3:</strong> <span style="color: #FFFFFF;">Aggregations</span></div>
+                        <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS4:</strong> <span style="color: #FFFFFF;">WHERE filters</span></div>
+                        <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS5:</strong> <span style="color: #FFFFFF;">Multi-table JOINs</span></div>
                     </div>
                 </div>
             """)
@@ -108,8 +87,7 @@ def dataset_viewer(shared_instruction, shared_schema):
                 label="Row Number",
                 value=0,
                 minimum=0,
-                precision=0,
-                info="Pick a row to test"
+                precision=0
             )
             
             send_to_model_btn = gr.Button("Run in Model Demo", variant="primary")
@@ -132,106 +110,11 @@ def dataset_viewer(shared_instruction, shared_schema):
                     scale=1
                 )
             
-            gr.HTML("""
-                <style>
-                /* HuggingFace-style table - FORCE DARK MODE */
-                .dataframe-container, .dataframe-container * {
-                    color: #E0E0E0 !important;
-                    background: var(--martian-black) !important;
-                }
-                
-                .dataframe-container label {
-                    display: none !important;
-                }
-                
-                .dataframe-container {
-                    border-radius: 8px !important;
-                    overflow: hidden !important;
-                    border: 1px solid #374151 !important;
-                }
-
-                .dataframe table {
-                    border-collapse: collapse !important;
-                    width: 100% !important;
-                    font-size: 0.875rem !important;
-                    background: #111827 !important;
-                }
-
-                .dataframe thead {
-                    background: #1f2937 !important;
-                }
-
-                .dataframe thead th {
-                    color: #9ca3af !important;
-                    font-weight: 600 !important;
-                    text-align: left !important;
-                    padding: 0.75rem 1rem !important;
-                    border-bottom: 1px solid #374151 !important;
-                    font-size: 0.75rem !important;
-                    text-transform: uppercase !important;
-                    letter-spacing: 0.05em !important;
-                    background: #1f2937 !important;
-                }
-
-                .dataframe tbody tr {
-                    background: #111827 !important;
-                    border-bottom: 1px solid #1f2937 !important;
-                    transition: all 0.15s ease !important;
-                    position: relative !important;
-                }
-
-                .dataframe tbody tr:hover {
-                    background: #1f2937 !important;
-                    box-shadow: 0 2px 8px rgba(255, 107, 74, 0.1) !important;
-                }
-
-                .dataframe tbody tr:hover::before {
-                    content: "Row " attr(data-row-index);
-                    position: absolute;
-                    left: -60px;
-                    top: 50%;
-                    transform: translateY(-50%);
-                    background: #FF6B4A;
-                    color: white;
-                    padding: 0.25rem 0.5rem;
-                    border-radius: 4px;
-                    font-size: 0.75rem;
-                    font-weight: 600;
-                    white-space: nowrap;
-                    opacity: 0.9;
-                }
-
-                .dataframe tbody td {
-                    padding: 0.75rem 1rem !important;
-                    color: #d1d5db !important;
-                    font-size: 0.875rem !important;
-                    line-height: 1.5 !important;
-                    max-width: 400px !important;
-                    overflow: hidden !important;
-                    text-overflow: ellipsis !important;
-                    background: #111827 !important;
-                }
-
-                .dataframe tbody tr:last-child {
-                    border-bottom: none !important;
-                }
-                
-                .dataframe tbody td:first-child,
-                .dataframe thead th:first-child {
-                    width: 0 !important;
-                    padding: 0 !important;
-                    opacity: 0 !important;
-                    position: absolute !important;
-                }
-                </style>
-            """)
-            
             df_display = gr.Dataframe(
                 headers=["index"] + COLUMNS,
                 datatype=["number", "str", "str", "str"],
                 interactive=False,
-                wrap=True,
-                elem_classes="dataframe-container"
+                wrap=True
             )
             
             stats_display = gr.Markdown("Click **Load Dataset** to begin")

tinysql_model_demo.py

@@ -63,19 +63,9 @@ def generate_sql(model_name, instruction, schema, max_length=256, temperature=0.
 
 def model_demo(shared_instruction, shared_schema):
     gr.HTML("""
-        <div style="text-align: center; padding: 2rem; background: linear-gradient(135deg, #3A3A3A 0%, #4A4A4A 100%); border-radius: 16px; margin-bottom: 1.5rem;">
+        <div style="text-align: center; padding: 2rem; background: linear-gradient(135deg, #3A3A3A 0%, #4A4A4A 100%); border-radius: 16px; margin-bottom: 2rem;">
             <h2 style="font-size: 2rem; font-weight: 700; margin-bottom: 0.5rem; color: white;">Interactive SQL Generation</h2>
-            <p style="font-size: 1rem; opacity: 0.9; line-height: 1.6; color: #D0D0D0;">
-                Transform natural language into SQL using <span style="color: #FF6B4A; font-weight: 600;">mechanistically interpretable</span> models
-            </p>
-        </div>
-    """)
-    
-    gr.HTML("""
-        <div style="background: linear-gradient(135deg, #2A2A2A 0%, #3A3A3A 100%); border-radius: 12px; padding: 1.5rem; margin: 1rem 0; border-left: 4px solid #FF6B4A;">
-            <p style="color: #D0D0D0; margin: 0; line-height: 1.6;">
-                <strong style="color: #FF6B4A;">How it works:</strong> Select a model, describe your query in plain English, and watch the model generate SQL.
-            </p>
+            <p style="font-size: 1rem; color: #D0D0D0; margin: 0;">Transform natural language into SQL</p>
         </div>
     """)
     
@@ -85,37 +75,36 @@ def model_demo(shared_instruction, shared_schema):
             model_dropdown = gr.Dropdown(
                 choices=list(MODELS.keys()),
                 value="BM2_CS2_Syn (0.5B)",
-                label="Model Selection",
-                info="Larger models = better accuracy"
+                label="Model Selection"
             )
             
             gr.HTML("""
-                <div style="background: #2A2A2A; border-radius: 12px; padding: 1.5rem; margin: 1.5rem 0; border: 1px solid #3A3A3A;">
-                    <h4 style="color: #FF6B4A; font-size: 1rem; margin: 0 0 1.25rem 0; font-weight: 700; border-bottom: 2px solid #3A3A3A; padding-bottom: 0.75rem;">Model Guide</h4>
+                <div style="background: #2A2A2A; border-radius: 12px; padding: 1.5rem; margin: 1.5rem 0;">
+                    <h4 style="color: #FF6B4A; font-size: 1rem; margin: 0 0 1.5rem 0; font-weight: 700;">Model Guide</h4>
                     
                     <div style="margin-bottom: 1rem;">
-                        <div style="color: #E0E0E0; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.35rem;">BM1 (33M parameters)</div>
-                        <div style="color: #999; font-size: 0.85rem; margin-left: 1rem;">TinyStories 33M fine-tuned</div>
+                        <div style="font-weight: 600; font-size: 0.9rem; margin-bottom: 0.4rem; color: #FFFFFF;">BM1 (33M parameters)</div>
+                        <div style="font-size: 0.85rem; margin-left: 1rem; color: #FFFFFF;">TinyStories 33M fine-tuned</div>
                     </div>
                     
                     <div style="margin-bottom: 1rem;">
-                        <div style="color: #E0E0E0; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.35rem;">BM2 (0.5B parameters)</div>
-                        <div style="color: #999; font-size: 0.85rem; margin-left: 1rem;">Qwen 2.5 0.5B fine-tuned</div>
+                        <div style="font-weight: 600; font-size: 0.9rem; margin-bottom: 0.4rem; color: #FFFFFF;">BM2 (0.5B parameters)</div>
+                        <div style="font-size: 0.85rem; margin-left: 1rem; color: #FFFFFF;">Qwen 2.5 0.5B fine-tuned</div>
                     </div>
                     
-                    <div style="margin-bottom: 1.25rem;">
-                        <div style="color: #E0E0E0; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.35rem;">BM3 (1B parameters)</div>
-                        <div style="color: #999; font-size: 0.85rem; margin-left: 1rem;">Llama 3.2 1B fine-tuned</div>
+                    <div style="margin-bottom: 1.5rem;">
+                        <div style="font-weight: 600; font-size: 0.9rem; margin-bottom: 0.4rem; color: #FFFFFF;">BM3 (1B parameters)</div>
+                        <div style="font-size: 0.85rem; margin-left: 1rem; color: #FFFFFF;">Llama 3.2 1B fine-tuned</div>
                     </div>
                     
                     <div style="padding-top: 1rem; border-top: 1px solid #3A3A3A;">
                         <div style="color: #FF6B4A; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.75rem;">Dataset Complexity</div>
-                        <div style="color: #999; font-size: 0.85rem; line-height: 1.8;">
-                            <div><strong style="color: #D0D0D0;">CS1:</strong> Basic SELECT-FROM</div>
-                            <div><strong style="color: #D0D0D0;">CS2:</strong> Adds ORDER BY</div>
-                            <div><strong style="color: #D0D0D0;">CS3:</strong> Aggregations</div>
-                            <div><strong style="color: #D0D0D0;">CS4:</strong> WHERE filters</div>
-                            <div><strong style="color: #D0D0D0;">CS5:</strong> Multi-table JOINs</div>
+                        <div style="font-size: 0.9rem; line-height: 2;">
+                            <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS1:</strong> <span style="color: #FFFFFF;">Basic SELECT-FROM</span></div>
+                            <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS2:</strong> <span style="color: #FFFFFF;">Adds ORDER BY</span></div>
+                            <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS3:</strong> <span style="color: #FFFFFF;">Aggregations</span></div>
+                            <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS4:</strong> <span style="color: #FFFFFF;">WHERE filters</span></div>
+                            <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS5:</strong> <span style="color: #FFFFFF;">Multi-table JOINs</span></div>
                         </div>
                     </div>
                 </div>
@@ -132,45 +121,23 @@ def model_demo(shared_instruction, shared_schema):
             )
             
             schema = gr.Code(
-                label="Database Schema",
+                label="Database Schema (SQL)",
                 language="sql",
-                value="CREATE TABLE employees (name VARCHAR(100), salary INT, department VARCHAR(100))",
-                lines=3
+                value="CREATE TABLE employees (\n    name VARCHAR(100),\n    salary INT,\n    department VARCHAR(100)\n);",
+                lines=6
             )
             
             with gr.Row():
-                max_length = gr.Slider(64, 512, value=256, step=32, label="Max Length", info="Token limit")
-                temperature = gr.Slider(0.0, 1.0, value=0.0, step=0.1, label="Temperature", info="Creativity level (0 = deterministic)")
+                max_length = gr.Slider(64, 512, value=256, step=32, label="Max Length")
+                temperature = gr.Slider(0.0, 1.0, value=0.0, 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=6)
-
-    gr.Markdown("### Example Query")
-    
-    # Better styled example with proper font
-    gr.HTML("""
-        <div style="background: #2A2A2A; border-radius: 12px; padding: 1.5rem; margin: 1rem 0; border: 1px solid #3A3A3A;">
-            <div style="margin-bottom: 1rem; padding-bottom: 1rem; border-bottom: 1px solid #3A3A3A;">
-                <div style="color: #FF6B4A; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.5rem;">Model</div>
-                <div style="background: #1A1A1A; padding: 0.75rem; border-radius: 6px; color: #D0D0D0; font-family: 'Inter', sans-serif; font-size: 0.9rem;">BM2_CS2_Syn (0.5B)</div>
-            </div>
-            
-            <div style="margin-bottom: 1rem; padding-bottom: 1rem; border-bottom: 1px solid #3A3A3A;">
-                <div style="color: #FF6B4A; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.5rem;">Natural Language Query</div>
-                <div style="background: #1A1A1A; padding: 0.75rem; border-radius: 6px; color: #D0D0D0; font-family: 'Inter', sans-serif; font-size: 0.9rem; line-height: 1.5;">List worker earnings from highest to lowest</div>
-            </div>
-            
-            <div>
-                <div style="color: #FF6B4A; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.5rem;">Database Schema</div>
-                <pre style="background: #1A1A1A; padding: 0.75rem; border-radius: 6px; margin: 0; color: #D0D0D0; font-family: 'Monaco', 'Courier New', monospace; font-size: 0.85rem; line-height: 1.5; overflow-x: auto;">CREATE TABLE employees (
-    name VARCHAR(100), 
-    salary INT, 
-    department VARCHAR(100)
-)</pre>
-            </div>
-        </div>
-    """)
+            output = gr.Code(
+                label="Generated SQL Query",
+                language="sql",
+                lines=8
+            )
 
     shared_instruction.change(
         fn=lambda x: x,