use concurrent futures instead of signal

duckdb-nsql/eval/evaluate.py

@@ -12,6 +12,7 @@ import click
 import pandas as pd
 from rich.console import Console
 from tqdm.auto import tqdm
+from concurrent.futures import ThreadPoolExecutor, TimeoutError
 
 sys.path.append(os.path.join(os.path.dirname(__file__), "."))
 # from metrics.spider import evaluation as spider_evaluation  # type: ignore # noqa: E402
@@ -113,15 +114,24 @@ def compute_exact_match_metric(
     return exact_match
 
 
+def evaluate_with_timeout(evaluator, *args, timeout):
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        future = executor.submit(evaluator.evaluate_one, *args)
+        try:
+            result = future.result(timeout=timeout)
+        except TimeoutError:
+            result = None
+    return result
+
 def compute_test_suite_metric(
-    predictions: list,
-    references: list,
-    gold_dbs: list,
-    setup_sqls: list,
-    validate_sqls: list,
-    kmaps: dict,
-    db_dir: str,
-    categories: list[str] = None,
+        predictions: list,
+        references: list,
+        gold_dbs: list,
+        setup_sqls: list,
+        validate_sqls: list,
+        kmaps: dict,
+        db_dir: str,
+        categories: list[str] = None,
 ) -> tuple[Any, list[int | None]]:
     """Compute test suite execution metric."""
     evaluator = test_suite_evaluation.Evaluator(
@@ -135,37 +145,27 @@ def compute_test_suite_metric(
     # Only used for Sparc/CoSQL
     turn_scores: dict[str, list] = {"exec": [], "exact": []}
     by_row_metrics: list[int | None] = []
+
     for prediction, reference, gold_db, setup_sql, validate_sql, category in tqdm(
-        zip(predictions, references, gold_dbs, setup_sqls, validate_sqls, categories),
-        total=len(predictions),
+            zip(predictions, references, gold_dbs, setup_sqls, validate_sqls, categories),
+            total=len(predictions),
     ):
         turn_idx = 0
         # skip final utterance-query pairs
         if turn_idx < 0:
             continue
 
-        # Register the timeout handler function
-        signal.signal(signal.SIGALRM, timeout_handler)
-        signal.alarm(TIMEOUT_SECONDS)
-
-        try:
-            ex_metrics = evaluator.evaluate_one(
-                gold_db,
-                reference,
-                prediction,
-                setup_sql,
-                validate_sql,
-                turn_scores,
-                idx=turn_idx,
-                category=category,
-            )
-            signal.alarm(0)
+        # Use the new function to evaluate with timeout
+        ex_metrics = evaluate_with_timeout(
+            evaluator, gold_db, reference, prediction, setup_sql, validate_sql,
+            turn_scores, timeout=TIMEOUT_SECONDS
+        )
 
+        if ex_metrics:
             by_row_metrics.append(int(ex_metrics["exec"]))
-        except Exception as e:
-            raise e
+        else:
             by_row_metrics.append(None)
-            pass
+
     evaluator.finalize()
     return evaluator.scores, by_row_metrics
 

evaluation_logic.py

@@ -60,8 +60,8 @@ def run_prediction(model_name, prompt_format, output_file):
             else:
                 table_params = []
 
-            if len(table_params) == 0:
-                yield f"[red] WARNING: No tables found for {db_id} [/red]"
+            #if len(table_params) == 0:
+                #yield f"[red] WARNING: No tables found for {db_id} [/red]"
 
             text_to_sql_inputs.append(TextToSQLParams(
                 instruction=question,