src/common/schema.py

from typing import Any

import pandera.pandas as pa


class DatasetSchema(pa.DataFrameModel):
    id_: pa.typing.Series[int] = pa.Field(alias="id")
    task_text: pa.typing.Series[str]
    answer_text: pa.typing.Series[str]
    correct_answer: pa.typing.Series[Any]
    date: pa.typing.Series[str]
    olymp_name: pa.typing.Series[str]
    grade: pa.typing.Series[str]
    description: pa.typing.Series[str]
    source: pa.typing.Series[str]
    answer_type: pa.typing.Series[str]
    check_type: pa.typing.Series[str]
    check_function: pa.typing.Series[str] = pa.Field(nullable=True)
    task_type: pa.typing.Series[str]
    task_note: pa.typing.Series[str]


class LeaderBoardSchema(pa.DataFrameModel):
    model_name: pa.typing.Series[str]
    model_size: pa.typing.Series[float] = pa.Field(nullable=True)
    model_url: pa.typing.Series[str] = pa.Field(nullable=True)
    pass1: pa.typing.Series[float]
    weighted_pass1: pa.typing.Series[float]
    arith_pass1: pa.typing.Series[float]
    geometry_pass1: pa.typing.Series[float]
    logic_pass1: pa.typing.Series[float]
    config: pa.typing.Series[str] = pa.Field(nullable=True, default={})