graphrag-eval 4.0.0__py3-none-any.whl

graphrag_eval/__init__.py

@@ -0,0 +1,4 @@
+from .aggregation import *
+from .evaluation import *
+from .steps import *
+from .steps.sparql import *

graphrag_eval/aggregation.py

@@ -0,0 +1,151 @@
+import json
+from collections import defaultdict
+from statistics import mean, median
+from typing import Any, Iterable
+
+
+METRICS = [
+    "answer_recall",
+    "answer_precision",
+    "answer_relevance",
+    "answer_relevance_cost",
+    "answer_f1",
+    "steps_score",
+    "input_tokens",
+    "output_tokens",
+    "total_tokens",
+    "elapsed_sec"
+]
+
+PROTECTED_METRICS = [
+    "input_tokens",
+    "output_tokens",
+    "total_tokens",
+    "elapsed_sec"
+]
+
+
+def stats_for_series(values: Iterable[int | float]) -> dict[str, float]:
+    return {
+        "sum": sum(values),
+        "mean": mean(values) if values else 0,
+        "median": median(values) if values else 0,
+        "min": min(values) if values else 0,
+        "max": max(values) if values else 0,
+    }
+
+
+def update_stats_per_template(
+    sample: dict,
+    stats_per_template: dict,
+    template_id: str
+):
+    for metric in METRICS:
+        value = sample.get(metric)
+        if value is not None:
+            stats_per_template[template_id][metric].append(value)
+
+
+def update_steps_summary_per_template(
+    sample: dict,
+    steps_summary_per_template: dict,
+    template_id: str
+):
+    seen = set()
+    for step in sample.get("actual_steps", []):
+        name = step["name"]
+        template_steps_summary = steps_summary_per_template[template_id]
+        template_steps_summary["total"][name] += 1
+        if step["status"] == "error":
+            template_steps_summary["errors"][name] += 1
+        if name not in seen:
+            seen.add(name)
+            template_steps_summary["once_per_sample"][name] += 1
+
+        if step["status"] != "error":
+            try:
+                res = json.loads(step["output"])
+                if "results" in res and "bindings" in res["results"]:
+                    if not res["results"]["bindings"]:
+                        template_steps_summary["empty_results"][name] += 1
+            except json.decoder.JSONDecodeError:
+                pass
+
+
+def compute_aggregates(samples: list[dict]) -> dict:
+    number_of_samples_per_template_by_status = defaultdict(lambda: defaultdict(int))
+    stats_per_template = defaultdict(lambda: defaultdict(list))
+    steps_summary_per_template = defaultdict(lambda: defaultdict(lambda: defaultdict(int)))
+
+    # Compute per-template stats
+    templates_ids = set()
+    for sample in samples:
+        template_id = sample["template_id"]
+        templates_ids.add(template_id)
+
+        if "error" in sample:
+            number_of_samples_per_template_by_status[template_id]["error"] += 1
+            continue
+        number_of_samples_per_template_by_status[template_id]["success"] += 1
+
+        update_stats_per_template(sample, stats_per_template, template_id)
+        update_steps_summary_per_template(
+            sample,
+            steps_summary_per_template,
+            template_id
+        )
+
+    summary = {"per_template": {}}
+
+    # Add per-template stats
+    for template_id in templates_ids:
+        template_summary: dict[str, Any] = {
+            "number_of_error_samples": number_of_samples_per_template_by_status[template_id]["error"],
+            "number_of_success_samples": number_of_samples_per_template_by_status[template_id]["success"],
+        }
+        steps_summary = {
+            k1: {k2: v2 for k2, v2 in v1.items()}
+            for k1, v1 in steps_summary_per_template[template_id].items()
+        }
+        if steps_summary:
+            template_summary.update({"steps": steps_summary})
+        for metric in METRICS:
+            results_for_template = stats_per_template[template_id]
+            series = results_for_template.get(metric, [])
+            if series or metric in PROTECTED_METRICS:
+                template_summary[metric] = stats_for_series(series)
+
+        summary["per_template"][template_id] = template_summary
+
+    # Add micro stats
+    values_ = number_of_samples_per_template_by_status.values()
+    summary["micro"] = {
+        "number_of_error_samples": sum(
+            values["error"] for values in values_
+        ),
+        "number_of_success_samples": sum(
+            values["success"] for values in values_
+        ),
+    }
+    for metric in METRICS:
+        series = [
+            i
+            for values in stats_per_template.values()
+            for i in values[metric]
+            if values.get(metric) is not None
+        ]
+        if series or metric in PROTECTED_METRICS:
+            summary["micro"][metric] = stats_for_series(series)
+
+    # Add macro stats
+    summary["macro"] = {}
+    for metric in METRICS:
+        means = [
+            values[metric]["mean"]
+            for template_id, values in summary["per_template"].items()
+            if values.get(metric) is not None
+        ]
+        if means or metric in PROTECTED_METRICS:
+            summary["macro"][metric] = {"mean": mean(means) if means else 0}
+
+    return summary

