app.py

from tinysql_model_demo import model_demo
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 DARK MODE - But preserve styled boxes */
.gradio-container, 
body, 
html,
.contain,
.app,
#root,
.main {
    background-color: var(--martian-black) !important;
    color-scheme: dark !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;
}

.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;
    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;
    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;
    color: white !important;
    padding: 0.5rem 1rem !important;
    border-radius: 8px !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;
    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;
    color: #E0E0E0 !important;
}

/* Citation box */
.citation-content {
    background: var(--martian-gray-medium) !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;
}

/* Code blocks */
code, pre {
    background: var(--martian-gray-dark) !important;
    color: #D0D0D0 !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:
    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>
    """)
    
    shared_instruction = gr.State("")
    shared_schema = gr.State("")
    
    with gr.Tabs():
        with gr.Tab("Dataset Viewer"):
            viewer_components = dataset_viewer(shared_instruction, shared_schema)
        
        with gr.Tab("Model Demo"):
            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,
    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>
        </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>
            </p>
        </div>
    """)

if __name__ == "__main__":
    demo.launch(ssr_mode=False)