Fine-tuning/backend/app/schemas/dataset.py

from enum import Enum

from pydantic import BaseModel, Field, Field


class DatasetFormat(str, Enum):
    JSONL = "jsonl"
    CSV = "csv"
    PARQUET = "parquet"
    JSON = "json"


class DatasetUploadResponse(BaseModel):
    id: str
    name: str
    format: DatasetFormat
    record_count: int
    file_path: str
    created_at: str


class DatasetDownloadRequest(BaseModel):
    dataset_id: str = Field(..., description="HuggingFace or ModelScope dataset ID, e.g. 'glue', 'MRPC'")
    use_modelscope: bool = Field(default=False, description="Use ModelScope instead of HuggingFace")


class DatasetDownloadResponse(BaseModel):
    dataset_id: str
    status: str  # "downloading" | "completed" | "failed"
    task_id: str
    path: str | None = None
    error: str | None = None


class DatasetPreviewRow(BaseModel):
    row_index: int
    data: dict


class DatasetPreviewResponse(BaseModel):
    total_records: int
    preview_rows: list[DatasetPreviewRow]
    columns: list[str]


class DatasetValidationResult(BaseModel):
    is_valid: bool
    errors: list[str] = []
    warnings: list[str] = []