evaluate.py

import argparse
import datetime
import enum
import json
import os
import re
import time
from typing import Optional

assert os.path.exists(".env"), "Please create a .env file with your API keys!"
with open(".env", "r") as f:
    for line in f.read().splitlines():
        key, value = line.split("=")
        os.environ[key] = value

from systems import (
    CanvasRetrieval,
    CanvasRetrievalContext,
    Gemma3Reasoner,
    GPT4Reasoner,
    GPT35Reasoner,
    Llama3_1_8bReasoner,
    Llama3_1_70bReasoner,
    Llama3_3Reasoner,
    MixtralReasoner,
    ContextualRAG,
    VanillaRAG,
)
from systems.evaluators import (
    BreakdownEvaluator,
    F1Evaluator,
    LLMEvaluator,
    PrecisionEvaluator,
    RecallEvaluator,
)


class MethodEnum(enum.Enum):
    RANDOM_EVIDENCE = "random_evidence"
    CANVAS_RETRIEVE = "canvas_retrieve"
    CANVAS_RETRIEVE_CONTEXT = "canvas_retrieve_context"
    GPT3 = "gpt3"
    GPT4 = "gpt4"
    LLAMA3_1_8B = "llama3_1_8b"
    LLAMA3_1_70B = "llama3_1_70b"
    LLAMA3_3 = "llama3_3"
    GEMMA3 = "gemma3"
    MIXTRAL = "mixtral"
    VANILLA_RAG = "vanilla_rag"

    def __str__(self):
        return self.name.lower()

    def __repr__(self):
        return str(self)

    @staticmethod
    def argparse(s):
        try:
            return MethodEnum[s.upper()]
        except KeyError:
            return s


REASONERS = {
    "vanilla_rag": VanillaRAG,
    "gpt3": GPT35Reasoner,
    "gpt4": GPT4Reasoner,
    "llama3_1_8b": Llama3_1_8bReasoner,
    "llama3_1_70b": Llama3_1_70bReasoner,
    "llama3_3": Llama3_3Reasoner,
    "gemma3": Gemma3Reasoner,
    "mixtral": MixtralReasoner,
    "rag_systems": ContextualRAG,
    "canvas_retrieve": CanvasRetrieval,
    "canvas_retrieve_context": CanvasRetrievalContext,
}

EVALUATORS = {
    "recall": RecallEvaluator,
    "breakdown": BreakdownEvaluator,
    "precision": PrecisionEvaluator,
    "f1": F1Evaluator,
    "llm": LLMEvaluator,
}

RESULTS_DIRECTORY = "output/results/"
EXPERIMENT_HISTORY_DIRECTORY = "output/experiments/"


USECASE_DATAPATH = {
    "hotpot-scrambled": {
        "data_path": "data/hotpot-scrambled",
        "qa_path": "workloads/hotpot-scrambled.json",
    },
    "stirpot-300": {
        "data_path": "data/stirpot-300",
        "qa_path": "workloads/stirpot-300.json",
    },
}


def _parse_result_cache(
    results_directory: str, system_name: str, usecase: str
) -> Optional[str]:
    """
    Returns the filename for the most recent cached experiment result for a system name and usecase
    and None if matching cache files are not found.
    """
    parsed_filenames = []
    filenames = [f for f in os.listdir(results_directory) if f.endswith(".json")]
    for fname in filenames:
        if fname.startswith(f"{usecase}_{system_name}"):
            timestamp = fname.split("_")[-1].split(".")[0]
            parsed_filenames.append((timestamp, fname))

    if len(parsed_filenames) == 0:
        return None
    sorted_filenames = sorted(parsed_filenames, key=lambda x: x[0])
    latest = sorted_filenames[-1][1]  # these are tuples of (timestamp, filename)
    return latest


