118 lines
3.9 KiB
Python
118 lines
3.9 KiB
Python
import argparse
|
|
import json
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
from datasets import Dataset
|
|
from swebench.harness.utils import load_swebench_dataset
|
|
from swebench.inference.make_datasets.create_instance import PATCH_EXAMPLE
|
|
|
|
import cognee
|
|
from cognee.api.v1.cognify.code_graph_pipeline import code_graph_pipeline
|
|
from cognee.api.v1.search import SearchType
|
|
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
|
|
|
|
|
|
async def generate_patch_with_cognee(instance, search_type=SearchType.CHUNKS):
|
|
|
|
await cognee.prune.prune_data()
|
|
await cognee.prune.prune_system(metadata=True)
|
|
|
|
dataset_name = "SWE_test_data"
|
|
code_text = instance["text"]
|
|
await cognee.add([code_text], dataset_name)
|
|
await code_graph_pipeline([dataset_name])
|
|
graph_engine = await get_graph_engine()
|
|
with open(graph_engine.filename, "r") as f:
|
|
graph_str = f.read()
|
|
|
|
problem_statement = instance['problem_statement']
|
|
instructions = read_query_prompt("patch_gen_instructions.txt")
|
|
|
|
prompt = "\n".join([
|
|
instructions,
|
|
"<patch>",
|
|
PATCH_EXAMPLE,
|
|
"</patch>",
|
|
"This is the knowledge graph:",
|
|
graph_str
|
|
])
|
|
|
|
llm_client = get_llm_client()
|
|
answer_prediction = await llm_client.acreate_structured_output(
|
|
text_input=problem_statement,
|
|
system_prompt=prompt,
|
|
response_model=str,
|
|
)
|
|
return answer_prediction
|
|
|
|
|
|
async def generate_patch_without_cognee(instance):
|
|
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,
|
|
response_model=str,
|
|
)
|
|
return answer_prediction
|
|
|
|
|
|
async def get_preds(dataset, with_cognee=True):
|
|
if with_cognee:
|
|
model_name = "with_cognee"
|
|
pred_func = generate_patch_with_cognee
|
|
else:
|
|
model_name = "without_cognee"
|
|
pred_func = generate_patch_without_cognee
|
|
|
|
preds = [{"instance_id": instance["instance_id"],
|
|
"model_patch": await pred_func(instance),
|
|
"model_name_or_path": model_name} for instance in dataset]
|
|
|
|
return preds
|
|
|
|
|
|
async def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Run LLM predictions on SWE-bench dataset")
|
|
parser.add_argument('--cognee_off', action='store_true')
|
|
args = parser.parse_args()
|
|
|
|
if args.cognee_off:
|
|
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]
|
|
filepath = Path("SWE-bench_testsample")
|
|
if filepath.exists():
|
|
dataset = Dataset.load_from_disk(filepath)
|
|
else:
|
|
dataset = download_instances(swe_dataset, filepath)
|
|
predictions_path = "preds.json"
|
|
preds = await get_preds(dataset, with_cognee=not args.cognee_off)
|
|
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"])
|
|
|
|
if __name__ == "__main__":
|
|
import asyncio
|
|
asyncio.run(main(), debug=True)
|