| |
| 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: |
| 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) |
| 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) |
|
|
