diff --git a/.gitignore b/.gitignore index ee29e42852..8c79c64263 100755 --- a/.gitignore +++ b/.gitignore @@ -20,7 +20,6 @@ unsloth_compiled_cache/ outputs/ *.gguf *.safetensors -/models/ # IDE / Editors .vscode/ @@ -35,3 +34,4 @@ Thumbs.db # Other resources/ +tmp/ diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 9561bf3288..0960d35599 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -5,36 +5,8 @@ from pydantic import BaseModel, Field from typing import Optional, List, Dict, Any -class ModelSearchRequest(BaseModel): - """Request schema for searching HuggingFace models""" - query: str = Field(..., description="Search query") - hf_token: Optional[str] = Field(None, description="HuggingFace token for authenticated searches") - - -class ModelInfo(BaseModel): - """Model information""" - id: str = Field(..., description="Model identifier") - name: Optional[str] = Field(None, description="Display name") - description: Optional[str] = Field(None, description="Model description") - size: Optional[str] = Field(None, description="Model size") - is_vision: bool = Field(False, description="Whether model is a vision model") - is_lora: bool = Field(False, description="Whether model is a LoRA adapter") - - -class ModelSearchResponse(BaseModel): - """Response schema for model search""" - models: List[ModelInfo] = Field(default_factory=list, description="List of matching models") - total: int = Field(0, description="Total number of results") - - -class ModelListResponse(BaseModel): - """Response schema for listing available models""" - models: List[ModelInfo] = Field(default_factory=list, description="List of available models") - default_models: List[str] = Field(default_factory=list, description="List of default model IDs") - - -class ModelConfigResponse(BaseModel): - """Response schema for model configuration""" +class ModelDetails(BaseModel): + """Detailed model configuration and metadata""" model_name: str = Field(..., description="Model identifier") config: Dict[str, Any] = Field(..., description="Model configuration dictionary") is_vision: bool = Field(False, description="Whether model is a vision model") diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 9105eaa7d3..2afee72e89 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -2,7 +2,7 @@ Pydantic schemas for Training API """ from pydantic import BaseModel, Field -from typing import Optional, List +from typing import Optional, List, Literal class TrainingStartRequest(BaseModel): @@ -59,38 +59,44 @@ class TrainingStartRequest(BaseModel): tensorboard_dir: Optional[str] = Field(None, description="TensorBoard directory") -class TrainingStartResponse(BaseModel): - """Response schema for training start""" - status: str = Field(..., description="Status: 'started' or 'error'") - job_id: Optional[str] = Field(None, description="Training job ID") - message: str = Field(..., description="Status message") - error: Optional[str] = Field(None, description="Error message if status is 'error'") +class TrainingJobResponse(BaseModel): + """Immediate response when training is initiated""" + job_id: str = Field(..., description="Unique training job identifier") + status: Literal["queued", "error"] = Field(..., description="Initial job status") + message: str = Field(..., description="Human-readable status message") + error: Optional[str] = Field(None, description="Error details if status is 'error'") -class TrainingStatusResponse(BaseModel): - """Response schema for training status""" - status: str = Field(..., description="Status: 'idle', 'preparing', 'training', 'stopping', 'error'") - is_active: bool = Field(..., description="Whether training is currently active (actual training running)") - message: str = Field(..., description="Status message") - current_step: Optional[int] = Field(None, description="Current training step") - total_steps: Optional[int] = Field(None, description="Total training steps") +class TrainingStatus(BaseModel): + """Current training job status - works for streaming or polling""" + job_id: str = Field(..., description="Training job identifier") + phase: Literal[ + "idle", + "loading_model", + "loading_dataset", + "configuring", + "training", + "completed", + "error", + "stopped" + ] = Field(..., description="Current phase of training pipeline") + is_training_running: bool = Field(..., description="True if training loop is actively running") + message: str = Field(..., description="Human-readable status message") + error: Optional[str] = Field(None, description="Error details if phase is 'error'") + details: Optional[dict] = Field(None, description="Phase-specific info, e.g. {'model_size': '8B'}") -class TrainingMetricsResponse(BaseModel): - """Response schema for training metrics""" - loss_history: List[float] = Field(default_factory=list, description="Loss values") - lr_history: List[float] = Field(default_factory=list, description="Learning rate values") - step_history: List[int] = Field(default_factory=list, description="Step numbers") - current_loss: Optional[float] = Field(None, description="Current loss value") - current_lr: Optional[float] = Field(None, description="Current learning rate") - current_step: Optional[int] = Field(None, description="Current step") - - -class TrainingProgressResponse(BaseModel): - """Response schema for training progress updates""" - step: int = Field(..., description="Current step") - loss: float = Field(..., description="Current loss") +class TrainingProgress(BaseModel): + """Training progress metrics - for streaming or polling""" + job_id: str = Field(..., description="Training job identifier") + step: int = Field(..., description="Current training step") + total_steps: int = Field(..., description="Total training steps") + loss: float = Field(..., description="Current loss value") learning_rate: float = Field(..., description="Current learning rate") - status_message: str = Field(..., description="Status message") - progress_percent: Optional[float] = Field(None, description="Progress percentage") + progress_percent: float = Field(..., description="Progress percentage (0.0 to 100.0)") + epoch: Optional[int] = Field(None, description="Current epoch") + elapsed_seconds: Optional[float] = Field(None, description="Time elapsed since training started") + eta_seconds: Optional[float] = Field(None, description="Estimated time remaining") + grad_norm: Optional[float] = Field(None, description="L2 norm of gradients, computed before gradient clipping") + num_tokens: Optional[int] = Field(None, description="Total number of tokens processed so far")