graphrag_eval/answer_correctness.py

@@ -0,0 +1,162 @@
+import csv
+from pathlib import Path
+
+from openai import OpenAI
+from tqdm import tqdm
+
+
+IN_FILE_PATH = "../data/data-1.tsv"
+PROMPT_FILE_PATH = "prompts/template.md"
+OUT_FILE_PATH = "results/data-1.tsv"
+OUT_FIELDS = ["#Reference", "#PTarget", "#Matching", "Reasoning", "Error"]
+LLM_MODEL = "gpt-4o-mini"
+TEMPERATURE = 0.0
+
+
+
+def compute_recall_precision_f1(
+    n_pos: int | None,
+    n_pred_pos: int | None,
+    n_true_pos: int | None,
+) -> tuple[float | None, float | None, float | None]:
+    recall = None
+    precision = None
+    f1 = None
+    if n_true_pos is not None and n_pos:
+        recall = n_true_pos / n_pos
+    if n_true_pos is not None and n_pred_pos:
+        precision = n_true_pos / n_pred_pos
+    if precision is not None and recall is not None and precision + recall > 0:
+        f1 = 2 * (precision * recall) / (precision + recall)
+    return recall, precision, f1
+
+
+def extract_response_values(
+    response: str
+) -> tuple[int | None, int | None, int | None, str, str]:
+    vals = response.split("\t")
+    n = len(vals)
+    if n < 4:
+        msg = f"Expected 4 tab-separated values: {response}"
+        return None, None, None, "", msg
+    vals = vals[:4]
+    try:
+        n_ref, n_target, n_matching = map(int, vals[:3])
+    except ValueError:
+        msg = f"Non-int value: {response}"
+        return None, None, None, vals[3], msg
+    if any([
+        n_ref < 1,
+        n_target < 1,
+        n_matching < 0,
+        n_matching > n_ref,
+        n_matching > n_target
+    ]):
+        msg = f"Invalid int values: {n_ref}\t{n_target}\t{n_matching}"
+        return None, None, None, vals[3], msg
+    return n_ref, n_target, n_matching, vals[3], ""
+
+
+class AnswerCorrectnessEvaluator:
+    def __init__(
+        self,
+        prompt_file_path: str | Path = PROMPT_FILE_PATH,
+        temperature : float = TEMPERATURE
+    ):
+        with open(prompt_file_path, encoding="utf-8") as f:
+            self.prompt_template = f.read()
+        self.openai_client = OpenAI()
+        self.temperature = temperature
+
+    def call_llm(self, prompt: str) -> str:
+        try:
+            response = self.openai_client.chat.completions.create(
+                model=LLM_MODEL,
+                messages=[{"role": "user", "content": prompt}],
+                temperature=self.temperature
+            )
+            return response.choices[0].message.content.strip("\n")
+        except Exception as e:
+            return str(e).replace("\n", "    ")
+
+    def evaluate_answer(
+        self,
+        question: str,
+        reference_answer: str,
+        actual_answer: str
+    ):
+        prompt = self.prompt_template.format(
+            question=question,
+            reference_answer=reference_answer,
+            candidate_answer=actual_answer,
+        )
+        response_str = self.call_llm(prompt)
+        return extract_response_values(response_str)
+
+    def get_correctness_dict(
+        self,
+        reference: dict,
+        target: dict,
+    ):
+        result = {}
+        result["reference_answer"] = reference["reference_answer"]
+        num_ref_claims, num_actual_claims, num_matching_claims, reason, error = \
+        self.evaluate_answer(
+            reference["question_text"],
+            reference["reference_answer"],
+            target["actual_answer"],
+        )
+        if error:
+            result["answer_eval_error"] = error
+        else:
+            result.update({
+                "answer_reference_claims_count": num_ref_claims,
+                "answer_actual_claims_count": num_actual_claims,
+                "answer_matching_claims_count": num_matching_claims,
+                "answer_correctness_reason": reason,
+            })
+            recall, precision, f1 = compute_recall_precision_f1(
+                num_ref_claims, num_actual_claims, num_matching_claims
+            )
+            if recall is not None:
+                result["answer_recall"] = recall
+            if precision is not None:
+                result["answer_precision"] = precision
+            if f1 is not None:
+                result["answer_f1"] = f1
+        return result
+
+
+def evaluate_and_write(
+    in_file_path: str | Path,
+    out_file_path: str | Path,
+) -> None:
+    evaluator = AnswerCorrectnessEvaluator(PROMPT_FILE_PATH)
+    with open(in_file_path, encoding="utf-8") as f:
+        reader = csv.DictReader(f, delimiter="\t")
+        rows = [row for row in reader]
+    print(f"Writing results to {out_file_path}")
+    Path(out_file_path).parent.mkdir(parents=True, exist_ok=True)
+    with open(out_file_path, "w", encoding="utf-8") as f:
+        writer = csv.writer(f, delimiter="\t")
+        writer.writerow(OUT_FIELDS)
+        for row in tqdm(rows):
+            vals = evaluator.evaluate_answer(
+                row["Question"],
+                row["Reference answer"],
+                row["Actual answer"]
+            )
+            writer.writerow(vals)
+            f.flush()
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--in-file", type=str, default=IN_FILE_PATH)
+    parser.add_argument("-o", "--out-file", type=str, default=OUT_FILE_PATH)
+    args = parser.parse_args()
+    evaluate_and_write(
+        in_file_path=args.in_file,
+        out_file_path=args.out_file,
+    )

