simple_evals.py

import json
import argparse
import pandas as pd
from common import make_report
from drop_eval import DropEval
from gpqa_eval import GPQAEval
from humaneval_eval import HumanEval
from math_eval import MathEval
from mgsm_eval import MGSMEval
from mmlu_eval import MMLUEval
from simpleqa_eval import SimpleQAEval
from sampler.chat_completion_sampler import (
    OPENAI_SYSTEM_MESSAGE_API,
    OPENAI_SYSTEM_MESSAGE_CHATGPT,
    ChatCompletionSampler,
)
from sampler.o1_chat_completion_sampler import O1ChatCompletionSampler
from sampler.claude_sampler import ClaudeCompletionSampler, CLAUDE_SYSTEM_MESSAGE_LMSYS
from sampler.exa_sampler import ExaSampler
from sampler.perplexity_sampler import PerplexitySampler
from sampler.you_sampler import YouSampler
from sampler.brave_sampler import BraveSampler
from sampler.bing_sampler import BingSampler
from sampler.tavily_sampler import TavilySampler
from sampler.serper_sampler import SerperSampler
from sampler.result_sampler import ResultSampler
from constants import (EXA_API_KEY, PERPLEXITY_API_KEY, YOU_API_KEY, BRAVE_API_KEY, BING_API_KEY, TAVILY_API_KEY, SERPER_API_KEY)

