8ec1e48ff6
* 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 * Run eval on a set of parameters and save results as json and png * script for running all param combinations * bugfix in simple rag * potential fix: single asyncio run * temp fix: exclude insights * Remove insights, have single asyncio run, refactor --------- Co-authored-by: lxobr <122801072+lxobr@users.noreply.github.com>
65 lines
2.2 KiB
Python
65 lines
2.2 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
|
|
|
|
|
|
async def get_raw_context(instance: dict) -> str:
|
|
return instance["context"]
|
|
|
|
|
|
async def cognify_instance(instance: dict):
|
|
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")
|
|
await cognee.cognify("QA")
|
|
|
|
|
|
async def get_context_with_cognee(instance: dict) -> str:
|
|
await cognify_instance(instance)
|
|
|
|
# TODO: Fix insights
|
|
# insights = await cognee.search(SearchType.INSIGHTS, query_text=instance["question"])
|
|
summaries = await cognee.search(SearchType.SUMMARIES, query_text=instance["question"])
|
|
# search_results = insights + summaries
|
|
search_results = summaries
|
|
|
|
search_results_str = "\n".join([context_item["text"] for context_item in search_results])
|
|
|
|
return search_results_str
|
|
|
|
|
|
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
|
|
|
|
|
|
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,
|
|
}
|