minor compatibility changes to make new eval work with hf_runner

defog-ai · Aug 15, 2023 · 9aedeb3 · 9aedeb3
1 parent c9a365b
commit 9aedeb3
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 23 deletions.
diff --git a/eval/eval.py b/eval/eval.py
@@ -203,9 +203,9 @@ def compare_query_results(
     query_gen: str,
     db_name: str,
     db_creds: dict,
-    timeout: float,
     question: str,
     query_category: str,
+    timeout: float = 10.0,
 ) -> "tuple[bool, bool]":
     """
     Compares the results of two queries and returns a tuple of booleans, where the first element is

diff --git a/eval/hf_runner.py b/eval/hf_runner.py
@@ -1,4 +1,4 @@
-from eval.eval import compare_df, query_postgres_db, subset_df
+from eval.eval import compare_query_results
 import pandas as pd
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
@@ -7,7 +7,6 @@
 from psycopg2.extensions import QueryCanceledError
 from time import time
 import gc
-import traceback
 
 
 def prepare_questions_df(questions_file, num_questions):
@@ -128,29 +127,25 @@ def run_hf_eval(
             question = row["question"]
             query_category = row["query_category"]
             exact_match = correct = 0
-            generated_result = expected_result = None
+            db_creds = {
+                "host": "localhost",
+                "port": 5432,
+                "user": "postgres",
+                "password": "postgres",
+                "database": db_name,
+            }
 
             try:
-                expected_result = query_postgres_db(golden_query, db_name).rename(
-                    columns=str.lower
+                exact_match, correct = compare_query_results(
+                    query_gold=golden_query,
+                    query_gen=generated_query,
+                    db_name=db_name,
+                    db_creds=db_creds,
+                    question=question,
+                    query_category=query_category,
                 )
-
-                generated_result = query_postgres_db(generated_query, db_name).rename(
-                    columns=str.lower
-                )
-
-                exact_match = correct = int(
-                    compare_df(
-                        expected_result, generated_result, query_category, question
-                    )
-                )
-                if not exact_match:
-                    correct = subset_df(
-                        df_sub=expected_result,
-                        df_super=generated_result,
-                        query_category=query_category,
-                        question=question,
-                    )
+                row["exact_match"] = int(exact_match)
+                row["correct"] = int(correct)
                 row["exact_match"] = int(exact_match)
                 row["correct"] = int(correct)
                 row["error_msg"] = ""