use smaller dataset, improve graph

app.py

@@ -86,7 +86,7 @@ def generate_query(dataset_id: str, query: str) -> str:
     ddl = get_dataset_ddl(dataset_id)
 
     system_prompt = f"""
-    You are an expert SQL assistant with access to the following DuckDB Table:
+    You are an expert SQL assistant with access to the following PostgreSQL Table:
     
     ```sql
     {ddl}
@@ -116,18 +116,24 @@ def generate_query(dataset_id: str, query: str) -> str:
 
 def query_dataset(dataset_id: str, query: str) -> Tuple[pd.DataFrame, str, plt.Figure]:
     response: SQLResponse = generate_query(dataset_id, query)
+
+    print("Querying Parquet...")
     df = conn.execute(response.sql).fetchdf()
 
     plot = None
 
     if response.visualization_type == OutputTypes.LINECHART:
         plot = df.plot(
-            kind="line", x=response.data_key, y=response.label_key
+            kind="line", x=response.label_key, y=response.data_key
         ).get_figure()
+        plt.xticks(rotation=45, ha="right")
+        plt.tight_layout()
     elif response.visualization_type == OutputTypes.BARCHART:
         plot = df.plot(
-            kind="bar", x=response.data_key, y=response.label_key
+            kind="bar", x=response.label_key, y=response.data_key
         ).get_figure()
+        plt.xticks(rotation=45, ha="right")
+        plt.tight_layout()
 
     markdown_output = f"""```sql\n{response.sql}\n```"""
     return df, markdown_output, plot
@@ -139,7 +145,7 @@ with gr.Blocks() as demo:
         label="Hub Dataset ID",
         placeholder="Find your favorite dataset...",
         search_type="dataset",
-        value="teknium/OpenHermes-2.5",
+        value="gretelai/synthetic_text_to_sql",
     )
     user_query = gr.Textbox("", label="Ask anything...")