updates

app.py

@@ -10,9 +10,25 @@ custom_css = """
     --martian-gray-medium: #2A2A2A;
 }
 
+/* FORCE DARK MODE - Override Gradio's light mode */
+.gradio-container, .gradio-container * {
+    color-scheme: dark !important;
+}
+
 .gradio-container {
     font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
     background-color: var(--martian-black) !important;
+    color: #E0E0E0 !important;
+}
+
+/* Force dark backgrounds everywhere */
+.block, .panel, div, section {
+    background-color: var(--martian-black) !important;
+}
+
+/* Force text colors */
+h1, h2, h3, h4, h5, h6, p, span, label, div {
+    color: #E0E0E0 !important;
 }
 
 /* Better tabs - bold, bigger, no icons, ORANGE selection */
@@ -32,6 +48,7 @@ custom_css = """
     transition: all 0.3s ease !important;
     color: #888 !important;
     border: none !important;
+    background: transparent !important;
 }
 
 .tab-nav button:hover {
@@ -45,7 +62,7 @@ custom_css = """
     box-shadow: 0 2px 8px rgba(255, 107, 74, 0.3) !important;
 }
 
-/* REMOVE ALL PURPLE - Replace with orange */
+/* Orange labels */
 label, .label-wrap {
     color: #E0E0E0 !important;
 }
@@ -59,7 +76,7 @@ label, .label-wrap {
     font-size: 0.9rem !important;
 }
 
-/* Orange dropdown and inputs - NO PURPLE */
+/* Orange inputs */
 select, input, textarea {
     background: var(--martian-gray-medium) !important;
     border: 1px solid #3A3A3A !important;
@@ -72,7 +89,7 @@ select:focus, input:focus, textarea:focus {
     box-shadow: 0 0 0 2px rgba(255, 107, 74, 0.2) !important;
 }
 
-/* Orange sliders - properly following thumb position */
+/* Orange sliders */
 input[type="range"] {
     -webkit-appearance: none;
     appearance: none;
@@ -120,25 +137,42 @@ input[type="range"]::-moz-range-progress {
     border-radius: 3px;
 }
 
-/* For Webkit browsers, use a gradient workaround */
-input[type="range"]::-webkit-slider-runnable-track {
-    background: linear-gradient(to right, 
-        var(--martian-orange) 0%, 
-        var(--martian-orange) var(--value, 0%), 
-        #3A3A3A var(--value, 0%), 
-        #3A3A3A 100%);
-    height: 6px;
-    border-radius: 3px;
-}
-
-/* Remove any purple/blue colors */
+/* Orange buttons */
 .primary, .secondary {
     background: var(--martian-orange) !important;
     border-color: var(--martian-orange) !important;
+    color: white !important;
 }
 
 button {
     border: none !important;
+    color: #E0E0E0 !important;
+}
+
+/* Compact header */
+.compact-header {
+    text-align: center;
+    padding: 1.5rem 0 1rem 0;
+    background: var(--martian-black);
+}
+
+.compact-header img {
+    height: 35px;
+    margin-bottom: 0.5rem;
+}
+
+.compact-header h1 {
+    font-size: 2rem;
+    font-weight: 700;
+    color: #FF6B4A;
+    margin: 0;
+    line-height: 1.2;
+}
+
+.compact-header p {
+    color: #999;
+    font-size: 0.95rem;
+    margin-top: 0.25rem;
 }
 
 /* Footer */
@@ -165,7 +199,7 @@ button {
 }
 
 .footer-text {
-    color: #999;
+    color: #999 !important;
     font-size: 0.95rem;
     margin: 0.5rem 0;
 }
@@ -175,13 +209,17 @@ button {
     animation: heartbeat 1.5s ease-in-out infinite;
 }
 
+.footer-text a {
+    color: #FF6B4A !important;
+}
+
 @keyframes heartbeat {
     0%, 100% { transform: scale(1); }
     50% { transform: scale(1.1); }
 }
 
 .citation-box {
-    background: #2A2A2A;
+    background: #2A2A2A !important;
     border: 1px solid #3A3A3A;
     border-radius: 12px;
     padding: 1.5rem;
@@ -193,15 +231,15 @@ button {
 
 .citation-header {
     font-weight: 700;
-    color: #FF6B4A;
+    color: #FF6B4A !important;
     margin-bottom: 1rem;
     font-size: 1.1rem;
     text-align: center;
 }
 
 .citation-box pre {
-    color: #D0D0D0;
-    background: #1A1A1A;
+    color: #D0D0D0 !important;
+    background: #1A1A1A !important;
     padding: 1rem;
     border-radius: 8px;
     overflow-x: auto;
@@ -210,12 +248,12 @@ button {
 """
 
 with gr.Blocks(css=custom_css, title="TinySQL Demo", theme=gr.themes.Soft()) as demo:
-    # Martian Logo Header
+    # Compact Header
     gr.HTML("""
-        <div style="text-align: center; padding: 2rem 0 1rem 0;">
-            <img src="https://withmartian.com/logo.png" alt="Martian" style="height: 60px; margin-bottom: 1rem;" onerror="this.style.display='none'">
-            <h1 style="font-size: 2.5rem; font-weight: 700; color: #FF6B4A; margin: 0;">TinySQL</h1>
-            <p style="color: #999; font-size: 1.1rem; margin-top: 0.5rem;">Mechanistic Interpretability for Text-to-SQL</p>
+        <div class="compact-header">
+            <img src="file/static/images/martian_logo.jpg" alt="Martian">
+            <h1>TinySQL</h1>
+            <p>Mechanistic Interpretability for Text-to-SQL</p>
         </div>
     """)
     
@@ -252,11 +290,11 @@ with gr.Blocks(css=custom_css, title="TinySQL Demo", theme=gr.themes.Soft()) as
             <p class="footer-text">
                 Brought to you with <span class="heart">❤️</span> from the Martian science team
             </p>
-            <img src="static/images/martian_logo.jpg" alt="Martian" class="footer-logo" style="width: 80px; margin: 1rem auto; display: block;">
+            <img src="/static/images/martian_logo.jpg" alt="Martian" class="footer-logo">
             <p class="footer-text" style="font-size: 0.85rem; margin-top: 1rem;">
-                <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>
+                <a href="https://arxiv.org/abs/2503.12730">Paper</a> • 
+                <a href="https://github.com/withmartian/TinySQL">Code</a> • 
+                <a href="https://huggingface.co/collections/withmartian/tinysql-6760e92748b63fa56a6ffc9f">Dataset</a>
             </p>
         </div>
     """)

tinysql_dataset_viewer.py

@@ -13,42 +13,60 @@ 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")
+            df = pd.DataFrame(ds).head(500)
+            if all(col in df.columns for col in COLUMNS):
+                df = df[COLUMNS]
+            df.insert(0, 'index', range(len(df)))
+            dataset_cache[name] = df
+            print(f"✓ Cached {name}")
+        except Exception as e:
+            print(f"✗ Failed to cache {name}: {e}")
+
+# Preload on import
+preload_datasets()
+
 def load_preview(dataset_name):
-    try:
-        ds = load_dataset(DATASETS[dataset_name], split="train")
-        df = pd.DataFrame(ds).head(500)
-        if all(col in df.columns for col in COLUMNS):
-            df = df[COLUMNS]
-        # Add index column
-        df.insert(0, 'index', range(len(df)))
-        return df
-    except Exception as e:
-        return pd.DataFrame({"Error": [str(e)]})
+    """Load from cache instantly"""
+    if dataset_name in dataset_cache:
+        return dataset_cache[dataset_name]
+    return pd.DataFrame({"Error": ["Dataset not found in cache"]})
 
-def filter_dataframe(df, search_query):
+def filter_dataframe(df, search_query, search_column):
     if not search_query or df.empty or "Error" in df.columns:
         return df
     
-    mask = df.astype(str).apply(
-        lambda row: row.str.contains(search_query, case=False, na=False).any(), 
-        axis=1
-    )
+    if search_column == "All Columns":
+        mask = df.astype(str).apply(
+            lambda row: row.str.contains(search_query, case=False, na=False).any(), 
+            axis=1
+        )
+    else:
+        mask = df[search_column].astype(str).str.contains(search_query, case=False, na=False)
+    
     return df[mask]
 
 def dataset_viewer(shared_instruction, shared_schema):
     gr.HTML("""
-        <div style="text-align: center; padding: 2rem 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: 2rem; font-weight: 700; margin-bottom: 0.5rem; color: #FF6B4A;">Dataset Explorer</h2>
-            <p style="font-size: 1rem; opacity: 0.9; line-height: 1.6; color: #D0D0D0;">
+        <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.5rem; margin: 1rem 0; border-left: 4px solid #FF6B4A;">
+        <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, click Load Dataset, then use search to filter. Pick any row and send it to the Model Demo tab.
+                <strong style="color: #FF6B4A;">Quick Start:</strong> Select a dataset and click Load Dataset. Use search to filter results.
             </p>
         </div>
     """)
@@ -64,28 +82,27 @@ def dataset_viewer(shared_instruction, shared_schema):
                 info="Select complexity level"
             )
             
-            # Simpler dataset guide - no colors, no beginner/intermediate
+            load_btn = gr.Button("Load Dataset", variant="primary", size="lg")
+            
             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;">Dataset Complexity Levels</h4>
+                <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.9rem; line-height: 2;">
-                        <div><strong>CS1:</strong> Basic SELECT-FROM queries</div>
-                        <div><strong>CS2:</strong> Adds ORDER BY clauses</div>
-                        <div><strong>CS3:</strong> Aggregations (COUNT, SUM, AVG)</div>
-                        <div><strong>CS4:</strong> Adds WHERE filters</div>
+                    <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: 1.5rem; padding-top: 1.25rem; border-top: 1px solid #3A3A3A;">
-                        <div style="color: #FF6B4A; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.5rem;">Synonym Variants</div>
-                        <div style="color: #999; font-size: 0.85rem; line-height: 1.6;">Natural language variations with semantic mappings</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>
                 </div>
             """)
             
