From 4fc9bbf0f1b6a150f1c56737226a2e1a0f0f881a Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Mon, 2 Feb 2026 20:00:04 +0000
Subject: [PATCH 1/4] update pydantic models for Models and Training routes

---
 .gitignore                        |   1 -
 studio/backend/models/models.py   |  32 +--------
 studio/backend/models/training.py |  64 +++++++++--------
 tmp/models.py                     |  27 +++++++
 tmp/pr_description.md             |  40 +++++++++++
 tmp/training.py                   |  99 ++++++++++++++++++++++++++
 tmp/training_pydantic.md          | 113 ++++++++++++++++++++++++++++++
 7 files changed, 315 insertions(+), 61 deletions(-)
 create mode 100644 tmp/models.py
 create mode 100644 tmp/pr_description.md
 create mode 100644 tmp/training.py
 create mode 100644 tmp/training_pydantic.md

diff --git a/.gitignore b/.gitignore
index ee29e42852..451ba79df7 100755
--- a/.gitignore
+++ b/.gitignore
@@ -20,7 +20,6 @@ unsloth_compiled_cache/
 outputs/
 *.gguf
 *.safetensors
-/models/
 
 # IDE / Editors
 .vscode/
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..86aead6462 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,42 @@ 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")
 
diff --git a/tmp/models.py b/tmp/models.py
new file mode 100644
index 0000000000..fd92a169b2
--- /dev/null
+++ b/tmp/models.py
@@ -0,0 +1,27 @@
+"""
+Pydantic schemas for Model Management API
+"""
+from pydantic import BaseModel, Field
+from typing import Optional, List, Dict, Any
+
+
+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")
+    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/tmp/pr_description.md b/tmp/pr_description.md
new file mode 100644
index 0000000000..012a3f03a4
--- /dev/null
+++ b/tmp/pr_description.md
@@ -0,0 +1,40 @@
+# Refactor Pydantic API Models
+
+## Summary
+Cleans up and restructures the Pydantic models for the training and model management APIs to reduce redundancy, improve naming clarity, and align with the frontend architecture.
+
+---
+
+## Training Models (`studio/backend/models/training.py`)
+
+| Before | After | Change |
+|--------|-------|--------|
+| `TrainingStartResponse` | `TrainingJobResponse` | Rename - clearer that it represents a created job |
+| `TrainingStatusResponse` | `TrainingStatus` | Rename + add `phase` field with explicit pipeline stages |
+| `TrainingProgressResponse` | `TrainingProgress` | Rename + add `epoch`, `elapsed_seconds`, `eta_seconds` |
+| `TrainingMetricsResponse` | — | **Removed** - frontend stores history in IndexedDB |
+
+**Key improvements:**
+- `TrainingStatus` unifies status polling and streaming with explicit `phase` literals: `idle`, `loading_model`, `loading_dataset`, `configuring`, `training`, `completed`, `error`, `stopped`
+- Renamed `is_active` → `is_training_running` for clarity
+- `TrainingProgress` now includes timing info (`elapsed_seconds`, `eta_seconds`) for better UX
+
+---
+
+## Model Management (`studio/backend/models/models.py`)
+
+| Before | After | Change |
+|--------|-------|--------|
+| `ModelSearchRequest` | — | **Removed** - search handled via query params |
+| `ModelSearchResponse` | — | **Removed** |
+| `ModelListResponse` | — | **Removed** |
+| `ModelInfo` | — | **Removed** - redundant with ModelDetails |
+| `ModelConfigResponse` | `ModelDetails` | Rename |
+
+**Remaining models:** `ModelDetails`, `LoRAInfo`, `LoRAScanResponse`
+
+---
+
+## Breaking Changes
+- All renamed/removed models will require updates in any code referencing them
+- Frontend TypeScript types should be regenerated
diff --git a/tmp/training.py b/tmp/training.py
new file mode 100644
index 0000000000..fed2076802
--- /dev/null
+++ b/tmp/training.py
@@ -0,0 +1,99 @@
+"""
+Pydantic schemas for Training API
+"""
+from pydantic import BaseModel, Field
+from typing import Optional, List, Literal
+
+
+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 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 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 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")
+    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")
diff --git a/tmp/training_pydantic.md b/tmp/training_pydantic.md
new file mode 100644
index 0000000000..63e5289ab0
--- /dev/null
+++ b/tmp/training_pydantic.md
@@ -0,0 +1,113 @@
+# Training Pydantic Models Recommendation
+
+## Current State
+
+| Model | Purpose |
+|-------|---------|
+| `TrainingStartRequest` | Request payload to start training |
+| `TrainingStartResponse` | Immediate response after `/start` |
+| `TrainingStatusResponse` | Response for `/status` polling |
+| `TrainingMetricsResponse` | Historical metrics |
+| `TrainingProgressResponse` | Per-step progress data |
+
+---
+
+## Proposed Models
+
+### 1. `TrainingStartRequest` — Keep as-is
+Well-structured, no changes needed.
+
+---
+
+### 2. `TrainingJobResponse` (rename from `TrainingStartResponse`)
+Returned **once** when `/train/start` is called.
+
+```python
+class TrainingJobResponse(BaseModel):
+    """Immediate response when training is initiated"""
+    job_id: str
+    status: Literal["queued", "error"]
+    message: str
+    error: Optional[str] = None
+```
+
+---
+
+### 3. `TrainingStatus` (unifies `TrainingStatusResponse` + `TrainingPhaseUpdate`)
+Single model for **both streaming and polling**.
+
+```python
+class TrainingStatus(BaseModel):
+    """Current training job status - works for streaming or polling"""
+    job_id: str
+    phase: Literal[
+        "idle",
+        "loading_model",
+        "loading_dataset",
+        "configuring",
+        "training",
+        "completed",
+        "error",
+        "stopped"
+    ]
+    is_training_running: bool    # True if training loop is actively running
+    message: str
+    error: Optional[str] = None
+    details: Optional[dict] = None  # Phase-specific info, e.g. {"model_size": "8B"}
+```
+
+**Usage:**
+- **Streaming**: Push when phase changes
+- **Polling**: Return from `GET /train/status/{job_id}`
+
+---
+
+### 4. `TrainingProgress` (rename from `TrainingProgressResponse`)
+Per-step metrics during active training.
+
+```python
+class TrainingProgress(BaseModel):
+    """Training progress metrics - for streaming or polling"""
+    step: int
+    total_steps: int
+    loss: float
+    learning_rate: float
+    progress_percent: float
+    epoch: Optional[int] = None
+    elapsed_seconds: Optional[float] = None
+    eta_seconds: Optional[float] = None
+```
+
+---
+
+### 5. `TrainingMetricsResponse` — Remove
+Frontend stores history in IndexedDB.
+
+---
+
+## Summary
+
+| Current | Proposed | Action |
+|---------|----------|--------|
+| `TrainingStartRequest` | `TrainingStartRequest` | Keep |
+| `TrainingStartResponse` | `TrainingJobResponse` | Rename |
+| `TrainingStatusResponse` | `TrainingStatus` | Rename + enhance |
+| `TrainingMetricsResponse` | — | **Remove** |
+| `TrainingProgressResponse` | `TrainingProgress` | Rename + enhance |
+
+---
+
+## API Flow
+
+```
+POST /train/start
+  └─► TrainingJobResponse { job_id, status: "queued" }
+
+StreamingResponse / Polling:
+  └─► TrainingStatus { phase: "loading_model", is_training_running: false }
+  └─► TrainingStatus { phase: "loading_dataset", is_training_running: false }
+  └─► TrainingStatus { phase: "training", is_training_running: true }
+  └─► TrainingProgress { step: 1, loss: 2.5, ... }
+  └─► TrainingProgress { step: 2, loss: 2.3, ... }
+  └─► TrainingStatus { phase: "completed", is_training_running: false }
+```

