diff --git a/eval/dailsql/preprocessed_data/beaver_preprocess.py b/eval/dailsql/preprocessed_data/beaver_preprocess.py index 85c4036..1d4547e 100644 --- a/eval/dailsql/preprocessed_data/beaver_preprocess.py +++ b/eval/dailsql/preprocessed_data/beaver_preprocess.py @@ -281,7 +281,9 @@ def convert_beaver_questions_to_dailsql_format(dataset, beaver_questions_path, o # Load spacy for tokenization nlp = spacy.load("en_core_web_sm") - retrieved_tables = get_retrieved_tables(dataset) + retrieved_tables = None + if option == 1: + retrieved_tables = get_retrieved_tables(dataset) dailsql_format = [] for question_info in tqdm(beaver_questions): @@ -289,7 +291,7 @@ def convert_beaver_questions_to_dailsql_format(dataset, beaver_questions_path, o # Option 1+: Add gold tables to question if option >= 1: - if option == 1: + if option == 1 and retrieved_tables is not None: gold_tables = retrieved_tables[question_info['id']] else: gold_tables = question_info.get('tables', []) diff --git a/eval/dinsql/preprocessed_data/beaver_preprocess_v2.py b/eval/dinsql/preprocessed_data/beaver_preprocess_v2.py index 096bb3c..7c6fd3b 100644 --- a/eval/dinsql/preprocessed_data/beaver_preprocess_v2.py +++ b/eval/dinsql/preprocessed_data/beaver_preprocess_v2.py @@ -234,7 +234,9 @@ def convert_beaver_questions_to_dinsql_format(dataset, beaver_questions_path, ou with open(beaver_questions_path, 'r') as f: beaver_questions = json.load(f) - retrieved_tables = get_retrieved_tables(dataset) + retrieved_tables = None + if option == 1: + retrieved_tables = get_retrieved_tables(dataset) dinsql_format = [] for question_info in beaver_questions: @@ -247,7 +249,7 @@ def convert_beaver_questions_to_dinsql_format(dataset, beaver_questions_path, ou # Option 1+: Store gold_tables for filtering if option >= 1: - if option == 1: + if option == 1 and retrieved_tables is not None: base_item['tables'] = retrieved_tables[question_info['id']] else: base_item['tables'] = question_info.get('tables', []) diff --git a/eval/reforce/convert_beaver_to_reforce.py b/eval/reforce/convert_beaver_to_reforce.py index 72d661f..1dce934 100644 --- a/eval/reforce/convert_beaver_to_reforce.py +++ b/eval/reforce/convert_beaver_to_reforce.py @@ -325,7 +325,9 @@ def get_table_key(t_name): new_join_keys.append([c.lower() for c in jk]) item['join_keys'] = new_join_keys - retrieved_tables = get_retrieved_tables(dataset) + retrieved_tables = None + if preprocessing_option == 1: + retrieved_tables = get_retrieved_tables(dataset) for idx, item in enumerate(beaver_questions): if (idx + 1) % 10 == 0: diff --git a/test_bug/test_retrieval_conditional.py b/test_bug/test_retrieval_conditional.py new file mode 100644 index 0000000..f89eacb --- /dev/null +++ b/test_bug/test_retrieval_conditional.py @@ -0,0 +1,125 @@ +""" +Test: get_retrieved_tables is only called when retrieval is actually needed. + +Affected files: + - eval/reforce/convert_beaver_to_reforce.py + - eval/dinsql/preprocessed_data/beaver_preprocess_v2.py + - eval/dailsql/preprocessed_data/beaver_preprocess.py + +Before the fix: + get_retrieved_tables(dataset) was called unconditionally before + the option/ setting check, causing AssertionError for Setting 1 (option=2) + and Setting 2 (option=3) even though they use gold tables from data. + +After the fix: + get_retrieved_tables(dataset) is only called when option == 1 + (Setting 0 / end-to-end mode requiring retrieval). +""" +import sys +import os +import json + +REPO_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO_DIR) + +DATA_DIR = os.path.join(REPO_DIR, "data") +OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__)) + + +def test_reforce_setting_1(): + """ReFoRCE: Setting 1 (option=2) without retrieval.""" + from eval.reforce.convert_beaver_to_reforce import convert_beaver_to_reforce + + out = os.path.join(OUTPUT_DIR, "reforce_opt2.json") + sampled = out.replace(".json", "_sampled.json") + convert_beaver_to_reforce( + dataset="dw", + beaver_questions_path=os.path.join(DATA_DIR, "dw", "dev_sampled.json"), + beaver_tables_path=os.path.join(DATA_DIR, "dw", "dev_tables.json"), + output_path=out, + sampled_questions_path=sampled, + preprocessing_option=2, + db_id="dw", + ) + assert os.path.exists(out) + with open(out) as f: + data = json.load(f) + assert len(data) > 0 + os.remove(out) + if os.path.exists(sampled): + os.remove(sampled) + print(f" [PASS] ReFoRCE option=2: {len(data)} entries") + + +def test_dinsql_setting_1(): + """DIN-SQL: Setting 1 (option=2) without retrieval.""" + from eval.dinsql.preprocessed_data.beaver_preprocess_v2 import convert_beaver_questions_to_dinsql_format + + out = os.path.join(OUTPUT_DIR, "dinsql_opt2.json") + convert_beaver_questions_to_dinsql_format( + dataset="dw", + beaver_questions_path=os.path.join(DATA_DIR, "dw", "dev_sampled.json"), + output_path=out, + option=2, + ) + assert os.path.exists(out) + with open(out) as f: + data = json.load(f) + assert len(data) > 0 + os.remove(out) + print(f" [PASS] DIN-SQL option=2: {len(data)} entries") + + +def test_dailsql_setting_1(): + """DAIL-SQL: Setting 1 (option=2) without retrieval.""" + from eval.dailsql.preprocessed_data.beaver_preprocess import convert_beaver_questions_to_dailsql_format + + out = os.path.join(OUTPUT_DIR, "dailsql_opt2.json") + convert_beaver_questions_to_dailsql_format( + dataset="dw", + beaver_questions_path=os.path.join(DATA_DIR, "dw", "dev_sampled.json"), + output_path=out, + option=2, + ) + assert os.path.exists(out) + with open(out) as f: + data = json.load(f) + assert len(data) > 0 + os.remove(out) + print(f" [PASS] DAIL-SQL option=2: {len(data)} entries") + + +def test_setting_0_still_requires_retrieval(): + """Setting 0 (option=1) should still fail if retrieval files are missing.""" + from eval.reforce.convert_beaver_to_reforce import convert_beaver_to_reforce + + out = os.path.join(OUTPUT_DIR, "setting0_should_fail.json") + sampled = out.replace(".json", "_sampled.json") + try: + convert_beaver_to_reforce( + dataset="dw", + beaver_questions_path=os.path.join(DATA_DIR, "dw", "dev_sampled.json"), + beaver_tables_path=os.path.join(DATA_DIR, "dw", "dev_tables.json"), + output_path=out, + sampled_questions_path=sampled, + preprocessing_option=1, + db_id="dw", + ) + print(" [FAIL] Setting 0 should have raised an error (no retrieval files)") + except AssertionError: + print(" [PASS] Setting 0 raises AssertionError as expected (no retrieval files)") + + if os.path.exists(out): + os.remove(out) + if os.path.exists(sampled): + os.remove(sampled) + + +if __name__ == "__main__": + print("=== Test: Non-retrieval settings without retrieval files ===\n") + print("ReFoRCE:"); test_reforce_setting_1() + print("DIN-SQL:"); test_dinsql_setting_1() + print("DAIL-SQL:"); test_dailsql_setting_1() + print("\n=== Test: Setting 0 still requires retrieval ===") + test_setting_0_still_requires_retrieval() + print("\n=== All tests completed! ===")