graphrag_eval/answer_relevance.py

@@ -0,0 +1,37 @@
+from langevals_ragas.response_relevancy import (
+    RagasResponseRelevancyEvaluator,
+    RagasResponseRelevancyEntry
+)
+
+
+def get_relevance_dict(
+    question_text: str,
+    actual_answer: str,
+    model_name : str = 'openai/gpt-4o-mini',
+    max_tokens: int = 65_536
+) -> dict:
+    settings_dict = {
+        'model': model_name,
+        'max_tokens': max_tokens
+    }
+    entry = RagasResponseRelevancyEntry(
+        input=question_text,
+        output=actual_answer
+    )
+    evaluator = RagasResponseRelevancyEvaluator(settings=settings_dict)
+    try:
+        result = evaluator.evaluate(entry)
+        if result.status == "processed":
+            return {
+                "answer_relevance": result.score,
+                "answer_relevance_cost": result.cost.amount,
+                "answer_relevance_reason": result.details,
+            }
+        else:
+            return {
+                "answer_relevance_error": result.details
+            }
+    except Exception as e:
+        return {
+            "answer_relevance_error": str(e),
+        }

graphrag_eval/evaluation.py

@@ -0,0 +1,62 @@
+from .steps import get_steps_evaluation_result_dict
+
+
+def run_evaluation(
+        qa_dataset: list[dict],
+        responses_dict: dict,
+) -> list[dict]:
+    # Output metrics are not nested, for simpler aggregation
+    answer_correctess_evaluator = None
+    evaluation_results = []
+    for template in qa_dataset:
+        template_id = template["template_id"]
+        for question in template["questions"]:
+            actual_result = responses_dict[question["id"]]
+            eval_result = {
+                "template_id": template_id,
+                "question_id": actual_result["question_id"],
+                "question_text": question["question_text"]
+            }
+            if "reference_answer" in question:
+                eval_result["reference_answer"] = question["reference_answer"]
+            if "reference_steps" in question:
+                eval_result["reference_steps"] = question["reference_steps"]
+            if "error" in actual_result:
+                eval_result.update({
+                    "status": "error",
+                    "error": actual_result["error"],
+                })
+                evaluation_results.append(eval_result)
+                continue
+            eval_result["status"] = "success"
+            if "actual_answer" in actual_result:
+                eval_result["actual_answer"] = actual_result["actual_answer"]
+                from graphrag_eval import answer_relevance
+                eval_result.update(
+                    answer_relevance.get_relevance_dict(
+                        question["question_text"],
+                        actual_result["actual_answer"],
+                    )
+                )
+            if "reference_answer" in question and "actual_answer" in actual_result:
+                from graphrag_eval.answer_correctness import AnswerCorrectnessEvaluator
+                if not answer_correctess_evaluator:
+                    answer_correctess_evaluator = AnswerCorrectnessEvaluator()
+                eval_result.update(
+                    answer_correctess_evaluator.get_correctness_dict(
+                        question,
+                        actual_result,
+                    )
+                )
+            if "steps" in actual_result:
+                eval_result.update(
+                    get_steps_evaluation_result_dict(question, actual_result)
+                )
+            eval_result.update({
+                "input_tokens": actual_result["input_tokens"],
+                "output_tokens": actual_result["output_tokens"],
+                "total_tokens": actual_result["total_tokens"],
+                "elapsed_sec": actual_result["elapsed_sec"],
+            })
+            evaluation_results.append(eval_result)
+    return evaluation_results

