evaluate script

evaluate_with_db.py

@@ -0,0 +1,67 @@
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import json
+import sqlite3
+from tqdm import tqdm
+from typing import List
+import os
+from pathlib import Path
+
+db_schemas_path = "db_schemas.json"
+model_path = "gaussalgo/T5-LM-Large-text2sql-spider"
+
+model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+
+def query_db(question: str, db_path: str) -> dict:
+    try:
+        # assert db_path.endswith('.sqlite')
+        con = sqlite3.connect(db_path)
+        cur = con.cursor()
+        cur.execute(question)
+        data = cur.fetchall()
+        return json.dumps(data)
+    except Exception as e:
+        print(question, " ", e)
+        pass
+
+
+def evaluate(eval_dataset: List[dict]):
+    reference = []
+    gen_queries = []
+
+    with open(db_schemas_path, "r") as schemas:
+        db_schema_dict = json.load(schemas)
+
+    for data in tqdm(eval_dataset, total=len(eval_dataset), desc="Executing queries"):
+        question = data["question"]
+        schema = data["db_id"]
+
+        filenames = [
+            i for i in os.listdir(Path(DB_PATH, schema)) if i.endswith(SQLITE_SUFFIX)
+        ]
+        path_to_db = Path(DB_PATH, schema, filenames[0])
+
+        input_text = " ".join(
+            ["Question: ", question, "Schema:", db_schema_dict[schema]]
+        )
+        model_inputs = tokenizer(input_text, return_tensors="pt")
+        outputs = model.generate(**model_inputs, max_length=512)
+
+        output_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+        reference.append(query_db(data["query"], path_to_db))
+        gen_queries.append(query_db(output_text, path_to_db))
+
+    equal_results = [ref == q for ref, q in zip(reference, gen_queries)]
+    eq_results_when_reference_works = [
+        ref == q for ref, q in zip(reference, gen_queries) if ref is not None
+    ]
+    num_of_working_ref = len([ref for ref in reference if ref is not None])
+    print("Length of eval dataset: ", len(eval_dataset))
+    print("Working references: ", num_of_working_ref)
+    print("Correct queries in labels: ", num_of_working_ref / len(eval_dataset))
+    print("Accuracy with whole dataset: ", sum(equal_results) / len(eval_dataset))
+    print(
+        "Accuracy with only working references: ",
+        sum(eq_results_when_reference_works) / num_of_working_ref,
+    )