tinysql_dataset_viewer.py

import gradio as gr
from datasets import load_dataset
import pandas as pd

DATASETS = {
    "CS1": "withmartian/cs1_dataset",
    "CS2": "withmartian/cs2_dataset",
    "CS3": "withmartian/cs3_dataset",
    "CS2 Synonyms": "withmartian/cs2_dataset_synonyms",
    "CS3 Synonyms": "withmartian/cs3_dataset_synonyms",
    "CS4 Synonyms": "withmartian/cs4_dataset_synonyms",
}

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):
    """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, search_column):
    if not search_query or df.empty or "Error" in df.columns:
        return df
    
    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: 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>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Controls")
            
            dataset_dropdown = gr.Dropdown(
                choices=list(DATASETS.keys()),
                value="CS1",
                label="Choose Dataset",
                info="Select complexity level"
            )
            
            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>
                </div>
            """)
            
            gr.Markdown("### Test Example")
            row_selector = gr.Number(
                label="Row Number",
                value=0,
                minimum=0,
                precision=0,
                info="Pick a row to test"
            )
            
            send_to_model_btn = gr.Button("Run in Model Demo", variant="primary")
        
        with gr.Column(scale=3):
            gr.Markdown("### Dataset Preview")
            
            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
                )
            
            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"
            )
            
            stats_display = gr.Markdown("Click **Load Dataset** to begin")
    
    df_state = gr.State(value=pd.DataFrame())
    
    def load_and_display(dataset_name):
        df = load_preview(dataset_name)
        if "Error" in df.columns:
            return df, df, "Error loading dataset"
        stats = f"**Loaded {len(df)} rows** • Columns: {', '.join(COLUMNS)}"
        return df, df, stats
    
    load_btn.click(
        fn=load_and_display,
        inputs=dataset_dropdown,
        outputs=[df_state, df_display, stats_display]
    )
    
    def search_and_display(df, query, column):
        if df.empty:
            return df, "Load a dataset first"
        
        filtered_df = filter_dataframe(df, query, column)
        stats = f"**Showing {len(filtered_df)} of {len(df)} rows**"
        if query:
            stats += f" • Search: '{query}' in {column}"
        return filtered_df, stats
    
    search_box.change(
        fn=search_and_display,
        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]
    )
    
    def send_to_model(df, row_num):
        if df.empty or row_num >= len(df):
            return "", "", "Invalid row or no data loaded"
        
        row = df.iloc[int(row_num)]
        instruction = row['english_prompt'] if 'english_prompt' in row else ""
        schema = row['create_statement'] if 'create_statement' in row else ""
        
        return instruction, schema, f"**Row {row_num} loaded!** Switch to Model Demo tab"
    
    send_to_model_btn.click(
        fn=send_to_model,
        inputs=[df_state, row_selector],
        outputs=[shared_instruction, shared_schema, stats_display]
    )
    
    return {'df_state': df_state, 'df_display': df_display}