llama-stack-mirror/llama_stack/apis/evals/client.py
2024-10-10 11:35:26 -07:00

85 lines
2.2 KiB
Python

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
import asyncio
import json
import fire
import httpx
from termcolor import cprint
from .evals import * # noqa: F403
class EvaluationClient(Evals):
def __init__(self, base_url: str):
self.base_url = base_url
async def initialize(self) -> None:
pass
async def shutdown(self) -> None:
pass
async def run_evals(
self,
model: str,
task: str,
dataset: Optional[str] = None,
eval_task_config: Optional[EvaluateTaskConfig] = None,
) -> EvaluateResponse:
async with httpx.AsyncClient() as client:
response = await client.post(
f"{self.base_url}/evals/run",
json={
"model": model,
"task": task,
"dataset": dataset,
"eval_task_config": (
json.loads(eval_task_config.json())
if eval_task_config
else None
),
},
headers={"Content-Type": "application/json"},
timeout=3600,
)
response.raise_for_status()
return EvaluateResponse(**response.json())
async def run_main(host: str, port: int):
client = EvaluationClient(f"http://{host}:{port}")
# CustomDataset
response = await client.run_evals(
model="Llama3.1-8B-Instruct",
dataset="mmlu-simple-eval-en",
task="mmlu",
eval_task_config=EvaluateTaskConfig(
n_samples=2,
),
)
cprint(f"evaluate response={response}", "green")
# Eleuther Eval Task
# response = await client.run_evals(
# model="Llama3.1-8B-Instruct",
# task="meta_mmlu_pro_instruct",
# # task="meta_ifeval",
# eval_task_config=EvaluateTaskConfig(
# n_samples=2,
# )
# )
# cprint(response.metrics["metrics_table"], "red")
def main(host: str, port: int):
asyncio.run(run_main(host, port))
if __name__ == "__main__":
fire.Fire(main)