from pydantic import BaseModel
from typing import List, Optional, Union, Any

class LMHarnessTaskRequest(BaseModel):
    model: str
    model_args: Optional[Union[str, dict]] = None
    tasks: Optional[List[Union[str, dict, object]]] = None
    num_fewshot: Optional[int] = None
    batch_size: Optional[int] = None
    device: Optional[str] = None
    limit: Optional[Union[int, float]] = None
    write_out: bool = False # Whether to write out an example document and model input, for checking task integrity


class LMHarnessTaskResponse(BaseModel):
    results: dict[Any]