def main():
    parser = argparse.ArgumentParser(
        description="Run sampling and evaluations using different samplers and evaluations."
    )
    parser.add_argument(
        "--list-models", action="store_true", help="List available models"
    )
    parser.add_argument("--model", type=str, help="Select a model by name")
    parser.add_argument("--debug", action="store_true", help="Run in debug mode")
    parser.add_argument(
        "--examples", type=int, help="Number of examples to use (overrides default)"
    )

    args = parser.parse_args()

    """
    models = {
        # chatgpt models:
        "gpt-4o-2024-11-20_assistant": ChatCompletionSampler(
            model="gpt-4o-2024-11-20",
            system_message=OPENAI_SYSTEM_MESSAGE_API,
            max_tokens=2048,
        ),
        "gpt-4o-2024-11-20_chatgpt": ChatCompletionSampler(
            model="gpt-4o-2024-11-20",
            system_message=OPENAI_SYSTEM_MESSAGE_CHATGPT,
            max_tokens=2048,
        ),
        "o1": O1ChatCompletionSampler(
            model="o1",
        ), 
        "o1-preview": O1ChatCompletionSampler(
            model="o1-preview",
        ),
        "o1-mini": O1ChatCompletionSampler(
            model="o1-mini",
        ),
        "gpt-4-turbo-2024-04-09_assistant": ChatCompletionSampler(
            model="gpt-4-turbo-2024-04-09",
            system_message=OPENAI_SYSTEM_MESSAGE_API,
        ),
        "gpt-4-turbo-2024-04-09_chatgpt": ChatCompletionSampler(
            model="gpt-4-turbo-2024-04-09",
            system_message=OPENAI_SYSTEM_MESSAGE_CHATGPT,
        ),
        "gpt-4o_assistant": ChatCompletionSampler(
            model="gpt-4o",
            system_message=OPENAI_SYSTEM_MESSAGE_API,
            max_tokens=2048,
        ),
        "gpt-4o_chatgpt": ChatCompletionSampler(
            model="gpt-4o",
            system_message=OPENAI_SYSTEM_MESSAGE_CHATGPT,
            max_tokens=2048,
        ),
        "gpt-4o-mini-2024-07-18": ChatCompletionSampler(
            model="gpt-4o-mini-2024-07-18",
            system_message=OPENAI_SYSTEM_MESSAGE_API,
            max_tokens=2048,
        ),
        # claude models:
        "claude-3-opus-20240229_empty": ClaudeCompletionSampler(
            model="claude-3-opus-20240229",
            system_message=CLAUDE_SYSTEM_MESSAGE_LMSYS,
        ),
    }
    """
    
    providers = {
        # New models:
        "exa": ExaSampler(
            api_key=EXA_API_KEY,
        ),
        "exa-fast": ExaSampler(
            api_key=EXA_API_KEY,
            mode="fast"
        ),
        # "perplexity-pro": PerplexitySampler(
        #     api_key=PERPLEXITY_API_KEY,
        #     model="sonar-pro",
        # ),
        # "perplexity": PerplexitySampler(
        #     api_key=PERPLEXITY_API_KEY,
        #     model="sonar",
        # ),
        # "you": YouSampler(
        #     api_key=YOU_API_KEY,
        # ),
        # Result-based models using different search providers
        # "brave-rag": ResultSampler(
        #     provider=BraveSampler(api_key=BRAVE_API_KEY),
        # ),
        # "bing-rag": ResultSampler(
        #     provider=BingSampler(api_key=BING_API_KEY),
        # ),
        # "tavily-rag": ResultSampler(
        #     provider=TavilySampler(api_key=TAVILY_API_KEY),
        # ),
        # "serper-rag": ResultSampler(
        #     provider=SerperSampler(api_key=SERPER_API_KEY),
        # ),
    }
    all_models = { **providers }

    if args.list_models:
        print("Available models:")
        for model_name in all_models.keys():
            print(f" - {model_name}")
        return

    if args.model:
        if args.model not in all_models:
            print(f"Error: Model '{args.model}' not found.")
            return
        all_models = {args.model: all_models[args.model]}

    grading_sampler = ChatCompletionSampler(model="gpt-4o")
    equality_checker = ChatCompletionSampler(model="gpt-4-turbo-preview")
    # ^^^ used for fuzzy matching, just for math

    def get_evals(eval_name, debug_mode):
        num_examples = (
            args.examples if args.examples is not None else (5 if debug_mode else None)
        )
        # Set num_examples = None to reproduce full evals
        match eval_name:
            case "mmlu":
                return MMLUEval(num_examples=1 if debug_mode else num_examples)
            case "math":
                return MathEval(
                    equality_checker=equality_checker,
                    num_examples=num_examples,
                    n_repeats=1 if debug_mode else 10,
                )
            case "gpqa":
                return GPQAEval(
                    n_repeats=1 if debug_mode else 10, num_examples=num_examples
                )
            case "mgsm":
                return MGSMEval(num_examples_per_lang=10 if debug_mode else 250)
            case "drop":
                return DropEval(
                    num_examples=10 if debug_mode else num_examples,
                    train_samples_per_prompt=3,
                )
            case "humaneval":
                return HumanEval(num_examples=10 if debug_mode else num_examples)
            case "simpleqa":
                return SimpleQAEval(
                    grader_model=grading_sampler,
                    num_examples=10 if debug_mode else num_examples,
                )
            case _:
                raise Exception(f"Unrecognized eval type: {eval_name}")

    evals = {
        eval_name: get_evals(eval_name, args.debug)
        # Excluded are:  "mmlu", "gpqa", "mgsm", "drop", "math", "humaneval"
        for eval_name in ["simpleqa"]
    }
    print(evals)
    debug_suffix = "_DEBUG" if args.debug else ""
    print(debug_suffix)
    mergekey2resultpath = {}
    for eval_name, eval_obj in evals.items():
        print(f"\nRunning {eval_name} evaluation:")
        for model_name, sampler in all_models.items():
            print(f"  Testing model: {model_name}")
            result = eval_obj(sampler)
            file_stem = f"{eval_name}_{model_name}"
            report_filename = f"/tmp/{file_stem}{debug_suffix}.html"
            print(f"  Writing report to {report_filename}")
            with open(report_filename, "w") as fh:
                fh.write(make_report(result))
            metrics = result.metrics | {"score": result.score}
            print(f"  Results: {metrics}")
            result_filename = f"/tmp/{file_stem}{debug_suffix}.json"
            with open(result_filename, "w") as f:
                f.write(json.dumps(metrics, indent=2))
            print(f"  Writing results to {result_filename}")
            mergekey2resultpath[f"{file_stem}"] = result_filename

    merge_metrics = []
    for eval_model_name, result_filename in mergekey2resultpath.items():
        try:
            result = json.load(open(result_filename, "r+"))
        except Exception as e:
            print(e, result_filename)
            continue
        result = result.get("f1_score", result.get("score", None))
        eval_name = eval_model_name[: eval_model_name.find("_")]
        model_name = eval_model_name[eval_model_name.find("_") + 1 :]
        merge_metrics.append(
            {"eval_name": eval_name, "model_name": model_name, "metric": result}
        )
    merge_metrics_df = pd.DataFrame(merge_metrics).pivot(
        index=["model_name"], columns="eval_name"
    )
    print("\nAll results: ")
    print(merge_metrics_df.to_markdown())
    return merge_metrics


if __name__ == "__main__":
    main()