112 lines
3.7 KiB
Python
112 lines
3.7 KiB
Python
import json
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
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 evals.eval_utils import download_instances
|
|
|
|
|
|
async def cognee_and_llm(dataset, search_type=SearchType.CHUNKS):
|
|
await cognee.prune.prune_data()
|
|
await cognee.prune.prune_system(metadata=True)
|
|
|
|
dataset_name = "SWE_test_data"
|
|
code_text = dataset[0]["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 = dataset[0]['problem_statement']
|
|
instructions = (
|
|
"I need you to solve this issue by looking at the provided knowledge graph and 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."
|
|
)
|
|
|
|
prompt = "\n".join([
|
|
instructions,
|
|
"<patch>",
|
|
PATCH_EXAMPLE,
|
|
"</patch>",
|
|
"This is the knowledge graph:",
|
|
graph_str
|
|
])
|
|
|
|
llm_client = get_llm_client()
|
|
answer_prediction = llm_client.create_structured_output(
|
|
text_input=problem_statement,
|
|
system_prompt=prompt,
|
|
response_model=str,
|
|
)
|
|
return answer_prediction
|
|
|
|
|
|
async def llm_on_preprocessed_data(dataset):
|
|
problem_statement = dataset[0]['problem_statement']
|
|
prompt = dataset[0]["text"]
|
|
|
|
llm_client = get_llm_client()
|
|
answer_prediction = llm_client.create_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:
|
|
text_output = await cognee_and_llm(dataset)
|
|
model_name = "with_cognee"
|
|
else:
|
|
text_output = await llm_on_preprocessed_data(dataset)
|
|
model_name = "without_cognee"
|
|
|
|
preds = [{"instance_id": dataset[0]["instance_id"],
|
|
"model_patch": text_output,
|
|
"model_name_or_path": model_name}]
|
|
|
|
return preds
|
|
|
|
|
|
async def main():
|
|
swe_dataset = load_swebench_dataset(
|
|
'princeton-nlp/SWE-bench', split='test')
|
|
swe_dataset_preprocessed = load_swebench_dataset(
|
|
'princeton-nlp/SWE-bench_bm25_13K', split='test')
|
|
test_data = swe_dataset[:1]
|
|
test_data_preprocessed = swe_dataset_preprocessed[:1]
|
|
assert test_data[0]["instance_id"] == test_data_preprocessed[0]["instance_id"]
|
|
filepath = Path("SWE-bench_testsample")
|
|
if filepath.exists():
|
|
from datasets import Dataset
|
|
dataset = Dataset.load_from_disk(filepath)
|
|
else:
|
|
dataset = download_instances(test_data, filepath)
|
|
|
|
cognee_preds = await get_preds(dataset, with_cognee=True)
|
|
# nocognee_preds = await get_preds(dataset, with_cognee=False)
|
|
with open("withcognee.json", "w") as file:
|
|
json.dump(cognee_preds, file)
|
|
|
|
subprocess.run(["python", "-m", "swebench.harness.run_evaluation",
|
|
"--dataset_name", 'princeton-nlp/SWE-bench',
|
|
"--split", "test",
|
|
"--predictions_path", "withcognee.json",
|
|
"--max_workers", "1",
|
|
"--instance_ids", test_data[0]["instance_id"],
|
|
"--run_id", "with_cognee"])
|
|
|
|
if __name__ == "__main__":
|
|
import asyncio
|
|
asyncio.run(main(), debug=True)
|