cognee/evals/run_qa_eval.py
2025-01-20 18:42:09 +01:00

59 lines
1.8 KiB
Python

import asyncio
from evals.eval_on_hotpot import eval_on_QA_dataset, incremental_eval_on_QA_dataset
from evals.qa_eval_utils import get_combinations, save_results_as_image
import argparse
from pathlib import Path
import json
async def run_evals_on_paramset(paramset: dict, out_path: str):
combinations = get_combinations(paramset)
json_path = Path(out_path) / Path("results.json")
results = {}
for params in combinations:
dataset = params["dataset"]
num_samples = params["num_samples"]
rag_option = params["rag_option"]
if dataset not in results:
results[dataset] = {}
if num_samples not in results[dataset]:
results[dataset][num_samples] = {}
if rag_option == "cognee_incremental":
result = await incremental_eval_on_QA_dataset(
dataset, num_samples, paramset["metric_names"], out_path
)
results[dataset][num_samples] |= result
else:
result = await eval_on_QA_dataset(
dataset, rag_option, num_samples, paramset["metric_names"], out_path
)
results[dataset][num_samples][rag_option] = result
with open(json_path, "w") as file:
json.dump(results, file, indent=1)
save_results_as_image(results, out_path)
return results
async def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--params_file", type=str, required=True, help="Which dataset to evaluate on"
)
parser.add_argument("--out_dir", type=str, help="Dir to save eval results")
args = parser.parse_args()
with open(args.params_file, "r") as file:
parameters = json.load(file)
await run_evals_on_paramset(parameters, args.out_dir)
if __name__ == "__main__":
asyncio.run(main())