Fix prompts buttons, and NL2SQL bug (#24)

- Fix prompts buttons, and NL2SQL bug (b8f53f4140ce72bf889c039fa072989834ee8d73)


Co-authored-by: Francesco Giannuzzo <[email protected]>

app.py

@@ -509,9 +509,13 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
             selected_models = [model for model, selected in zip(model_list, model_selections) if selected]
             input_data['models'] = selected_models
             button_state = bool(selected_models and '{db_schema}' in input_data["prompt"] and '{question}' in input_data["prompt"])
-            return selected_models, gr.update(open=True, visible=True), gr.update(interactive=button_state), gr.update(interactive=button_state)
+            return selected_models, gr.update(open=True, visible=True), gr.update(interactive=button_state)
 
         # Add the Textbox to the interface
+        with gr.Row():
+            button_prompt_nlsql = gr.Button("Choose NL2SQL task")
+            button_prompt_tqa = gr.Button("Choose TQA task")
+
         prompt = gr.TextArea(
             label="Customise the prompt for selected models here or leave the default one.",
             placeholder=prompt_default,
@@ -522,17 +526,20 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
 
         # Submit button (initially disabled)
         with gr.Row():
-            submit_models_button = gr.Button("Submit Models for NL2SQL task", interactive=False)
-            submit_models_button_tqa = gr.Button("Submit Models for TQA task", interactive=False)
+            submit_models_button = gr.Button("Submit Models", interactive=False)
 
         def check_prompt(prompt):
             #TODO
             missing_elements = []
             if(prompt==""):
-                input_data["prompt"] = prompt_default
+                global flag_TQA
+                if not flag_TQA:
+                    input_data["prompt"] = prompt_default
+                else:
+                    input_data["prompt"] = prompt_default_tqa
                 button_state = bool(len(input_data['models']) > 0 and '{db_schema}' in input_data["prompt"] and '{question}' in input_data["prompt"])
             else:
-                input_data["prompt"]=prompt
+                input_data["prompt"] = prompt
                 if "{db_schema}" not in prompt:
                     missing_elements.append("{db_schema}")
                 if "{question}" not in prompt:
@@ -543,21 +550,21 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
                         value=f"<div style='text-align: center; font-size: 18px; font-weight: bold;'>"
                             f"❌ Missing {', '.join(missing_elements)} in the prompt ❌</div>",
                         visible=True
-                    ), gr.update(interactive=button_state)
-            return gr.update(visible=False),  gr.update(interactive=button_state)
+                    ), gr.update(interactive=button_state), gr.TextArea(placeholder=input_data["prompt"])
+            return gr.update(visible=False), gr.update(interactive=button_state), gr.TextArea(placeholder=input_data["prompt"])
 
-        prompt.change(fn=check_prompt, inputs=[prompt], outputs=[warning_prompt, submit_models_button, submit_models_button_tqa])
+        prompt.change(fn=check_prompt, inputs=[prompt], outputs=[warning_prompt, submit_models_button])
         # Link checkboxes to selection events
         for checkbox in model_checkboxes:
             checkbox.change(
                 fn=get_selected_models, 
                 inputs=model_checkboxes, 
-                outputs=[selected_models_output, select_model_acc, submit_models_button, submit_models_button_tqa]
+                outputs=[selected_models_output, select_model_acc, submit_models_button]
             )
         prompt.change(
             fn=get_selected_models, 
             inputs=model_checkboxes, 
-            outputs=[selected_models_output, select_model_acc, submit_models_button, submit_models_button_tqa]
+            outputs=[selected_models_output, select_model_acc, submit_models_button]
         )
 
         submit_models_button.click(
@@ -566,11 +573,6 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
             outputs=[selected_models_output, select_model_acc, qatch_acc]
         )
 
-        submit_models_button_tqa.click(
-            fn=lambda *args: (get_selected_models(*args), gr.update(open=False, visible=True), gr.update(open=True, visible=True)),
-            inputs=model_checkboxes,
-            outputs=[selected_models_output, select_model_acc, qatch_acc]
-        )
         def change_flag():
             global flag_TQA
             flag_TQA = True
@@ -579,8 +581,14 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
             global flag_TQA
             flag_TQA = False
 
-        submit_models_button.click(fn = dis_flag, inputs=[], outputs=[])
-        submit_models_button_tqa.click(fn = change_flag, inputs=[], outputs=[])
+        button_prompt_tqa.click(fn = change_flag, inputs=[], outputs=[])
+
+        button_prompt_nlsql.click(fn = dis_flag, inputs=[], outputs=[])
+
+        button_prompt_tqa.click(fn=check_prompt, inputs=[prompt], outputs=[warning_prompt, submit_models_button, prompt])
+
+        button_prompt_nlsql.click(fn=check_prompt, inputs=[prompt], outputs=[warning_prompt, submit_models_button, prompt])      
+
 
         def enable_disable(enable):
             return (
@@ -592,7 +600,6 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
                 gr.update(interactive=enable), 
                 gr.update(interactive=enable),
                 *[gr.update(interactive=enable) for _ in table_outputs],
-                gr.update(interactive=enable),
                 gr.update(interactive=enable)
             )
         
