merge nightly
This commit is contained in:
Shine1i
commit
ba09bbfaba
64 changed files with 1391 additions and 130 deletions
145
README.md
145
README.md
|
|
@ -1 +1,144 @@
|
|||
# new-ui-prototype
|
||||
<p align="center">
|
||||
<img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20logo%20white%20text.png" alt="Unsloth Studio" width="400"/>
|
||||
</p>
|
||||
|
||||
<h3 align="center">🦥 Unsloth Studio</h3>
|
||||
|
||||
<p align="center">
|
||||
A modern, full-stack web interface for fine-tuning, managing, and chatting with large language models — locally or in the cloud.
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="#features">Features</a> •
|
||||
<a href="#quick-start">Quick Start</a> •
|
||||
<a href="#api-reference">API Reference</a> •
|
||||
<a href="#project-structure">Project Structure</a>
|
||||
</p>
|
||||
|
||||
---
|
||||
|
||||
## Features
|
||||
|
||||
| Area | Capabilities |
|
||||
|---|---|
|
||||
| **Training** | Configure and launch LoRA / QLoRA fine-tuning jobs with real-time SSE progress streaming, live loss charts, and one-click stop / resume |
|
||||
| **Model Management** | Browse, load, and manage Hugging Face hub models and local checkpoints |
|
||||
| **Inference** | Interactive chat playground for testing fine-tuned models |
|
||||
| **Dataset Tools** | Upload, preview, and prepare datasets (JSON, CSV, Parquet, PDF, DOCX) |
|
||||
| **Export** | Export & push trained adapters to the Hugging Face Hub |
|
||||
| **Auth** | Token-based authentication with JWT access / refresh flow and first-time setup token |
|
||||
|
||||
## Quick Start
|
||||
|
||||
### One-command setup
|
||||
|
||||
```bash
|
||||
bash setup.sh
|
||||
```
|
||||
|
||||
This script will:
|
||||
1. Install **Node.js ≥ 20** via nvm (if needed)
|
||||
2. Build the frontend to `studio/frontend/dist`
|
||||
3. Create a Python virtual environment and install all dependencies (including `unsloth`)
|
||||
4. Register a convenient `unsloth-ui` shell alias
|
||||
|
||||
### Launch the studio
|
||||
|
||||
```bash
|
||||
# After setup, open a new terminal (or source ~/.bashrc), then:
|
||||
unsloth-ui -H 0.0.0.0 -p 8000
|
||||
```
|
||||
|
||||
On **first launch**, a one-time setup token is printed to the console. Use it in the browser to create your admin account.
|
||||
|
||||
## API Reference
|
||||
|
||||
All endpoints require a valid JWT `Authorization: Bearer <token>` header (except `/api/auth/*` and `/api/health`).
|
||||
|
||||
| Method | Endpoint | Description |
|
||||
|---|---|---|
|
||||
| `GET` | `/api/health` | Health check |
|
||||
| `GET` | `/api/system` | System info (GPU, CPU, memory) |
|
||||
| `POST` | `/api/auth/signup` | Create account (requires setup token on first run) |
|
||||
| `POST` | `/api/auth/login` | Login and receive JWT tokens |
|
||||
| `POST` | `/api/auth/refresh` | Refresh an expired access token |
|
||||
| `GET` | `/api/auth/status` | Check if auth is initialized |
|
||||
| `POST` | `/api/train/start` | Start a training job |
|
||||
| `POST` | `/api/train/stop` | Stop a running training job |
|
||||
| `POST` | `/api/train/reset` | Reset training state |
|
||||
| `GET` | `/api/train/status` | Get current training status |
|
||||
| `GET` | `/api/train/metrics` | Get training metrics (loss, LR, steps) |
|
||||
| `GET` | `/api/train/stream` | SSE stream of real-time training progress |
|
||||
| `GET` | `/api/models/` | List available models |
|
||||
| `POST` | `/api/inference/chat` | Send a chat message for inference |
|
||||
| `GET` | `/api/datasets/` | List / manage datasets |
|
||||
|
||||
> Full interactive docs are available at `/docs` (Swagger UI) and `/redoc` when the server is running.
|
||||
|
||||
## CLI Commands
|
||||
|
||||
The Unsloth CLI (`cli.py`) provides the following commands:
|
||||
|
||||
```
|
||||
Usage: cli.py [COMMAND]
|
||||
|
||||
Commands:
|
||||
train Fine-tune a model
|
||||
inference Run inference on a trained model
|
||||
export Export a trained adapter
|
||||
list-checkpoints List saved checkpoints
|
||||
ui Launch the Unsloth Studio web UI
|
||||
studio Launch the studio (alias)
|
||||
```
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
new-ui-prototype/
|
||||
├── cli.py # CLI entry point
|
||||
├── cli/ # Typer CLI commands
|
||||
│ └── commands/
|
||||
│ ├── train.py
|
||||
│ ├── inference.py
|
||||
│ ├── export.py
|
||||
│ ├── ui.py
|
||||
│ └── studio.py
|
||||
├── setup.sh # One-command bootstrap script
|
||||
└── studio/
|
||||
├── backend/
|
||||
│ ├── main.py # FastAPI app & middleware
|
||||
│ ├── run.py # Server launcher (uvicorn)
|
||||
│ ├── auth/ # Auth storage & JWT logic
|
||||
│ ├── routes/ # API route handlers
|
||||
│ │ ├── training.py
|
||||
│ │ ├── models.py
|
||||
│ │ ├── inference.py
|
||||
│ │ ├── datasets.py
|
||||
│ │ └── auth.py
|
||||
│ ├── models/ # Pydantic request/response schemas
|
||||
│ ├── core/ # Training engine & config
|
||||
│ ├── utils/ # Hardware detection, helpers
|
||||
│ └── requirements.txt
|
||||
├── frontend/
|
||||
│ ├── src/
|
||||
│ │ ├── features/ # Feature modules
|
||||
│ │ │ ├── auth/ # Login / signup flow
|
||||
│ │ │ ├── training/ # Training config & monitoring
|
||||
│ │ │ ├── studio/ # Main studio workspace
|
||||
│ │ │ ├── chat/ # Inference chat UI
|
||||
│ │ │ ├── export/ # Model export flow
|
||||
│ │ │ └── onboarding/# Onboarding wizard
|
||||
│ │ ├── components/ # Shared UI components (shadcn)
|
||||
│ │ ├── hooks/ # Custom React hooks
|
||||
│ │ ├── stores/ # Zustand state stores
|
||||
│ │ └── types/ # TypeScript type definitions
|
||||
│ ├── package.json
|
||||
│ └── vite.config.ts
|
||||
└── tests/ # Backend test suite
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
This project is licensed under the [GNU Affero General Public License v3.0 (AGPL-3.0)](https://www.gnu.org/licenses/agpl-3.0.html).
|
||||
|
||||
Copyright © 2026 Unsloth AI.
|
||||
|
|
|
|||
|
|
@ -45,3 +45,9 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.7
|
||||
top_p: 0.95
|
||||
top_k: -1
|
||||
min_p: 0.01
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/ERNIE-4.5-VL-28B-A3B-PT
|
||||
# Based on ERNIE_4_5_VL_28B_A3B_PT_Vision.ipynb
|
||||
# Also applies to: unsloth/ERNIE-4.5-VL-28B-A3B-PT
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.5
|
||||
min_p: 0.1
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/codegemma-7b-bnb-4bit
|
||||
# Based on CodeGemma_(7B)-Conversational.ipynb
|
||||
# Also applies to: unsloth/codegemma-7b, google/codegemma-7b
|
||||
# added inference parameters from Ollama
|
||||
|
||||
training:
|
||||
max_seq_length: 4096
|
||||
|
|
@ -44,3 +45,7 @@ logging:
|
|||
enable_tensorboard: false
|
||||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0
|
||||
top_p: 0.9
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/functiongemma-270m-it
|
||||
# Based on FunctionGemma_(270M).ipynb
|
||||
# Also applies to: unsloth/functiongemma-270m-it-unsloth-bnb-4bit, google/functiongemma-270m-it, unsloth/functiongemma-270m-it-unsloth-bnb-4bit
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 4096
|
||||
|
|
@ -45,3 +46,9 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.0
|
||||
top_k: 64
|
||||
top_p: 0.95
|
||||
min_p: 0.0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/gemma-3-270m-it
|
||||
# Based on Gemma3_(270M).ipynb
|
||||
# Also applies to: unsloth/gemma-3-270m-it-unsloth-bnb-4bit, google/gemma-3-270m-it, unsloth/gemma-3-270m-it-bnb-4bit
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,9 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.0
|
||||
top_k: 64
|
||||
top_p: 0.95
|
||||
min_p: 0.0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/gemma-3-27b-it
|
||||
# Based on Gemma3_(27B)_A100-Conversational.ipynb
|
||||
# Also applies to: unsloth/gemma-3-27b-it-unsloth-bnb-4bit, google/gemma-3-27b-it, unsloth/gemma-3-27b-it-bnb-4bit
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -39,3 +40,9 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.0
|
||||
top_k: 64
|
||||
top_p: 0.95
|
||||
min_p: 0.0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/gemma-3-4b-it
|
||||
# Based on Gemma3_(4B).ipynb
|
||||
# Also applies to: unsloth/gemma-3-4b-it-unsloth-bnb-4bit, google/gemma-3-4b-it, unsloth/gemma-3-4b-it-bnb-4bit
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -39,3 +40,9 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.0
|
||||
top_k: 64
|
||||
top_p: 0.95
|
||||
min_p: 0.0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/gemma-3-4b-pt
|
||||
# Based on Gemma3_(4B)-Vision.ipynb
|
||||
# Also applies to: unsloth/gemma-3-4b-pt-unsloth-bnb-4bit, google/gemma-3-4b-pt, unsloth/gemma-3-4b-pt-bnb-4bit
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -39,3 +40,9 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.0
|
||||
top_k: 64
|
||||
top_p: 0.95
|
||||
min_p: 0.0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/gemma-3n-E4B-it
|
||||
# Based on Gemma3N_(4B)-Conversational.ipynb
|
||||
# Also applies to: unsloth/gemma-3n-E4B-it-unsloth-bnb-4bit, google/gemma-3n-E4B-it, unsloth/gemma-3n-E4B-it-unsloth-bnb-4bit
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 1024
|
||||
|
|
@ -39,3 +40,9 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.0
|
||||
top_k: 64
|
||||
top_p: 0.95
|
||||
min_p: 0.0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/gemma-3n-E4B
|
||||
# Based on Gemma3N_(4B)-Vision.ipynb
|
||||
# Also applies to: unsloth/gemma-3n-E4B-unsloth-bnb-4bit, google/gemma-3n-E4B
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -39,3 +40,9 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.0
|
||||
top_k: 64
|
||||
top_p: 0.95
|
||||
min_p: 0.0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/gpt-oss-120b
|
||||
# Based on gpt-oss-(120B)_A100-Fine-tuning.ipynb
|
||||
# Also applies to: openai/gpt-oss-120b, unsloth/gpt-oss-120b-unsloth-bnb-4bit
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 4096
|
||||
|
|
@ -45,3 +46,8 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.0
|
||||
top_p: 1.0
|
||||
top_k: 0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/gpt-oss-20b
|
||||
# Based on gpt-oss-(20B)-Fine-tuning.ipynb
|
||||
# Also applies to: openai/gpt-oss-20b, unsloth/gpt-oss-20b-unsloth-bnb-4bit, unsloth/gpt-oss-20b-BF16
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 1024
|
||||
|
|
@ -45,3 +46,8 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.0
|
||||
top_p: 1.0
|
||||
top_k: 0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/granite-4.0-350m
|
||||
# Based on Granite4.0_350M.ipynb
|
||||
# Also applies to: ibm-granite/granite-4.0-350m, unsloth/granite-4.0-350m-bnb-4bit
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -47,3 +48,8 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: 0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/granite-4.0-h-micro
|
||||
# Based on Granite4.0.ipynb
|
||||
# Also applies to: ibm-granite/granite-4.0-h-micro, unsloth/granite-4.0-h-micro-bnb-4bit, unsloth/granite-4.0-h-micro-unsloth-bnb-4bit
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -47,3 +48,8 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: 0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Llama-3.2-11B-Vision-Instruct
|
||||
# Based on Llama3.2_(11B)-Vision.ipynb
|
||||
# Also applies to: unsloth/Llama-3.2-11B-Vision-Instruct-unsloth-bnb-4bit, meta-llama/Llama-3.2-11B-Vision-Instruct, unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -39,3 +40,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.5
|
||||
min_p: 0.1
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Llama-3.2-3B-Instruct
|
||||
# Based on Llama3.2_(1B_and_3B)-Conversational.ipynb
|
||||
# Also applies to: unsloth/Llama-3.2-3B-Instruct-unsloth-bnb-4bit, meta-llama/Llama-3.2-3B-Instruct, unsloth/Llama-3.2-3B-Instruct-bnb-4bit, RedHatAI/Llama-3.2-3B-Instruct-FP8, unsloth/Llama-3.2-3B-Instruct-FP8-Block, unsloth/Llama-3.2-3B-Instruct-FP8-Dynamic
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.5
|
||||
min_p: 0.1
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Llama-3.3-70B-Instruct
|
||||
# Based on Llama3.3_(70B)_A100-Conversational.ipynb
|
||||
# Also applies to: unsloth/Llama-3.3-70B-Instruct-unsloth-bnb-4bit, meta-llama/Llama-3.3-70B-Instruct, unsloth/Llama-3.3-70B-Instruct-bnb-4bit, RedHatAI/Llama-3.3-70B-Instruct-FP8, unsloth/Llama-3.3-70B-Instruct-FP8-Block, unsloth/Llama-3.3-70B-Instruct-FP8-Dynamic
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.5
|
||||
min_p: 0.1
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Llasa-3B
|
||||
# Based on Llasa_TTS_(3B).ipynb and Llasa_TTS_(1B).ipynb
|
||||
# Also applies to: HKUSTAudio/Llasa-1B
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -40,3 +41,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.2
|
||||
top_p: 1.2
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Magistral-Small-2509
|
||||
# Based on Magistral_(24B)-Reasoning-Conversational.ipynb
|
||||
# Also applies to: mistralai/Magistral-Small-2509, unsloth/Magistral-Small-2509-bnb-4bit
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,8 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.7
|
||||
min_p: 0.01
|
||||
top_p: 0.95
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Ministral-3-3B-Instruct-2512
|
||||
# Based on Ministral_3_VL_(3B)_Vision.ipynb
|
||||
# Also applies to: unsloth/Ministral-3-3B-Instruct-2512
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.15
|
||||
top_p: default
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Pixtral-12B-2409
|
||||
# Based on Pixtral_(12B)-Vision.ipynb
|
||||
# Also applies to: unsloth/Pixtral-12B-2409-unsloth-bnb-4bit, mistralai/Pixtral-12B-2409, unsloth/Pixtral-12B-2409-bnb-4bit
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -39,3 +40,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.5
|
||||
min_p: 0.1
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for OuteAI/Llama-OuteTTS-1.0-1B
|
||||
# Based on Oute_TTS_(1B).ipynb
|
||||
# Also applies to: OuteAI/Llama-OuteTTS-1.0-1B
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -40,3 +41,9 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.4
|
||||
top_k: 40
|
||||
top_p: 0.9
|
||||
min_p: 0.05
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for Spark-TTS-0.5B/LLM
|
||||
# Based on Spark_TTS_(0_5B).ipynb
|
||||
# Also applies to: Spark-TTS-0.5B/LLM
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,8 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.8
|
||||
top_k: 50
|
||||
top_p: 1.0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/LFM2-1.2B
|
||||
# Based on Liquid_LFM2_(1.2B)-Conversational.ipynb
|
||||
# Also applies to: unsloth/LFM2-1.2B
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -39,3 +40,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.3
|
||||
min_p: 0.15
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Nemotron-3-Nano-30B-A3B
|
||||
# Based on Nemotron-3-Nano-30B-A3B_A100.ipynb
|
||||
# Also applies to: unsloth/Nemotron-3-Nano-30B-A3B
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -47,3 +48,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.0
|
||||
top_p: 1.0
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/PaddleOCR-VL
|
||||
# Based on Paddle_OCR_(1B)_Vision.ipynb
|
||||
# Also applies to: unsloth/PaddleOCR-VL
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.5
|
||||
min_p: 0.1
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/orpheus-3b-0.1-ft
|
||||
# Based on Orpheus_(3B)-TTS.ipynb
|
||||
# Also applies to: unsloth/orpheus-3b-0.1-ft-unsloth-bnb-4bit, canopylabs/orpheus-3b-0.1-ft, unsloth/orpheus-3b-0.1-ft-bnb-4bit
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.6
|
||||
top_p: 0.95
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Phi-4
|
||||
# Based on Phi_4-Conversational.ipynb
|
||||
# Also applies to: unsloth/phi-4-unsloth-bnb-4bit, microsoft/phi-4, unsloth/phi-4-bnb-4bit
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.8
|
||||
top_p: 0.95
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Qwen2-VL-7B-Instruct
|
||||
# Based on Qwen2_VL_(7B)-Vision.ipynb
|
||||
# Also applies to: unsloth/Qwen2-VL-7B-Instruct-unsloth-bnb-4bit, Qwen/Qwen2-VL-7B-Instruct, unsloth/Qwen2-VL-7B-Instruct-bnb-4bit
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -39,3 +40,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.5
|
||||
min_p: 0.1
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Qwen2.5-Coder-14B-Instruct
|
||||
# Based on Qwen2.5_Coder_(14B)-Conversational.ipynb
|
||||
# Also applies to: unsloth/Qwen2.5-Coder-14B-Instruct-bnb-4bit, Qwen/Qwen2.5-Coder-14B-Instruct
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.5
|
||||
min_p: 0.1
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Qwen2.5-VL-7B-Instruct-bnb-4bit
|
||||
# Based on Qwen2.5_VL_(7B)-Vision.ipynb
|
||||
# Also applies to: unsloth/Qwen2.5-VL-7B-Instruct, Qwen/Qwen2.5-VL-7B-Instruct, unsloth/Qwen2.5-VL-7B-Instruct-unsloth-bnb-4bit
|
||||
# added inference parameters from unsloth notebook
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -39,3 +40,7 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 1.5
|
||||
min_p: 0.1
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Qwen3-0.6B
|
||||
# Based on Qwen3_(0_6B)-Phone_Deployment.ipynb
|
||||
# Also applies to: unsloth/Qwen3-0.6B-unsloth-bnb-4bit, Qwen/Qwen3-0.6B, unsloth/Qwen3-0.6B-bnb-4bit, Qwen/Qwen3-0.6B-FP8, unsloth/Qwen3-0.6B-FP8
|
||||
# added inference parameters from Ollama
|
||||
|
||||
training:
|
||||
max_seq_length: 1024
|
||||
|
|
@ -45,3 +46,8 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.6
|
||||
top_k: 20
|
||||
top_p: 0.95
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Qwen3-14B-Base
|
||||
# Based on Qwen3_(14B)-Alpaca.ipynb
|
||||
# Also applies to: unsloth/Qwen3-14B-Base, Qwen/Qwen3-14B-Base, unsloth/Qwen3-14B-Base-bnb-4bit
|
||||
# added inference parameters from Ollama
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,8 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.6
|
||||
top_k: 20
|
||||
top_p: 0.95
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Qwen3-14B
|
||||
# Based on Qwen3_(14B).ipynb
|
||||
# Also applies to: unsloth/Qwen3-14B-unsloth-bnb-4bit, Qwen/Qwen3-14B, unsloth/Qwen3-14B-bnb-4bit, Qwen/Qwen3-14B-FP8, unsloth/Qwen3-14B-FP8
|
||||
# added inference parameters from Ollama
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,8 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.6
|
||||
top_k: 20
|
||||
top_p: 0.95
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Qwen3-32B
|
||||
# Based on Qwen3_(32B)_A100-Reasoning-Conversational.ipynb
|
||||
# Also applies to: unsloth/Qwen3-32B-unsloth-bnb-4bit, Qwen/Qwen3-32B, unsloth/Qwen3-32B-bnb-4bit, Qwen/Qwen3-32B-FP8, unsloth/Qwen3-32B-FP8
|
||||
# added inference parameters from Ollama
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,8 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.6
|
||||
top_k: 20
|
||||
top_p: 0.95
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Qwen3-4B-Instruct-2507
|
||||
# Based on Qwen3_(4B)-Instruct.ipynb
|
||||
# Also applies to: unsloth/Qwen3-4B-Instruct-2507-unsloth-bnb-4bit, Qwen/Qwen3-4B-Instruct-2507, unsloth/Qwen3-4B-Instruct-2507-bnb-4bit, Qwen/Qwen3-4B-Instruct-2507-FP8, unsloth/Qwen3-4B-Instruct-2507-FP8
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,9 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.7
|
||||
top_p: 0.80
|
||||
top_k: 20
|
||||
min_p: 0.00
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Qwen3-4B-Thinking-2507
|
||||
# Based on Qwen3_(4B)-Thinking.ipynb
|
||||
# Also applies to: unsloth/Qwen3-4B-Thinking-2507-unsloth-bnb-4bit, Qwen/Qwen3-4B-Thinking-2507, unsloth/Qwen3-4B-Thinking-2507-bnb-4bit, Qwen/Qwen3-4B-Thinking-2507-FP8, unsloth/Qwen3-4B-Thinking-2507-FP8
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -45,3 +46,9 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.6
|
||||
top_p: 0.95
|
||||
top_k: 20
|
||||
min_p: 0.00
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Model defaults for unsloth/Qwen3-VL-8B-Instruct
|
||||
# Based on Qwen3_VL_(8B)-Vision.ipynb
|
||||
# Also applies to: Qwen/Qwen3-VL-8B-Instruct-FP8, unsloth/Qwen3-VL-8B-Instruct-FP8, unsloth/Qwen3-VL-8B-Instruct, Qwen/Qwen3-VL-8B-Instruct, unsloth/Qwen3-VL-8B-Instruct-bnb-4bit
|
||||
# added inference parameters from unsloth guides
|
||||
|
||||
training:
|
||||
max_seq_length: 2048
|
||||
|
|
@ -39,3 +40,8 @@ logging:
|
|||
tensorboard_dir: "runs"
|
||||
log_frequency: 10
|
||||
|
||||
inference:
|
||||
temperature: 0.7
|
||||
top_p: 0.8
|
||||
top_k: 20
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from peft import PeftModel, PeftModelForCausalLM
|
|||
|
||||
import sys
|
||||
import torch
|
||||
from typing import Optional, Generator, Tuple
|
||||
from typing import Optional, Union, Generator, Tuple
|
||||
from utils.models import ModelConfig, get_base_model_from_lora
|
||||
from utils.paths import is_model_cached
|
||||
from utils.utils import format_error_message
|
||||
|
|
@ -179,24 +179,21 @@ class InferenceBackend:
|
|||
model = self.models[base_model_name].get("model")
|
||||
|
||||
try:
|
||||
# Step 1: Unload the adapter weights. This returns the base model object.
|
||||
# This step is only necessary if the model is currently a PeftModel instance.
|
||||
# Step 1: Unload the adapter weights if model is a PeftModel.
|
||||
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
|
||||
logger.info("Model is a PeftModel. Unloading adapters...")
|
||||
logger.info(f"Unloading LoRA adapters from '{base_model_name}'...")
|
||||
unwrapped_base_model = model.unload()
|
||||
self.models[base_model_name]["model"] = unwrapped_base_model
|
||||
model = unwrapped_base_model # Continue with the unwrapped model
|
||||
model = unwrapped_base_model
|
||||
|
||||
# Step 2: Delete any lingering adapter configurations from the object.
|
||||
# This is the crucial step you identified.
|
||||
if hasattr(model, 'peft_config') and model.peft_config:
|
||||
logger.info("Found lingering adapter configurations. Deleting them now...")
|
||||
# Create a static list of keys before iterating and deleting
|
||||
for name in list(model.peft_config.keys()):
|
||||
logger.info(f"Deleting adapter config: '{name}'")
|
||||
model.delete_adapter(name)
|
||||
# Step 2: Clear any lingering peft_config from the unwrapped model.
|
||||
# After model.unload(), the base model may still carry a peft_config
|
||||
# attribute. Removing it ensures PeftModel.from_pretrained() gets
|
||||
# a clean base model without "multiple adapters" warnings.
|
||||
if hasattr(model, 'peft_config'):
|
||||
del model.peft_config
|
||||
|
||||
logger.info("Model has been successfully reverted to a clean base state.")
|
||||
logger.info(f"Model '{base_model_name}' reverted to clean base state.")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
|
|
@ -204,35 +201,29 @@ class InferenceBackend:
|
|||
import traceback
|
||||
logger.error(traceback.format_exc())
|
||||
return False
|
||||
pass
|
||||
|
||||
def activate_lora_adapter(self, base_model_name: str, lora_path: str) -> Tuple[bool, Optional[str]]:
|
||||
"""
|
||||
Activates a specific LoRA adapter on what is assumed to be a clean base model.
|
||||
Uses PeftModel.from_pretrained() which correctly wraps the base model.
|
||||
"""
|
||||
model = self.models[base_model_name].get("model")
|
||||
adapter_name_to_load = lora_path.split("/")[-1].replace(".", "_")
|
||||
|
||||
try:
|
||||
# At this point, the model should be clean thanks to revert_to_base_model.
|
||||
# We can now safely load and set the new adapter.
|
||||
|
||||
# Step 3: Load the new adapter.
|
||||
logger.info(f"Loading adapter '{adapter_name_to_load}' from '{lora_path}'")
|
||||
model.load_adapter(lora_path, adapter_name=adapter_name_to_load)
|
||||
|
||||
# Step 4: Set the new adapter as active.
|
||||
logger.info(f"Setting '{adapter_name_to_load}' as the active adapter.")
|
||||
model.set_adapter(adapter_name_to_load)
|
||||
# Use PeftModel.from_pretrained to wrap the clean base model with the adapter.
|
||||
# This is the correct approach after model.unload() + del peft_config.
|
||||
logger.info(f"Loading LoRA adapter '{adapter_name_to_load}' from '{lora_path}'...")
|
||||
model = PeftModel.from_pretrained(model, lora_path, adapter_name=adapter_name_to_load)
|
||||
self.models[base_model_name]["model"] = model
|
||||
logger.info(f"LoRA adapter '{adapter_name_to_load}' activated successfully.")
|
||||
|
||||
return True, adapter_name_to_load
|
||||
except Exception as e:
|
||||
# This will catch the "already exists" error if revert_to_base_model failed.
|
||||
logger.error(f"Failed to activate LoRA adapter '{adapter_name_to_load}': {e}")
|
||||
import traceback
|
||||
logger.error(traceback.format_exc())
|
||||
return False, None
|
||||
pass
|
||||
|
||||
def load_adapter(self, base_model_name: str, adapter_path: str, adapter_name: str = None) -> bool:
|
||||
"""
|
||||
|
|
@ -448,6 +439,74 @@ class InferenceBackend:
|
|||
return False
|
||||
pass
|
||||
|
||||
def _apply_adapter_state(self, use_adapter: Optional[Union[bool, str]]) -> None:
|
||||
"""
|
||||
Apply adapter state before generation. Must be called under _generation_lock.
|
||||
|
||||
Uses revert_to_base_model() / activate_lora_adapter() which work correctly
|
||||
for models loaded by Unsloth as complete PeftModels (via model.unload() /
|
||||
model.load_adapter()), matching the proven pattern from the Gradio eval page.
|
||||
|
||||
Args:
|
||||
use_adapter: None = no change, False = disable (base model),
|
||||
True = enable current adapter, str = enable specific adapter.
|
||||
"""
|
||||
if use_adapter is None:
|
||||
return
|
||||
|
||||
base = self.active_model_name
|
||||
if not base or base not in self.models:
|
||||
return
|
||||
|
||||
model_info = self.models[base]
|
||||
|
||||
if use_adapter is False:
|
||||
# Revert to pure base model by unloading adapter weights
|
||||
logger.info(f"Compare mode: reverting '{base}' to base model for generation")
|
||||
self.revert_to_base_model(base)
|
||||
|
||||
elif use_adapter is True:
|
||||
# Activate the LoRA adapter from the original model path
|
||||
lora_path = model_info.get("model_path")
|
||||
if lora_path and model_info.get("is_lora"):
|
||||
logger.info(f"Compare mode: activating LoRA adapter from '{lora_path}' on '{base}'")
|
||||
self.activate_lora_adapter(base, lora_path)
|
||||
else:
|
||||
# Fallback for dynamically attached adapters
|
||||
loaded = model_info.get("loaded_adapters", {})
|
||||
if loaded:
|
||||
adapter_name = list(loaded.keys())[-1]
|
||||
logger.info(f"Compare mode: enabling adapter '{adapter_name}' on '{base}'")
|
||||
self.set_active_adapter(base, adapter_name)
|
||||
else:
|
||||
logger.warning("use_adapter=true but no adapter path/adapters on model")
|
||||
|
||||
elif isinstance(use_adapter, str):
|
||||
# Activate a specific adapter by path
|
||||
logger.info(f"Compare mode: activating specific adapter '{use_adapter}' on '{base}'")
|
||||
self.activate_lora_adapter(base, use_adapter)
|
||||
|
||||
def generate_with_adapter_control(
|
||||
self,
|
||||
use_adapter: Optional[Union[bool, str]] = None,
|
||||
**gen_kwargs,
|
||||
) -> Generator[str, None, None]:
|
||||
"""
|
||||
Thread-safe generation with optional adapter toggling.
|
||||
|
||||
Acquires the generation lock, applies adapter state, then generates.
|
||||
This ensures adapter toggle + generation are atomic — critical for
|
||||
compare mode where base and LoRA panes fire concurrently.
|
||||
|
||||
Args:
|
||||
use_adapter: Adapter control (None/False/True/str). See _apply_adapter_state.
|
||||
**gen_kwargs: Forwarded to generate_chat_response.
|
||||
"""
|
||||
with self._generation_lock:
|
||||
self._apply_adapter_state(use_adapter)
|
||||
# Delegate to the lock-free generation path
|
||||
yield from self._generate_chat_response_inner(**gen_kwargs)
|
||||
|
||||
def generate_chat_response(self,
|
||||
messages: list,
|
||||
system_prompt: str,
|
||||
|
|
@ -459,11 +518,33 @@ class InferenceBackend:
|
|||
repetition_penalty: float = 1.1) -> Generator[str, None, None]:
|
||||
"""
|
||||
Generate response for text or vision models.
|
||||
Acquires the generation lock. For adapter-controlled generation,
|
||||
use generate_with_adapter_control() instead.
|
||||
"""
|
||||
with self._generation_lock:
|
||||
yield from self._generate_chat_response_inner(
|
||||
messages=messages,
|
||||
system_prompt=system_prompt,
|
||||
image=image,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
top_k=top_k,
|
||||
max_new_tokens=max_new_tokens,
|
||||
repetition_penalty=repetition_penalty,
|
||||
)
|
||||
|
||||
1. Messages are already in ChatML format (role/content)
|
||||
2. Apply get_chat_template() if model in mapper
|
||||
3. Apply tokenizer.apply_chat_template()
|
||||
4. Generate
|
||||
def _generate_chat_response_inner(self,
|
||||
messages: list,
|
||||
system_prompt: str = "",
|
||||
image=None,
|
||||
temperature: float = 0.7,
|
||||
top_p: float = 0.9,
|
||||
top_k: int = 40,
|
||||
max_new_tokens: int = 256,
|
||||
repetition_penalty: float = 1.1) -> Generator[str, None, None]:
|
||||
"""
|
||||
Inner generation logic (no lock). Called by both generate_chat_response
|
||||
and generate_with_adapter_control.
|
||||
"""
|
||||
if not self.active_model_name:
|
||||
yield "Error: No active model"
|
||||
|
|
@ -473,55 +554,54 @@ class InferenceBackend:
|
|||
is_vision = model_info.get("is_vision", False)
|
||||
tokenizer = model_info.get("tokenizer") or model_info.get("processor")
|
||||
|
||||
with self._generation_lock:
|
||||
if is_vision:
|
||||
# Vision model generation
|
||||
yield from self._generate_vision_response(
|
||||
messages, system_prompt, image,
|
||||
temperature, top_p, top_k, max_new_tokens, repetition_penalty
|
||||
)
|
||||
else:
|
||||
# Text model: Use training pipeline approach
|
||||
# Messages are already in ChatML format from eval.py
|
||||
if is_vision:
|
||||
# Vision model generation
|
||||
yield from self._generate_vision_response(
|
||||
messages, system_prompt, image,
|
||||
temperature, top_p, top_k, max_new_tokens, repetition_penalty
|
||||
)
|
||||
else:
|
||||
# Text model: Use training pipeline approach
|
||||
# Messages are already in ChatML format from eval.py
|
||||
|
||||
# Step 1: Apply get_chat_template if model is in mapper
|
||||
try:
|
||||
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, get_tokenizer_chat_template
|
||||
# Step 1: Apply get_chat_template if model is in mapper
|
||||
try:
|
||||
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, get_tokenizer_chat_template
|
||||
|
||||
model_name_lower = self.active_model_name.lower()
|
||||
model_name_lower = self.active_model_name.lower()
|
||||
|
||||
# Check if model has a registered template
|
||||
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
|
||||
template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
|
||||
logger.info(f"Applying chat template '{template_name}' for {self.active_model_name}")
|
||||
# Check if model has a registered template
|
||||
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
|
||||
template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
|
||||
logger.info(f"Applying chat template '{template_name}' for {self.active_model_name}")
|
||||
|
||||
# This modifies the tokenizer with the correct template
|
||||
tokenizer = get_chat_template(
|
||||
tokenizer,
|
||||
self.active_model_name
|
||||
)
|
||||
else:
|
||||
logger.info(f"No registered template for {self.active_model_name}, using tokenizer default")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not apply get_chat_template: {e}")
|
||||
|
||||
# Step 2: Format with tokenizer.apply_chat_template()
|
||||
try:
|
||||
formatted_prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
# This modifies the tokenizer with the correct template
|
||||
tokenizer = get_chat_template(
|
||||
tokenizer,
|
||||
self.active_model_name
|
||||
)
|
||||
logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...")
|
||||
except Exception as e:
|
||||
logger.error(f"Error applying chat template: {e}")
|
||||
# Fallback to manual formatting
|
||||
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
|
||||
else:
|
||||
logger.info(f"No registered template for {self.active_model_name}, using tokenizer default")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not apply get_chat_template: {e}")
|
||||
|
||||
# Step 3: Generate
|
||||
yield from self.generate_stream(
|
||||
formatted_prompt, temperature, top_p, top_k, max_new_tokens, repetition_penalty
|
||||
# Step 2: Format with tokenizer.apply_chat_template()
|
||||
try:
|
||||
formatted_prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)
|
||||
logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...")
|
||||
except Exception as e:
|
||||
logger.error(f"Error applying chat template: {e}")
|
||||
# Fallback to manual formatting
|
||||
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
|
||||
|
||||
# Step 3: Generate
|
||||
yield from self.generate_stream(
|
||||
formatted_prompt, temperature, top_p, top_k, max_new_tokens, repetition_penalty
|
||||
)
|
||||
|
||||
def _generate_vision_response(self, messages, system_prompt, image,
|
||||
temperature, top_p, top_k, max_new_tokens,
|
||||
|
|
|
|||
|
|
@ -2,13 +2,16 @@
|
|||
Unsloth Training Backend
|
||||
Integrates Unsloth training capabilities with the FastAPI backend
|
||||
"""
|
||||
import os
|
||||
# Prevent tokenizer parallelism deadlocks when datasets uses multiprocessing fork
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
import torch
|
||||
from utils.hardware import clear_gpu_cache
|
||||
torch._dynamo.config.recompile_limit = 64
|
||||
from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
|
||||
import os
|
||||
import json
|
||||
import threading
|
||||
import math
|
||||
|
|
@ -58,6 +61,7 @@ class UnslothTrainer:
|
|||
self.progress_callbacks = []
|
||||
self.is_training = False
|
||||
self.should_stop = False
|
||||
self.save_on_stop = True
|
||||
|
||||
# Model state tracking
|
||||
self.is_vlm = False
|
||||
|
|
@ -558,6 +562,12 @@ class UnslothTrainer:
|
|||
config_args["warmup_steps"] = 5
|
||||
print(f"Using default warmup_steps: 5\n")
|
||||
|
||||
# Add save_steps if specified
|
||||
save_steps_val = training_args.get('save_steps', 0)
|
||||
if save_steps_val and save_steps_val > 0:
|
||||
config_args["save_steps"] = save_steps_val
|
||||
config_args["save_strategy"] = "steps"
|
||||
|
||||
# If max_steps is specified, use it instead of epochs
|
||||
max_steps_val = training_args.get('max_steps', 0)
|
||||
if max_steps_val and max_steps_val > 0:
|
||||
|
|
@ -756,16 +766,32 @@ class UnslothTrainer:
|
|||
self.trainer.train()
|
||||
|
||||
# ========== SAVE MODEL ==========
|
||||
self.trainer.save_model()
|
||||
self.tokenizer.save_pretrained(output_dir)
|
||||
print(f"\nTraining completed! Model saved to {output_dir}\n")
|
||||
|
||||
self._update_progress(
|
||||
is_training=False,
|
||||
is_completed=True,
|
||||
#status_message=status_msg
|
||||
status_message=f"Training completed! Model saved to {output_dir}",
|
||||
)
|
||||
if self.should_stop and self.save_on_stop:
|
||||
# Stopped by user — save model at current checkpoint
|
||||
self.trainer.save_model()
|
||||
self.tokenizer.save_pretrained(output_dir)
|
||||
print(f"\nTraining stopped. Model saved to {output_dir}\n")
|
||||
self._update_progress(
|
||||
is_training=False,
|
||||
status_message=f"Training stopped. Model saved to {output_dir}",
|
||||
)
|
||||
elif self.should_stop:
|
||||
# Cancelled by user — don't save
|
||||
print("\nTraining cancelled.\n")
|
||||
self._update_progress(
|
||||
is_training=False,
|
||||
status_message="Training cancelled.",
|
||||
)
|
||||
else:
|
||||
# Normal completion
|
||||
self.trainer.save_model()
|
||||
self.tokenizer.save_pretrained(output_dir)
|
||||
print(f"\nTraining completed! Model saved to {output_dir}\n")
|
||||
self._update_progress(
|
||||
is_training=False,
|
||||
is_completed=True,
|
||||
status_message=f"Training completed! Model saved to {output_dir}",
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Training error: {e}")
|
||||
|
|
@ -774,10 +800,11 @@ class UnslothTrainer:
|
|||
finally:
|
||||
self.is_training = False
|
||||
|
||||
def stop_training(self):
|
||||
def stop_training(self, save: bool = True):
|
||||
"""Stop ongoing training"""
|
||||
print("\nStopping training...")
|
||||
print(f"\nStopping training (save={save})...")
|
||||
self.should_stop = True
|
||||
self.save_on_stop = save
|
||||
self.is_training = False
|
||||
# Clear the status message so timer doesn't show stale status
|
||||
self._update_progress(is_training=False, status_message="")
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from typing import Any, Generator, Tuple
|
|||
import logging
|
||||
|
||||
from .trainer import get_trainer, TrainingProgress
|
||||
from utils.hardware import clear_gpu_cache
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -101,8 +102,34 @@ class TrainingBackend:
|
|||
True if training started successfully, False otherwise.
|
||||
"""
|
||||
try:
|
||||
# Wait for any previous training thread to finish
|
||||
old_thread = getattr(self.trainer, "training_thread", None)
|
||||
if old_thread and old_thread.is_alive():
|
||||
logger.info("Waiting for previous training thread to finish...")
|
||||
old_thread.join(timeout=30)
|
||||
|
||||
# Explicitly free old SFTTrainer and CUDA resources before loading new model.
|
||||
# Without this, forked multiprocessing workers (num_proc tokenization) inherit
|
||||
# stale CUDA state from the previous run, causing extreme slowdowns or crashes.
|
||||
if self.trainer.trainer is not None:
|
||||
logger.info("Cleaning up previous SFTTrainer...")
|
||||
self.trainer.trainer = None
|
||||
if self.trainer.model is not None:
|
||||
self.trainer.model = None
|
||||
if self.trainer.tokenizer is not None:
|
||||
self.trainer.tokenizer = None
|
||||
# Flush all pending async CUDA ops so forked tokenization processes
|
||||
# don't inherit stale async state that causes pool join to hang.
|
||||
import torch as _torch
|
||||
if _torch.cuda.is_available():
|
||||
_torch.cuda.synchronize()
|
||||
import gc
|
||||
gc.collect()
|
||||
clear_gpu_cache()
|
||||
|
||||
# Reset stop flag and clear history
|
||||
self.trainer.should_stop = False
|
||||
self.trainer.save_on_stop = True
|
||||
self.loss_history = []
|
||||
self.lr_history = []
|
||||
self.step_history = []
|
||||
|
|
@ -224,16 +251,19 @@ class TrainingBackend:
|
|||
)
|
||||
return False
|
||||
|
||||
def stop_training(self) -> bool:
|
||||
def stop_training(self, save: bool = True) -> bool:
|
||||
"""
|
||||
Stop ongoing training.
|
||||
|
||||
Args:
|
||||
save: If True, save the model at the current checkpoint.
|
||||
|
||||
Returns:
|
||||
True if training was successfully stopped.
|
||||
"""
|
||||
try:
|
||||
logger.info("Stopping training...")
|
||||
self.trainer.stop_training()
|
||||
logger.info(f"Stopping training (save={save})...")
|
||||
self.trainer.stop_training(save=save)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Error stopping training: {e}")
|
||||
|
|
@ -293,6 +323,10 @@ class TrainingBackend:
|
|||
True if training is in progress, False otherwise
|
||||
"""
|
||||
try:
|
||||
# If user requested stop, training is no longer considered active
|
||||
if self.trainer.should_stop:
|
||||
return False
|
||||
|
||||
progress = self.trainer.get_training_progress()
|
||||
# Training is active if is_training is True
|
||||
# Also check if we're in loading/preparation phase (status_message indicates activity)
|
||||
|
|
|
|||
|
|
@ -14,12 +14,13 @@ from datetime import datetime
|
|||
|
||||
# Import routers
|
||||
from routes import (
|
||||
training_router,
|
||||
models_router,
|
||||
inference_router,
|
||||
datasets_router,
|
||||
auth_router,
|
||||
data_recipe_router,
|
||||
datasets_router,
|
||||
export_router,
|
||||
inference_router,
|
||||
models_router,
|
||||
training_router,
|
||||
)
|
||||
from auth import storage
|
||||
from utils.hardware import detect_hardware
|
||||
|
|
@ -75,6 +76,7 @@ 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"])
|
||||
app.include_router(data_recipe_router, prefix="/api/data-recipe", tags=["data-recipe"])
|
||||
app.include_router(export_router, prefix="/api/export", tags=["export"])
|
||||
|
||||
|
||||
# ============ Health and System Endpoints ============
|
||||
|
|
@ -136,7 +138,7 @@ def setup_frontend(app: FastAPI, build_path: Path):
|
|||
|
||||
@app.get("/")
|
||||
async def serve_root():
|
||||
return FileResponse(build_path / "index.html")
|
||||
return FileResponse(build_path / "index.html", headers={"Cache-Control": "no-cache, no-store, must-revalidate"})
|
||||
|
||||
@app.get("/{full_path:path}")
|
||||
async def serve_frontend(full_path: str):
|
||||
|
|
@ -147,7 +149,7 @@ def setup_frontend(app: FastAPI, build_path: Path):
|
|||
if file_path.is_file():
|
||||
return FileResponse(file_path)
|
||||
|
||||
return FileResponse(build_path / "index.html")
|
||||
return FileResponse(build_path / "index.html", headers={"Cache-Control": "no-cache, no-store, must-revalidate"})
|
||||
|
||||
return True
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -19,6 +19,17 @@ from .auth import (
|
|||
RefreshTokenRequest,
|
||||
AuthStatusResponse,
|
||||
)
|
||||
from .export import (
|
||||
CheckpointInfo,
|
||||
CheckpointListResponse,
|
||||
LoadCheckpointRequest,
|
||||
ExportStatusResponse,
|
||||
ExportOperationResponse,
|
||||
ExportMergedModelRequest,
|
||||
ExportBaseModelRequest,
|
||||
ExportGGUFRequest,
|
||||
ExportLoRAAdapterRequest,
|
||||
)
|
||||
from .users import Token
|
||||
from .datasets import (
|
||||
CheckFormatRequest,
|
||||
|
|
@ -62,6 +73,16 @@ __all__ = [
|
|||
"AuthLoginRequest",
|
||||
"RefreshTokenRequest",
|
||||
"AuthStatusResponse",
|
||||
# Export schemas
|
||||
"CheckpointInfo",
|
||||
"CheckpointListResponse",
|
||||
"LoadCheckpointRequest",
|
||||
"ExportStatusResponse",
|
||||
"ExportOperationResponse",
|
||||
"ExportMergedModelRequest",
|
||||
"ExportBaseModelRequest",
|
||||
"ExportGGUFRequest",
|
||||
"ExportLoRAAdapterRequest",
|
||||
"Token",
|
||||
# Dataset schemas
|
||||
"CheckFormatRequest",
|
||||
|
|
|
|||
141
studio/backend/models/export.py
Normal file
141
studio/backend/models/export.py
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
"""
|
||||
Pydantic schemas for Export API.
|
||||
"""
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Literal, Dict, Any
|
||||
|
||||
|
||||
class CheckpointInfo(BaseModel):
|
||||
"""Information about a discovered checkpoint directory."""
|
||||
|
||||
display_name: str = Field(..., description="User-friendly checkpoint name (folder name)")
|
||||
path: str = Field(..., description="Full path to the checkpoint directory")
|
||||
|
||||
|
||||
class CheckpointListResponse(BaseModel):
|
||||
"""Response for listing available checkpoints in an outputs directory."""
|
||||
|
||||
outputs_dir: str = Field(..., description="Directory that was scanned")
|
||||
checkpoints: List[CheckpointInfo] = Field(
|
||||
default_factory=list,
|
||||
description="List of discovered checkpoints",
|
||||
)
|
||||
|
||||
|
||||
class LoadCheckpointRequest(BaseModel):
|
||||
"""Request for loading a checkpoint into the export backend."""
|
||||
|
||||
checkpoint_path: str = Field(..., description="Path to the checkpoint directory")
|
||||
max_seq_length: int = Field(
|
||||
2048,
|
||||
ge=128,
|
||||
le=32768,
|
||||
description="Maximum sequence length for loading the model",
|
||||
)
|
||||
load_in_4bit: bool = Field(
|
||||
True,
|
||||
description="Whether to load the model in 4-bit quantization",
|
||||
)
|
||||
|
||||
|
||||
class ExportStatusResponse(BaseModel):
|
||||
"""Current export backend status."""
|
||||
|
||||
current_checkpoint: Optional[str] = Field(
|
||||
None,
|
||||
description="Path to the currently loaded checkpoint, if any",
|
||||
)
|
||||
is_vision: bool = Field(
|
||||
False,
|
||||
description="True if the loaded checkpoint is a vision model",
|
||||
)
|
||||
is_peft: bool = Field(
|
||||
False,
|
||||
description="True if the loaded checkpoint is a PEFT (LoRA) model",
|
||||
)
|
||||
|
||||
|
||||
class ExportOperationResponse(BaseModel):
|
||||
"""Generic response for export operations."""
|
||||
|
||||
success: bool = Field(..., description="True if the operation succeeded")
|
||||
message: str = Field(..., description="Human-readable status or error message")
|
||||
details: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description="Optional extra details about the operation",
|
||||
)
|
||||
|
||||
|
||||
class ExportCommonOptions(BaseModel):
|
||||
"""Common options for export operations that save locally and/or push to Hub."""
|
||||
|
||||
save_directory: str = Field(
|
||||
...,
|
||||
description="Local directory where the exported artifacts will be written",
|
||||
)
|
||||
push_to_hub: bool = Field(
|
||||
False,
|
||||
description="If True, also push the exported model to the Hugging Face Hub",
|
||||
)
|
||||
repo_id: Optional[str] = Field(
|
||||
None,
|
||||
description="Hugging Face Hub repository ID (username/model-name)",
|
||||
)
|
||||
hf_token: Optional[str] = Field(
|
||||
None,
|
||||
description="Hugging Face access token used for Hub operations",
|
||||
)
|
||||
private: bool = Field(
|
||||
False,
|
||||
description="If True, create a private repository on the Hub (where applicable)",
|
||||
)
|
||||
|
||||
|
||||
class ExportMergedModelRequest(ExportCommonOptions):
|
||||
"""Request for exporting a merged PEFT model."""
|
||||
|
||||
format_type: Literal["16-bit (FP16)", "4-bit (FP4)"] = Field(
|
||||
"16-bit (FP16)",
|
||||
description="Export precision / format for the merged model",
|
||||
)
|
||||
|
||||
|
||||
class ExportBaseModelRequest(ExportCommonOptions):
|
||||
"""Request for exporting a non-PEFT (base) model."""
|
||||
|
||||
# Uses fields from ExportCommonOptions only
|
||||
pass
|
||||
|
||||
|
||||
class ExportGGUFRequest(BaseModel):
|
||||
"""Request for exporting the current model to GGUF format."""
|
||||
|
||||
save_directory: str = Field(
|
||||
...,
|
||||
description="Directory where GGUF files will be saved",
|
||||
)
|
||||
quantization_method: str = Field(
|
||||
"Q4_K_M",
|
||||
description='GGUF quantization method (e.g. "Q4_K_M")',
|
||||
)
|
||||
push_to_hub: bool = Field(
|
||||
False,
|
||||
description="If True, also push GGUF artifacts to the Hugging Face Hub",
|
||||
)
|
||||
repo_id: Optional[str] = Field(
|
||||
None,
|
||||
description="Hugging Face Hub repository ID for GGUF upload",
|
||||
)
|
||||
hf_token: Optional[str] = Field(
|
||||
None,
|
||||
description="Hugging Face token for GGUF upload",
|
||||
)
|
||||
|
||||
|
||||
class ExportLoRAAdapterRequest(ExportCommonOptions):
|
||||
"""Request for exporting only the LoRA adapter (not merged)."""
|
||||
|
||||
# Uses fields from ExportCommonOptions only
|
||||
pass
|
||||
|
||||
|
||||
|
|
@ -43,6 +43,7 @@ class LoadResponse(BaseModel):
|
|||
display_name: str = Field(..., description="Display name of the model")
|
||||
is_vision: bool = Field(False, description="Whether model is a vision model")
|
||||
is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
|
||||
inference: dict = Field(..., description="Inference parameters (temperature, top_p, top_k, min_p)")
|
||||
|
||||
|
||||
class UnloadResponse(BaseModel):
|
||||
|
|
@ -130,6 +131,16 @@ class ChatCompletionRequest(BaseModel):
|
|||
top_k: int = Field(40, ge=1, le=100, description="[x-unsloth] Top-k sampling")
|
||||
repetition_penalty: float = Field(1.1, ge=1.0, le=2.0, description="[x-unsloth] Repetition penalty")
|
||||
image_base64: Optional[str] = Field(None, description="[x-unsloth] Base64-encoded image for vision models")
|
||||
use_adapter: Optional[Union[bool, str]] = Field(
|
||||
None,
|
||||
description=(
|
||||
"[x-unsloth] Adapter control for compare mode. "
|
||||
"null = no change (default), "
|
||||
"false = disable adapters (base model), "
|
||||
"true = enable the current adapter, "
|
||||
"string = enable a specific adapter by name."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# ── Streaming response chunks ────────────────────────────────────
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from routes.inference import router as inference_router
|
|||
from routes.datasets import router as datasets_router
|
||||
from routes.auth import router as auth_router
|
||||
from routes.data_recipe import router as data_recipe_router
|
||||
from routes.export import router as export_router
|
||||
|
||||
__all__ = [
|
||||
"training_router",
|
||||
|
|
@ -16,4 +17,5 @@ __all__ = [
|
|||
"datasets_router",
|
||||
"auth_router",
|
||||
"data_recipe_router",
|
||||
"export_router",
|
||||
]
|
||||
|
|
|
|||
310
studio/backend/routes/export.py
Normal file
310
studio/backend/routes/export.py
Normal file
|
|
@ -0,0 +1,310 @@
|
|||
"""
|
||||
Export API routes: checkpoint discovery and model export operations.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
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))
|
||||
|
||||
# Auth
|
||||
from auth.authentication import get_current_subject
|
||||
|
||||
# Import backend functions
|
||||
try:
|
||||
from core.export import get_export_backend
|
||||
except ImportError:
|
||||
parent_backend = backend_path.parent / "backend"
|
||||
if str(parent_backend) not in sys.path:
|
||||
sys.path.insert(0, str(parent_backend))
|
||||
from core.export import get_export_backend
|
||||
|
||||
# Import Pydantic models
|
||||
from models import (
|
||||
CheckpointInfo,
|
||||
CheckpointListResponse,
|
||||
LoadCheckpointRequest,
|
||||
ExportStatusResponse,
|
||||
ExportOperationResponse,
|
||||
ExportMergedModelRequest,
|
||||
ExportBaseModelRequest,
|
||||
ExportGGUFRequest,
|
||||
ExportLoRAAdapterRequest,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
@router.get("/checkpoints", response_model=CheckpointListResponse)
|
||||
async def list_checkpoints(
|
||||
outputs_dir: str = Query(
|
||||
default="./outputs",
|
||||
description="Directory to scan for checkpoints",
|
||||
),
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
List available checkpoints in the outputs directory.
|
||||
|
||||
Wraps ExportBackend.scan_checkpoints.
|
||||
"""
|
||||
try:
|
||||
backend = get_export_backend()
|
||||
raw_checkpoints = backend.scan_checkpoints(outputs_dir=outputs_dir)
|
||||
|
||||
checkpoints = [
|
||||
CheckpointInfo(display_name=display_name, path=path)
|
||||
for display_name, path in raw_checkpoints
|
||||
]
|
||||
|
||||
return CheckpointListResponse(
|
||||
outputs_dir=outputs_dir,
|
||||
checkpoints=checkpoints,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing checkpoints: {e}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to list checkpoints: {str(e)}",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/load-checkpoint", response_model=ExportOperationResponse)
|
||||
async def load_checkpoint(
|
||||
request: LoadCheckpointRequest,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
Load a checkpoint into the export backend.
|
||||
|
||||
Wraps ExportBackend.load_checkpoint.
|
||||
"""
|
||||
try:
|
||||
backend = get_export_backend()
|
||||
success, message = backend.load_checkpoint(
|
||||
checkpoint_path=request.checkpoint_path,
|
||||
max_seq_length=request.max_seq_length,
|
||||
load_in_4bit=request.load_in_4bit,
|
||||
)
|
||||
|
||||
if not success:
|
||||
raise HTTPException(status_code=400, detail=message)
|
||||
|
||||
return ExportOperationResponse(success=True, message=message)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading checkpoint: {e}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to load checkpoint: {str(e)}",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/cleanup", response_model=ExportOperationResponse)
|
||||
async def cleanup_export_memory(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
Cleanup export-related models from memory (GPU/CPU).
|
||||
|
||||
Wraps ExportBackend.cleanup_memory.
|
||||
"""
|
||||
try:
|
||||
backend = get_export_backend()
|
||||
success = backend.cleanup_memory()
|
||||
|
||||
if not success:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail="Memory cleanup failed. See server logs for details.",
|
||||
)
|
||||
|
||||
return ExportOperationResponse(
|
||||
success=True,
|
||||
message="Memory cleanup completed successfully",
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error during export memory cleanup: {e}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to cleanup export memory: {str(e)}",
|
||||
)
|
||||
|
||||
|
||||
@router.get("/status", response_model=ExportStatusResponse)
|
||||
async def get_export_status(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
Get current export backend status (loaded checkpoint, model type, PEFT flag).
|
||||
"""
|
||||
try:
|
||||
backend = get_export_backend()
|
||||
return ExportStatusResponse(
|
||||
current_checkpoint=backend.current_checkpoint,
|
||||
is_vision=bool(getattr(backend, "is_vision", False)),
|
||||
is_peft=bool(getattr(backend, "is_peft", False)),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting export status: {e}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to get export status: {str(e)}",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/export/merged", response_model=ExportOperationResponse)
|
||||
async def export_merged_model(
|
||||
request: ExportMergedModelRequest,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
Export a merged PEFT model (e.g., 16-bit or 4-bit) and optionally push to Hub.
|
||||
|
||||
Wraps ExportBackend.export_merged_model.
|
||||
"""
|
||||
try:
|
||||
backend = get_export_backend()
|
||||
success, message = backend.export_merged_model(
|
||||
save_directory=request.save_directory,
|
||||
format_type=request.format_type,
|
||||
push_to_hub=request.push_to_hub,
|
||||
repo_id=request.repo_id,
|
||||
hf_token=request.hf_token,
|
||||
private=request.private,
|
||||
)
|
||||
|
||||
if not success:
|
||||
raise HTTPException(status_code=400, detail=message)
|
||||
|
||||
return ExportOperationResponse(success=True, message=message)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error exporting merged model: {e}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to export merged model: {str(e)}",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/export/base", response_model=ExportOperationResponse)
|
||||
async def export_base_model(
|
||||
request: ExportBaseModelRequest,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
Export a non-PEFT base model and optionally push to Hub.
|
||||
|
||||
Wraps ExportBackend.export_base_model.
|
||||
"""
|
||||
try:
|
||||
backend = get_export_backend()
|
||||
success, message = backend.export_base_model(
|
||||
save_directory=request.save_directory,
|
||||
push_to_hub=request.push_to_hub,
|
||||
repo_id=request.repo_id,
|
||||
hf_token=request.hf_token,
|
||||
private=request.private,
|
||||
)
|
||||
|
||||
if not success:
|
||||
raise HTTPException(status_code=400, detail=message)
|
||||
|
||||
return ExportOperationResponse(success=True, message=message)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error exporting base model: {e}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to export base model: {str(e)}",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/export/gguf", response_model=ExportOperationResponse)
|
||||
async def export_gguf(
|
||||
request: ExportGGUFRequest,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
Export the current model to GGUF format and optionally push to Hub.
|
||||
|
||||
Wraps ExportBackend.export_gguf.
|
||||
"""
|
||||
try:
|
||||
backend = get_export_backend()
|
||||
success, message = backend.export_gguf(
|
||||
save_directory=request.save_directory,
|
||||
quantization_method=request.quantization_method,
|
||||
push_to_hub=request.push_to_hub,
|
||||
repo_id=request.repo_id,
|
||||
hf_token=request.hf_token,
|
||||
)
|
||||
|
||||
if not success:
|
||||
raise HTTPException(status_code=400, detail=message)
|
||||
|
||||
return ExportOperationResponse(success=True, message=message)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error exporting GGUF model: {e}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to export GGUF model: {str(e)}",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/export/lora", response_model=ExportOperationResponse)
|
||||
async def export_lora_adapter(
|
||||
request: ExportLoRAAdapterRequest,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
Export only the LoRA adapter (if the loaded model is PEFT).
|
||||
|
||||
Wraps ExportBackend.export_lora_adapter.
|
||||
"""
|
||||
try:
|
||||
backend = get_export_backend()
|
||||
success, message = backend.export_lora_adapter(
|
||||
save_directory=request.save_directory,
|
||||
push_to_hub=request.push_to_hub,
|
||||
repo_id=request.repo_id,
|
||||
hf_token=request.hf_token,
|
||||
private=request.private,
|
||||
)
|
||||
|
||||
if not success:
|
||||
raise HTTPException(status_code=400, detail=message)
|
||||
|
||||
return ExportOperationResponse(success=True, message=message)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error exporting LoRA adapter: {e}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to export LoRA adapter: {str(e)}",
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -22,12 +22,14 @@ if str(backend_path) not in sys.path:
|
|||
try:
|
||||
from core.inference import get_inference_backend
|
||||
from utils.models import ModelConfig
|
||||
from utils.inference import load_inference_config
|
||||
except ImportError:
|
||||
parent_backend = backend_path.parent / "backend"
|
||||
if str(parent_backend) not in sys.path:
|
||||
sys.path.insert(0, str(parent_backend))
|
||||
from core.inference import get_inference_backend
|
||||
from utils.models import ModelConfig
|
||||
from utils.inference import load_inference_config
|
||||
|
||||
from models.inference import (
|
||||
LoadRequest,
|
||||
|
|
@ -64,15 +66,17 @@ async def load_model(request: LoadRequest):
|
|||
Load a model for inference.
|
||||
|
||||
The model_path should be a clean identifier from GET /models/list.
|
||||
Returns inference configuration parameters (temperature, top_p, top_k, min_p)
|
||||
from the model's YAML config, falling back to default.yaml for missing values.
|
||||
"""
|
||||
try:
|
||||
backend = get_inference_backend()
|
||||
|
||||
# Create config using clean factory method
|
||||
# is_lora is auto-detected from adapter_config.json on disk/HF
|
||||
config = ModelConfig.from_identifier(
|
||||
model_id=request.model_path,
|
||||
hf_token=request.hf_token,
|
||||
is_lora=request.is_lora,
|
||||
)
|
||||
|
||||
if not config:
|
||||
|
|
@ -97,12 +101,16 @@ async def load_model(request: LoadRequest):
|
|||
|
||||
logger.info(f"Loaded model: {config.identifier}")
|
||||
|
||||
# Load inference configuration parameters
|
||||
inference_config = load_inference_config(config.identifier)
|
||||
|
||||
return LoadResponse(
|
||||
status="loaded",
|
||||
model=config.identifier,
|
||||
display_name=config.display_name,
|
||||
is_vision=config.is_vision,
|
||||
is_lora=config.is_lora,
|
||||
inference=inference_config,
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
|
|
@ -365,6 +373,18 @@ async def openai_chat_completions(request: ChatCompletionRequest):
|
|||
repetition_penalty=request.repetition_penalty,
|
||||
)
|
||||
|
||||
# ── Choose generation path (adapter-controlled or standard) ──
|
||||
if request.use_adapter is not None:
|
||||
# Compare mode: toggle adapter state atomically with generation
|
||||
def generate():
|
||||
return backend.generate_with_adapter_control(
|
||||
use_adapter=request.use_adapter, **gen_kwargs
|
||||
)
|
||||
else:
|
||||
# Standard path: no adapter toggling
|
||||
def generate():
|
||||
return backend.generate_chat_response(**gen_kwargs)
|
||||
|
||||
model_name = backend.active_model_name or request.model
|
||||
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
|
||||
created = int(time.time())
|
||||
|
|
@ -388,7 +408,7 @@ async def openai_chat_completions(request: ChatCompletionRequest):
|
|||
# Content chunks — generate_chat_response yields cumulative
|
||||
# text, so we diff to get incremental deltas.
|
||||
prev_text = ""
|
||||
for cumulative in backend.generate_chat_response(**gen_kwargs):
|
||||
for cumulative in generate():
|
||||
new_text = cumulative[len(prev_text):]
|
||||
prev_text = cumulative
|
||||
if not new_text:
|
||||
|
|
@ -439,7 +459,7 @@ async def openai_chat_completions(request: ChatCompletionRequest):
|
|||
else:
|
||||
try:
|
||||
full_text = ""
|
||||
for token in backend.generate_chat_response(**gen_kwargs):
|
||||
for token in generate():
|
||||
full_text = token # generate_stream yields cumulative text
|
||||
|
||||
response = ChatCompletion(
|
||||
|
|
|
|||
|
|
@ -37,6 +37,11 @@ from models import (
|
|||
TrainingProgress,
|
||||
)
|
||||
from models.responses import TrainingStopResponse, TrainingMetricsResponse
|
||||
from pydantic import BaseModel as PydanticBaseModel
|
||||
|
||||
|
||||
class TrainingStopRequest(PydanticBaseModel):
|
||||
save: bool = True
|
||||
|
||||
router = APIRouter()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -251,10 +256,14 @@ async def start_training(
|
|||
|
||||
@router.post("/stop", response_model=TrainingStopResponse)
|
||||
async def stop_training(
|
||||
body: TrainingStopRequest = TrainingStopRequest(),
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
Stop the currently running training job.
|
||||
|
||||
Body:
|
||||
save (bool): If True (default), save the model at the current checkpoint.
|
||||
"""
|
||||
try:
|
||||
backend = get_training_backend()
|
||||
|
|
@ -266,7 +275,7 @@ async def stop_training(
|
|||
)
|
||||
|
||||
# Call backend stop method
|
||||
backend.stop_training()
|
||||
backend.stop_training(save=body.save)
|
||||
|
||||
return TrainingStopResponse(
|
||||
status="stopped",
|
||||
|
|
@ -281,6 +290,29 @@ async def stop_training(
|
|||
)
|
||||
|
||||
|
||||
@router.post("/reset")
|
||||
async def reset_training(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
Reset training state so the user can return to configuration.
|
||||
"""
|
||||
try:
|
||||
backend = get_training_backend()
|
||||
backend.trainer.should_stop = False
|
||||
backend.trainer.training_progress = backend.trainer.training_progress.__class__()
|
||||
backend.loss_history = []
|
||||
backend.lr_history = []
|
||||
backend.step_history = []
|
||||
return {"status": "ok"}
|
||||
except Exception as e:
|
||||
logger.error(f"Error resetting training: {e}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to reset training: {str(e)}",
|
||||
)
|
||||
|
||||
|
||||
@router.get("/status")
|
||||
async def get_training_status(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
|
|
@ -313,6 +345,9 @@ async def get_training_status(
|
|||
) or "Ready to train"
|
||||
error_message = getattr(progress, "error", None) if progress else None
|
||||
|
||||
# Check if training was stopped by user
|
||||
trainer_stopped = getattr(backend.trainer, "should_stop", False)
|
||||
|
||||
# Derive high-level phase
|
||||
if error_message:
|
||||
phase = "error"
|
||||
|
|
@ -326,6 +361,8 @@ async def get_training_status(
|
|||
phase = "configuring"
|
||||
else:
|
||||
phase = "training"
|
||||
elif trainer_stopped:
|
||||
phase = "stopped"
|
||||
elif progress and getattr(progress, "is_completed", False):
|
||||
phase = "completed"
|
||||
elif has_thread:
|
||||
|
|
|
|||
7
studio/backend/utils/inference/__init__.py
Normal file
7
studio/backend/utils/inference/__init__.py
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
"""
|
||||
Inference utility functions
|
||||
"""
|
||||
from utils.inference.inference_config import load_inference_config
|
||||
|
||||
__all__ = ["load_inference_config"]
|
||||
|
||||
65
studio/backend/utils/inference/inference_config.py
Normal file
65
studio/backend/utils/inference/inference_config.py
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
"""
|
||||
Inference configuration loading utilities.
|
||||
|
||||
This module provides functions to load inference parameters (temperature, top_p, top_k, min_p)
|
||||
from model YAML configuration files, with fallback to default.yaml.
|
||||
"""
|
||||
from pathlib import Path
|
||||
from typing import Dict, Any
|
||||
import yaml
|
||||
import logging
|
||||
|
||||
from utils.models.model_config import load_model_defaults
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def load_inference_config(model_identifier: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Load inference configuration parameters for a model.
|
||||
|
||||
This function loads inference parameters (temperature, top_p, top_k, min_p) from the
|
||||
model's YAML configuration file using the same mapping logic as the /config endpoint.
|
||||
If a parameter is missing from the model's config, it falls back to the value in
|
||||
default.yaml.
|
||||
|
||||
Args:
|
||||
model_identifier: Model identifier (e.g., "unsloth/llama-3-8b-bnb-4bit")
|
||||
|
||||
Returns:
|
||||
Dictionary containing inference parameters:
|
||||
{
|
||||
"temperature": float,
|
||||
"top_p": float,
|
||||
"top_k": int,
|
||||
"min_p": float
|
||||
}
|
||||
"""
|
||||
# Load model defaults to get inference parameters
|
||||
model_defaults = load_model_defaults(model_identifier)
|
||||
|
||||
# Load default.yaml for fallback values
|
||||
script_dir = Path(__file__).parent.parent.parent
|
||||
defaults_dir = script_dir / "assets" / "configs" / "model_defaults"
|
||||
default_config_path = defaults_dir / "default.yaml"
|
||||
|
||||
default_inference = {}
|
||||
if default_config_path.exists():
|
||||
try:
|
||||
with open(default_config_path, 'r', encoding='utf-8') as f:
|
||||
default_config = yaml.safe_load(f) or {}
|
||||
default_inference = default_config.get("inference", {})
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load default.yaml: {e}")
|
||||
|
||||
# Extract inference parameters from model config, fallback to defaults
|
||||
model_inference = model_defaults.get("inference", {})
|
||||
inference_config = {
|
||||
"temperature": model_inference.get("temperature", default_inference.get("temperature", 0.7)),
|
||||
"top_p": model_inference.get("top_p", default_inference.get("top_p", 0.95)),
|
||||
"top_k": model_inference.get("top_k", default_inference.get("top_k", -1)),
|
||||
"min_p": model_inference.get("min_p", default_inference.get("min_p", 0.01)),
|
||||
}
|
||||
|
||||
return inference_config
|
||||
|
||||
|
|
@ -12,9 +12,21 @@ async function hasActiveSession(): Promise<boolean> {
|
|||
return refreshSession();
|
||||
}
|
||||
|
||||
async function checkAuthInitialized(): Promise<boolean> {
|
||||
try {
|
||||
const res = await fetch("/api/auth/status");
|
||||
if (!res.ok) return true; // fallback to login on error
|
||||
const data = (await res.json()) as { initialized: boolean };
|
||||
return data.initialized;
|
||||
} catch {
|
||||
return true; // fallback to login on error
|
||||
}
|
||||
}
|
||||
|
||||
export async function requireAuth(): Promise<void> {
|
||||
if (await hasActiveSession()) return;
|
||||
throw redirect({ to: "/login" });
|
||||
const initialized = await checkAuthInitialized();
|
||||
throw redirect({ to: initialized ? "/login" : "/signup" });
|
||||
}
|
||||
|
||||
export async function requireGuest(): Promise<void> {
|
||||
|
|
|
|||
|
|
@ -10,6 +10,26 @@ type RefreshResponse = {
|
|||
refresh_token: string;
|
||||
};
|
||||
|
||||
let isRedirecting = false;
|
||||
|
||||
async function redirectToAuth(): Promise<void> {
|
||||
if (isRedirecting) return;
|
||||
isRedirecting = true;
|
||||
|
||||
let target = "/login";
|
||||
try {
|
||||
const res = await fetch("/api/auth/status");
|
||||
if (res.ok) {
|
||||
const data = (await res.json()) as { initialized: boolean };
|
||||
if (!data.initialized) target = "/signup";
|
||||
}
|
||||
} catch {
|
||||
// Fall through to /login on error
|
||||
}
|
||||
|
||||
window.location.href = target;
|
||||
}
|
||||
|
||||
export async function refreshSession(): Promise<boolean> {
|
||||
const refreshToken = getRefreshToken();
|
||||
if (!refreshToken) return false;
|
||||
|
|
@ -48,7 +68,11 @@ export async function authFetch(
|
|||
if (response.status !== 401) return response;
|
||||
|
||||
const refreshed = await refreshSession();
|
||||
if (!refreshed) return response;
|
||||
if (!refreshed) {
|
||||
clearAuthTokens();
|
||||
void redirectToAuth();
|
||||
return response;
|
||||
}
|
||||
|
||||
const retryHeaders = new Headers(init?.headers);
|
||||
const newToken = getAuthToken();
|
||||
|
|
|
|||
|
|
@ -63,7 +63,18 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null {
|
|||
const response = await fetch("/api/auth/status");
|
||||
if (!response.ok) throw new Error("Failed to load auth status.");
|
||||
const result = (await response.json()) as AuthStatusResponse;
|
||||
if (!canceled) setInitialized(result.initialized);
|
||||
if (!canceled) {
|
||||
setInitialized(result.initialized);
|
||||
// Auto-redirect to the correct page based on init state
|
||||
if (mode === "login" && result.initialized === false) {
|
||||
navigate({ to: "/signup" });
|
||||
return;
|
||||
}
|
||||
if (mode === "signup" && result.initialized === true) {
|
||||
navigate({ to: "/login" });
|
||||
return;
|
||||
}
|
||||
}
|
||||
} catch (err: unknown) {
|
||||
if (!canceled) {
|
||||
setError(err instanceof Error ? err.message : "Failed to load.");
|
||||
|
|
@ -199,6 +210,9 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null {
|
|||
)}
|
||||
</Button>
|
||||
</div>
|
||||
{!isLoginMode && (
|
||||
<p className="text-xs text-muted-foreground">Must be at least 8 characters</p>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{!isLoginMode && (
|
||||
|
|
@ -223,7 +237,7 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null {
|
|||
<Button
|
||||
type="submit"
|
||||
className="w-full"
|
||||
disabled={loading || statusLoading || blockedByState}
|
||||
disabled={loading || statusLoading || blockedByState || (!isLoginMode && password.length < 8)}
|
||||
>
|
||||
{loading ? "Please wait..." : submitLabel}
|
||||
</Button>
|
||||
|
|
|
|||
|
|
@ -223,7 +223,7 @@ export function DatasetSection() {
|
|||
|
||||
<div className="flex flex-col gap-2">
|
||||
<span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
|
||||
Dataset Format
|
||||
Target Format
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild={true}>
|
||||
<button
|
||||
|
|
@ -258,7 +258,7 @@ export function DatasetSection() {
|
|||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="auto">Auto-detect</SelectItem>
|
||||
<SelectItem value="auto">Auto</SelectItem>
|
||||
<SelectItem value="alpaca">Alpaca</SelectItem>
|
||||
<SelectItem value="chatml">ChatML</SelectItem>
|
||||
<SelectItem value="sharegpt">ShareGPT</SelectItem>
|
||||
|
|
|
|||
|
|
@ -1,4 +1,14 @@
|
|||
import { SectionCard } from "@/components/section-card";
|
||||
import {
|
||||
AlertDialog,
|
||||
AlertDialogAction,
|
||||
AlertDialogCancel,
|
||||
AlertDialogContent,
|
||||
AlertDialogDescription,
|
||||
AlertDialogFooter,
|
||||
AlertDialogHeader,
|
||||
AlertDialogTitle,
|
||||
} from "@/components/ui/alert-dialog";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import {
|
||||
Popover,
|
||||
|
|
@ -104,6 +114,7 @@ export function ProgressSection(): ReactElement {
|
|||
);
|
||||
|
||||
const { stopTrainingRun } = useTrainingActions();
|
||||
const [stopDialogOpen, setStopDialogOpen] = useState(false);
|
||||
const localStartAtRef = useRef<number | null>(null);
|
||||
const [, setLocalTick] = useState(0);
|
||||
|
||||
|
|
@ -225,15 +236,39 @@ export function ProgressSection(): ReactElement {
|
|||
</div>
|
||||
</PopoverContent>
|
||||
</Popover>
|
||||
<Button
|
||||
variant="destructive"
|
||||
size="sm"
|
||||
className="h-7 cursor-pointer px-3 text-xs"
|
||||
onClick={() => void stopTrainingRun()}
|
||||
disabled={!runtime.isTrainingRunning}
|
||||
>
|
||||
<HugeiconsIcon icon={StopIcon} className="size-3" /> Stop
|
||||
</Button>
|
||||
<AlertDialog open={stopDialogOpen} onOpenChange={setStopDialogOpen}>
|
||||
<Button
|
||||
variant="destructive"
|
||||
size="sm"
|
||||
className="h-7 cursor-pointer px-3 text-xs"
|
||||
onClick={() => setStopDialogOpen(true)}
|
||||
disabled={!runtime.isTrainingRunning}
|
||||
>
|
||||
<HugeiconsIcon icon={StopIcon} className="size-3" /> Stop
|
||||
</Button>
|
||||
<AlertDialogContent>
|
||||
<AlertDialogHeader>
|
||||
<AlertDialogTitle>Stop Training</AlertDialogTitle>
|
||||
<AlertDialogDescription>
|
||||
Choose how you want to stop the current training run.
|
||||
</AlertDialogDescription>
|
||||
</AlertDialogHeader>
|
||||
<AlertDialogFooter>
|
||||
<AlertDialogCancel>Continue Training</AlertDialogCancel>
|
||||
<AlertDialogAction
|
||||
variant="destructive"
|
||||
onClick={() => void stopTrainingRun(false)}
|
||||
>
|
||||
Cancel Training
|
||||
</AlertDialogAction>
|
||||
<AlertDialogAction
|
||||
onClick={() => void stopTrainingRun(true)}
|
||||
>
|
||||
Stop and Save
|
||||
</AlertDialogAction>
|
||||
</AlertDialogFooter>
|
||||
</AlertDialogContent>
|
||||
</AlertDialog>
|
||||
</div>
|
||||
}
|
||||
>
|
||||
|
|
|
|||
|
|
@ -1,8 +1,12 @@
|
|||
import { Button } from "@/components/ui/button";
|
||||
import {
|
||||
shouldShowTrainingView,
|
||||
useTrainingActions,
|
||||
useTrainingRuntimeLifecycle,
|
||||
useTrainingRuntimeStore,
|
||||
} from "@/features/training";
|
||||
import { ArrowLeft01Icon } from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import type { ReactElement } from "react";
|
||||
import { DatasetSection } from "./sections/dataset-section";
|
||||
import { ModelSection } from "./sections/model-section";
|
||||
|
|
@ -14,12 +18,28 @@ export function StudioPage(): ReactElement {
|
|||
useTrainingRuntimeLifecycle();
|
||||
const showTrainingView = useTrainingRuntimeStore(shouldShowTrainingView);
|
||||
const runtimeMessage = useTrainingRuntimeStore((state) => state.message);
|
||||
const runtimePhase = useTrainingRuntimeStore((state) => state.phase);
|
||||
const isHydratingRuntime = useTrainingRuntimeStore((state) => state.isHydrating);
|
||||
const hasHydratedRuntime = useTrainingRuntimeStore((state) => state.hasHydrated);
|
||||
const { dismissTrainingRun } = useTrainingActions();
|
||||
|
||||
const canGoBack = runtimePhase === "stopped" || runtimePhase === "error";
|
||||
|
||||
return (
|
||||
<div className="min-h-screen bg-background">
|
||||
<main className="mx-auto max-w-7xl px-6 py-4">
|
||||
{canGoBack && (
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="sm"
|
||||
className="mb-2 cursor-pointer gap-1.5 text-muted-foreground"
|
||||
onClick={() => void dismissTrainingRun()}
|
||||
>
|
||||
<HugeiconsIcon icon={ArrowLeft01Icon} className="size-4" />
|
||||
Back to configuration
|
||||
</Button>
|
||||
)}
|
||||
|
||||
{/* Header */}
|
||||
<div className="mb-8 flex flex-col gap-0.5">
|
||||
<h1 className="text-2xl font-semibold tracking-tight">
|
||||
|
|
|
|||
|
|
@ -41,11 +41,22 @@ export async function startTraining(
|
|||
return parseJson<TrainingStartResponse>(response);
|
||||
}
|
||||
|
||||
export async function stopTraining(): Promise<TrainingStopResponse> {
|
||||
const response = await authFetch("/api/train/stop", { method: "POST" });
|
||||
export async function stopTraining(save = true): Promise<TrainingStopResponse> {
|
||||
const response = await authFetch("/api/train/stop", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ save }),
|
||||
});
|
||||
return parseJson<TrainingStopResponse>(response);
|
||||
}
|
||||
|
||||
export async function resetTraining(): Promise<void> {
|
||||
const response = await authFetch("/api/train/reset", { method: "POST" });
|
||||
if (!response.ok) {
|
||||
throw new Error(await readError(response));
|
||||
}
|
||||
}
|
||||
|
||||
export async function getTrainingStatus(): Promise<TrainingStatusResponse> {
|
||||
const response = await authFetch("/api/train/status");
|
||||
return parseJson<TrainingStatusResponse>(response);
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import { useCallback } from "react";
|
||||
import { useTrainingConfigStore } from "../stores/training-config-store";
|
||||
import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
|
||||
import { startTraining, stopTraining } from "../api/train-api";
|
||||
import { startTraining, stopTraining, resetTraining } from "../api/train-api";
|
||||
import { buildTrainingStartPayload } from "../api/mappers";
|
||||
import { syncTrainingRuntimeFromBackend } from "../lib/sync-runtime";
|
||||
import { validateTrainingConfig } from "../lib/validation";
|
||||
|
|
@ -45,12 +45,12 @@ export function useTrainingActions() {
|
|||
}
|
||||
}, []);
|
||||
|
||||
const stopTrainingRun = useCallback(async (): Promise<boolean> => {
|
||||
const stopTrainingRun = useCallback(async (save = true): Promise<boolean> => {
|
||||
const runtimeStore = useTrainingRuntimeStore.getState();
|
||||
runtimeStore.setStartError(null);
|
||||
|
||||
try {
|
||||
await stopTraining();
|
||||
await stopTraining(save);
|
||||
await syncTrainingRuntimeFromBackend();
|
||||
return true;
|
||||
} catch (error) {
|
||||
|
|
@ -61,10 +61,20 @@ export function useTrainingActions() {
|
|||
}
|
||||
}, []);
|
||||
|
||||
const dismissTrainingRun = useCallback(async (): Promise<void> => {
|
||||
useTrainingRuntimeStore.getState().resetRuntime();
|
||||
try {
|
||||
await resetTraining();
|
||||
} catch {
|
||||
// Frontend already reset; backend will catch up on next poll
|
||||
}
|
||||
}, []);
|
||||
|
||||
return {
|
||||
isStarting,
|
||||
startError,
|
||||
startTrainingRun,
|
||||
stopTrainingRun,
|
||||
dismissTrainingRun,
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import { hasAuthToken } from "@/features/auth";
|
||||
import { useEffect } from "react";
|
||||
import {
|
||||
getTrainingMetrics,
|
||||
|
|
@ -48,23 +49,27 @@ export function useTrainingRuntimeLifecycle(): void {
|
|||
};
|
||||
|
||||
const pollMetrics = async () => {
|
||||
if (!hasAuthToken()) return;
|
||||
const gen = runtimeStore.getState().resetGeneration;
|
||||
try {
|
||||
const metrics = await getTrainingMetrics();
|
||||
if (disposed) {
|
||||
if (disposed || runtimeStore.getState().resetGeneration !== gen) {
|
||||
return;
|
||||
}
|
||||
runtimeStore.getState().applyMetrics(metrics);
|
||||
} catch (error) {
|
||||
if (!isAbortError(error) && !disposed) {
|
||||
if (!isAbortError(error) && !disposed && hasAuthToken()) {
|
||||
runtimeStore.getState().setSseConnected(false);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const pollStatus = async () => {
|
||||
if (!hasAuthToken()) return;
|
||||
const gen = runtimeStore.getState().resetGeneration;
|
||||
try {
|
||||
const status = await getTrainingStatus();
|
||||
if (disposed) {
|
||||
if (disposed || runtimeStore.getState().resetGeneration !== gen) {
|
||||
return;
|
||||
}
|
||||
|
||||
|
|
@ -77,7 +82,7 @@ export function useTrainingRuntimeLifecycle(): void {
|
|||
stopStream();
|
||||
}
|
||||
} catch (error) {
|
||||
if (!isAbortError(error) && !disposed) {
|
||||
if (!isAbortError(error) && !disposed && hasAuthToken()) {
|
||||
runtimeStore.getState().setSseConnected(false);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,12 +6,17 @@ import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
|
|||
import type { TrainingStatusResponse } from "../types/runtime";
|
||||
|
||||
export async function syncTrainingRuntimeFromBackend(): Promise<TrainingStatusResponse> {
|
||||
const gen = useTrainingRuntimeStore.getState().resetGeneration;
|
||||
|
||||
const [status, metrics] = await Promise.all([
|
||||
getTrainingStatus(),
|
||||
getTrainingMetrics(),
|
||||
]);
|
||||
|
||||
const runtimeStore = useTrainingRuntimeStore.getState();
|
||||
if (runtimeStore.resetGeneration !== gen) {
|
||||
return status;
|
||||
}
|
||||
runtimeStore.applyStatus(status);
|
||||
runtimeStore.applyMetrics(metrics);
|
||||
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ const initialState: TrainingRuntimeState = {
|
|||
lossHistory: [],
|
||||
lrHistory: [],
|
||||
gradNormHistory: [],
|
||||
resetGeneration: 0,
|
||||
};
|
||||
|
||||
function sortSeries(points: TrainingSeriesPoint[]): TrainingSeriesPoint[] {
|
||||
|
|
@ -95,12 +96,13 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => (
|
|||
setLastEventId: (value) => set({ lastEventId: value }),
|
||||
|
||||
resetRuntime: () =>
|
||||
set({
|
||||
set((state) => ({
|
||||
...initialState,
|
||||
lossHistory: [],
|
||||
lrHistory: [],
|
||||
gradNormHistory: [],
|
||||
}),
|
||||
resetGeneration: state.resetGeneration + 1,
|
||||
})),
|
||||
|
||||
setStartQueued: (jobId, message) =>
|
||||
set({
|
||||
|
|
|
|||
|
|
@ -82,6 +82,7 @@ export interface TrainingRuntimeState {
|
|||
lossHistory: TrainingSeriesPoint[];
|
||||
lrHistory: TrainingSeriesPoint[];
|
||||
gradNormHistory: TrainingSeriesPoint[];
|
||||
resetGeneration: number;
|
||||
}
|
||||
|
||||
export interface TrainingRuntimeActions {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue