2e010f8dd1
* QA eval dataset as argument, with hotpot and 2wikimultihop as options. Json schema validation for datasets. * Load dataset file by filename, outsource utilities * restructure metric selection * Add comprehensiveness, diversity and empowerment metrics * add promptfoo as an option * refactor RAG solution in eval;2C * LLM as a judge metrics implemented in a uniform way * Use requests.get instead of wget * clean up promptfoo config template * minor fixes * get promptfoo path instead of hardcoding * minor fixes * Add LLM as a judge prompts * Support 4 different rag options in eval * Minor refactor and logger usage * feat: make tasks a configurable argument in the cognify function * Run eval on a set of parameters and save results as json and png * fix: add data points task * script for running all param combinations * enable context provider to get tasks as param * bugfix in simple rag * Incremental eval of cognee pipeline * potential fix: single asyncio run * temp fix: exclude insights * Remove insights, have single asyncio run, refactor * minor fixes * handle pipeline slices in utils * include all options in params json --------- Co-authored-by: lxobr <122801072+lxobr@users.noreply.github.com> Co-authored-by: hajdul88 <52442977+hajdul88@users.noreply.github.com>
89 lines
3 KiB
Python
89 lines
3 KiB
Python
import cognee
|
|
from cognee.api.v1.search import SearchType
|
|
from cognee.infrastructure.databases.vector import get_vector_engine
|
|
from cognee.modules.retrieval.brute_force_triplet_search import brute_force_triplet_search
|
|
from cognee.tasks.completion.graph_query_completion import retrieved_edges_to_string
|
|
from functools import partial
|
|
from cognee.api.v1.cognify.cognify_v2 import get_default_tasks
|
|
|
|
|
|
async def get_raw_context(instance: dict) -> str:
|
|
return instance["context"]
|
|
|
|
|
|
async def cognify_instance(instance: dict, task_indices: list[int] = None):
|
|
await cognee.prune.prune_data()
|
|
await cognee.prune.prune_system(metadata=True)
|
|
for title, sentences in instance["context"]:
|
|
await cognee.add("\n".join(sentences), dataset_name="QA")
|
|
all_cognify_tasks = await get_default_tasks()
|
|
if task_indices:
|
|
selected_tasks = [all_cognify_tasks[ind] for ind in task_indices]
|
|
else:
|
|
selected_tasks = all_cognify_tasks
|
|
await cognee.cognify("QA", tasks=selected_tasks)
|
|
|
|
|
|
async def get_context_with_cognee(
|
|
instance: dict,
|
|
task_indices: list[int] = None,
|
|
search_types: list[SearchType] = [SearchType.SUMMARIES, SearchType.CHUNKS],
|
|
) -> str:
|
|
await cognify_instance(instance, task_indices)
|
|
|
|
search_results = []
|
|
for search_type in search_types:
|
|
search_results += await cognee.search(search_type, query_text=instance["question"])
|
|
|
|
search_results_str = "\n".join([context_item["text"] for context_item in search_results])
|
|
|
|
return search_results_str
|
|
|
|
|
|
def create_cognee_context_getter(
|
|
task_indices=None, search_types=[SearchType.SUMMARIES, SearchType.CHUNKS]
|
|
):
|
|
return partial(get_context_with_cognee, task_indices=task_indices, search_types=search_types)
|
|
|
|
|
|
async def get_context_with_simple_rag(instance: dict) -> str:
|
|
await cognee.prune.prune_data()
|
|
await cognee.prune.prune_system(metadata=True)
|
|
|
|
for title, sentences in instance["context"]:
|
|
await cognee.add("\n".join(sentences), dataset_name="QA")
|
|
|
|
vector_engine = get_vector_engine()
|
|
found_chunks = await vector_engine.search("document_chunk_text", instance["question"], limit=5)
|
|
|
|
search_results_str = "\n".join([context_item.payload["text"] for context_item in found_chunks])
|
|
|
|
return search_results_str
|
|
|
|
|
|
async def get_context_with_brute_force_triplet_search(instance: dict) -> str:
|
|
await cognify_instance(instance)
|
|
|
|
found_triplets = await brute_force_triplet_search(instance["question"], top_k=5)
|
|
|
|
search_results_str = retrieved_edges_to_string(found_triplets)
|
|
|
|
return search_results_str
|
|
|
|
|
|
valid_pipeline_slices = {
|
|
"base": [0, 1, 5],
|
|
"extract_chunks": [0, 1, 2, 5],
|
|
"extract_graph": [0, 1, 2, 3, 5],
|
|
"summarize": [0, 1, 2, 3, 4, 5],
|
|
}
|
|
|
|
qa_context_providers = {
|
|
"no_rag": get_raw_context,
|
|
"cognee": get_context_with_cognee,
|
|
"simple_rag": get_context_with_simple_rag,
|
|
"brute_force": get_context_with_brute_force_triplet_search,
|
|
} | {
|
|
name: create_cognee_context_getter(task_indices=slice)
|
|
for name, slice in valid_pipeline_slices.items()
|
|
}
|