@@ -610,24 +617,7 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
                 default_checkbox,
                 table_selector,
                 *table_outputs,
-                open_model_selection,
-                submit_models_button_tqa
-            ]
-        )
-        submit_models_button_tqa.click(
-            fn=enable_disable,
-            inputs=[gr.State(False)],
-            outputs=[
-                *model_checkboxes,
-                submit_models_button,
-                preview_output,
-                submit_button,
-                file_input,
-                default_checkbox,
-                table_selector,
-                *table_outputs,
-                open_model_selection,
-                submit_models_button_tqa
+                open_model_selection
             ]
         )
         
@@ -645,8 +635,7 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
                 default_checkbox,
                 table_selector,
                 *table_outputs,
-                open_model_selection,
-                submit_models_button_tqa
+                open_model_selection
             ]
         )
 
@@ -749,13 +738,28 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
                 yield gr.Markdown(eval_text, visible=True), gr.Image(), gr.Markdown(), gr.Markdown(), gr.Markdown(), metrics_conc, *[predictions_dict[model][columns_to_visulize] for model in model_list]
             
             else:
+                global flag_TQA
                 orchestrator_generator = OrchestratorGenerator()
                 target_df = orchestrator_generator.generate_dataset(connector=input_data['data']['db'], tables_to_include=input_data['data']['selected_tables'])
 
                 #create target_df[target_answer]
                 if flag_TQA : 
-                    if (input_data["prompt"] == prompt_default):
-                        input_data["prompt"] = prompt_default_tqa
+                    # if (input_data["prompt"] == prompt_default):
+                    #     input_data["prompt"] = prompt_default_tqa
+
+                    target_df['db_schema'] = target_df.apply(
+                        lambda row: utils_get_db_tables_info.utils_extract_db_schema_as_string(
+                            db_id=input_data["db_name"],
+                            base_path=input_data["data_path"],
+                            normalize=False,
+                            sql=row["query"],
+                            get_insert_into=True,
+                            model=None,
+                            prompt=input_data["prompt"].format(question=row["question"], db_schema="")
+                        ),
+                        axis=1
+                    )
+                    
                     target_df = us.extract_answer(target_df)
 
                 predictor = ModelPrediction()
@@ -766,6 +770,7 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
                     count=0
                     for index, row in target_df.iterrows():
                         if (reset_flag == False):
+                            global flag_TQA
                             percent_complete = round(((index+1) / len(target_df)) * 100, 2)
                             load_text = f"{generate_loading_text(percent_complete)}"
                             
@@ -780,7 +785,6 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
                             #samples = us.generate_some_samples(input_data["data_path"], row["tbl_name"])
                             model_to_send = None if not flag_TQA else model
 
-
                             db_schema_text = utils_get_db_tables_info.utils_extract_db_schema_as_string(
                                 db_id = input_data["db_name"], 
                                 base_path = input_data["data_path"], 
@@ -806,11 +810,11 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
                                 prompt=f"{prompt_to_send}",
                                 task=task 
                             )
+                            end_time = time.time()
                             prediction = response['response_parsed']
                             price = response['cost']
                             answer = response['response']
 
-                            end_time = time.time()
                             if flag_TQA:
                                 task_string = "Answer"
                             else:
@@ -857,6 +861,7 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
                 evaluator = OrchestratorEvaluator()
 
                 for model in input_data["models"]:
+                    global flag_TQA
                     if not flag_TQA:
                         metrics_df_model = evaluator.evaluate_df(
                             df=predictions_dict[model],
@@ -920,11 +925,6 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
             inputs=[],
             outputs=[tab_dict[model] for model in model_list]  # Update TabItem visibility
         )
-        submit_models_button_tqa.click(
-            change_tab,
-            inputs=[],
-            outputs=[tab_dict[model] for model in model_list]  # Update TabItem visibility
-        )
 
         selected_models_display = gr.JSON(label="Final input data", visible=False)
         metrics_df = gr.DataFrame(visible=False)
@@ -936,20 +936,10 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
             outputs=[evaluation_loading, model_logo, variable, question_display, prediction_display, metrics_df] + list(dataframe_per_model.values())
         )
 