From ecebd30eca68d730243da26db33d847ea8d119af Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Mon, 2 Feb 2026 20:00:23 +0000
Subject: [PATCH 2/4] update pydantic models for Models and Training routes

---
 .gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitignore b/.gitignore
index 451ba79df7..8c79c64263 100755
--- a/.gitignore
+++ b/.gitignore
@@ -34,3 +34,4 @@ Thumbs.db
 
 # Other
 resources/
+tmp/

From 6c017686d9b173f8d30b3ab16f11f974745da1e8 Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Mon, 2 Feb 2026 20:02:18 +0000
Subject: [PATCH 3/4] delete tmp directory

---
 tmp/models.py            |  27 ----------
 tmp/pr_description.md    |  40 --------------
 tmp/training.py          |  99 ----------------------------------
 tmp/training_pydantic.md | 113 ---------------------------------------
 4 files changed, 279 deletions(-)
 delete mode 100644 tmp/models.py
 delete mode 100644 tmp/pr_description.md
 delete mode 100644 tmp/training.py
 delete mode 100644 tmp/training_pydantic.md

diff --git a/tmp/models.py b/tmp/models.py
deleted file mode 100644
index fd92a169b2..0000000000
--- a/tmp/models.py
+++ /dev/null
@@ -1,27 +0,0 @@
-"""
-Pydantic schemas for Model Management API
-"""
-from pydantic import BaseModel, Field
-from typing import Optional, List, Dict, Any
-
-
-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")
-    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/tmp/pr_description.md b/tmp/pr_description.md
deleted file mode 100644
index 012a3f03a4..0000000000
--- a/tmp/pr_description.md
+++ /dev/null
@@ -1,40 +0,0 @@
-# Refactor Pydantic API Models
-
-## Summary
-Cleans up and restructures the Pydantic models for the training and model management APIs to reduce redundancy, improve naming clarity, and align with the frontend architecture.
-
----
-
-## Training Models (`studio/backend/models/training.py`)
-
-| Before | After | Change |
-|--------|-------|--------|
-| `TrainingStartResponse` | `TrainingJobResponse` | Rename - clearer that it represents a created job |
-| `TrainingStatusResponse` | `TrainingStatus` | Rename + add `phase` field with explicit pipeline stages |
-| `TrainingProgressResponse` | `TrainingProgress` | Rename + add `epoch`, `elapsed_seconds`, `eta_seconds` |
-| `TrainingMetricsResponse` | — | **Removed** - frontend stores history in IndexedDB |
-
-**Key improvements:**
-- `TrainingStatus` unifies status polling and streaming with explicit `phase` literals: `idle`, `loading_model`, `loading_dataset`, `configuring`, `training`, `completed`, `error`, `stopped`
-- Renamed `is_active` → `is_training_running` for clarity
-- `TrainingProgress` now includes timing info (`elapsed_seconds`, `eta_seconds`) for better UX
-
----
-
-## Model Management (`studio/backend/models/models.py`)
-
-| Before | After | Change |
-|--------|-------|--------|
-| `ModelSearchRequest` | — | **Removed** - search handled via query params |
-| `ModelSearchResponse` | — | **Removed** |
-| `ModelListResponse` | — | **Removed** |
-| `ModelInfo` | — | **Removed** - redundant with ModelDetails |
-| `ModelConfigResponse` | `ModelDetails` | Rename |
-
-**Remaining models:** `ModelDetails`, `LoRAInfo`, `LoRAScanResponse`
-
----
-
-## Breaking Changes
-- All renamed/removed models will require updates in any code referencing them
-- Frontend TypeScript types should be regenerated
diff --git a/tmp/training.py b/tmp/training.py
deleted file mode 100644
index fed2076802..0000000000
--- a/tmp/training.py
+++ /dev/null
@@ -1,99 +0,0 @@
-"""
-Pydantic schemas for Training API
-"""
-from pydantic import BaseModel, Field
-from typing import Optional, List, Literal
-
-
-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 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 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 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")
-    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")
diff --git a/tmp/training_pydantic.md b/tmp/training_pydantic.md
deleted file mode 100644
index 63e5289ab0..0000000000
--- a/tmp/training_pydantic.md
+++ /dev/null
@@ -1,113 +0,0 @@
-# Training Pydantic Models Recommendation
-
-## Current State
-
-| Model | Purpose |
-|-------|---------|
-| `TrainingStartRequest` | Request payload to start training |
-| `TrainingStartResponse` | Immediate response after `/start` |
-| `TrainingStatusResponse` | Response for `/status` polling |
-| `TrainingMetricsResponse` | Historical metrics |
-| `TrainingProgressResponse` | Per-step progress data |
-
----
-
-## Proposed Models
-
-### 1. `TrainingStartRequest` — Keep as-is
-Well-structured, no changes needed.
-
----
-
-### 2. `TrainingJobResponse` (rename from `TrainingStartResponse`)
-Returned **once** when `/train/start` is called.
-
-```python
-class TrainingJobResponse(BaseModel):
-    """Immediate response when training is initiated"""
-    job_id: str
-    status: Literal["queued", "error"]
-    message: str
-    error: Optional[str] = None
-```
-
----
-
-### 3. `TrainingStatus` (unifies `TrainingStatusResponse` + `TrainingPhaseUpdate`)
-Single model for **both streaming and polling**.
-
-```python
-class TrainingStatus(BaseModel):
-    """Current training job status - works for streaming or polling"""
-    job_id: str
-    phase: Literal[
-        "idle",
-        "loading_model",
-        "loading_dataset",
-        "configuring",
-        "training",
-        "completed",
-        "error",
-        "stopped"
-    ]
-    is_training_running: bool    # True if training loop is actively running
-    message: str
-    error: Optional[str] = None
-    details: Optional[dict] = None  # Phase-specific info, e.g. {"model_size": "8B"}
-```
-
-**Usage:**
-- **Streaming**: Push when phase changes
-- **Polling**: Return from `GET /train/status/{job_id}`
-
----
-
-### 4. `TrainingProgress` (rename from `TrainingProgressResponse`)
-Per-step metrics during active training.
-
-```python
-class TrainingProgress(BaseModel):
-    """Training progress metrics - for streaming or polling"""
-    step: int
-    total_steps: int
-    loss: float
-    learning_rate: float
-    progress_percent: float
-    epoch: Optional[int] = None
-    elapsed_seconds: Optional[float] = None
-    eta_seconds: Optional[float] = None
-```
-
----
-
-### 5. `TrainingMetricsResponse` — Remove
-Frontend stores history in IndexedDB.
-
----
-
-## Summary
-
-| Current | Proposed | Action |
-|---------|----------|--------|
-| `TrainingStartRequest` | `TrainingStartRequest` | Keep |
-| `TrainingStartResponse` | `TrainingJobResponse` | Rename |
-| `TrainingStatusResponse` | `TrainingStatus` | Rename + enhance |
-| `TrainingMetricsResponse` | — | **Remove** |
-| `TrainingProgressResponse` | `TrainingProgress` | Rename + enhance |
-
----
-
-## API Flow
-
-```
-POST /train/start
-  └─► TrainingJobResponse { job_id, status: "queued" }
-
-StreamingResponse / Polling:
-  └─► TrainingStatus { phase: "loading_model", is_training_running: false }
-  └─► TrainingStatus { phase: "loading_dataset", is_training_running: false }
-  └─► TrainingStatus { phase: "training", is_training_running: true }
-  └─► TrainingProgress { step: 1, loss: 2.5, ... }
-  └─► TrainingProgress { step: 2, loss: 2.3, ... }
-  └─► TrainingStatus { phase: "completed", is_training_running: false }
-```

From 7bb0aeb756bfb7485357ade17e66cdf980e476da Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Tue, 3 Feb 2026 10:35:58 +0000
Subject: [PATCH 4/4] add grad_norm and num_tokens to TrainingProgress response
 object

---
 studio/backend/models/training.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 86aead6462..2afee72e89 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -97,4 +97,6 @@ class TrainingProgress(BaseModel):
     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")