diff --git a/.gitignore b/.gitignore index 3bb5a62a82..ee29e42852 100755 --- a/.gitignore +++ b/.gitignore @@ -20,7 +20,7 @@ unsloth_compiled_cache/ outputs/ *.gguf *.safetensors -models/ +/models/ # IDE / Editors .vscode/ diff --git a/studio/backend/models/.gitkeep b/studio/backend/models/.gitkeep new file mode 100644 index 0000000000..e69de29bb2 diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py new file mode 100644 index 0000000000..07836f2168 --- /dev/null +++ b/studio/backend/models/__init__.py @@ -0,0 +1,37 @@ +""" +Pydantic models for API request/response schemas +""" +from .training import ( + TrainingStartRequest, + TrainingStartResponse, + TrainingStatusResponse, + TrainingMetricsResponse, + TrainingProgressResponse, +) +from .models import ( + ModelSearchRequest, + ModelSearchResponse, + ModelListResponse, + ModelConfigResponse, + LoRAScanResponse, + LoRAInfo, + ModelInfo, +) + +__all__ = [ + # Training schemas + "TrainingStartRequest", + "TrainingStartResponse", + "TrainingStatusResponse", + "TrainingMetricsResponse", + "TrainingProgressResponse", + # Model management schemas + "ModelSearchRequest", + "ModelSearchResponse", + "ModelListResponse", + "ModelConfigResponse", + "LoRAScanResponse", + "LoRAInfo", + "ModelInfo", +] + diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py new file mode 100644 index 0000000000..9561bf3288 --- /dev/null +++ b/studio/backend/models/models.py @@ -0,0 +1,56 @@ +""" +Pydantic schemas for Model Management API +""" +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""" + 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") + is_lora: bool = Field(False, description="Whether model is a LoRA adapter") + base_model: Optional[str] = Field(None, description="Base model if this is a LoRA adapter") + + +class LoRAInfo(BaseModel): + """LoRA adapter information""" + display_name: str = Field(..., description="Display name for the LoRA") + adapter_path: str = Field(..., description="Path to the LoRA adapter") + base_model: Optional[str] = Field(None, description="Base model identifier") + + +class LoRAScanResponse(BaseModel): + """Response schema for scanning trained LoRA adapters""" + loras: List[LoRAInfo] = Field(default_factory=list, description="List of found LoRA adapters") + outputs_dir: str = Field(..., description="Directory that was scanned") + diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py new file mode 100644 index 0000000000..9105eaa7d3 --- /dev/null +++ b/studio/backend/models/training.py @@ -0,0 +1,96 @@ +""" +Pydantic schemas for Training API +""" +from pydantic import BaseModel, Field +from typing import Optional, List + + +class TrainingStartRequest(BaseModel): + """Request schema for starting training""" + # Model parameters + model_name: str = Field(..., description="Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')") + training_type: str = Field(..., description="Training type: 'LoRA/QLoRA' or 'Full Finetuning'") + hf_token: Optional[str] = Field(None, description="HuggingFace token") + load_in_4bit: bool = Field(True, description="Load model in 4-bit quantization") + max_seq_length: int = Field(2048, description="Maximum sequence length") + + # Dataset parameters + hf_dataset: Optional[str] = Field(None, description="HuggingFace dataset identifier") + local_datasets: List[str] = Field(default_factory=list, description="List of local dataset paths") + format_type: str = Field(..., description="Dataset format type") + + # Training parameters + num_epochs: int = Field(1, description="Number of training epochs") + learning_rate: str = Field("2e-4", description="Learning rate") + batch_size: int = Field(1, description="Batch size") + gradient_accumulation_steps: int = Field(1, description="Gradient accumulation steps") + warmup_steps: Optional[int] = Field(None, description="Warmup steps") + warmup_ratio: Optional[float] = Field(None, description="Warmup ratio") + max_steps: Optional[int] = Field(None, description="Maximum training steps") + save_steps: int = Field(100, description="Steps between checkpoints") + weight_decay: float = Field(0.01, description="Weight decay") + random_seed: int = Field(42, description="Random seed") + packing: bool = Field(False, description="Enable sequence packing") + optim: str = Field("adamw_8bit", description="Optimizer") + lr_scheduler_type: str = Field("linear", description="Learning rate scheduler type") + + # LoRA parameters + use_lora: bool = Field(True, description="Use LoRA (derived from training_type)") + lora_r: int = Field(16, description="LoRA rank") + lora_alpha: int = Field(16, description="LoRA alpha") + lora_dropout: float = Field(0.0, description="LoRA dropout") + target_modules: List[str] = Field(default_factory=list, description="Target modules for LoRA") + gradient_checkpointing: str = Field("", description="Gradient checkpointing setting") + use_rslora: bool = Field(False, description="Use RSLoRA") + use_loftq: bool = Field(False, description="Use LoftQ") + train_on_completions: bool = Field(False, description="Train on completions only") + + # Vision-specific LoRA parameters + finetune_vision_layers: bool = Field(False, description="Finetune vision layers") + finetune_language_layers: bool = Field(False, description="Finetune language layers") + finetune_attention_modules: bool = Field(False, description="Finetune attention modules") + finetune_mlp_modules: bool = Field(False, description="Finetune MLP modules") + + # Logging parameters + enable_wandb: bool = Field(False, description="Enable Weights & Biases logging") + wandb_token: Optional[str] = Field(None, description="W&B token") + wandb_project: Optional[str] = Field(None, description="W&B project name") + enable_tensorboard: bool = Field(False, description="Enable TensorBoard logging") + 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 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 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") + learning_rate: float = Field(..., description="Current learning rate") + status_message: str = Field(..., description="Status message") + progress_percent: Optional[float] = Field(None, description="Progress percentage") + diff --git a/studio/backend/models/users.py b/studio/backend/models/users.py new file mode 100644 index 0000000000..ba6bb8f8e7 --- /dev/null +++ b/studio/backend/models/users.py @@ -0,0 +1,32 @@ +"""Pydantic models for user-related API endpoints. + +This module defines the data models used for user authentication and management +in the FastAPI application. +""" + +from pydantic import BaseModel + + +class User(BaseModel): + """Basic user model containing username.""" + + username: str + + +class UserInDB(BaseModel): + """User model with password for database storage.""" + + password: str + + +class Token(BaseModel): + """Authentication token model with access token and type.""" + + access_token: str + token_type: str + + +class TokenData(BaseModel): + """Token payload model containing username.""" + + username: str | None = None