feat: small ux updates

app.py

@@ -88,22 +88,26 @@ with gr.Blocks() as demo:
                 search_type="dataset",
                 sumbit_on_select=True,
             )
-
+            query = gr.Textbox(
+                label="Natural Language Query",
+                placeholder="Enter a natural language query to generate SQL",
+            )
+            sql_out = gr.Code(
+                label="SQL Query",
+                interactive=True,
+                language="sql",
+                lines=1,
+                visible=False,
+            )
+    with gr.Row():
+        with gr.Column():
             btn = gr.Button("Show Dataset")
+        with gr.Column():
+            btn2 = gr.Button("Query Dataset")
     with gr.Row():
         search_out = gr.HTML(label="Search Results")
     with gr.Row():
         features = gr.Code(label="Features", language="json", visible=False)
-    with gr.Row():
-        query = gr.Textbox(
-            label="Natural Language Query",
-            placeholder="Enter a natural language query to generate SQL",
-        )
-    with gr.Row():
-        sql_out = gr.Code(label="SQL Query", interactive=True, language="sql", lines=1)
-    with gr.Row():
-        btn2 = gr.Button("Query Dataset")
-
     gr.on(
         [btn.click, search_in.submit],
         fn=get_iframe,
@@ -114,8 +118,8 @@ with gr.Blocks() as demo:
         inputs=[search_in],
         outputs=[features],
     )
-
-    btn2.click(
+    gr.on(
+        [btn2.click, query.submit],
         fn=query_dataset,
         inputs=[search_in, features, query],
         outputs=[sql_out, search_out],