diff --git a/studio/backend/main.py b/studio/backend/main.py index 8f7fecaaf3..f335839f8d 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -9,7 +9,7 @@ from pathlib import Path from datetime import datetime # Import routers -from routes import training_router, models_router, inference_router +from routes import training_router, models_router, inference_router, datasets_router # Create FastAPI app app = FastAPI( @@ -33,6 +33,7 @@ app.add_middleware( app.include_router(training_router, prefix="/api/train", tags=["training"]) app.include_router(models_router, prefix="/api/models", tags=["models"]) app.include_router(inference_router, prefix="/api/inference", tags=["inference"]) +app.include_router(datasets_router, prefix="/api/datasets", tags=["datasets"]) # ============ Health and System Endpoints ============ diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py new file mode 100644 index 0000000000..81c2451ebe --- /dev/null +++ b/studio/backend/models/datasets.py @@ -0,0 +1,23 @@ +""" +Dataset-related Pydantic models for API requests and responses. +""" +from pydantic import BaseModel +from typing import Optional, Dict, List + + +class CheckFormatRequest(BaseModel): + """Request for dataset format check""" + dataset_name: str # HuggingFace dataset name or local path + is_vlm: bool = False + hf_token: Optional[str] = None + split: Optional[str] = "train" + + +class CheckFormatResponse(BaseModel): + """Response for dataset format check""" + requires_manual_mapping: bool + detected_format: str + columns: List[str] + suggested_mapping: Optional[Dict[str, str]] = None + detected_image_column: Optional[str] = None + detected_text_column: Optional[str] = None diff --git a/studio/backend/routes/__init__.py b/studio/backend/routes/__init__.py index a281c709dd..a865794d54 100644 --- a/studio/backend/routes/__init__.py +++ b/studio/backend/routes/__init__.py @@ -5,5 +5,6 @@ API Routes from routes.training import router as training_router from routes.models import router as models_router from routes.inference import router as inference_router +from routes.datasets import router as datasets_router -__all__ = ["training_router", "models_router", "inference_router"] \ No newline at end of file +__all__ = ["training_router", "models_router", "inference_router", "datasets_router"] \ No newline at end of file diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py new file mode 100644 index 0000000000..8e21acaa9e --- /dev/null +++ b/studio/backend/routes/datasets.py @@ -0,0 +1,93 @@ +""" +Datasets API routes +""" +import sys +from pathlib import Path +from fastapi import APIRouter, HTTPException +import logging + +# Add backend directory to path +backend_path = Path(__file__).parent.parent.parent +if str(backend_path) not in sys.path: + sys.path.insert(0, str(backend_path)) + +# Import dataset utilities +from utils.datasets import check_dataset_format + +router = APIRouter() +logger = logging.getLogger(__name__) + +# Configure logger +if not logger.handlers: + handler = logging.StreamHandler() + handler.setLevel(logging.INFO) + formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') + handler.setFormatter(formatter) + logger.addHandler(handler) + logger.setLevel(logging.INFO) + + +from models.datasets import CheckFormatRequest, CheckFormatResponse + + +# --- Endpoints --- + +@router.post("/check-format", response_model=CheckFormatResponse) +async def check_format(request: CheckFormatRequest): + """ + Check if a dataset requires manual column mapping. + + This is a lightweight check that only runs format detection, + not full processing. Use before starting training to determine + if the user needs to manually map columns. + """ + try: + from datasets import load_dataset + + logger.info(f"Checking format for dataset: {request.dataset_name}") + + # Load dataset + dataset_path = Path(request.dataset_name) + + if dataset_path.exists(): + # Local dataset + if dataset_path.suffix in ['.json', '.jsonl']: + dataset = load_dataset('json', data_files=str(dataset_path), split=request.split) + elif dataset_path.suffix == '.csv': + dataset = load_dataset('csv', data_files=str(dataset_path), split=request.split) + elif dataset_path.suffix == '.parquet': + dataset = load_dataset('parquet', data_files=str(dataset_path), split=request.split) + else: + raise HTTPException( + status_code=400, + detail=f"Unsupported file format: {dataset_path.suffix}" + ) + else: + # HuggingFace dataset + load_kwargs = {"path": request.dataset_name, "split": request.split} + if request.hf_token: + load_kwargs["token"] = request.hf_token + dataset = load_dataset(**load_kwargs) + + # Run lightweight format check + result = check_dataset_format(dataset, is_vlm=request.is_vlm) + + logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}") + + return CheckFormatResponse( + requires_manual_mapping=result["requires_manual_mapping"], + detected_format=result["detected_format"], + columns=result["columns"], + suggested_mapping=result.get("suggested_mapping"), + detected_image_column=result.get("detected_image_column"), + detected_text_column=result.get("detected_text_column"), + ) + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error checking dataset format: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to check dataset format: {str(e)}" + ) diff --git a/studio/backend/utils/datasets/__init__.py b/studio/backend/utils/datasets/__init__.py index 9c26f431fa..b47db737c3 100644 --- a/studio/backend/utils/datasets/__init__.py +++ b/studio/backend/utils/datasets/__init__.py @@ -59,6 +59,7 @@ from .model_mappings import ( # Legacy imports from the original dataset_utils.py for backward compatibility # These functions have not yet been refactored into separate modules from .dataset_utils import ( + check_dataset_format, format_and_template_dataset, format_dataset, ) @@ -90,7 +91,8 @@ __all__ = [ "TEMPLATE_TO_MODEL_MAPPER", "MODEL_TO_TEMPLATE_MAPPER", "TEMPLATE_TO_RESPONSES_MAPPER", - # Legacy (backward compat) + # Main entry points + "check_dataset_format", "format_and_template_dataset", "format_dataset", ] diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 46f3dd58fc..a1137ae1f6 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -2,6 +2,7 @@ Dataset utilities for format detection, conversion, and template application. This module provides the main entry points for dataset processing: +- check_dataset_format: Lightweight check if manual mapping is needed (for frontend) - format_dataset: Detects and normalizes dataset formats - format_and_template_dataset: End-to-end processing with chat template application @@ -39,6 +40,79 @@ from .data_collators import DeepSeekOCRDataCollator, VLMDataCollator from .model_mappings import TEMPLATE_TO_MODEL_MAPPER, RESPONSE_MARKERS +def check_dataset_format(dataset, is_vlm: bool = False) -> dict: + """ + Lightweight format check without processing - for frontend validation. + + Use this to quickly determine if user needs to manually map columns + before calling the full format_and_template_dataset(). + + Args: + dataset: HuggingFace dataset + is_vlm: Whether this is a Vision-Language Model dataset + + Returns: + dict: { + "requires_manual_mapping": bool - True if user must map columns, + "detected_format": str - The detected format, + "columns": list - Available column names for mapping UI, + "suggested_mapping": dict or None - Auto-detected mapping if available, + "detected_image_column": str or None - For VLM only, + "detected_text_column": str or None - For VLM only, + } + """ + columns = list(dataset.column_names) if hasattr(dataset, 'column_names') else list(next(iter(dataset)).keys()) + + if is_vlm: + vlm_structure = detect_vlm_dataset_structure(dataset) + requires_mapping = vlm_structure["format"] == "unknown" + + return { + "requires_manual_mapping": requires_mapping, + "detected_format": vlm_structure["format"], + "columns": columns, + "suggested_mapping": None, + "detected_image_column": vlm_structure.get("image_column"), + "detected_text_column": vlm_structure.get("text_column"), + } + else: + # LLM flow + detected = detect_dataset_format(dataset) + + # If format is unknown, try heuristic detection + if detected["format"] == "unknown": + heuristic_mapping = detect_custom_format_heuristic(dataset) + if heuristic_mapping: + # Heuristic succeeded - no manual mapping needed + return { + "requires_manual_mapping": False, + "detected_format": "custom_heuristic", + "columns": columns, + "suggested_mapping": heuristic_mapping, + "detected_image_column": None, + "detected_text_column": None, + } + else: + # Both detection and heuristic failed + return { + "requires_manual_mapping": True, + "detected_format": "unknown", + "columns": columns, + "suggested_mapping": None, + "detected_image_column": None, + "detected_text_column": None, + } + + # Known format detected + return { + "requires_manual_mapping": False, + "detected_format": detected["format"], + "columns": columns, + "suggested_mapping": None, + "detected_image_column": None, + "detected_text_column": None, + } + def _apply_user_mapping(dataset, mapping: dict, batch_size: int = 1000): """ Apply user-provided column mapping to convert dataset to conversations format.