src/models.py

from typing import Any, Optional

from pydantic import BaseModel


class Dataset(BaseModel):
    name: str
    repo: str
    num_tasks: int


class Harness(BaseModel):
    name: str
    skills: list[str]


class Model(BaseModel):
    name: str
    repo: str | None = None
    is_oss: bool
    num_params: int
    precision: str


class Environment(BaseModel):
    name: str
    config: Optional[dict[str, Any]] = None


class Metrics(BaseModel):
    score: float
    time: Optional[int] = None
    costUSD: Optional[float] = None


class Result(BaseModel):
    dataset: Dataset
    harness: Harness
    model: Model
    environment: Environment
    metrics: Metrics