svc/schemas.py

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


from pydantic import BaseModel


class OAuth2PasswordRequestForm(BaseModel):
    username: str
    hf_token: str

class User(BaseModel):
    username: str
    hf_token: str


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

class DeepEvalSuiteRequest(BaseModel):
    model_name: str
    # dtype: str
    # weight: str
    # base_model: str
    tasks: Optional[List[str]] = None


class TaskResponse(BaseModel):
    results: Json # dict[Any,Any]

class LoadModelRequest(BaseModel):
    model_name: str