new version

app.py

@@ -17,11 +17,9 @@ MODELS = {
     "BM3_CS3_Syn (1B)": "withmartian/sql_interp_bm3_cs3_experiment_9.3",
 }
 
-# Cache loaded models
 model_cache = {}
 
 def load_model(model_name):
-    """Load model and tokenizer with caching"""
     if model_name not in model_cache:
         model_id = MODELS[model_name]
         tokenizer = AutoTokenizer.from_pretrained(model_id)
@@ -34,19 +32,18 @@ def load_model(model_name):
     return model_cache[model_name]
 
 def generate_sql(model_name, instruction, schema, max_length=256, temperature=0.7):
-    """Generate SQL query from natural language"""
+    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)
         
-        # Format prompt
         prompt = f"""### Instruction: {instruction}
 ### Context: {schema}
 ### Response:"""
         
-        # Tokenize
         inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
         
-        # Generate
         outputs = model.generate(
             **inputs,
             max_length=max_length,
@@ -55,10 +52,8 @@ def generate_sql(model_name, instruction, schema, max_length=256, temperature=0.
             pad_token_id=tokenizer.eos_token_id
         )
         
-        # Decode
         generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Extract only the SQL response
         if "### Response:" in generated:
             sql = generated.split("### Response:")[-1].strip()
         else:
@@ -67,12 +62,12 @@ def generate_sql(model_name, instruction, schema, max_length=256, temperature=0.
         return sql
         
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"❌ Error: {str(e)}"
 
 # Example queries
 examples = [
     [
-        "BM1_CS1 (33M)",
+        "BM1_CS1_Syn (33M)",
         "Show me the name and salary from employees",
         "CREATE TABLE employees (name VARCHAR(100), salary INT, department VARCHAR(100))"
     ],
@@ -82,44 +77,167 @@ examples = [
         "CREATE TABLE employees (name VARCHAR(100), salary INT, department VARCHAR(100))"
     ],
     [
-        "BM3_CS3 (1B)",
+        "BM3_CS3_Syn (1B)",
         "Count how many employees in each department",
         "CREATE TABLE employees (name VARCHAR(100), salary INT, department VARCHAR(100))"
     ],
 ]
 
+# Custom CSS for beautiful styling
+custom_css = """
+.gradio-container {
+    font-family: 'Inter', sans-serif;
+}
+
+.header-section {
+    text-align: center;
+    padding: 2rem 0;
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    border-radius: 12px;
+    margin-bottom: 2rem;
+    color: white;
+}
+
+.logo-container {
+    display: flex;
+    justify-content: center;
+    align-items: center;
+    gap: 1rem;
+    margin-bottom: 1rem;
+}
+
+.martian-badge {
+    background: rgba(255, 255, 255, 0.2);
+    padding: 0.5rem 1rem;
+    border-radius: 20px;
+    font-size: 0.9rem;
+    backdrop-filter: blur(10px);
+}
+
+.info-box {
+    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
+    border-radius: 12px;
+    padding: 1.5rem;
+    margin: 1rem 0;
+    border-left: 4px solid #667eea;
+}
+
+.citation-box {
+    background: #f8f9fa;
+    border: 1px solid #dee2e6;
+    border-radius: 8px;
+    padding: 1.5rem;
+    margin: 2rem 0;
+    font-family: 'Monaco', 'Courier New', monospace;
+    font-size: 0.85rem;
+}
+
+.citation-header {
+    font-weight: bold;
+    color: #495057;
+    margin-bottom: 0.5rem;
+    display: flex;
+    align-items: center;
+    gap: 0.5rem;
+}
+
+.resource-links {
+    display: flex;
+    gap: 1rem;
+    justify-content: center;
+    margin: 1.5rem 0;
+    flex-wrap: wrap;
+}
+
+.resource-link {
+    background: white;
+    padding: 0.75rem 1.5rem;
+    border-radius: 8px;
+    text-decoration: none;
+    color: #667eea;
+    border: 2px solid #667eea;
+    font-weight: 500;
+    transition: all 0.3s ease;
+}
+
+.resource-link:hover {
+    background: #667eea;
+    color: white;
+}
+
+footer {
+    text-align: center;
+    padding: 2rem 0;
+    color: #6c757d;
+    border-top: 1px solid #dee2e6;
+    margin-top: 3rem;
+}
+"""
+
 # Create Gradio interface
-with gr.Blocks(title="TinySQL Demo") as demo:
-    gr.Markdown("""
-    # 🔍 TinySQL: Text-to-SQL Generation Demo
+with gr.Blocks(css=custom_css, title="TinySQL Demo | Martian", theme=gr.themes.Soft()) as demo:
     
-    Generate SQL queries from natural language using models trained on TinySQL.
-    Select a model, provide a natural language instruction and database schema, then click **Generate**.
+    # Header with Martian branding
+    gr.HTML("""
+        <div class="header-section">
+            <div class="logo-container">
+                <h1 style="margin: 0; font-size: 2.5rem;">🔮 TinySQL Interactive Demo</h1>
+            </div>
+            <div class="martian-badge">
+                ⚡ Powered by Martian
+            </div>
+            <p style="font-size: 1.1rem; margin-top: 1rem; opacity: 0.9;">
+                Transform natural language into SQL queries using mechanistically interpretable models
+            </p>
+        </div>
+    """)
     
-    **Model Types:**
-    - **BM1** (33M params): TinyStories-based, fastest
-    - **BM2** (0.5B params): Qwen2.5-based, balanced
-    - **BM3** (1B params): Llama-3.2-based, most accurate
-    - **Syn** variants: Trained on synonym dataset (handles semantic mappings)
+    # 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=2):
+        with gr.Column(scale=1):
+            gr.Markdown("### 🎛️ Configuration")
+            
             model_dropdown = gr.Dropdown(
                 choices=list(MODELS.keys()),
-                value="BM2_CS1_Syn (0.5B)",
-                label="Select Model",
-                info="Choose model size and training dataset"
+                value="BM2_CS2_Syn (0.5B)",
+                label="🤖 Model Selection",
+                info="Larger models = better accuracy, slower inference"
             )
             
+            gr.Markdown("""
+            **Model Guide:**
+            - 🟢 **BM1 (33M)**: Lightning fast, great for simple queries
+            - 🟡 **BM2 (0.5B)**: Balanced performance and speed
+            - 🔴 **BM3 (1B)**: Most accurate, handles complex queries
+            
+            **Dataset Complexity:**
+            - **CS1**: Basic SELECT-FROM queries
+            - **CS2**: Adds ORDER BY clauses
+            - **CS3**: Aggregations (COUNT, SUM, AVG)
+            - **CS4**: Adds WHERE filters
+            - **CS5**: Multi-table JOINs
+            """)
+            
+        with gr.Column(scale=2):
+            gr.Markdown("### 💬 Your Query")
+            
             instruction = gr.Textbox(
-                label="Natural Language Query",
-                placeholder="e.g., Show me all employees with salary greater than 50000",
+                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",
+                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))"
@@ -131,39 +249,80 @@ with gr.Blocks(title="TinySQL Demo") as demo:
                     maximum=512,
                     value=256,
                     step=32,
-                    label="Max Length"
+                    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"
+                    label="Temperature",
+                    info="Higher = more creative (use 0.1 for accuracy)"
                 )
             
-            generate_btn = gr.Button("Generate SQL", variant="primary")
-        
-        with gr.Column(scale=1):
-            output = gr.Textbox(
-                label="Generated SQL",
-                lines=10,
-                placeholder="SQL query will appear here..."
+            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.Markdown("### 💡 Try These Examples")
     gr.Examples(
         examples=examples,
         inputs=[model_dropdown, instruction, schema],
     )
     
-    gr.Markdown("""
-    ---
-    **Paper:** [TinySQL: A Progressive Text-to-SQL Dataset for Mechanistic Interpretability Research](https://arxiv.org/abs/2503.12730)
+    # 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">
+                🤗 Get Dataset & Models
+            </a>
+            <a href="https://withmartian.com" class="resource-link" target="_blank">
+                🚀 Visit Martian
+            </a>
+        </div>
+    """)
+    
+    # Citation box
+    gr.HTML("""
+        <div class="citation-box">
+            <div class="citation-header">
+                📚 Citation
+            </div>
+            <pre style="margin: 0; overflow-x: auto;">@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>
+    """)
     
-    **Resources:** [GitHub](https://github.com/withmartian/TinySQL) | [Dataset](https://huggingface.co/collections/withmartian/tinysql-6760e92748b63fa56a6ffc9f)
+    # Footer
+    gr.HTML("""
+        <footer>
+            <p style="margin: 0.5rem 0;">
+                Built with ❤️ by the Martian team
+            </p>
+            <p style="margin: 0; font-size: 0.9rem;">
+                Bridging the gap between toy tasks and real-world interpretability
+            </p>
+        </footer>
     """)
     
-    # Connect button
     generate_btn.click(
         fn=generate_sql,
         inputs=[model_dropdown, instruction, schema, max_length, temperature],