File size: 9,462 Bytes
e07e1fe
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a6c02ed
e07e1fe
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a6c02ed
e07e1fe
 
 
 
 
 
 
 
 
 
 
 
 
 
f35a40c
 
599e176
 
 
 
f35a40c
 
 
 
 
599e176
 
 
 
f35a40c
 
 
 
 
 
 
 
 
 
599e176
f35a40c
599e176
f35a40c
599e176
 
 
 
f72a9c9
 
599e176
 
 
f72a9c9
 
599e176
 
 
f72a9c9
 
599e176
 
 
 
 
 
 
 
 
 
 
 
f35a40c
 
 
 
 
f72a9c9
0f515d8
599e176
0f515d8
 
 
f35a40c
a6c02ed
599e176
f35a40c
599e176
 
 
f35a40c
 
 
599e176
a6c02ed
f35a40c
 
599e176
 
f35a40c
599e176
f35a40c
a6c02ed
599e176
 
a6c02ed
 
 
599e176
a6c02ed
 
 
 
599e176
a6c02ed
599e176
a6c02ed
 
 
 
 
599e176
 
 
 
 
f35a40c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
599e176
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

MODELS = {
    "BM1_CS1_Syn (33M)": "withmartian/sql_interp_bm1_cs1_experiment_1.10",
    "BM1_CS2_Syn (33M)": "withmartian/sql_interp_bm1_cs2_experiment_2.10",
    "BM1_CS3_Syn (33M)": "withmartian/sql_interp_bm1_cs3_experiment_3.10",
    "BM1_CS4_Syn (33M)": "withmartian/sql_interp_bm1_cs4_dataset_synonyms_experiment_1.1",
    "BM1_CS5_Syn (33M)": "withmartian/sql_interp_bm1_cs5_dataset_synonyms_experiment_1.2",
    "BM2_CS1_Syn (0.5B)": "withmartian/sql_interp_bm2_cs1_experiment_4.3",
    "BM2_CS2_Syn (0.5B)": "withmartian/sql_interp_bm2_cs2_experiment_5.3",
    "BM2_CS3_Syn (0.5B)": "withmartian/sql_interp_bm2_cs3_experiment_6.3",
    "BM3_CS1_Syn (1B)": "withmartian/sql_interp_bm3_cs1_experiment_7.3",
    "BM3_CS2_Syn (1B)": "withmartian/sql_interp_bm3_cs2_experiment_8.3",
    "BM3_CS3_Syn (1B)": "withmartian/sql_interp_bm3_cs3_experiment_9.3",
}

model_cache = {}

def load_model(model_name):
    if model_name not in model_cache:
        model_id = MODELS[model_name]
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        model_cache[model_name] = (tokenizer, model)
    return model_cache[model_name]