-        submit_models_button_tqa.click(
-            fn=qatch_flow_nl_sql,
-            inputs=[],
-            outputs=[evaluation_loading, model_logo, variable, question_display, prediction_display, metrics_df] + list(dataframe_per_model.values())
-        )
-
         submit_models_button.click(
             fn=lambda: gr.update(value=input_data), 
             outputs=[selected_models_display]
         )
-        submit_models_button_tqa.click(
-            fn=lambda: gr.update(value=input_data), 
-            outputs=[selected_models_display]
-        )
 
         # Works for METRICS
         metrics_df.change(fn=change_text, inputs=[metrics_df], outputs=[metrics_df_out])
@@ -972,10 +962,6 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
             fn=lambda: gr.update(visible=False),
             outputs=[download_metrics]
         )
-        submit_models_button_tqa.click(
-            fn=lambda: gr.update(visible=False),
-            outputs=[download_metrics]
-        )
 
         def refresh():
             global reset_flag
@@ -1007,8 +993,7 @@ with gr.Blocks(theme='shivi/calm_seafoam', css_paths='style.css', js=js_func) as
                 default_checkbox,
                 table_selector,
                 *table_outputs,
-                open_model_selection,
-                submit_models_button_tqa
+                open_model_selection
             ]
         )
 

utilities.py

@@ -8,6 +8,7 @@ import os
 from qatch.connectors.sqlite_connector import SqliteConnector
 from qatch.evaluate_dataset.metrics_evaluators import CellPrecision, CellRecall, ExecutionAccuracy, TupleCardinality, TupleConstraint, TupleOrder, ValidEfficiencyScore
 import qatch.evaluate_dataset.orchestrator_evaluator as eva
+import utils_get_db_tables_info
 #import tiktoken
 from transformers import AutoTokenizer
 
@@ -151,11 +152,16 @@ def extract_answer(df):
     answers = []
     for _, row in df.iterrows():
         query = row["query"]
-        db_path = row["db_path"]
-        try:
-            conn = SqliteConnector(relative_db_path = db_path , db_name= "db")
-            answer = eva._utils_run_query_if_str(query, conn)
+        db_schema = row["db_schema"]
+        #db_path = row["db_path"]
+        try: 
+            conn = utils_get_db_tables_info.create_db_temp(db_schema)
+
+            result = pd.read_sql_query(query, conn)
+            answer = result.values.tolist()  # Convert the DataFrame to a list of lists
+
             answers.append(answer)
+            conn.close()
         except Exception as e:
             answers.append(f"Error: {e}")
     

utils_get_db_tables_info.py

@@ -49,11 +49,15 @@ def _get_schema_entries(cursor, sql=None, get_insert_into=False, model: str | No
         tables = [tbl[0] for tbl in cursor.fetchall()]
 
     for table in tables:
+        entries_per_table = []
         # Retrieve the CREATE TABLE statement for each table
         cursor.execute(f"SELECT sql FROM sqlite_master WHERE type='table' AND name='{table}' AND sql IS NOT NULL;")
         create_table_stmt = cursor.fetchone()
         if create_table_stmt:
-            entries.append(create_table_stmt[0])
+            stmt = create_table_stmt[0].strip()
+            if not stmt.endswith(';'):
+                stmt += ';'
+            entries_per_table.append(stmt)
 
         if get_insert_into:
             # Retrieve all data from the table
@@ -70,9 +74,10 @@ def _get_schema_entries(cursor, sql=None, get_insert_into=False, model: str | No
             for row in rows[:max_len]:
                 values = ', '.join(f"'{str(value)}'" if isinstance(value, str) else str(value) for value in row)
                 insert_stmt = f"INSERT INTO {table} ({', '.join(column_names)}) VALUES ({values});"
-                entries.append(insert_stmt)
+                entries_per_table.append(insert_stmt)
     
-    if model != None : entries = us.crop_entries_per_token(entries, model, prompt)
+        if model != None : entries_per_table = us.crop_entries_per_token(entries_per_table, model, prompt)
+        entries.extend(entries_per_table)
 
     return entries
 
@@ -112,3 +117,26 @@ def _combine_schema_entries(schema_entries, normalize):
         )
         for entry in schema_entries
     )
+
+
+def create_db_temp(schema_sql: str) -> sqlite3.Connection:
+    """
+    Creates a temporary SQLite database in memory by executing the provided SQL schema.
+    
+    Args:
+        schema_sql (str): The SQL code containing CREATE TABLE and INSERT INTO.
+        
+    Returns:
+        sqlite3.Connection: Connection object to the temporary database.
+    """
+    conn = sqlite3.connect(':memory:')
+    cursor = conn.cursor()
+    
+    try:
+        cursor.executescript(schema_sql)
+        conn.commit()
+    except sqlite3.Error as e:
+        conn.close()
+        raise
+
+    return conn