Add datasets check-format endpoint

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 ============

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

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"]
+__all__ = ["training_router", "models_router", "inference_router", "datasets_router"]

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)}"
+        )

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",
 ]

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.