def generate_sql(model_name, instruction, schema, max_length=256, temperature=0.0):
    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)
        
        prompt = f"""### Instruction: {instruction}
### Context: {schema}
### Response:"""
        
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            temperature=temperature if temperature > 0 else 1.0,
            do_sample=temperature > 0,
            pad_token_id=tokenizer.eos_token_id
        )
        
        generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
        if "### Response:" in generated:
            sql = generated.split("### Response:")[-1].strip()
        else:
            sql = generated.strip()
        return sql
        
    except Exception as e:
        return f"Error: {str(e)}"

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;">
            <h2 style="font-size: 2rem; font-weight: 700; margin-bottom: 0.5rem; color: white;">Interactive SQL Generation</h2>
            <p style="font-size: 1rem; opacity: 0.9; line-height: 1.6; color: #D0D0D0;">
                Transform natural language into SQL using <span style="color: #FF6B4A; font-weight: 600;">mechanistically interpretable</span> models
            </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;">
            <p style="color: #D0D0D0; margin: 0; line-height: 1.6;">
                <strong style="color: #FF6B4A;">How it works:</strong> Select a model, describe your query in plain English, and watch the model generate SQL.
            </p>
        </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Configuration")
            model_dropdown = gr.Dropdown(
                choices=list(MODELS.keys()),
                value="BM2_CS2_Syn (0.5B)",
                label="Model Selection",
                info="Larger models = better accuracy"
            )
            
            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: #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: #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: #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;">
                        <div style="color: #FF6B4A; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.75rem;">Dataset Complexity</div>
                        <div style="color: #999; font-size: 0.85rem; line-height: 1.8;">
                            <div><strong style="color: #D0D0D0;">CS1:</strong> Basic SELECT-FROM</div>
                            <div><strong style="color: #D0D0D0;">CS2:</strong> Adds ORDER BY</div>
                            <div><strong style="color: #D0D0D0;">CS3:</strong> Aggregations</div>
                            <div><strong style="color: #D0D0D0;">CS4:</strong> WHERE filters</div>
                            <div><strong style="color: #D0D0D0;">CS5:</strong> Multi-table JOINs</div>
                        </div>
                    </div>
                </div>
            """)

        with gr.Column(scale=2):
            gr.Markdown("### Your Query")
            
            instruction = gr.Textbox(
                label="Natural Language Query",
                placeholder="e.g., Find all employees earning more than 50000 sorted by name",
                lines=2,
                value=""
            )
            
            schema = gr.Code(
                label="Database Schema",
                language="sql",
                value="CREATE TABLE employees (name VARCHAR(100), salary INT, department VARCHAR(100))",
                lines=3
            )
            
            with gr.Row():
                max_length = gr.Slider(64, 512, value=256, step=32, label="Max Length", info="Token limit")
                temperature = gr.Slider(0.0, 1.0, value=0.0, step=0.1, label="Temperature", info="Creativity level (0 = deterministic)")
            
            generate_btn = gr.Button("Generate SQL", variant="primary", size="lg")
            
            output = gr.Code(label="Generated SQL Query", language="sql", lines=6)

    gr.Markdown("### Example Query")
    
    # Better styled example with proper font
    gr.HTML("""
        <div style="background: #2A2A2A; border-radius: 12px; padding: 1.5rem; margin: 1rem 0; border: 1px solid #3A3A3A;">
            <div style="margin-bottom: 1rem; padding-bottom: 1rem; border-bottom: 1px solid #3A3A3A;">
                <div style="color: #FF6B4A; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.5rem;">Model</div>
                <div style="background: #1A1A1A; padding: 0.75rem; border-radius: 6px; color: #D0D0D0; font-family: 'Inter', sans-serif; font-size: 0.9rem;">BM2_CS2_Syn (0.5B)</div>
            </div>
            
            <div style="margin-bottom: 1rem; padding-bottom: 1rem; border-bottom: 1px solid #3A3A3A;">
                <div style="color: #FF6B4A; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.5rem;">Natural Language Query</div>
                <div style="background: #1A1A1A; padding: 0.75rem; border-radius: 6px; color: #D0D0D0; font-family: 'Inter', sans-serif; font-size: 0.9rem; line-height: 1.5;">List worker earnings from highest to lowest</div>
            </div>
            
            <div>
                <div style="color: #FF6B4A; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.5rem;">Database Schema</div>
                <pre style="background: #1A1A1A; padding: 0.75rem; border-radius: 6px; margin: 0; color: #D0D0D0; font-family: 'Monaco', 'Courier New', monospace; font-size: 0.85rem; line-height: 1.5; overflow-x: auto;">CREATE TABLE employees (
    name VARCHAR(100), 
    salary INT, 
    department VARCHAR(100)
)</pre>
            </div>
        </div>
    """)

    shared_instruction.change(
        fn=lambda x: x,
        inputs=shared_instruction,
        outputs=instruction
    )
    
    shared_schema.change(
        fn=lambda x: x,
        inputs=shared_schema,
        outputs=schema
    )

    generate_btn.click(
        fn=generate_sql,
        inputs=[model_dropdown, instruction, schema, max_length, temperature],
        outputs=output
    )
    
    return {'instruction': instruction, 'schema': schema, 'output': output}