def evaluate_system(
    system_name: str,
    usecase: Optional[str],
    data_path: Optional[str],
    workload_path: Optional[str],
    evaluator: str,
    topk: int,
    query_limit: int,
    use_system_cache: bool,
    use_result_cache: bool,
    verbose: bool,
):
    timestamp = time.strftime("%Y-%m-%d-%H%M")

    # Prepares the workload
    assert (usecase is None) ^ (data_path is None)
    assert (usecase is None) ^ (workload_path is None)
    if usecase is not None:
        data_path = USECASE_DATAPATH[usecase]["data_path"]
        workload_path = USECASE_DATAPATH[usecase]["qa_path"]
    with open(workload_path) as f:
        qa_dicts = json.load(f)
    queries = qa_dicts[:query_limit]

    print(f"Benchmark: {usecase} ({len(queries)} questions) on method {method_name}")

    # Prepares the system and evaluator
    reasoner = REASONERS[system_name](
        data_path=data_path,
        topk=topk,
        use_cache=use_system_cache,
        verbose=verbose,
        usecase=usecase,
        method_name=system_name,
    )
    evaluator = EVALUATORS[evaluator](use_cache=use_system_cache, verbose=verbose)

    # Attempt to load result cache
    result_filename = None
    if use_result_cache:
        result_filename = _parse_result_cache(
            results_directory=RESULTS_DIRECTORY,
            system_name=system_name,
            usecase=usecase,
        )

    if result_filename is not None:
        print(f"Using cached results from {result_filename}")
        with open(f"{RESULTS_DIRECTORY}/{method_name}/{result_filename}") as f:
            results = json.load(f)
        for idx, result in results.items():
            measures = evaluator.evaluate(result)
            results[idx] = result | measures

        with open(
            f"{EXPERIMENT_HISTORY_DIRECTORY}/{method_name}/{result_filename}"
        ) as f:
            experiment_history = json.load(f)
        overall_time = experiment_history["overall_time"]

    else:
        results = {}
        start_time = time.time()
        for idx, query in enumerate(queries):
            if verbose:
                print()
                print(f"Running query {idx}: {query['question']}")
            predicted = reasoner(query["question"])
            try:
                target = query["answer"]
            except KeyError:
                target = [answer["value"] for answer in query["answer_set"]]
            result = query | {"predicted": predicted, "target": target}
            del result["answer"]
            measures = evaluator.evaluate(result)
            results[idx] = result | measures
            end_time = time.time()
            if verbose:
                print(f"Target: {target} - Accuracy: {measures['accuracy']}")
                print()
            overall_time = end_time - start_time
    aggregate_scores = evaluator.aggregate_scores()

    print("Overall time", overall_time)
    for k in aggregate_scores:
        print(k, ":", aggregate_scores[k])

    if not use_result_cache:
        experiment_history = {
            "date": timestamp,
            "usecase": usecase,
            "method": method_name,
            "overall_time": overall_time,
        } | aggregate_scores

        experiment_string = f"{usecase}_{method_name}_{timestamp}.json"
        if "vanilla_rag" in method_name:
            experiment_string = f"{usecase}_{method_name}_topk{topk}_{timestamp}.json"
            
        result_path = f"{RESULTS_DIRECTORY}{method_name}/{experiment_string}"

        os.makedirs(f"{RESULTS_DIRECTORY}{method_name}", exist_ok=True)
        with open(result_path,"w+") as f:
            json.dump(results, f, indent=4, ensure_ascii=False)

        if "vanilla_rag" in method_name:
            experiment_history["rag_topk"] = topk

        history_path = f"{EXPERIMENT_HISTORY_DIRECTORY}{method_name}/{experiment_string}.json"
        os.makedirs(f"{EXPERIMENT_HISTORY_DIRECTORY}{method_name}", exist_ok=True)
        with open(history_path,"w+") as f:
            json.dump(experiment_history, f, indent=4, ensure_ascii=False)


if __name__ == "__main__":
    argparser = argparse.ArgumentParser()
    argparser.add_argument(
        "--usecase",
        help="The workload to be used for the evaluation, e.g., bdf-tiny,openmmqa-medium, ma-tiny",
        default="openmmqa-tiny",
    )
    argparser.add_argument(
        "--query-limit",
        type=int,
        help="Limit the number of queries to run",
        default=0,
    )
    argparser.add_argument(
        "--system-cache",
        help="Use system cache if it's available",
        default=False,
        action="store_true",
    )
    argparser.add_argument(
        "--result-cache",
        help="Use cached results if they are available",
        default=False,
        action="store_true",
    )
    argparser.add_argument(
        "--verbose",
        help="Print intermediate results",
        default=False,
        action="store_true",
    )
    argparser.add_argument(
        "--method",
        help="Choose method to evaluate for",
        type=MethodEnum.argparse,
        choices=list(MethodEnum),
        default="canvas_systems",
    )
    argparser.add_argument(
        "--evaluator",
        help="Choose evaluator class, e.g., recall, breakdown",
        default="llm",
    )
    argparser.add_argument(
        "--topk",
        help="The number of top-k documents to retrieve if using a RAG method (default 50)",
        default=50,
    )
    args = argparser.parse_args()
    ### Parsing arguments

    method = args.method
    method_name = method.name.lower()

    evaluate_system(
        system_name=method_name,
        usecase=args.usecase,
        data_path=None,
        workload_path=None,
        evaluator=args.evaluator,
        topk=int(args.topk),
        query_limit=int(args.query_limit) or None,
        use_system_cache=args.system_cache,
        use_result_cache=args.result_cache,
        verbose=args.verbose,
    )