Source code for chariot.bulk_inference.models

from pydantic import BaseModel, ConfigDict
from enum import StrEnum




class ProgressState(StrEnum):
    IN_PROGRESS = "in_progress"
    FAILED = "failed"
    SUCCEEDED = "succeeded"
    INITIALIZED = "initialized"
    RUNNING = "running"
    RUNNING_FAILED = "running_failed"
    ABORTED = "aborted"





class NewBulkInferenceJobRequest(BaseModel):
    model_project_id: str
    model_id: str

    dataset_project_id: str
    dataset_id: str
    dataset_snapshot_id: str
    dataset_snapshot_split: str

    inference_method: str = ""
    evaluate_metrics: bool = False
    batch_size: int = 1

    model_config = ConfigDict(protected_namespaces=())





class BulkInferenceJob(BaseModel):
    execution_id: str
    execution_status: str
    num_computed: int | None = None
    total_expected: int | None = None

    @property
    def job_id(self):
        return self.execution_id

    @property
    def job_status(self) -> ProgressState:
        return ProgressState(self.execution_status)