graphrag_eval/steps/__init__.py

@@ -0,0 +1,120 @@
+import json
+from collections import defaultdict
+
+from .retrieval import recall_at_k
+from .sparql import compare_sparql_results
+
+
+def compare_steps_outputs(reference: dict, actual: dict) -> float:
+    ref_output = reference["output"]
+    act_output = actual["output"]
+    if reference.get("output_media_type") == "application/sparql-results+json":
+        return compare_sparql_results(
+            json.loads(ref_output),
+            json.loads(act_output),
+            reference["required_columns"],
+            reference.get("ordered", False),
+        )
+    if reference.get("output_media_type") == "application/json":
+        return float(json.loads(ref_output) == json.loads(act_output))
+    if reference["name"] == "retrieval":
+        k = reference["args"]["k"]
+        return recall_at_k(ref_output, act_output, k)
+    return float(ref_output == act_output)
+
+
+def match_group_by_output(
+        reference_steps: list[list[dict]],
+        group_idx: int,
+        actual_steps: list[dict],
+        candidates_by_name: dict[str, list[int]],
+) -> list[tuple[int, int, int, float]]:
+    used_actual_indices = set()
+    matches = []
+
+    reference_group = reference_steps[group_idx]
+    for reference_idx, reference_step in enumerate(reference_group):
+        name = reference_step["name"]
+        candidates = reversed(candidates_by_name.get(name, []))
+        for actual_idx in candidates:
+            if actual_idx in used_actual_indices:
+                continue
+            actual_step = actual_steps[actual_idx]
+            score = compare_steps_outputs(reference_step, actual_step)
+            if score > 0.0:
+                matches.append((group_idx, reference_idx, actual_idx, score))
+                used_actual_indices.add(actual_idx)
+                break
+    return matches
+
+
+def collect_possible_matches_by_name_and_status(
+        group: list[dict],
+        actual_steps: list[dict],
+        search_upto: int,
+) -> dict[str, list[int]]:
+    group_by_name = defaultdict(list)
+
+    for j in range(search_upto):
+        name = actual_steps[j]["name"]
+        if actual_steps[j]["status"] == "success":
+            group_by_name[name].append(j)
+
+    reference_names = {item["name"] for item in group}
+    return {name: group_by_name[name] for name in reference_names if name in group_by_name}
+
+
+def get_steps_matches(
+        reference_steps: list[list[dict]],
+        actual_steps: list[dict],
+) -> list[tuple[int, int, int, float]]:
+    # when we have autocomplete
+    # matches = []
+    # search_upto = len(actual_steps)
+    # for group_idx in reversed(range(len(reference_steps))):
+    #     group = reference_steps[group_idx]
+    #     candidates = collect_possible_matches_by_name(group, actual_steps, search_upto)
+    #
+    #     matched = match_group_by_output(reference_steps, group_idx, actual_steps, candidates)
+    #     if len(matched) == len(group):
+    #         # update search_upto to just before the highest matched actual index
+    #         matches.extend(matched)
+    #         search_upto = min(j for (_, j) in matched)
+    #     elif len(matched) < len(group):
+    #         matches.extend(matched)
+    #         break # a step is not matched and missing, abort
+    #     else:
+    #         break  # a step is not matched and missing, abort
+    # return matches
+
+    # for now, we have only the last step(s)
+    last_group = reference_steps[-1]
+    candidates = collect_possible_matches_by_name_and_status(last_group, actual_steps, len(actual_steps))
+    return match_group_by_output(reference_steps, -1, actual_steps, candidates)
+
+
+def evaluate_steps(
+    reference_steps_groups: list[list[dict]],
+    actual_steps: list[dict]
+) -> float:
+    matches = get_steps_matches(reference_steps_groups, actual_steps)
+    matches_by_group = defaultdict(list)
+    scores_by_group = defaultdict(float)
+    for ref_group_idx, ref_match_idx, actual_idx, score in matches:
+        matches_by_group[ref_group_idx].append(ref_match_idx)
+        scores_by_group[ref_group_idx] += score
+        reference_steps_groups[ref_group_idx][ref_match_idx]["matches"] \
+            = actual_steps[actual_idx]["id"]
+    group_ix = -1  # For now, consider only the last reference group of steps
+    return scores_by_group[group_ix] / len(reference_steps_groups[group_ix])
+
+
+def get_steps_evaluation_result_dict(reference: dict, target: dict) -> dict:
+    act_steps = target["steps"]
+    eval_result = {}
+    eval_result["actual_steps"] = act_steps
+    if "reference_steps" in reference:
+        ref_steps = reference["reference_steps"]
+        steps_score = evaluate_steps(ref_steps, act_steps)
+        eval_result["steps_score"] = steps_score
+    return eval_result

