using the code graph pipeline instead of cognify

This commit is contained in:
Rita Aleksziev 2024-11-15 17:54:41 +01:00
parent 4d6229bf7e
commit ed08cdb9f9

View file

@ -20,7 +20,7 @@ async def cognee_and_llm(dataset, search_type = SearchType.CHUNKS):
dataset_name = "SWE_test_data" dataset_name = "SWE_test_data"
code_text = dataset[0]["text"][:100000] code_text = dataset[0]["text"][:100000]
await cognee.add([code_text], dataset_name) await cognee.add([code_text], dataset_name)
await cognee.cognify([dataset_name]) await code_graph_pipeline([dataset_name])
graph_engine = await get_graph_engine() graph_engine = await get_graph_engine()
with open(graph_engine.filename, "r") as f: with open(graph_engine.filename, "r") as f:
graph_str = f.read() graph_str = f.read()
@ -63,7 +63,7 @@ async def llm_on_preprocessed_data(dataset):
) )
return answer_prediction return answer_prediction
async def get_preds(dataset, with_cognee): async def get_preds(dataset, with_cognee=True):
if with_cognee: if with_cognee:
text_output = await cognee_and_llm(dataset) text_output = await cognee_and_llm(dataset)
model_name = "with_cognee" model_name = "with_cognee"