-            load_btn = gr.Button("Load Dataset", variant="primary", size="lg")
-            
             gr.Markdown("### Test Example")
             row_selector = gr.Number(
                 label="Row Number",
@@ -100,16 +117,29 @@ def dataset_viewer(shared_instruction, shared_schema):
         with gr.Column(scale=3):
             gr.Markdown("### Dataset Preview")
             
-            search_box = gr.Textbox(
-                label="Search",
-                placeholder="Search across all columns...",
-                lines=1
-            )
+            with gr.Row():
+                search_box = gr.Textbox(
+                    label="Search",
+                    placeholder="Enter search term...",
+                    lines=1,
+                    scale=3
+                )
+                
+                search_column = gr.Dropdown(
+                    choices=["All Columns", "create_statement", "english_prompt", "sql_statement"],
+                    value="All Columns",
+                    label="Search In",
+                    scale=1
+                )
             
-            # HuggingFace-style table with row index on hover
             gr.HTML("""
                 <style>
-                /* True HuggingFace-style table - NO "Results" label */
+                /* HuggingFace-style table - FORCE DARK MODE */
+                .dataframe-container, .dataframe-container * {
+                    color: #E0E0E0 !important;
+                    background: var(--martian-black) !important;
+                }
+                
                 .dataframe-container label {
                     display: none !important;
                 }
@@ -124,6 +154,7 @@ def dataset_viewer(shared_instruction, shared_schema):
                     border-collapse: collapse !important;
                     width: 100% !important;
                     font-size: 0.875rem !important;
+                    background: #111827 !important;
                 }
 
                 .dataframe thead {
@@ -139,6 +170,7 @@ def dataset_viewer(shared_instruction, shared_schema):
                     font-size: 0.75rem !important;
                     text-transform: uppercase !important;
                     letter-spacing: 0.05em !important;
+                    background: #1f2937 !important;
                 }
 
                 .dataframe tbody tr {
@@ -177,13 +209,13 @@ def dataset_viewer(shared_instruction, shared_schema):
                     max-width: 400px !important;
                     overflow: hidden !important;
                     text-overflow: ellipsis !important;
+                    background: #111827 !important;
                 }
 
                 .dataframe tbody tr:last-child {
                     border-bottom: none !important;
                 }
                 
-                /* Hide index column but keep it for reference */
                 .dataframe tbody td:first-child,
                 .dataframe thead th:first-child {
                     width: 0 !important;
@@ -202,7 +234,7 @@ def dataset_viewer(shared_instruction, shared_schema):
                 elem_classes="dataframe-container"
             )
             
-            stats_display = gr.Markdown("Click **Load Dataset** to begin exploring")
+            stats_display = gr.Markdown("Click **Load Dataset** to begin")
     
     df_state = gr.State(value=pd.DataFrame())
     
@@ -219,19 +251,25 @@ def dataset_viewer(shared_instruction, shared_schema):
         outputs=[df_state, df_display, stats_display]
     )
     
-    def search_and_display(df, query):
+    def search_and_display(df, query, column):
         if df.empty:
             return df, "Load a dataset first"
         
-        filtered_df = filter_dataframe(df, query)
+        filtered_df = filter_dataframe(df, query, column)
         stats = f"**Showing {len(filtered_df)} of {len(df)} rows**"
         if query:
-            stats += f" • Search: '{query}'"
+            stats += f" • Search: '{query}' in {column}"
         return filtered_df, stats
     
     search_box.change(
         fn=search_and_display,
-        inputs=[df_state, search_box],
+        inputs=[df_state, search_box, search_column],
+        outputs=[df_display, stats_display]
+    )
+    
+    search_column.change(
+        fn=search_and_display,
+        inputs=[df_state, search_box, search_column],
         outputs=[df_display, stats_display]
     )