graphrag_eval/steps/retrieval.py

@@ -0,0 +1,55 @@
+from typing import Iterable
+
+
+def recall_at_k(relevant_docs: Iterable, retrieved_docs: list, k: int = 10) -> float:
+    """
+    Calculates Recall@k.
+
+    Args:
+        relevant_docs (Iterable): A set of ground truth relevant document IDs.
+        retrieved_docs (list): A list of retrieved document IDs, ordered by rank.
+        k (int): The cutoff for the retrieval list.
+
+    Returns:
+        float: The Recall@k score.
+    """
+    retrieved_at_k = retrieved_docs[:k]
+
+    relevant_set = set(relevant_docs)
+    retrieved_set = set(retrieved_at_k)
+    true_positives = len(relevant_set.intersection(retrieved_set))
+
+    total_relevant = len(relevant_set)
+
+    if total_relevant == 0:
+        return 0.0
+
+    return true_positives / total_relevant
+
+
+def average_precision(relevant_docs: Iterable, retrieved_docs: Iterable) -> float:
+    """
+    Calculates Average Precision (AP) for a single query.
+
+    Args:
+        relevant_docs (Iterable): A set of ground truth relevant document IDs.
+        retrieved_docs (Iterable): A list of retrieved document IDs, ordered by rank.
+
+    Returns:
+        float: The Average Precision score.
+    """
+    relevant_set = set(relevant_docs)
+    hits = 0
+    sum_of_precisions = 0.0
+
+    for i, doc_id in enumerate(retrieved_docs):
+        if doc_id in relevant_set:
+            hits += 1
+            precision_at_k = hits / (i + 1)
+            sum_of_precisions += precision_at_k
+
+    total_relevant = len(relevant_set)
+    if total_relevant == 0:
+        return 0.0
+
+    return sum_of_precisions / total_relevant