From 2855dd7c1341c2fd61d7586d79edf47311b5f255 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 24 Jun 2025 20:02:07 +0300 Subject: [PATCH] Use LLM-as-a-judge to test QA --- tests/__init__.py | 0 tests/conftest.py | 7 ++++ tests/generate_benchmark_db.py | 46 +++++++++++++++++++++-- tests/llm_judge.py | 68 ++++++++++++++++++++++++++++++++++ tests/test_qa.py | 42 +++++++++++++++++++++ 5 files changed, 160 insertions(+), 3 deletions(-) create mode 100644 tests/__init__.py create mode 100644 tests/llm_judge.py create mode 100644 tests/test_qa.py diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/conftest.py b/tests/conftest.py index 31ed812d..2dcea549 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,6 +3,8 @@ from pathlib import Path import pytest from datasets import Dataset, load_dataset, load_from_disk +from .llm_judge import LLMJudge + @pytest.fixture(scope="session") def qa_corpus() -> Dataset: @@ -16,3 +18,8 @@ def qa_corpus() -> Dataset: corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories") corpus.save_to_disk(ds_path) return corpus + + +@pytest.fixture(scope="session") +def llm_judge() -> LLMJudge: + return LLMJudge() diff --git a/tests/generate_benchmark_db.py b/tests/generate_benchmark_db.py index 01ce7c20..b37519a8 100644 --- a/tests/generate_benchmark_db.py +++ b/tests/generate_benchmark_db.py @@ -1,13 +1,16 @@ +import asyncio from pathlib import Path from datasets import Dataset, load_dataset +from llm_judge import LLMJudge from tqdm import tqdm from haiku.rag.client import HaikuRAG +from haiku.rag.qa.ollama import QA async def populate_db(): - if (Path(__file__).parent / "benchmark.sqlite").exists(): + if (Path(__file__).parent / "data" / "benchmark.sqlite").exists(): print("Benchmark database already exists. Skipping creation.") return @@ -61,6 +64,7 @@ async def run_match_benchmark(): recall_at_2 = correct_at_2 / total_queries recall_at_3 = correct_at_3 / total_queries + print("\n=== Retrieval Benchmark Results ===") print(f"Total queries: {total_queries}") print(f"Recall@1: {recall_at_1:.4f}") print(f"Recall@2: {recall_at_2:.4f}") @@ -69,12 +73,48 @@ 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(): + """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") + + judge = LLMJudge() + correct_answers = 0 + total_questions = 0 + + async with HaikuRAG(Path(__file__).parent / "benchmark.sqlite") as rag: + qa = QA(rag) + + for i, doc in enumerate(tqdm(corpus, desc="QA Benchmarking")): + question = doc["question"] # type: ignore + expected_answer = doc["answer"] # type: ignore + + generated_answer = await qa.answer(question) + is_equivalent = await judge.judge_answers( + question, generated_answer, expected_answer + ) + + if is_equivalent: + correct_answers += 1 + total_questions += 1 + + accuracy = correct_answers / total_questions if total_questions > 0 else 0 + + print("\n=== QA Benchmark Results ===") + print(f"Total questions: {total_questions}") + print(f"Correct answers: {correct_answers}") + print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)") + + async def main(): await populate_db() + + print("Running retrieval benchmarks...") await run_match_benchmark() + print("\nRunning QA benchmarks...") + await run_qa_benchmark() + if __name__ == "__main__": - import asyncio - asyncio.run(main()) diff --git a/tests/llm_judge.py b/tests/llm_judge.py new file mode 100644 index 00000000..66bfd2cb --- /dev/null +++ b/tests/llm_judge.py @@ -0,0 +1,68 @@ +import json + +from ollama import AsyncClient +from pydantic import BaseModel + +from haiku.rag.config import Config + + +class LLMJudgeResponseSchema(BaseModel): + equivalent: bool + + +class LLMJudge: + """LLM-as-judge for evaluating answer equivalence using Ollama.""" + + def __init__(self, model: str = "qwen3"): + self.model = model + self.client = AsyncClient(host=Config.OLLAMA_BASE_URL) + + async def judge_answers( + self, question: str, answer: str, expected_answer: str + ) -> bool: + """ + Judge whether two answers are equivalent for a given question. + + Args: + question: The original question + answer: The generated answer to evaluate + expected_answer: The reference/expected answer + + Returns: + Dictionary with judgment result: + - equivalent: bool indicating if answers are equivalent + - explanation: str explaining the reasoning + - score: str rating from 1-5 + """ + + prompt = f""" + You are an expert judge evaluating the equivalence of two answers to the same question. + + Question: {question} + + Generated Answer: {answer} + + Expected Answer: {expected_answer} + + Your task is to determine if these two answers are equivalent in meaning and both correctly answer the question. Consider: + + 1. Do both answers provide the same answer? + 2. Do both answers directly address the question asked? + 3. Minor differences in wording or style are acceptable if the meaning of the answer is the same. + + Be strict but fair in your evaluation. Focus on factual correctness and whether both answers would satisfy someone asking the question.""" + + response = await self.client.chat( + model=self.model, + messages=[{"role": "user", "content": prompt}], + format=LLMJudgeResponseSchema.model_json_schema(), + think=False, + ) + + answer = response["message"]["content"].strip() + try: + res = json.loads(answer) + assert "equivalent" in res, "Response must contain 'equivalent' key" + return res["equivalent"] + except json.JSONDecodeError: + assert False, "Response is not valid JSON" diff --git a/tests/test_qa.py b/tests/test_qa.py new file mode 100644 index 00000000..18496df7 --- /dev/null +++ b/tests/test_qa.py @@ -0,0 +1,42 @@ +from typing import TYPE_CHECKING + +import pytest +from datasets import Dataset + +from haiku.rag.client import HaikuRAG +from haiku.rag.qa.ollama import QA + +if TYPE_CHECKING: + import sys + from pathlib import Path + + sys.path.append(str(Path(__file__).parent)) + from llm_judge import LLMJudge + + +@pytest.mark.asyncio +async def test_qa_with_dataset_question(qa_corpus: Dataset, llm_judge: "LLMJudge"): + """Test QA with actual question from the dataset using LLM judge.""" + client = HaikuRAG(":memory:") + qa = QA(client) + + # Use the first document from the corpus + doc = qa_corpus[1] + + # Add the document to database + await client.create_document( + content=doc["document_extracted"], uri=doc["document_id"] + ) + + question = doc["question"] + expected_answer = doc["answer"] + + answer = await qa.answer(question) + # Use LLM judge to evaluate answer equivalence + is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer) + + assert isinstance(answer, str) + assert len(answer) > 0 + assert is_equivalent, ( + f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}" + )