src/schemas/eval_queue.py

from typing import Annotated

from pydantic import BaseModel, ConfigDict, Field, computed_field


class EvalQueue(BaseModel):
    """Model of the eval queue json file."""

    model_config = ConfigDict(extra="allow", frozen=True)

    model_name: Annotated[str, Field(..., description="The model name. e.g. Qwen/Qwen2.5-3B")]
    base_model: Annotated[str, Field(..., description="The base model name. e.g. Qwen/Qwen2.5")]

    model_dtype: Annotated[str, Field(..., description="The model dtype. e.g. bfloat16")]
    model_sha: Annotated[
        str | None, Field(..., description="The model sha. e.g. 3aab1f1954e9cc14eb9509a215f9e5ca08227a9b")
    ]

    @computed_field
    @property
    def model_key(self) -> str:
        return self.model_name.lower().replace("/", "_")