File size: 2,360 Bytes
c351861 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | # Strict BIRD execution accuracy: set(pred rows) == set(gold rows). usage: python score.py preds.jsonl db_root
import json, os, sqlite3, sys, collections, time
from concurrent.futures import ProcessPoolExecutor, TimeoutError as TO
HERE = os.path.dirname(os.path.abspath(__file__))
preds_path, db_root = sys.argv[1], sys.argv[2]
gold = {d["question_id"]: d for d in (json.loads(l) for l in open(os.path.join(HERE, "mini_dev_sqlite.jsonl"), encoding="utf-8"))}
preds = {}
for l in open(preds_path, encoding="utf-8"):
p = json.loads(l); preds[p["question_id"]] = p
def run(db, sql, limit=None):
con = sqlite3.connect(f"file:{db}?mode=ro", uri=True)
if limit: # abort runaway queries inside the worker; a pool timeout alone leaves them running
end = time.time() + limit
con.set_progress_handler(lambda: time.time() > end, 100000)
try:
return set(con.execute(sql).fetchall())
finally:
con.close()
def judge(args):
db, psql, gsql = args
try:
g = run(db, gsql)
except Exception as e:
return "gold_error"
if not psql: return "no_sql"
try:
return "ok" if run(db, psql, 60) == g else "wrong"
except sqlite3.OperationalError as e:
return "timeout" if "interrupted" in str(e) else "exec_error"
except Exception:
return "exec_error"
if __name__ == "__main__":
jobs = []
for qid, g in gold.items():
db = os.path.join(db_root, g["db_id"], g["db_id"] + ".sqlite")
jobs.append((qid, (db, preds.get(qid, {}).get("sql", ""), g["SQL"])))
res = {}
with ProcessPoolExecutor(8) as ex:
futs = {qid: ex.submit(judge, a) for qid, a in jobs}
for qid, fu in futs.items():
try: res[qid] = fu.result(timeout=600) # the 60 s per-query limit is enforced inside the worker; this only guards a hung pool
except TO: res[qid] = "timeout"
c = collections.Counter(res.values())
by = collections.defaultdict(lambda: [0, 0])
for qid, v in res.items():
d = gold[qid]["difficulty"]; by[d][1] += 1; by[d][0] += v == "ok"
print(f"EX = {c['ok']}/{len(res)} = {100 * c['ok'] / len(res):.2f}% ", dict(c))
for d, (a, n) in sorted(by.items()): print(f" {d:12s} {a}/{n} = {100 * a / n:.1f}%")
json.dump(res, open(preds_path.replace(".jsonl", "_score.json"), "w"), indent=0)
|