diff --git a/src/haiku/rag/qa/ollama.py b/src/haiku/rag/qa/ollama.py index 273317b1..021190ca 100644 --- a/src/haiku/rag/qa/ollama.py +++ b/src/haiku/rag/qa/ollama.py @@ -4,6 +4,8 @@ from haiku.rag.client import HaikuRAG from haiku.rag.config import Config from haiku.rag.qa.base import QABase +OLLAMA_OPTIONS = {"temperature": 0.0, "seed": 42, "num_ctx": 64000} + class QA(QABase): def __init__(self, client: HaikuRAG, model: str = Config.QA_MODEL): @@ -48,7 +50,7 @@ class QA(QABase): model=self._model, messages=messages, tools=tools, - options={"temperature": 0.0, "seed": 42}, + options=OLLAMA_OPTIONS, think=False, ) @@ -82,7 +84,7 @@ class QA(QABase): model=self._model, messages=messages, think=False, - options={"temperature": 0.0, "seed": 42}, + options=OLLAMA_OPTIONS, ) return final_response["message"]["content"] else: diff --git a/tests/generate_benchmark_db.py b/tests/generate_benchmark_db.py index b37519a8..bbebb5eb 100644 --- a/tests/generate_benchmark_db.py +++ b/tests/generate_benchmark_db.py @@ -8,16 +8,18 @@ from tqdm import tqdm from haiku.rag.client import HaikuRAG from haiku.rag.qa.ollama import QA +db_path = Path(__file__).parent / "data" / "benchmark.sqlite" + async def populate_db(): - if (Path(__file__).parent / "data" / "benchmark.sqlite").exists(): + if (db_path).exists(): print("Benchmark database already exists. Skipping creation.") return ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"] # type: ignore corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories") - async with HaikuRAG(Path(__file__).parent / "benchmark.sqlite") as rag: + async with HaikuRAG(db_path) as rag: for i, doc in enumerate(tqdm(corpus)): await rag.create_document( content=doc["document_extracted"], # type: ignore @@ -34,7 +36,7 @@ async def run_match_benchmark(): correct_at_3 = 0 total_queries = 0 - async with HaikuRAG(Path(__file__).parent / "benchmark.sqlite") as rag: + async with HaikuRAG(db_path) as rag: for i, doc in enumerate(tqdm(corpus)): doc_id = doc["document_id"] # type: ignore matches = await rag.search( @@ -73,16 +75,19 @@ async def run_match_benchmark(): return {"recall@1": recall_at_1, "recall@2": recall_at_2, "recall@3": recall_at_3} -async def run_qa_benchmark(): +async def run_qa_benchmark(k: int | None = None): """Run QA benchmarking on the corpus.""" ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"] # type: ignore corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories") + if k is not None: + corpus = corpus.select(range(min(k, len(corpus)))) + judge = LLMJudge() correct_answers = 0 total_questions = 0 - async with HaikuRAG(Path(__file__).parent / "benchmark.sqlite") as rag: + async with HaikuRAG(db_path) as rag: qa = QA(rag) for i, doc in enumerate(tqdm(corpus, desc="QA Benchmarking")): @@ -93,6 +98,10 @@ async def run_qa_benchmark(): is_equivalent = await judge.judge_answers( question, generated_answer, expected_answer ) + print(f"Question: {question}") + print(f"Expected: {expected_answer}") + print(f"Generated: {generated_answer}") + print(f"Equivalent: {is_equivalent}\n") if is_equivalent: correct_answers += 1