updates

app.py

@@ -15,7 +15,7 @@ custom_css = """
     background-color: var(--martian-black) !important;
 }
 
-/* Better tabs - bold, bigger, no icons */
+/* Better tabs - bold, bigger, no icons, ORANGE selection */
 .tab-nav {
     border-bottom: 2px solid var(--martian-gray-dark) !important;
     background: var(--martian-gray-dark) !important;
@@ -31,6 +31,7 @@ custom_css = """
     margin: 0 0.25rem !important;
     transition: all 0.3s ease !important;
     color: #888 !important;
+    border: none !important;
 }
 
 .tab-nav button:hover {
@@ -44,7 +45,7 @@ custom_css = """
     box-shadow: 0 2px 8px rgba(255, 107, 74, 0.3) !important;
 }
 
-/* Orange labels - remove all purple */
+/* REMOVE ALL PURPLE - Replace with orange */
 label, .label-wrap {
     color: #E0E0E0 !important;
 }
@@ -58,23 +59,56 @@ label, .label-wrap {
     font-size: 0.9rem !important;
 }
 
-/* Orange sliders */
-input[type="range"] {
+/* Orange dropdown and inputs - NO PURPLE */
+select, input, textarea {
     background: var(--martian-gray-medium) !important;
+    border: 1px solid #3A3A3A !important;
+    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 - selected area is orange */
+input[type="range"] {
+    background: transparent !important;
 }
 
 input[type="range"]::-webkit-slider-thumb {
     background: var(--martian-orange) !important;
     cursor: pointer !important;
+    width: 16px !important;
+    height: 16px !important;
+    border-radius: 50% !important;
 }
 
 input[type="range"]::-moz-range-thumb {
     background: var(--martian-orange) !important;
     cursor: pointer !important;
+    width: 16px !important;
+    height: 16px !important;
+    border-radius: 50% !important;
 }
 
 input[type="range"]::-webkit-slider-runnable-track {
-    background: linear-gradient(to right, var(--martian-orange) 0%, var(--martian-gray-light) 100%) !important;
+    background: linear-gradient(to right, var(--martian-orange) 0%, var(--martian-orange) var(--value, 50%), #3A3A3A var(--value, 50%), #3A3A3A 100%) !important;
+    height: 4px !important;
+    border-radius: 2px !important;
+}
+
+input[type="range"]::-moz-range-track {
+    background: #3A3A3A !important;
+    height: 4px !important;
+    border-radius: 2px !important;
+}
+
+input[type="range"]::-moz-range-progress {
+    background: var(--martian-orange) !important;
+    height: 4px !important;
+    border-radius: 2px !important;
 }
 
 /* Remove any purple/blue colors */
@@ -83,6 +117,10 @@ input[type="range"]::-webkit-slider-runnable-track {
     border-color: var(--martian-orange) !important;
 }
 
+button {
+    border: none !important;
+}
+
 /* Footer */
 .footer {
     text-align: center;
@@ -119,6 +157,34 @@ input[type="range"]::-webkit-slider-runnable-track {
     0%, 100% { transform: scale(1); }
     50% { transform: scale(1.1); }
 }
+
+.citation-box {
+    background: #2A2A2A;
+    border: 1px solid #3A3A3A;
+    border-radius: 12px;
+    padding: 1.5rem;
+    margin: 2rem auto;
+    max-width: 900px;
+    font-family: 'Monaco', 'Courier New', monospace;
+    font-size: 0.85rem;
+}
+
+.citation-header {
+    font-weight: 700;
+    color: #FF6B4A;
+    margin-bottom: 1rem;
+    font-size: 1.1rem;
+    text-align: center;
+}
+
+.citation-box pre {
+    color: #D0D0D0;
+    background: #1A1A1A;
+    padding: 1rem;
+    border-radius: 8px;
+    overflow-x: auto;
+    margin: 0;
+}
 """
 
 with gr.Blocks(css=custom_css, title="TinySQL Demo", theme=gr.themes.Soft()) as demo:
@@ -142,6 +208,22 @@ with gr.Blocks(css=custom_css, title="TinySQL Demo", theme=gr.themes.Soft()) as
         with gr.Tab("Model Demo"):
             model_components = model_demo(shared_instruction, shared_schema)
     
+    # Citation section
+    gr.HTML("""
+        <div class="citation-box">
+            <div class="citation-header">Cite Our Work</div>
+            <pre>@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 with Martian branding
     gr.HTML("""
         <div class="footer">
@@ -152,7 +234,7 @@ with gr.Blocks(css=custom_css, title="TinySQL Demo", theme=gr.themes.Soft()) as
             <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://huggingface.co/collections/withmartian/tinysql-tinysql-6760e92748b63fa56a6ffc9f" style="color: #FF6B4A; text-decoration: none;">Dataset</a>
             </p>
         </div>
     """)

tinysql_dataset_viewer.py

@@ -19,6 +19,8 @@ def load_preview(dataset_name):
         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)]})
@@ -62,32 +64,17 @@ def dataset_viewer(shared_instruction, shared_schema):
                 info="Select complexity level"
             )
             
-            # Better formatted model guide card
+            # Simpler dataset guide - no colors, no beginner/intermediate
             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="margin-bottom: 1.5rem;">
-                        <div style="color: #4CAF50; font-weight: 600; font-size: 0.95rem; margin-bottom: 0.5rem;">Basic Level</div>
-                        <div style="margin-left: 1rem; color: #999; font-size: 0.85rem; line-height: 1.8;">
-                            <div><strong style="color: #D0D0D0;">CS1:</strong> Basic SELECT-FROM queries</div>
-                            <div><strong style="color: #D0D0D0;">CS2:</strong> Adds ORDER BY clauses</div>
-                        </div>
-                    </div>
-                    
-                    <div style="margin-bottom: 1.5rem;">
-                        <div style="color: #FF9800; font-weight: 600; font-size: 0.95rem; margin-bottom: 0.5rem;">Intermediate Level</div>
-                        <div style="margin-left: 1rem; color: #999; font-size: 0.85rem; line-height: 1.8;">
-                            <div><strong style="color: #D0D0D0;">CS3:</strong> Aggregations (COUNT, SUM, AVG)</div>
-                            <div><strong style="color: #D0D0D0;">CS4:</strong> Adds WHERE filters</div>
-                        </div>
-                    </div>
-                    
-                    <div>
-                        <div style="color: #f44336; font-weight: 600; font-size: 0.95rem; margin-bottom: 0.5rem;">Advanced Level</div>
-                        <div style="margin-left: 1rem; color: #999; font-size: 0.85rem; line-height: 1.8;">
-                            <div><strong style="color: #D0D0D0;">CS5:</strong> Multi-table JOINs</div>
-                        </div>
+                    <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><strong>CS5:</strong> Multi-table JOINs</div>
                     </div>
                     
                     <div style="margin-top: 1.5rem; padding-top: 1.25rem; border-top: 1px solid #3A3A3A;">
@@ -119,10 +106,14 @@ def dataset_viewer(shared_instruction, shared_schema):
                 lines=1
             )
             
-            # HuggingFace-style table
+            # HuggingFace-style table with row index on hover
             gr.HTML("""
                 <style>
-                /* True HuggingFace-style table */
+                /* True HuggingFace-style table - NO "Results" label */
+                .dataframe-container label {
+                    display: none !important;
+                }
+                
                 .dataframe-container {
                     border-radius: 8px !important;
                     overflow: hidden !important;
@@ -153,11 +144,29 @@ def dataset_viewer(shared_instruction, shared_schema):
                 .dataframe tbody tr {
                     background: #111827 !important;
                     border-bottom: 1px solid #1f2937 !important;
-                    transition: background-color 0.15s ease !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 {
@@ -173,15 +182,23 @@ def dataset_viewer(shared_instruction, shared_schema):
                 .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;
+                    padding: 0 !important;
+                    opacity: 0 !important;
+                    position: absolute !important;
+                }
                 </style>
             """)
             
             df_display = gr.Dataframe(
-                headers=COLUMNS,
-                datatype=["str", "str", "str"],
+                headers=["index"] + COLUMNS,
+                datatype=["number", "str", "str", "str"],
                 interactive=False,
                 wrap=True,
-                label="Results",
                 elem_classes="dataframe-container"
             )
             

tinysql_model_demo.py

@@ -61,15 +61,6 @@ def generate_sql(model_name, instruction, schema, max_length=256, temperature=0.
     except Exception as e:
         return f"Error: {str(e)}"
 
-# Only ONE example
-examples = [
-    [
-        "BM2_CS2_Syn (0.5B)",
-        "List worker earnings from highest to lowest",
-        "CREATE TABLE employees (name VARCHAR(100), salary INT, department VARCHAR(100))"
-    ],
-]
-
 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;">
@@ -98,24 +89,24 @@ def model_demo(shared_instruction, shared_schema):
                 info="Larger models = better accuracy"
             )
             
-            # Better formatted guide
+            # Better formatted guide with base models
             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="margin-bottom: 1rem;">
-                        <div style="color: #4CAF50; 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;">Lightning fast, ideal for simple queries</div>
+                        <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>
                     
                     <div style="margin-bottom: 1rem;">
-                        <div style="color: #FF9800; 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;">Balanced performance and speed</div>
+                        <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>
                     
                     <div style="margin-bottom: 1.25rem;">
-                        <div style="color: #f44336; 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;">Most accurate for complex queries</div>
+                        <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>
                     
                     <div style="padding-top: 1rem; border-top: 1px solid #3A3A3A;">
@@ -133,9 +124,12 @@ def model_demo(shared_instruction, shared_schema):
 
         with gr.Column(scale=2):
             gr.Markdown("### Your Query")
-            instruction = gr.Textbox(
+            
+            # Query in code box like schema
+            instruction = gr.Code(
                 label="Natural Language Query",
-                placeholder="e.g., Find all employees earning more than 50000 sorted by name",
+                language="text",
+                value="",
                 lines=2
             )
             
@@ -158,14 +152,15 @@ def model_demo(shared_instruction, shared_schema):
 
     gr.Markdown("### Example Query")
     
-    # Custom examples with 2-line schema display
+    # Custom example with query in code box
     gr.HTML("""
         <div style="background: #2A2A2A; border-radius: 12px; padding: 1.5rem; margin: 1rem 0; border: 1px solid #3A3A3A;">
             <div style="margin-bottom: 0.75rem;">
                 <strong style="color: #FF6B4A;">Model:</strong> <span style="color: #D0D0D0;">BM2_CS2_Syn (0.5B)</span>
             </div>
-            <div style="margin-bottom: 0.75rem;">
-                <strong style="color: #FF6B4A;">Query:</strong> <span style="color: #D0D0D0;">List worker earnings from highest to lowest</span>
+<div style="margin-bottom: 0.75rem;">
+                <strong style="color: #FF6B4A;">Query:</strong>
+                <pre style="background: #1A1A1A; padding: 0.75rem; border-radius: 6px; margin-top: 0.5rem; color: #D0D0D0; font-size: 0.85rem; line-height: 1.5;">List worker earnings from highest to lowest</pre>
             </div>
             <div>
                 <strong style="color: #FF6B4A;">Schema:</strong>