From 1bfa3a0ea317c79f36fe817adcf96abf0a24f88e Mon Sep 17 00:00:00 2001 From: Leon Luithlen Date: Fri, 29 Nov 2024 11:30:30 +0100 Subject: [PATCH 1/6] Rebase onto code-graph --- evals/EC2_README.md | 34 +++++++ evals/cloud/setup_ubuntu_instance.sh | 43 +++++++++ evals/eval_swe_bench.py | 127 ++++++++++++++++++--------- 3 files changed, 163 insertions(+), 41 deletions(-) create mode 100644 evals/EC2_README.md create mode 100644 evals/cloud/setup_ubuntu_instance.sh diff --git a/evals/EC2_README.md b/evals/EC2_README.md new file mode 100644 index 000000000..50a92bc27 --- /dev/null +++ b/evals/EC2_README.md @@ -0,0 +1,34 @@ +Create an EC2 Instance with the + +`Ubuntu Image` + +Many instance types will work, we used: + +`m7a.2xlarge` # more than 8 parallel processes doesn't seem to speed up overall process. Maybe to do with docker parallelism? + +DON'T FORGET TO ADD + +`500 GB storage` + +Or the evaluation run will run out of space + +-------------------------------------------------------- + +Then ssh into the instance, run + +source evals/cloud/setup_ubuntu_instance.sh + +sudo usermod -aG docker $USER + +disconnect, and reconnect. + +Then enter a `screen` and activate the virtual env + +screen +source venv/bin/activate + +then, from cognee, you can run swe_bench: + +python evals/eval_swe_bench --cognee_off --max_workers=N_CPUS + +Building the environment images takes roughly 17 minutes \ No newline at end of file diff --git a/evals/cloud/setup_ubuntu_instance.sh b/evals/cloud/setup_ubuntu_instance.sh new file mode 100644 index 000000000..e5386c372 --- /dev/null +++ b/evals/cloud/setup_ubuntu_instance.sh @@ -0,0 +1,43 @@ + +sudo apt-get update +sudo apt-get install ca-certificates curl +sudo install -m 0755 -d /etc/apt/keyrings +sudo curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc +sudo chmod a+r /etc/apt/keyrings/docker.asc + +# Add the repository to Apt sources: +echo \ + "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu \ + $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \ + sudo tee /etc/apt/sources.list.d/docker.list > /dev/null +sudo apt-get update + +sudo apt-get install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin + +sudo docker run hello-world + +sudo apt install unzip + +sudo apt-get install python3-virtualenv + +sudo add-apt-repository ppa:deadsnakes/ppa +sudo apt update + +sudo apt install python3.11 + +virtualenv venv --python=python3.11 + +source venv/bin/activate + +pip install poetry + +poetry install + +pip install swebench transformers sentencepiece + +groups | grep docker + +python evals/eval_swe_bench.py --cognee_off + +sudo usermod -aG docker $USER + diff --git a/evals/eval_swe_bench.py b/evals/eval_swe_bench.py index 1dd0e58ab..5cbea58ee 100644 --- a/evals/eval_swe_bench.py +++ b/evals/eval_swe_bench.py @@ -1,6 +1,7 @@ import argparse import json import subprocess +import sys from pathlib import Path from datasets import Dataset @@ -29,7 +30,28 @@ from evals.eval_utils import ingest_repos from evals.eval_utils import download_github_repo from evals.eval_utils import delete_repo -async def generate_patch_with_cognee(instance): + +def check_install_package(package_name): + """ + Check if a pip package is installed and install it if not. + Returns True if package is/was installed successfully, False otherwise. + """ + try: + __import__(package_name) + return True + except ImportError: + try: + subprocess.check_call( + [sys.executable, "-m", "pip", "install", package_name] + ) + return True + except subprocess.CalledProcessError: + return False + + + +async def generate_patch_with_cognee(instance, search_type=SearchType.CHUNKS): + await cognee.prune.prune_data() await cognee.prune.prune_system() @@ -59,23 +81,22 @@ async def generate_patch_with_cognee(instance): await render_graph(None, include_labels = True, include_nodes = True) - problem_statement = instance['problem_statement'] + problem_statement = instance["problem_statement"] instructions = read_query_prompt("patch_gen_instructions.txt") graph_str = 'HERE WE SHOULD PASS THE TRIPLETS FROM GRAPHRAG' - prompt = "\n".join([ - instructions, - "", - PATCH_EXAMPLE, - "", - "This is the knowledge graph:", - graph_str - ]) + prompt = "\n".join( + [ + instructions, + "", + PATCH_EXAMPLE, + "", + "This is the knowledge graph:", + graph_str, + ] + ) - return 0 - - ''' :TODO: We have to find out how do we do the generation llm_client = get_llm_client() answer_prediction = await llm_client.acreate_structured_output( text_input=problem_statement, @@ -84,13 +105,11 @@ async def generate_patch_with_cognee(instance): ) return answer_prediction - ''' -async def generate_patch_without_cognee(instance): - problem_statement = instance['problem_statement'] +async def generate_patch_without_cognee(instance, llm_client): + problem_statement = instance["problem_statement"] prompt = instance["text"] - llm_client = get_llm_client() answer_prediction = await llm_client.acreate_structured_output( text_input=problem_statement, system_prompt=prompt, @@ -100,43 +119,56 @@ async def generate_patch_without_cognee(instance): async def get_preds(dataset, with_cognee=True): + llm_client = get_llm_client() + if with_cognee: model_name = "with_cognee" - pred_func = generate_patch_with_cognee + futures = [ + (instance["instance_id"], generate_patch_with_cognee(instance)) + for instance in dataset + ] else: model_name = "without_cognee" - pred_func = generate_patch_without_cognee + futures = [ + (instance["instance_id"], generate_patch_without_cognee(instance, llm_client)) + for instance in dataset + ] + model_patches = await asyncio.gather(*[x[1] for x in futures]) + preds = [ + { + "instance_id": instance_id, + "model_patch": model_patch, + "model_name_or_path": model_name, + } + for (instance_id, _), model_patch in zip(futures, model_patches) + ] - for instance in dataset: - await pred_func(instance) - - ''' - preds = [{"instance_id": instance["instance_id"], - "model_patch": await pred_func(instance), - "model_name_or_path": model_name} for instance in dataset] - ''' - return 0 + return preds async def main(): parser = argparse.ArgumentParser( - description="Run LLM predictions on SWE-bench dataset") - parser.add_argument('--cognee_off', action='store_true') + description="Run LLM predictions on SWE-bench dataset" + ) + parser.add_argument("--cognee_off", action="store_true") + parser.add_argument("--max_workers", type=int, required=True) args = parser.parse_args() + for dependency in ["transformers", "sentencepiece", "swebench"]: + check_install_package(dependency) + if args.cognee_off: - dataset_name = 'princeton-nlp/SWE-bench_Lite_bm25_13K' - dataset = load_swebench_dataset(dataset_name, split='test') + dataset_name = "princeton-nlp/SWE-bench_Lite_bm25_13K" + dataset = load_swebench_dataset(dataset_name, split="test") predictions_path = "preds_nocognee.json" if not Path(predictions_path).exists(): preds = await get_preds(dataset, with_cognee=False) with open(predictions_path, "w") as file: json.dump(preds, file) else: - dataset_name = 'princeton-nlp/SWE-bench_Lite' - swe_dataset = load_swebench_dataset( - dataset_name, split='test')[:1] + dataset_name = "princeton-nlp/SWE-bench_Lite" + swe_dataset = load_swebench_dataset(dataset_name, split="test")[:1] filepath = Path("SWE-bench_testsample") if filepath.exists(): dataset = Dataset.load_from_disk(filepath) @@ -147,12 +179,25 @@ async def main(): with open(predictions_path, "w") as file: json.dump(preds, file) - subprocess.run(["python", "-m", "swebench.harness.run_evaluation", - "--dataset_name", dataset_name, - "--split", "test", - "--predictions_path", predictions_path, - "--max_workers", "1", - "--run_id", "test_run"]) + + subprocess.run( + [ + "python", + "-m", + "swebench.harness.run_evaluation", + "--dataset_name", + dataset_name, + "--split", + "test", + "--predictions_path", + predictions_path, + "--max_workers", + str(args.max_workers), + "--run_id", + "test_run", + ] + ) + if __name__ == "__main__": import asyncio From 5036f3a85f22b84bf2dd9c3ee0aa4ba56de0ff32 Mon Sep 17 00:00:00 2001 From: Leon Luithlen Date: Tue, 26 Nov 2024 11:18:36 +0100 Subject: [PATCH 2/6] Add -y to setup_ubuntu_instance.sh commands and update EC2_README --- evals/EC2_README.md | 44 +++++++++++++++++++++++----- evals/cloud/setup_ubuntu_instance.sh | 38 +++++++++--------------- 2 files changed, 51 insertions(+), 31 deletions(-) diff --git a/evals/EC2_README.md b/evals/EC2_README.md index 50a92bc27..8e3dccb11 100644 --- a/evals/EC2_README.md +++ b/evals/EC2_README.md @@ -1,3 +1,5 @@ +## Creating the EC2 Instance + Create an EC2 Instance with the `Ubuntu Image` @@ -12,23 +14,51 @@ DON'T FORGET TO ADD Or the evaluation run will run out of space --------------------------------------------------------- +Add a key pair login where you have access to the corresponding key file (*.pem) -Then ssh into the instance, run +## Accessing your instance and setup -source evals/cloud/setup_ubuntu_instance.sh +To ssh into the instance, you have to save your key pair file (*.pem) to an appropriate location, such as ~/.aws. After launching the instance, you can access the Instance Summary, and retrieve "Public IPv4 DNS" address. Then run + +`ssh -i PATH_TO_KEY ubuntu@IPv4ADDRESS` + +to gain command line access to the instance. + +To copy your current state of cognee, go to the folder that contains "cognee" on your local machine, zip it to cognee.zip and run: + +`zip -r cognee.zip cognee` +`scp -i PATH_TO_KEY cognee.zip ubuntu@IPv4ADDRESS:cognee.zip` + +And unzip cognee.zip in your SSH session: + +`sudo apt install unzip` +`unzip cognee.zip` + +Then run: +`cd cognee` +`source evals/cloud/setup_ubuntu_instance.sh` sudo usermod -aG docker $USER disconnect, and reconnect. +Confirm that `ubuntu` has been added to the docker user group with + +`groups | grep docker` + +## Running SWE-bench + Then enter a `screen` and activate the virtual env -screen -source venv/bin/activate +`screen` +`source venv/bin/activate` then, from cognee, you can run swe_bench: -python evals/eval_swe_bench --cognee_off --max_workers=N_CPUS +`cd cognee` -Building the environment images takes roughly 17 minutes \ No newline at end of file +`python evals/eval_swe_bench.py --cognee_off --max_workers=N_CPUS` + +Building the environment images should take roughly 17 minutes + +If the virtual env wasn't set up correctly for some reason, just run the last few lines of `setup_ubuntu_instance.sh` manually \ No newline at end of file diff --git a/evals/cloud/setup_ubuntu_instance.sh b/evals/cloud/setup_ubuntu_instance.sh index e5386c372..e05b761e2 100644 --- a/evals/cloud/setup_ubuntu_instance.sh +++ b/evals/cloud/setup_ubuntu_instance.sh @@ -1,43 +1,33 @@ - -sudo apt-get update -sudo apt-get install ca-certificates curl +sudo apt-get update -y +sudo apt-get install -y ca-certificates curl sudo install -m 0755 -d /etc/apt/keyrings sudo curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc sudo chmod a+r /etc/apt/keyrings/docker.asc # Add the repository to Apt sources: echo \ - "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu \ - $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \ - sudo tee /etc/apt/sources.list.d/docker.list > /dev/null -sudo apt-get update +"deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu \ + $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \ +sudo tee /etc/apt/sources.list.d/docker.list > /dev/null +sudo apt-get update -y -sudo apt-get install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin +sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin sudo docker run hello-world -sudo apt install unzip +sudo apt install -y unzip -sudo apt-get install python3-virtualenv +sudo apt-get install -y python3-virtualenv -sudo add-apt-repository ppa:deadsnakes/ppa -sudo apt update +sudo add-apt-repository -y ppa:deadsnakes/ppa -sudo apt install python3.11 +sudo apt update -y + +sudo apt install -y python3.11 virtualenv venv --python=python3.11 source venv/bin/activate - pip install poetry - poetry install - -pip install swebench transformers sentencepiece - -groups | grep docker - -python evals/eval_swe_bench.py --cognee_off - -sudo usermod -aG docker $USER - +pip install swebench transformers sentencepiece datasets tiktoken protobuf From 618d476c301e1c30bf0f7b505b117d1c918e1473 Mon Sep 17 00:00:00 2001 From: Leon Luithlen Date: Tue, 26 Nov 2024 11:20:07 +0100 Subject: [PATCH 3/6] Add code formating to usermod command --- evals/EC2_README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/evals/EC2_README.md b/evals/EC2_README.md index 8e3dccb11..d6a937ed7 100644 --- a/evals/EC2_README.md +++ b/evals/EC2_README.md @@ -38,7 +38,7 @@ Then run: `cd cognee` `source evals/cloud/setup_ubuntu_instance.sh` -sudo usermod -aG docker $USER +`sudo usermod -aG docker $USER` disconnect, and reconnect. From b46af5a6f6d049fde2ffc66992ec9e9a530d8fff Mon Sep 17 00:00:00 2001 From: Leon Luithlen Date: Wed, 27 Nov 2024 12:12:12 +0100 Subject: [PATCH 4/6] Update eval_swe_bench --- evals/eval_swe_bench.py | 41 ++++++++++++++++++++++++++++------------- 1 file changed, 28 insertions(+), 13 deletions(-) diff --git a/evals/eval_swe_bench.py b/evals/eval_swe_bench.py index 5cbea58ee..980be4bd2 100644 --- a/evals/eval_swe_bench.py +++ b/evals/eval_swe_bench.py @@ -26,9 +26,6 @@ from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.llm.get_llm_client import get_llm_client from cognee.infrastructure.llm.prompts import read_query_prompt from evals.eval_utils import download_instances -from evals.eval_utils import ingest_repos -from evals.eval_utils import download_github_repo -from evals.eval_utils import delete_repo def check_install_package(package_name): @@ -49,8 +46,14 @@ def check_install_package(package_name): return False +<<<<<<< HEAD async def generate_patch_with_cognee(instance, search_type=SearchType.CHUNKS): +======= +async def generate_patch_with_cognee( + instance, search_type=SearchType.CHUNKS +): +>>>>>>> c4e3634 (Update eval_swe_bench) await cognee.prune.prune_data() await cognee.prune.prune_system() @@ -81,7 +84,7 @@ async def generate_patch_with_cognee(instance, search_type=SearchType.CHUNKS): await render_graph(None, include_labels = True, include_nodes = True) - problem_statement = instance["problem_statement"] + problem_statement = instance['problem_statement'] instructions = read_query_prompt("patch_gen_instructions.txt") graph_str = 'HERE WE SHOULD PASS THE TRIPLETS FROM GRAPHRAG' @@ -97,7 +100,6 @@ async def generate_patch_with_cognee(instance, search_type=SearchType.CHUNKS): ] ) - llm_client = get_llm_client() answer_prediction = await llm_client.acreate_structured_output( text_input=problem_statement, system_prompt=prompt, @@ -106,8 +108,9 @@ async def generate_patch_with_cognee(instance, search_type=SearchType.CHUNKS): return answer_prediction + async def generate_patch_without_cognee(instance, llm_client): - problem_statement = instance["problem_statement"] + problem_statement = instance['problem_statement'] prompt = instance["text"] answer_prediction = await llm_client.acreate_structured_output( @@ -134,7 +137,10 @@ async def get_preds(dataset, with_cognee=True): for instance in dataset ] model_patches = await asyncio.gather(*[x[1] for x in futures]) +<<<<<<< HEAD +======= +>>>>>>> c4e3634 (Update eval_swe_bench) preds = [ { "instance_id": instance_id, @@ -149,9 +155,8 @@ async def get_preds(dataset, with_cognee=True): async def main(): parser = argparse.ArgumentParser( - description="Run LLM predictions on SWE-bench dataset" - ) - parser.add_argument("--cognee_off", action="store_true") + description="Run LLM predictions on SWE-bench dataset") + parser.add_argument('--cognee_off', action='store_true') parser.add_argument("--max_workers", type=int, required=True) args = parser.parse_args() @@ -159,16 +164,17 @@ async def main(): check_install_package(dependency) if args.cognee_off: - dataset_name = "princeton-nlp/SWE-bench_Lite_bm25_13K" - dataset = load_swebench_dataset(dataset_name, split="test") + dataset_name = 'princeton-nlp/SWE-bench_Lite_bm25_13K' + dataset = load_swebench_dataset(dataset_name, split='test') predictions_path = "preds_nocognee.json" if not Path(predictions_path).exists(): preds = await get_preds(dataset, with_cognee=False) with open(predictions_path, "w") as file: json.dump(preds, file) else: - dataset_name = "princeton-nlp/SWE-bench_Lite" - swe_dataset = load_swebench_dataset(dataset_name, split="test")[:1] + dataset_name = 'princeton-nlp/SWE-bench_Lite' + swe_dataset = load_swebench_dataset( + dataset_name, split='test')[:1] filepath = Path("SWE-bench_testsample") if filepath.exists(): dataset = Dataset.load_from_disk(filepath) @@ -179,6 +185,7 @@ async def main(): with open(predictions_path, "w") as file: json.dump(preds, file) +<<<<<<< HEAD subprocess.run( [ @@ -198,6 +205,14 @@ async def main(): ] ) +======= + subprocess.run(["python", "-m", "swebench.harness.run_evaluation", + "--dataset_name", dataset_name, + "--split", "test", + "--predictions_path", predictions_path, + "--max_workers", "1", + "--run_id", "test_run"]) +>>>>>>> c4e3634 (Update eval_swe_bench) if __name__ == "__main__": import asyncio From d9fc740ec0a8f9d84d1e9a3fabc5fa3edfa84e40 Mon Sep 17 00:00:00 2001 From: Leon Luithlen Date: Fri, 29 Nov 2024 11:33:05 +0100 Subject: [PATCH 5/6] Fix merge conflicts --- evals/eval_swe_bench.py | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/evals/eval_swe_bench.py b/evals/eval_swe_bench.py index 980be4bd2..45b83970a 100644 --- a/evals/eval_swe_bench.py +++ b/evals/eval_swe_bench.py @@ -45,15 +45,7 @@ def check_install_package(package_name): except subprocess.CalledProcessError: return False - -<<<<<<< HEAD - async def generate_patch_with_cognee(instance, search_type=SearchType.CHUNKS): -======= -async def generate_patch_with_cognee( - instance, search_type=SearchType.CHUNKS -): ->>>>>>> c4e3634 (Update eval_swe_bench) await cognee.prune.prune_data() await cognee.prune.prune_system() @@ -137,10 +129,7 @@ async def get_preds(dataset, with_cognee=True): for instance in dataset ] model_patches = await asyncio.gather(*[x[1] for x in futures]) -<<<<<<< HEAD -======= ->>>>>>> c4e3634 (Update eval_swe_bench) preds = [ { "instance_id": instance_id, @@ -185,7 +174,6 @@ async def main(): with open(predictions_path, "w") as file: json.dump(preds, file) -<<<<<<< HEAD subprocess.run( [ @@ -205,14 +193,6 @@ async def main(): ] ) -======= - subprocess.run(["python", "-m", "swebench.harness.run_evaluation", - "--dataset_name", dataset_name, - "--split", "test", - "--predictions_path", predictions_path, - "--max_workers", "1", - "--run_id", "test_run"]) ->>>>>>> c4e3634 (Update eval_swe_bench) if __name__ == "__main__": import asyncio From a5ae9185cd397f9526627023bf7f98453b35df8f Mon Sep 17 00:00:00 2001 From: Leon Luithlen Date: Fri, 29 Nov 2024 11:40:51 +0100 Subject: [PATCH 6/6] Replicate PR 33 --- .../llm/prompts/patch_gen_instructions.txt | 3 +- .../llm/prompts/patch_gen_kg_instructions.txt | 3 ++ evals/eval_swe_bench.py | 32 +++++++++---------- 3 files changed, 19 insertions(+), 19 deletions(-) create mode 100644 cognee/infrastructure/llm/prompts/patch_gen_kg_instructions.txt diff --git a/cognee/infrastructure/llm/prompts/patch_gen_instructions.txt b/cognee/infrastructure/llm/prompts/patch_gen_instructions.txt index 1553753ab..5e7e48dda 100644 --- a/cognee/infrastructure/llm/prompts/patch_gen_instructions.txt +++ b/cognee/infrastructure/llm/prompts/patch_gen_instructions.txt @@ -1,3 +1,2 @@ -I need you to solve this issue by looking at the provided knowledge graph and -generating a single patch file that I can apply directly to this repository using git apply. +I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format. \ No newline at end of file diff --git a/cognee/infrastructure/llm/prompts/patch_gen_kg_instructions.txt b/cognee/infrastructure/llm/prompts/patch_gen_kg_instructions.txt new file mode 100644 index 000000000..1553753ab --- /dev/null +++ b/cognee/infrastructure/llm/prompts/patch_gen_kg_instructions.txt @@ -0,0 +1,3 @@ +I need you to solve this issue by looking at the provided knowledge graph and +generating a single patch file that I can apply directly to this repository using git apply. +Please respond with a single patch file in the following format. \ No newline at end of file diff --git a/evals/eval_swe_bench.py b/evals/eval_swe_bench.py index 45b83970a..4a59457e1 100644 --- a/evals/eval_swe_bench.py +++ b/evals/eval_swe_bench.py @@ -45,7 +45,7 @@ def check_install_package(package_name): except subprocess.CalledProcessError: return False -async def generate_patch_with_cognee(instance, search_type=SearchType.CHUNKS): +async def generate_patch_with_cognee(instance, llm_client, search_type=SearchType.CHUNKS): await cognee.prune.prune_data() await cognee.prune.prune_system() @@ -77,13 +77,13 @@ async def generate_patch_with_cognee(instance, search_type=SearchType.CHUNKS): await render_graph(None, include_labels = True, include_nodes = True) problem_statement = instance['problem_statement'] - instructions = read_query_prompt("patch_gen_instructions.txt") + instructions = read_query_prompt("patch_gen_kg_instructions.txt") graph_str = 'HERE WE SHOULD PASS THE TRIPLETS FROM GRAPHRAG' prompt = "\n".join( [ - instructions, + problem_statement, "", PATCH_EXAMPLE, "", @@ -93,8 +93,8 @@ async def generate_patch_with_cognee(instance, search_type=SearchType.CHUNKS): ) answer_prediction = await llm_client.acreate_structured_output( - text_input=problem_statement, - system_prompt=prompt, + text_input=prompt, + system_prompt=instructions, response_model=str, ) @@ -102,12 +102,11 @@ async def generate_patch_with_cognee(instance, search_type=SearchType.CHUNKS): async def generate_patch_without_cognee(instance, llm_client): - problem_statement = instance['problem_statement'] - prompt = instance["text"] + instructions = read_query_prompt("patch_gen_instructions.txt") answer_prediction = await llm_client.acreate_structured_output( - text_input=problem_statement, - system_prompt=prompt, + text_input=instance["text"], + system_prompt=instructions, response_model=str, ) return answer_prediction @@ -118,16 +117,15 @@ async def get_preds(dataset, with_cognee=True): if with_cognee: model_name = "with_cognee" - futures = [ - (instance["instance_id"], generate_patch_with_cognee(instance)) - for instance in dataset - ] + pred_func = generate_patch_with_cognee else: model_name = "without_cognee" - futures = [ - (instance["instance_id"], generate_patch_without_cognee(instance, llm_client)) - for instance in dataset - ] + pred_func = generate_patch_without_cognee + + futures = [ + (instance["instance_id"], pred_func(instance, llm_client)) + for instance in dataset + ] model_patches = await asyncio.gather(*[x[1] for x in futures]) preds = [