ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Final cleanup

Browse source
This commit is contained in:
Roland Tannous 2026-03-12 18:28:04 +00:00
commit 47654cb91c
123 changed files with 7474 additions and 5805 deletions
Download patch file Download diff file Expand all files Collapse all files

240
README.md
View file

@ -1,240 +0,0 @@
<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
### Prerequisites
| Requirement | Linux / WSL | Windows |
|---|---|---|
| **GPU** | NVIDIA GPU with working driver | NVIDIA GPU with working driver |
| **Python** | 3.11 – 3.13 | 3.11 – 3.13 |
| **Git** | Pre-installed on most distros | Auto-installed by setup script (via `winget`) |
| **CMake** | Pre-installed or `sudo apt install cmake` | Auto-installed by setup script (via `winget`) |
| **C++ compiler** | `build-essential` (auto-detected) | Visual Studio Build Tools 2022 (auto-installed by setup script) |
| **CUDA Toolkit** | Optional — setup auto-detects `nvcc` | Auto-installed by setup script (version matched to driver) |
> [!NOTE]
> On **WSL**, the setup script will also run `sudo apt-get install build-essential cmake curl git libcurl4-openssl-dev` so that GGUF export works in non-interactive subprocesses. You may be prompted for your password during setup.
---
### Linux / Windows WSL
```bash
# 1. Clone the repo
git clone https://github.com/unslothai/unsloth-studio.git
cd unsloth-studio
# 2. Run setup (installs Node, builds frontend, creates .venv, builds llama.cpp)
bash setup.sh
# 3. Open a new terminal (or source your shell rc), then launch:
unsloth-studio -H 0.0.0.0 -p 8000
```
<details>
<summary><b>What does <code>setup.sh</code> do?</b></summary>
1. Installs **Node.js ≥ 20** via nvm (if needed)
2. Runs `npm install && npm run build` for the React frontend
3. Detects the best **Python 3.11 – 3.13** on your system and creates a `.venv`
4. Installs all Python dependencies (unsloth, PyTorch with CUDA, triton kernels, etc.)
5. On **WSL**: pre-installs build dependencies via `apt-get`
6. Clones and builds **llama.cpp** at `~/.unsloth/llama.cpp` (GPU-accelerated if CUDA is found)
7. Registers `unsloth-studio` and `unsloth-ui` shell aliases in your shell rc (bash, zsh, fish, or ksh)
</details>
---
### Windows (Native)
> [!IMPORTANT]
> Requires an **NVIDIA GPU** — CPU-only machines are not supported on Windows.
```powershell
# 1. Clone the repo
git clone https://github.com/unslothai/unsloth-studio.git
cd unsloth-studio
# 2. Run setup (Right-click → "Run with PowerShell", or from a terminal):
.\setup.bat
# Or directly:
powershell -ExecutionPolicy Bypass -File setup.ps1
```
After setup completes, **open a new terminal** and run:
```powershell
# PowerShell
unsloth-studio -H 0.0.0.0 -p 8000
# Or cmd.exe
unsloth-studio -H 0.0.0.0 -p 8000
```
<details>
<summary><b>What does <code>setup.ps1</code> do?</b></summary>
1. Enables **Windows Long Paths** (required for deep dependency trees — prompts for UAC)
2. Auto-installs missing system tools via `winget`: **Git**, **CMake**, **Visual Studio Build Tools 2022**, **CUDA Toolkit** (version-matched to your driver), **Node.js LTS**, **Python 3.12**, **OpenSSL dev**
3. Builds the React frontend (`npm install && npm run build`)
4. Creates a `.venv` and installs all Python dependencies (including CUDA-enabled PyTorch from the official index)
5. Sets `TORCHINDUCTOR_CACHE_DIR=C:\tc` to avoid Windows MAX_PATH issues with Triton
6. Clones and builds **llama.cpp** at `%USERPROFILE%\.unsloth\llama.cpp` with CUDA + Visual Studio
7. Registers `unsloth-studio` and `unsloth-ui` commands in both PowerShell profile and `cmd.exe` (via batch files on PATH)
</details>
---
### Google Colab
The setup script auto-detects Colab and installs everything into the existing system Python (no venv):
```python
!bash setup.sh
```
---
### Launching the Studio
After setup on any platform, the command is the same:
```bash
unsloth-studio -H 0.0.0.0 -p 8000
```
| Flag | Description |
|---|---|
| `-H` / `--host` | Bind address (`0.0.0.0` for all interfaces, `127.0.0.1` for local only) |
| `-p` / `--port` | Port number (default: `8000`) |
On **first launch**, a one-time setup token is printed to the console. Open the URL shown in your browser and use this token to create your admin account.
> [!TIP]
> This repo is in active development. After pulling new changes, **always re-run the setup script** (`bash setup.sh` or `.\setup.bat`) to pick up dependency and build updates.
## 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 # Bootstrap script (Linux / WSL / Colab)
├── setup.ps1 # Bootstrap script (Windows native)
├── setup.bat # Wrapper to launch setup.ps1 via double-click
├── install_python_stack.py # Cross-platform Python dependency installer
└── 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.

6
cli/__init__.py
View file

@ -10,8 +10,8 @@ from cli.commands.ui import ui
from cli.commands.studio import studio_app
app = typer.Typer(
help="Command-line interface for Unsloth training, inference, and export.",
context_settings={"help_option_names": ["-h", "--help"]},
help = "Command-line interface for Unsloth training, inference, and export.",
context_settings = {"help_option_names": ["-h", "--help"]},
)
app.command()(train)
@ -19,4 +19,4 @@ app.command()(inference)
app.command()(export)
app.command("list-checkpoints")(list_checkpoints)
app.command()(ui)
app.add_typer(studio_app, name="studio", help="Unsloth Studio commands.")
app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.")

2
cli/commands/__init__.py
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

90
cli/commands/export.py
View file

@ -13,14 +13,14 @@ GGUF_QUANTS = ["q4_k_m", "q5_k_m", "q8_0", "f16"]
def list_checkpoints(
outputs_dir: Path = typer.Option(
Path("./outputs"), "--outputs-dir", help="Directory that holds training runs."
Path("./outputs"), "--outputs-dir", help = "Directory that holds training runs."
),
):
"""List checkpoints detected in the outputs directory."""
from studio.backend.core.export import ExportBackend
backend = ExportBackend()
checkpoints = backend.scan_checkpoints(outputs_dir=str(outputs_dir))
checkpoints = backend.scan_checkpoints(outputs_dir = str(outputs_dir))
if not checkpoints:
typer.echo("No checkpoints found.")
raise typer.Exit()
@ -33,96 +33,100 @@ def list_checkpoints(
def export(
checkpoint: Path = typer.Argument(..., help="Path to checkpoint directory."),
output_dir: Path = typer.Argument(..., help="Directory to save exported model."),
checkpoint: Path = typer.Argument(..., help = "Path to checkpoint directory."),
output_dir: Path = typer.Argument(..., help = "Directory to save exported model."),
format: str = typer.Option(
"merged-16bit",
"--format",
"-f",
help=f"Export format: {', '.join(EXPORT_FORMATS)}",
help = f"Export format: {', '.join(EXPORT_FORMATS)}",
),
quantization: str = typer.Option(
"q4_k_m",
"--quantization",
"-q",
help=f"GGUF quantization method: {', '.join(GGUF_QUANTS)}",
help = f"GGUF quantization method: {', '.join(GGUF_QUANTS)}",
),
push_to_hub: bool = typer.Option(
False, "--push-to-hub", help="Push exported model to HuggingFace Hub."
False, "--push-to-hub", help = "Push exported model to HuggingFace Hub."
),
repo_id: Optional[str] = typer.Option(
None, "--repo-id", help="HuggingFace repo ID (username/model-name)."
None, "--repo-id", help = "HuggingFace repo ID (username/model-name)."
),
hf_token: Optional[str] = typer.Option(
None, "--hf-token", envvar="HF_TOKEN", help="HuggingFace token."
None, "--hf-token", envvar = "HF_TOKEN", help = "HuggingFace token."
),
private: bool = typer.Option(
False, "--private", help="Make the HuggingFace repo private."
False, "--private", help = "Make the HuggingFace repo private."
),
max_seq_length: int = typer.Option(2048, "--max-seq-length"),
load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
):
"""Export a checkpoint to various formats (merged, GGUF, LoRA adapter)."""
if format not in EXPORT_FORMATS:
typer.echo(f"Error: Invalid format '{format}'. Choose from: {', '.join(EXPORT_FORMATS)}", err=True)
raise typer.Exit(code=2)
typer.echo(
f"Error: Invalid format '{format}'. Choose from: {', '.join(EXPORT_FORMATS)}",
err = True,
)
raise typer.Exit(code = 2)
if push_to_hub and not repo_id:
typer.echo("Error: --repo-id required when using --push-to-hub", err=True)
raise typer.Exit(code=2)
typer.echo("Error: --repo-id required when using --push-to-hub", err = True)
raise typer.Exit(code = 2)
from studio.backend.core.export import ExportBackend
backend = ExportBackend()
typer.echo(f"Loading checkpoint: {checkpoint}")
success, message = backend.load_checkpoint(
checkpoint_path=str(checkpoint),
max_seq_length=max_seq_length,
load_in_4bit=load_in_4bit,
checkpoint_path = str(checkpoint),
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
)
if not success:
typer.echo(f"Error: {message}", err=True)
raise typer.Exit(code=1)
typer.echo(f"Error: {message}", err = True)
raise typer.Exit(code = 1)
typer.echo(message)
typer.echo(f"Exporting as {format}...")
if format == "merged-16bit":
success, message = backend.export_merged_model(
save_directory=str(output_dir),
format_type="16-bit (FP16)",
push_to_hub=push_to_hub,
repo_id=repo_id,
hf_token=hf_token,
private=private,
save_directory = str(output_dir),
format_type = "16-bit (FP16)",
push_to_hub = push_to_hub,
repo_id = repo_id,
hf_token = hf_token,
private = private,
)
elif format == "merged-4bit":
success, message = backend.export_merged_model(
save_directory=str(output_dir),
format_type="4-bit (FP4)",
push_to_hub=push_to_hub,
repo_id=repo_id,
hf_token=hf_token,
private=private,
save_directory = str(output_dir),
format_type = "4-bit (FP4)",
push_to_hub = push_to_hub,
repo_id = repo_id,
hf_token = hf_token,
private = private,
)
elif format == "gguf":
success, message = backend.export_gguf(
save_directory=str(output_dir),
quantization_method=quantization.upper(),
push_to_hub=push_to_hub,
repo_id=repo_id,
hf_token=hf_token,
save_directory = str(output_dir),
quantization_method = quantization.upper(),
push_to_hub = push_to_hub,
repo_id = repo_id,
hf_token = hf_token,
)
elif format == "lora":
success, message = backend.export_lora_adapter(
save_directory=str(output_dir),
push_to_hub=push_to_hub,
repo_id=repo_id,
hf_token=hf_token,
private=private,
save_directory = str(output_dir),
push_to_hub = push_to_hub,
repo_id = repo_id,
hf_token = hf_token,
private = private,
)
if not success:
typer.echo(f"Error: {message}", err=True)
raise typer.Exit(code=1)
typer.echo(f"Error: {message}", err = True)
raise typer.Exit(code = 1)
typer.echo(message)

44
cli/commands/inference.py
View file

@ -8,10 +8,10 @@ import typer
def inference(
model: str = typer.Argument(..., help="HF model id or local path."),
prompt: str = typer.Argument(..., help="Prompt to send to the model."),
model: str = typer.Argument(..., help = "HF model id or local path."),
prompt: str = typer.Argument(..., help = "Prompt to send to the model."),
hf_token: Optional[str] = typer.Option(
None, "--hf-token", envvar="HF_TOKEN", help="Hugging Face token if needed."
None, "--hf-token", envvar = "HF_TOKEN", help = "Hugging Face token if needed."
),
temperature: float = typer.Option(0.7, "--temperature"),
top_p: float = typer.Option(0.9, "--top-p"),
@ -21,7 +21,7 @@ def inference(
system_prompt: str = typer.Option(
"",
"--system-prompt",
help="Optional system prompt to prepend.",
help = "Optional system prompt to prepend.",
),
max_seq_length: int = typer.Option(2048, "--max-seq-length"),
load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
@ -31,36 +31,36 @@ def inference(
inference_backend = get_inference_backend()
model_config = ModelConfig.from_ui_selection(
dropdown_value=model, search_value=None, hf_token=hf_token, is_lora=False
dropdown_value = model, search_value = None, hf_token = hf_token, is_lora = False
)
if not model_config:
typer.echo("Could not resolve model config", err=True)
raise typer.Exit(code=1)
typer.echo("Could not resolve model config", err = True)
raise typer.Exit(code = 1)
if not inference_backend.load_model(
config=model_config,
max_seq_length=max_seq_length,
load_in_4bit=load_in_4bit,
hf_token=hf_token,
config = model_config,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
hf_token = hf_token,
):
typer.echo("Model load failed", err=True)
raise typer.Exit(code=1)
typer.echo("Model load failed", err = True)
raise typer.Exit(code = 1)
messages = [{"role": "user", "content": prompt}]
stream = inference_backend.generate_chat_response(
messages=messages,
system_prompt=system_prompt,
temperature=temperature,
top_p=top_p,
top_k=top_k,
max_new_tokens=max_new_tokens,
repetition_penalty=repetition_penalty,
messages = messages,
system_prompt = system_prompt,
temperature = temperature,
top_p = top_p,
top_k = top_k,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
)
typer.echo("Assistant:", nl=True)
typer.echo("Assistant:", nl = True)
previous = ""
for chunk in stream:
delta = chunk[len(previous):]
delta = chunk[len(previous) :]
if delta:
sys.stdout.write(delta)
sys.stdout.flush()

249
cli/commands/studio.py
View file

@ -10,28 +10,44 @@ from pathlib import Path
from typing import Optional
import typer
studio_app = typer.Typer(help="Unsloth Studio commands.")
studio_app = typer.Typer(help = "Unsloth Studio commands.")
STUDIO_HOME = Path.home() / ".unsloth" / "studio"
# __file__ is cli/commands/studio.py — two parents up is the package root
# (either site-packages or the repo root for editable installs).
_PACKAGE_ROOT = Path(__file__).resolve().parent.parent.parent
def _is_repo_root(path: Path) -> bool:
"""Check if a directory looks like the repo root (actual git clone, not site-packages)."""
return (
(path / ".git").exists()
and (path / "pyproject.toml").is_file()
and (
(path / "studio" / "setup.sh").is_file()
or (path / "studio" / "setup.ps1").is_file()
)
)
def _get_repo_root() -> Optional[Path]:
"""Find the git clone repo root, or None if pure pip install."""
"""Find the git clone repo root.
Used only by setup() — checks __file__ first (editable install),
then walks CWD parents (wheel install, user is inside the clone).
"""
# Check 1: __file__ is in the repo (editable install)
candidate = Path(__file__).resolve().parent.parent.parent
if (candidate / "pyproject.toml").is_file() and (candidate / "studio" / "setup.sh").is_file():
return candidate
# Check 2: CWD is the repo (non-editable wheel, running from repo dir)
cwd = Path.cwd()
if (cwd / "pyproject.toml").is_file() and (cwd / "studio" / "setup.sh").is_file():
return cwd
if _is_repo_root(_PACKAGE_ROOT):
return _PACKAGE_ROOT
# Check 2: CWD or any parent is the repo
cwd = Path.cwd().resolve()
for parent in (cwd, *cwd.parents):
if _is_repo_root(parent):
return parent
return None
def _is_git_clone() -> bool:
return _get_repo_root() is not None
def _studio_venv_python() -> Optional[Path]:
"""Return the studio venv Python binary, or None if not set up."""
if platform.system() == "Windows":
@ -42,37 +58,69 @@ def _studio_venv_python() -> Optional[Path]:
def _find_run_py() -> Optional[Path]:
"""Find studio/backend/run.py."""
# 1. Repo root (git clone / editable)
repo = _get_repo_root()
if repo:
run_py = repo / "studio" / "backend" / "run.py"
if run_py.is_file():
return run_py
# 2. Studio venv's site-packages
for match in (STUDIO_HOME / ".venv").glob("lib/python*/site-packages/studio/backend/run.py"):
return match
# 3. Current package's site-packages
run_py = Path(__file__).resolve().parent.parent.parent / "studio" / "backend" / "run.py"
return run_py if run_py.is_file() else None
"""Find studio/backend/run.py.
No CWD dependency — works from any directory.
Since studio/ is now a proper package (has __init__.py), it lives in
site-packages after pip install, right next to cli/.
"""
# 1. Relative to __file__ (site-packages or editable repo root)
run_py = _PACKAGE_ROOT / "studio" / "backend" / "run.py"
if run_py.is_file():
return run_py
# 2. Studio venv's site-packages (Linux + Windows layouts)
for pattern in (
"lib/python*/site-packages/studio/backend/run.py",
"Lib/site-packages/studio/backend/run.py",
):
for match in (STUDIO_HOME / ".venv").glob(pattern):
return match
return None
def _find_install_script() -> Optional[Path]:
"""Find studio/install_python_stack.py."""
# 1. Repo root
repo = _get_repo_root()
if repo:
s = repo / "studio" / "install_python_stack.py"
if s.is_file():
return s
# 2. Relative to __file__ (in site-packages)
s = Path(__file__).resolve().parent.parent.parent / "studio" / "install_python_stack.py"
return s if s.is_file() else None
"""Find studio/install_python_stack.py.
No CWD dependency — works from any directory.
"""
# 1. Relative to __file__ (site-packages or editable repo root)
s = _PACKAGE_ROOT / "studio" / "install_python_stack.py"
if s.is_file():
return s
# 2. Studio venv's site-packages
for pattern in (
"lib/python*/site-packages/studio/install_python_stack.py",
"Lib/site-packages/studio/install_python_stack.py",
):
for match in (STUDIO_HOME / ".venv").glob(pattern):
return match
return None
def _find_setup_script() -> Optional[Path]:
"""Find studio/setup.sh or studio/setup.ps1.
No CWD dependency — works from any directory.
"""
name = "setup.ps1" if platform.system() == "Windows" else "setup.sh"
# 1. Relative to __file__ (site-packages or editable repo root)
s = _PACKAGE_ROOT / "studio" / name
if s.is_file():
return s
# 2. Studio venv's site-packages
for pattern in (
f"lib/python*/site-packages/studio/{name}",
f"Lib/site-packages/studio/{name}",
):
for match in (STUDIO_HOME / ".venv").glob(pattern):
return match
return None
# ── unsloth studio (server) ──────────────────────────────────────────
@studio_app.callback(invoke_without_command=True)
@studio_app.callback(invoke_without_command = True)
def studio_default(
ctx: typer.Context,
port: int = typer.Option(8000, "--port", "-p"),
@ -94,7 +142,14 @@ def studio_default(
if studio_python and run_py:
if not silent:
typer.echo("Launching with studio venv...")
args = [str(studio_python), str(run_py), "--host", host, "--port", str(port)]
args = [
str(studio_python),
str(run_py),
"--host",
host,
"--port",
str(port),
]
if frontend:
args.extend(["--frontend", str(frontend)])
if silent:
@ -108,14 +163,15 @@ def studio_default(
if not silent:
from studio.backend.run import _resolve_external_ip
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}")
run_server(
host=host,
port=port,
frontend_path=frontend,
silent=silent,
host = host,
port = port,
frontend_path = frontend,
silent = silent,
)
try:
@ -127,25 +183,30 @@ def studio_default(
# ── unsloth studio setup ─────────────────────────────────────────────
@studio_app.command()
def setup():
"""Run one-time Studio environment setup."""
if _is_git_clone():
_dev_setup()
# If we're inside a git clone, use the full setup script (builds frontend, etc.)
repo = _get_repo_root()
if repo:
_dev_setup(repo)
else:
_pip_setup()
def _dev_setup():
def _dev_setup(repo_root: Path):
"""Git-clone: run setup.sh / setup.ps1."""
repo_root = _get_repo_root()
studio_dir = repo_root / "studio"
if platform.system() == "Windows":
script = studio_dir / "setup.ps1"
subprocess.run(["powershell", "-ExecutionPolicy", "Bypass", "-File", str(script)], check=True)
subprocess.run(
["powershell", "-ExecutionPolicy", "Bypass", "-File", str(script)],
check = True,
)
else:
script = studio_dir / "setup.sh"
subprocess.run(["bash", str(script)], check=True)
subprocess.run(["bash", str(script)], check = True)
def _pip_setup():
@ -167,14 +228,14 @@ def _pip_setup():
# 1. Create venv
if not venv_python.is_file():
typer.echo(f" Creating venv at {venv_dir}...")
STUDIO_HOME.mkdir(parents=True, exist_ok=True)
_venv.create(str(venv_dir), with_pip=True)
STUDIO_HOME.mkdir(parents = True, exist_ok = True)
_venv.create(str(venv_dir), with_pip = True)
# 2. Install all Python deps via install_python_stack.py
install_script = _find_install_script()
if install_script:
typer.echo(" Installing Python dependencies...")
subprocess.run([str(venv_python), str(install_script)], check=True)
subprocess.run([str(venv_python), str(install_script)], check = True)
else:
typer.echo("Error: Could not find install_python_stack.py")
raise typer.Exit(1)
@ -184,16 +245,28 @@ def _pip_setup():
typer.echo(f" Transformers 5.x overlay already at {venv_t5_dir}")
else:
typer.echo(" Installing transformers 5.x overlay...")
venv_t5_dir.mkdir(parents=True, exist_ok=True)
venv_t5_dir.mkdir(parents = True, exist_ok = True)
subprocess.run(
[str(venv_pip), "install",
"--target", str(venv_t5_dir), "--no-deps", "transformers==5.2.0"],
check=True,
[
str(venv_pip),
"install",
"--target",
str(venv_t5_dir),
"--no-deps",
"transformers==5.2.0",
],
check = True,
)
subprocess.run(
[str(venv_pip), "install",
"--target", str(venv_t5_dir), "--no-deps", "huggingface_hub==1.3.0"],
check=True,
[
str(venv_pip),
"install",
"--target",
str(venv_t5_dir),
"--no-deps",
"huggingface_hub==1.3.0",
],
check = True,
)
typer.echo(f" Installed to {venv_t5_dir}")
@ -221,12 +294,26 @@ def _build_llama_cpp():
typer.echo(" Building llama.cpp for GGUF inference...")
if llama_dir.exists():
shutil.rmtree(llama_dir)
unsloth_home.mkdir(parents=True, exist_ok=True)
# necessary because shutil.rmtree fails on Windows because .git pack files are read-only
def _force_remove_readonly(func, path, exc_info):
"""Clear read-only flag and retry — needed on Windows for .git pack files."""
import stat
os.chmod(path, stat.S_IWRITE)
func(path)
shutil.rmtree(llama_dir, onerror = _force_remove_readonly)
unsloth_home.mkdir(parents = True, exist_ok = True)
result = subprocess.run(
["git", "clone", "--depth", "1", "https://github.com/ggml-org/llama.cpp.git", str(llama_dir)],
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
[
"git",
"clone",
"--depth",
"1",
"https://github.com/ggml-org/llama.cpp.git",
str(llama_dir),
],
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
)
if result.returncode != 0:
typer.echo(" Failed to clone llama.cpp")
@ -245,7 +332,8 @@ def _build_llama_cpp():
build_dir = llama_dir / "build"
result = subprocess.run(
["cmake", "-S", str(llama_dir), "-B", str(build_dir)] + cmake_args,
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
)
if result.returncode != 0:
typer.echo(" cmake configure failed")
@ -253,21 +341,42 @@ def _build_llama_cpp():
ncpu = str(os.cpu_count() or 4)
result = subprocess.run(
["cmake", "--build", str(build_dir), "--config", "Release",
"--target", "llama-server", f"-j{ncpu}"],
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
[
"cmake",
"--build",
str(build_dir),
"--config",
"Release",
"--target",
"llama-server",
f"-j{ncpu}",
],
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
)
if result.returncode != 0:
typer.echo(" llama-server build failed")
return
subprocess.run(
["cmake", "--build", str(build_dir), "--config", "Release",
"--target", "llama-quantize", f"-j{ncpu}"],
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
[
"cmake",
"--build",
str(build_dir),
"--config",
"Release",
"--target",
"llama-quantize",
f"-j{ncpu}",
],
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
)
server_bin = build_dir / "bin" / "llama-server"
if sys.platform == "win32":
server_bin = build_dir / "bin" / "Release" / "llama-server.exe"
else:
server_bin = build_dir / "bin" / "llama-server"
if server_bin.is_file():
typer.echo(f" llama-server built at {server_bin}")
else:

74
cli/commands/train.py
View file

@ -17,18 +17,18 @@ def train(
None,
"--config",
"-c",
help="Path to YAML/JSON config file. CLI flags override config values.",
help = "Path to YAML/JSON config file. CLI flags override config values.",
),
hf_token: Optional[str] = typer.Option(
None, "--hf-token", envvar="HF_TOKEN", help="Hugging Face token if needed."
None, "--hf-token", envvar = "HF_TOKEN", help = "Hugging Face token if needed."
),
wandb_token: Optional[str] = typer.Option(
None, "--wandb-token", envvar="WANDB_API_KEY", help="Weights & Biases API key."
None, "--wandb-token", envvar = "WANDB_API_KEY", help = "Weights & Biases API key."
),
dry_run: bool = typer.Option(
False,
"--dry-run",
help="Show resolved config and exit without training.",
help = "Show resolved config and exit without training.",
),
config_overrides: dict = None,
):
@ -36,14 +36,15 @@ def train(
try:
cfg = load_config(config)
except FileNotFoundError as e:
typer.echo(f"Error: {e}", err=True)
raise typer.Exit(code=2)
typer.echo(f"Error: {e}", err = True)
raise typer.Exit(code = 2)
cfg.apply_overrides(**config_overrides)
# CLI/env tokens take precedence over config
# Handle case where typer.Option isn't resolved (decorator interaction)
from typer.models import OptionInfo
if isinstance(hf_token, OptionInfo):
hf_token = None
if isinstance(wandb_token, OptionInfo):
@ -53,33 +54,38 @@ def train(
if dry_run:
import yaml
data = cfg.model_dump()
data["training"]["output_dir"] = str(data["training"]["output_dir"])
typer.echo(yaml.dump(data, default_flow_style=False, sort_keys=False))
raise typer.Exit(code=0)
typer.echo(yaml.dump(data, default_flow_style = False, sort_keys = False))
raise typer.Exit(code = 0)
if not cfg.model:
typer.echo("Error: provide --model or set model in --config", err=True)
raise typer.Exit(code=2)
typer.echo("Error: provide --model or set model in --config", err = True)
raise typer.Exit(code = 2)
if not cfg.data.dataset and not cfg.data.local_dataset:
typer.echo(
"Error: provide --dataset or --local-dataset (or via --config)", err=True
"Error: provide --dataset or --local-dataset (or via --config)", err = True
)
raise typer.Exit(code=2)
raise typer.Exit(code = 2)
# Check if the model path is a LoRA adapter (has adapter_config.json)
model_path = Path(cfg.model) if cfg.model else None
model_is_lora = model_path and model_path.is_dir() and (model_path / "adapter_config.json").exists()
model_is_lora = (
model_path
and model_path.is_dir()
and (model_path / "adapter_config.json").exists()
)
use_lora = cfg.training.training_type.lower() == "lora"
if model_is_lora and not use_lora:
typer.echo(
"Error: Cannot do full finetuning on a LoRA adapter. "
"Use --training-type lora or provide a base model.",
err=True,
err = True,
)
raise typer.Exit(code=2)
raise typer.Exit(code = 2)
from studio.backend.core.training.trainer import UnslothTrainer
@ -87,38 +93,40 @@ def train(
# Load model (trainer.is_vlm is set after this)
if not trainer.load_model(
model_name=cfg.model,
max_seq_length=cfg.training.max_seq_length,
load_in_4bit=cfg.training.load_in_4bit if use_lora else False,
hf_token=hf_token,
model_name = cfg.model,
max_seq_length = cfg.training.max_seq_length,
load_in_4bit = cfg.training.load_in_4bit if use_lora else False,
hf_token = hf_token,
):
typer.echo("Model load failed", err=True)
raise typer.Exit(code=1)
typer.echo("Model load failed", err = True)
raise typer.Exit(code = 1)
is_vision = trainer.is_vlm
if not trainer.prepare_model_for_training(**cfg.model_kwargs(use_lora, is_vision)):
typer.echo("Model preparation failed", err=True)
raise typer.Exit(code=1)
typer.echo("Model preparation failed", err = True)
raise typer.Exit(code = 1)
result = trainer.load_and_format_dataset(
dataset_source=cfg.data.dataset or "",
format_type=cfg.data.format_type,
local_datasets=cfg.data.local_dataset,
dataset_source = cfg.data.dataset or "",
format_type = cfg.data.format_type,
local_datasets = cfg.data.local_dataset,
)
if result is None:
typer.echo("Dataset load failed", err=True)
raise typer.Exit(code=1)
typer.echo("Dataset load failed", err = True)
raise typer.Exit(code = 1)
ds, eval_ds = result
training_kwargs = cfg.training_kwargs()
training_kwargs["wandb_token"] = wandb_token # CLI/env takes precedence
started = trainer.start_training(dataset=ds, eval_dataset=eval_ds, **training_kwargs)
started = trainer.start_training(
dataset = ds, eval_dataset = eval_ds, **training_kwargs
)
if not started:
typer.echo("Training failed to start", err=True)
raise typer.Exit(code=1)
typer.echo("Training failed to start", err = True)
raise typer.Exit(code = 1)
try:
while trainer.training_thread and trainer.training_thread.is_alive():
@ -132,5 +140,5 @@ def train(
final = trainer.get_training_progress()
if getattr(final, "error", None):
typer.echo(f"Training error: {final.error}", err=True)
raise typer.Exit(code=1)
typer.echo(f"Training error: {final.error}", err = True)
raise typer.Exit(code = 1)

59
cli/commands/ui.py
View file

@ -1,6 +1,8 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import os
import sys
import time
from pathlib import Path
from typing import Optional
@ -9,27 +11,64 @@ import typer
def ui(
port: int = typer.Option(8000, "--port", "-p", help="Port to run the UI server on."),
host: str = typer.Option("0.0.0.0", "--host", "-H", help="Host address to bind to."),
frontend: Optional[Path] = typer.Option(None, "--frontend", "-f", help="Path to frontend build directory."),
silent: bool = typer.Option(False, "--silent", "-q", help="Suppress startup messages."),
port: int = typer.Option(
8000, "--port", "-p", help = "Port to run the UI server on."
),
host: str = typer.Option(
"0.0.0.0", "--host", "-H", help = "Host address to bind to."
),
frontend: Optional[Path] = typer.Option(
None, "--frontend", "-f", help = "Path to frontend build directory."
),
silent: bool = typer.Option(
False, "--silent", "-q", help = "Suppress startup messages."
),
):
"""Launch the Unsloth web UI backend server."""
"""Launch the Unsloth web UI backend server (alias for 'unsloth studio')."""
from cli.commands.studio import _studio_venv_python, _find_run_py, STUDIO_HOME
# Re-execute in studio venv if available and not already inside it
studio_venv_dir = STUDIO_HOME / ".venv"
in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
if not in_studio_venv:
studio_python = _studio_venv_python()
run_py = _find_run_py()
if studio_python and run_py:
if not silent:
typer.echo("Launching with studio venv...")
args = [
str(studio_python),
str(run_py),
"--host",
host,
"--port",
str(port),
]
if frontend:
args.extend(["--frontend", str(frontend)])
if silent:
args.append("--silent")
os.execvp(str(studio_python), args)
else:
typer.echo("Studio not set up. Run 'unsloth studio setup' first.")
raise typer.Exit(1)
from studio.backend.run import run_server
if not silent:
from studio.backend.run import _resolve_external_ip
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}")
run_server(
host=host,
port=port,
frontend_path=frontend,
silent=silent,
host = host,
port = port,
frontend_path = frontend,
silent = silent,
)
# Keep running until interrupted
try:
while True:
time.sleep(1)

17
cli/config.py
View file

@ -58,10 +58,10 @@ class LoggingConfig(BaseModel):
class Config(BaseModel):
model: Optional[str] = None
data: DataConfig = Field(default_factory=DataConfig)
training: TrainingConfig = Field(default_factory=TrainingConfig)
lora: LoraConfig = Field(default_factory=LoraConfig)
logging: LoggingConfig = Field(default_factory=LoggingConfig)
data: DataConfig = Field(default_factory = DataConfig)
training: TrainingConfig = Field(default_factory = TrainingConfig)
lora: LoraConfig = Field(default_factory = LoraConfig)
logging: LoggingConfig = Field(default_factory = LoggingConfig)
def apply_overrides(self, **kwargs):
"""Apply CLI overrides by matching arg names to config fields."""
@ -83,7 +83,11 @@ class Config(BaseModel):
# Vision models expect a string (e.g., "all-linear"); fall back to None to use trainer defaults
target_modules = "all-linear" if self.lora.vision_all_linear else None
else:
parsed = [m.strip() for m in str(self.lora.target_modules).split(",") if m and m.strip()]
parsed = [
m.strip()
for m in str(self.lora.target_modules).split(",")
if m and m.strip()
]
target_modules = parsed or None
return {
@ -134,11 +138,12 @@ def load_config(path: Optional[Path]) -> Config:
if not path.exists():
raise FileNotFoundError(f"Config file not found: {path}")
text = path.read_text(encoding="utf-8")
text = path.read_text(encoding = "utf-8")
if path.suffix.lower() in {".yaml", ".yml"}:
data = yaml.safe_load(text) or {}
else:
import json
data = json.loads(text or "{}")
return Config(**data)

18
cli/options.py
View file

@ -80,7 +80,9 @@ def add_options_from_config(config_class: type[BaseModel]) -> Callable:
which will receive a dict of all CLI-provided config values.
"""
fields = _collect_config_fields(config_class)
field_names = {name for name, field_info in fields if not _is_list_type(field_info.annotation)}
field_names = {
name for name, field_info in fields if not _is_list_type(field_info.annotation)
}
def decorator(func: Callable) -> Callable:
sig = inspect.signature(func)
@ -105,22 +107,22 @@ def add_options_from_config(config_class: type[BaseModel]) -> Callable:
default = typer.Option(
None,
f"{flag_name}/--no-{field_name.replace('_', '-')}",
help=help_text,
help = help_text,
)
param = inspect.Parameter(
field_name,
inspect.Parameter.POSITIONAL_OR_KEYWORD,
default=default,
annotation=Optional[bool],
default = default,
annotation = Optional[bool],
)
else:
py_type = _get_python_type(annotation)
default = typer.Option(None, flag_name, help=help_text)
default = typer.Option(None, flag_name, help = help_text)
param = inspect.Parameter(
field_name,
inspect.Parameter.POSITIONAL_OR_KEYWORD,
default=default,
annotation=Optional[py_type],
default = default,
annotation = Optional[py_type],
)
new_params.append(param)
@ -129,7 +131,7 @@ def add_options_from_config(config_class: type[BaseModel]) -> Callable:
if param.name != "config_overrides":
new_params.append(param)
new_sig = sig.replace(parameters=new_params)
new_sig = sig.replace(parameters = new_params)
@functools.wraps(func)
def wrapper(*args, **kwargs):

12
Unsloth_Studio_Colab.ipynb → studio/Unsloth_Studio_Colab.ipynb
View file

@ -1,5 +1,17 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "f2b0c6a1",
"metadata": {},
"source": [
"**License Notice**\n",
"\n",
"SPDX-License-Identifier: AGPL-3.0-only\n",
"\n",
"Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0\n"
]
},
{
"cell_type": "code",
"execution_count": null,

2
studio/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

2
studio/backend/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

2
studio/backend/assets/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

2
studio/backend/assets/configs/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

4
studio/backend/auth/__init__.py
View file

@ -4,6 +4,7 @@
"""
Authentication module for JWT-based auth with SQLite storage.
"""
from .authentication import (
create_access_token,
create_refresh_token,
@ -44,6 +45,3 @@ __all__ = [
"hash_password",
"verify_password",
]

20
studio/backend/auth/authentication.py
View file

@ -38,10 +38,10 @@ def create_access_token(
"""
to_encode = {"sub": subject}
expire = datetime.now(timezone.utc) + (
expires_delta or timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
expires_delta or timedelta(minutes = ACCESS_TOKEN_EXPIRE_MINUTES)
)
to_encode.update({"exp": expire})
return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
return jwt.encode(to_encode, SECRET_KEY, algorithm = ALGORITHM)
def create_refresh_token(subject: str) -> str:
@ -51,7 +51,7 @@ def create_refresh_token(subject: str) -> str:
Refresh tokens are opaque (not JWTs) and expire after REFRESH_TOKEN_EXPIRE_DAYS.
"""
token = secrets.token_urlsafe(48)
expires_at = datetime.now(timezone.utc) + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
expires_at = datetime.now(timezone.utc) + timedelta(days = REFRESH_TOKEN_EXPIRE_DAYS)
save_refresh_token(token, subject, expires_at.isoformat())
return token
@ -66,13 +66,13 @@ def refresh_access_token(refresh_token: str) -> Optional[str]:
username = verify_refresh_token(refresh_token)
if username is None:
return None
return create_access_token(subject=username)
return create_access_token(subject = username)
def reload_secret() -> None:
"""
Reload the JWT secret from SQLite.
Call this after setup to ensure new tokens use the persistent secret.
"""
global SECRET_KEY
@ -93,16 +93,16 @@ async def get_current_subject(
"""
token = credentials.credentials
try:
payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
payload = jwt.decode(token, SECRET_KEY, algorithms = [ALGORITHM])
subject: Optional[str] = payload.get("sub")
if subject is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid token payload",
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Invalid token payload",
)
return subject
except jwt.InvalidTokenError:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid or expired token",
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Invalid or expired token",
)

6
studio/backend/auth/hashing.py
View file

@ -4,6 +4,7 @@
"""
Password hashing utilities using PBKDF2.
"""
import hashlib
import hmac
import secrets
@ -13,7 +14,7 @@ from typing import Tuple
def hash_password(password: str, salt: str | None = None) -> Tuple[str, str]:
"""
Hash a password using PBKDF2-HMAC-SHA256.
Returns (salt, hex_hash) tuple.
"""
if salt is None:
@ -30,7 +31,7 @@ def hash_password(password: str, salt: str | None = None) -> Tuple[str, str]:
def verify_password(password: str, salt: str, hashed: str) -> bool:
"""
Verify a password against a stored salt and hash.
Uses constant-time comparison to prevent timing attacks.
"""
dk = hashlib.pbkdf2_hmac(
@ -40,4 +41,3 @@ def verify_password(password: str, salt: str, hashed: str) -> bool:
100_000,
)
return hmac.compare_digest(dk.hex(), hashed)

17
studio/backend/auth/storage.py
View file

@ -4,6 +4,7 @@
"""
SQLite storage for authentication data (user credentials + JWT secret).
"""
import hashlib
import sqlite3
from datetime import datetime, timezone
@ -69,11 +70,11 @@ def is_initialized() -> bool:
def create_initial_user(username: str, password: str, jwt_secret: str) -> None:
"""
Create the initial admin user in the database.
Raises sqlite3.IntegrityError if username already exists.
"""
from .hashing import hash_password
salt, pwd_hash = hash_password(password)
conn = get_connection()
try:
@ -92,7 +93,7 @@ def create_initial_user(username: str, password: str, jwt_secret: str) -> None:
def delete_user(username: str) -> None:
"""
Delete a user from the database.
Used for rollback when setup fails after user creation.
"""
conn = get_connection()
@ -106,7 +107,7 @@ def delete_user(username: str) -> None:
def get_user_and_secret(username: str) -> Optional[Tuple[str, str, str]]:
"""
Get user's password salt, hash, and JWT secret.
Returns (password_salt, password_hash, jwt_secret) or None if user not found.
"""
conn = get_connection()
@ -130,7 +131,7 @@ def get_user_and_secret(username: str) -> Optional[Tuple[str, str, str]]:
def load_jwt_secret() -> str:
"""
Load the JWT secret from the database.
Raises RuntimeError if auth is not initialized.
"""
conn = get_connection()
@ -138,7 +139,9 @@ def load_jwt_secret() -> str:
cur = conn.execute("SELECT jwt_secret FROM auth_user LIMIT 1")
row = cur.fetchone()
if not row:
raise RuntimeError("Auth is not initialized. Please set up a password first.")
raise RuntimeError(
"Auth is not initialized. Please set up a password first."
)
return row["jwt_secret"]
finally:
conn.close()
@ -161,7 +164,7 @@ def save_setup_token(token: str) -> None:
def consume_setup_token(token: str) -> bool:
"""
Verify a setup token and delete it if valid.
Returns True if the token was valid (and is now consumed), False otherwise.
"""
token_hash = _hash_token(token)

30
studio/backend/colab.py
View file

@ -5,6 +5,7 @@
Colab-specific helpers for running Unsloth Studio.
Uses Colab's built-in proxy - no external tunneling needed!
"""
from pathlib import Path
import sys
@ -14,20 +15,19 @@ if backend_path not in sys.path:
sys.path.insert(0, backend_path)
from loggers import get_logger
logger = get_logger(__name__)
def get_colab_url(port: int = 8000) -> str:
"""
Get the actual Colab proxy URL for a port.
"""
try:
from google.colab.output import eval_js
# Use Colab's proxy mechanism
url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec=5)
url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 5)
return url if url else f"http://localhost:{port}"
except Exception as e:
logger.info(f"Note: Could not get Colab URL ({e})")
@ -37,10 +37,10 @@ def get_colab_url(port: int = 8000) -> str:
def show_link(port: int = 8000):
"""Display a styled clickable link to the UI."""
from IPython.display import display, HTML
# Get real Colab proxy URL
url = get_colab_url(port)
html = f"""
<div style="padding: 20px; background: linear-gradient(135deg, #22c55e 0%, #16a34a 100%);
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
@ -65,32 +65,32 @@ def show_link(port: int = 8000):
def start(port: int = 8000):
"""
Start Unsloth Studio server in Colab and display the URL.
Usage:
from colab import start
start()
"""
import sys
logger.info("🦥 Starting Unsloth Studio...")
logger.info(" Loading backend...")
from run import run_server
# Auto-detect frontend path
repo_root = Path(__file__).parent.parent
frontend_path = repo_root / "frontend" / "dist"
if not frontend_path.exists():
logger.info("❌ Frontend not built! Please run the setup cell first.")
return
logger.info(" Starting server...")
# Start server silently
run_server(host="0.0.0.0", port=port, frontend_path=frontend_path, silent=True)
run_server(host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True)
logger.info(" Server started!")
# Show the clickable link with real URL
show_link(port)

142
studio/backend/core/__init__.py
View file

@ -12,101 +12,123 @@ code has a chance to run.
__all__ = [
# Inference
'InferenceBackend',
'get_inference_backend',
"InferenceBackend",
"get_inference_backend",
# Training
'get_training_backend',
'TrainingBackend',
'TrainingProgress',
"get_training_backend",
"TrainingBackend",
"TrainingProgress",
# Config
'ModelConfig',
'is_vision_model',
'scan_trained_loras',
'load_model_defaults',
'get_base_model_from_lora',
"ModelConfig",
"is_vision_model",
"scan_trained_loras",
"load_model_defaults",
"get_base_model_from_lora",
# Utils
'format_and_template_dataset',
'normalize_path',
'is_local_path',
'is_model_cached',
'without_hf_auth',
'format_error_message',
'get_gpu_memory_info',
'log_gpu_memory',
'get_device',
'is_apple_silicon',
'clear_gpu_cache',
'DeviceType',
"format_and_template_dataset",
"normalize_path",
"is_local_path",
"is_model_cached",
"without_hf_auth",
"format_error_message",
"get_gpu_memory_info",
"log_gpu_memory",
"get_device",
"is_apple_silicon",
"clear_gpu_cache",
"DeviceType",
]
def __getattr__(name):
# Inference
if name in ('InferenceBackend', 'get_inference_backend'):
if name in ("InferenceBackend", "get_inference_backend"):
from .inference import InferenceBackend, get_inference_backend
globals()['InferenceBackend'] = InferenceBackend
globals()['get_inference_backend'] = get_inference_backend
globals()["InferenceBackend"] = InferenceBackend
globals()["get_inference_backend"] = get_inference_backend
return globals()[name]
# Training
if name in ('TrainingBackend', 'get_training_backend', 'TrainingProgress'):
if name in ("TrainingBackend", "get_training_backend", "TrainingProgress"):
from .training import TrainingBackend, get_training_backend, TrainingProgress
globals()['TrainingBackend'] = TrainingBackend
globals()['get_training_backend'] = get_training_backend
globals()['TrainingProgress'] = TrainingProgress
globals()["TrainingBackend"] = TrainingBackend
globals()["get_training_backend"] = get_training_backend
globals()["TrainingProgress"] = TrainingProgress
return globals()[name]
# Config (from utils.models)
if name in ('is_vision_model', 'ModelConfig', 'scan_trained_loras',
'load_model_defaults', 'get_base_model_from_lora'):
if name in (
"is_vision_model",
"ModelConfig",
"scan_trained_loras",
"load_model_defaults",
"get_base_model_from_lora",
):
from utils.models import (
is_vision_model, ModelConfig, scan_trained_loras,
load_model_defaults, get_base_model_from_lora,
is_vision_model,
ModelConfig,
scan_trained_loras,
load_model_defaults,
get_base_model_from_lora,
)
globals()['is_vision_model'] = is_vision_model
globals()['ModelConfig'] = ModelConfig
globals()['scan_trained_loras'] = scan_trained_loras
globals()['load_model_defaults'] = load_model_defaults
globals()['get_base_model_from_lora'] = get_base_model_from_lora
globals()["is_vision_model"] = is_vision_model
globals()["ModelConfig"] = ModelConfig
globals()["scan_trained_loras"] = scan_trained_loras
globals()["load_model_defaults"] = load_model_defaults
globals()["get_base_model_from_lora"] = get_base_model_from_lora
return globals()[name]
# Paths
if name in ('normalize_path', 'is_local_path', 'is_model_cached'):
if name in ("normalize_path", "is_local_path", "is_model_cached"):
from utils.paths import normalize_path, is_local_path, is_model_cached
globals()['normalize_path'] = normalize_path
globals()['is_local_path'] = is_local_path
globals()['is_model_cached'] = is_model_cached
globals()["normalize_path"] = normalize_path
globals()["is_local_path"] = is_local_path
globals()["is_model_cached"] = is_model_cached
return globals()[name]
# Utils
if name in ('without_hf_auth', 'format_error_message'):
if name in ("without_hf_auth", "format_error_message"):
from utils.utils import without_hf_auth, format_error_message
globals()['without_hf_auth'] = without_hf_auth
globals()['format_error_message'] = format_error_message
globals()["without_hf_auth"] = without_hf_auth
globals()["format_error_message"] = format_error_message
return globals()[name]
# Hardware
if name in ('get_device', 'is_apple_silicon', 'clear_gpu_cache',
'get_gpu_memory_info', 'log_gpu_memory', 'DeviceType'):
if name in (
"get_device",
"is_apple_silicon",
"clear_gpu_cache",
"get_gpu_memory_info",
"log_gpu_memory",
"DeviceType",
):
from utils.hardware import (
get_device, is_apple_silicon, clear_gpu_cache,
get_gpu_memory_info, log_gpu_memory, DeviceType,
get_device,
is_apple_silicon,
clear_gpu_cache,
get_gpu_memory_info,
log_gpu_memory,
DeviceType,
)
globals()['get_device'] = get_device
globals()['is_apple_silicon'] = is_apple_silicon
globals()['clear_gpu_cache'] = clear_gpu_cache
globals()['get_gpu_memory_info'] = get_gpu_memory_info
globals()['log_gpu_memory'] = log_gpu_memory
globals()['DeviceType'] = DeviceType
globals()["get_device"] = get_device
globals()["is_apple_silicon"] = is_apple_silicon
globals()["clear_gpu_cache"] = clear_gpu_cache
globals()["get_gpu_memory_info"] = get_gpu_memory_info
globals()["log_gpu_memory"] = log_gpu_memory
globals()["DeviceType"] = DeviceType
return globals()[name]
# Datasets
if name == 'format_and_template_dataset':
if name == "format_and_template_dataset":
from utils.datasets import format_and_template_dataset
globals()['format_and_template_dataset'] = format_and_template_dataset
globals()["format_and_template_dataset"] = format_and_template_dataset
return format_and_template_dataset
raise AttributeError(f"module 'core' has no attribute {name!r}")

1
studio/backend/core/data_recipe/jobs/__init__.py
View file

@ -4,4 +4,3 @@
from .manager import JobManager, get_job_manager
__all__ = ["JobManager", "get_job_manager"]

80
studio/backend/core/data_recipe/jobs/manager.py
View file

@ -52,12 +52,10 @@ class Subscription:
if event_id is None:
self._next_id += 1
event_id = self._next_id
body = json.dumps(event, separators=(",", ":"), ensure_ascii=False)
body = json.dumps(event, separators = (",", ":"), ensure_ascii = False)
event_type = event.get("type") or "message"
return (
f"id: {event_id}\n"
f"event: {event_type}\n"
f"data: {body}\n\n"
f"id: {event_id}\n" f"event: {event_type}\n" f"data: {body}\n\n"
).encode("utf-8")
@ -68,7 +66,7 @@ class JobManager:
self._job: Job | None = None
self._proc: mp.Process | None = None
self._mp_q: Any | None = None
self._events: deque[dict] = deque(maxlen=5000)
self._events: deque[dict] = deque(maxlen = 5000)
self._subs: list[queue.Queue] = []
self._pump_thread: threading.Thread | None = None
self._seq: int = 0
@ -92,7 +90,7 @@ class JobManager:
raise RuntimeError("job already running")
job_id = uuid.uuid4().hex
self._job = Job(job_id=job_id, status="pending", started_at=time.time())
self._job = Job(job_id = job_id, status = "pending", started_at = time.time())
self._job.progress_columns_total = llm_column_count
self._events.clear()
self._seq = 0
@ -101,18 +99,20 @@ class JobManager:
run_payload["_job_id"] = job_id
mp_q = _CTX.Queue()
proc = _CTX.Process(
target=run_job_process,
kwargs={"event_queue": mp_q, "recipe": recipe, "run": run_payload},
daemon=True,
target = run_job_process,
kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload},
daemon = True,
)
proc.start()
self._mp_q = mp_q
self._proc = proc
self._pump_thread = threading.Thread(target=self._pump_loop, daemon=True)
self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread.start()
self._emit({"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id})
self._emit(
{"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id}
)
return job_id
def cancel(self, job_id: str) -> bool:
@ -123,7 +123,9 @@ class JobManager:
if self._proc is None or not self._proc.is_alive():
return True
self._job.status = "cancelling"
self._emit({"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id})
self._emit(
{"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id}
)
try:
self._proc.terminate()
except (AttributeError, OSError):
@ -225,7 +227,7 @@ class JobManager:
if in_memory_dataset is not None:
total = len(in_memory_dataset)
rows = in_memory_dataset[offset:offset + limit]
rows = in_memory_dataset[offset : offset + limit]
return {"dataset": rows, "total": total}
if not artifact_path:
if job_status in {"completed", "error", "cancelled"}:
@ -238,7 +240,9 @@ class JobManager:
if not parquet_dir.exists():
return {"error": f"dataset path missing: {parquet_dir}"}
return self._load_dataset_page(parquet_dir=parquet_dir, limit=limit, offset=offset)
return self._load_dataset_page(
parquet_dir = parquet_dir, limit = limit, offset = offset
)
except Exception as exc:
return {"error": f"dataset load failed: {exc}"}
@ -250,16 +254,16 @@ class JobManager:
offset: int,
) -> dict[str, Any]:
dataset_page = JobManager._load_dataset_page_with_duckdb(
parquet_dir=parquet_dir,
limit=limit,
offset=offset,
parquet_dir = parquet_dir,
limit = limit,
offset = offset,
)
if dataset_page is not None:
return dataset_page
return JobManager._load_dataset_page_with_data_designer(
parquet_dir=parquet_dir,
limit=limit,
offset=offset,
parquet_dir = parquet_dir,
limit = limit,
offset = offset,
)
@staticmethod
@ -299,9 +303,9 @@ class JobManager:
for helper_col in ("filename", "__row_num__"):
if helper_col in dataframe.columns:
dataframe = dataframe.drop(columns=[helper_col])
dataframe = dataframe.drop(columns = [helper_col])
rows = dataframe.to_dict(orient="records")
rows = dataframe.to_dict(orient = "records")
return {"dataset": to_preview_jsonable(rows), "total": total}
@staticmethod
@ -315,21 +319,23 @@ class JobManager:
dataframe = read_parquet_dataset(parquet_dir)
total = int(len(dataframe.index))
rows = dataframe.iloc[offset:offset + limit].to_dict(orient="records")
rows = dataframe.iloc[offset : offset + limit].to_dict(orient = "records")
return {"dataset": to_preview_jsonable(rows), "total": total}
def subscribe(self, job_id: str, *, after_seq: int | None = None) -> Subscription | None:
def subscribe(
self, job_id: str, *, after_seq: int | None = None
) -> Subscription | None:
"""SSE subscribe: get replay buffer + live events stream."""
with self._lock:
if self._job is None or self._job.job_id != job_id:
return None
q: queue.Queue = queue.Queue(maxsize=2000)
q: queue.Queue = queue.Queue(maxsize = 2000)
self._subs.append(q)
if after_seq is None:
replay = list(self._events)
else:
replay = [e for e in self._events if int(e.get("seq") or 0) > after_seq]
return Subscription(replay=replay, _q=q)
return Subscription(replay = replay, _q = q)
def unsubscribe(self, sub: Subscription) -> None:
"""Drop SSE subscriber (client disconnected)."""
@ -361,7 +367,7 @@ class JobManager:
def _read_queue_with_timeout(q: Any, *, timeout_sec: float) -> dict | None:
"""Try read 1 event from mp queue. Timeout = pump stays responsive."""
try:
return coerce_event(q.get(timeout=timeout_sec))
return coerce_event(q.get(timeout = timeout_sec))
except queue.Empty:
return None
except (EOFError, OSError, ValueError):
@ -387,7 +393,7 @@ class JobManager:
return
job, proc, mp_q = snap
event = self._read_queue_with_timeout(mp_q, timeout_sec=0.25)
event = self._read_queue_with_timeout(mp_q, timeout_sec = 0.25)
if event is not None:
self._handle_event(job, event)
continue
@ -399,7 +405,11 @@ class JobManager:
self._handle_event(job, e)
with self._lock:
if self._job and self._job.status in {"pending", "active", "cancelling"}:
if self._job and self._job.status in {
"pending",
"active",
"cancelling",
}:
if self._job.status == "cancelling":
self._job.status = "cancelled"
else:
@ -407,9 +417,17 @@ class JobManager:
self._job.error = self._job.error or "process exited"
self._job.finished_at = time.time()
event_type = (
EVENT_JOB_CANCELLED if self._job.status == "cancelled" else EVENT_JOB_ERROR
EVENT_JOB_CANCELLED
if self._job.status == "cancelled"
else EVENT_JOB_ERROR
)
self._emit(
{
"type": event_type,
"ts": time.time(),
"job_id": self._job.job_id,
}
)
self._emit({"type": event_type, "ts": time.time(), "job_id": self._job.job_id})
return
def _handle_event(self, job: Job, event: dict) -> None:

90
studio/backend/core/data_recipe/jobs/parse.py
View file

@ -22,7 +22,7 @@ from .constants import (
from .types import Job, ModelUsage, Progress
@dataclass(frozen=True)
@dataclass(frozen = True)
class ParsedUpdate:
stage: str | None = None
current_column: str | None = None
@ -42,6 +42,7 @@ class ParsedUpdate:
usage_rpm: float | None = None
usage_section_start: bool | None = None
# kinda of a bummber but currently only option, Best effort parser from data-designer logs -> structured status for UI.
_RE_SAMPLERS = re.compile(
r"Preparing samplers to generate (?P<rows>\d+) records across (?P<cols>\d+) columns"
@ -66,76 +67,76 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
m = _RE_SAMPLERS.search(msg)
if m:
return ParsedUpdate(
stage=STAGE_SAMPLING,
rows=int(m.group("rows")),
cols=int(m.group("cols")),
stage = STAGE_SAMPLING,
rows = int(m.group("rows")),
cols = int(m.group("cols")),
)
if "Sorting column configs into a Directed Acyclic Graph" in msg:
return ParsedUpdate(stage=STAGE_DAG)
return ParsedUpdate(stage = STAGE_DAG)
if "Running health checks for models" in msg:
return ParsedUpdate(stage=STAGE_HEALTHCHECK)
return ParsedUpdate(stage = STAGE_HEALTHCHECK)
if "Preview generation in progress" in msg:
return ParsedUpdate(stage=STAGE_PREVIEW)
return ParsedUpdate(stage = STAGE_PREVIEW)
if "Creating Data Designer dataset" in msg:
return ParsedUpdate(stage=STAGE_CREATE)
return ParsedUpdate(stage = STAGE_CREATE)
if "Measuring dataset column statistics" in msg:
return ParsedUpdate(stage=STAGE_PROFILING)
return ParsedUpdate(stage = STAGE_PROFILING)
m = _RE_COLCFG.search(msg)
if m:
col = m.group("col")
return ParsedUpdate(stage=STAGE_COLUMN_CONFIG, current_column=col)
return ParsedUpdate(stage = STAGE_COLUMN_CONFIG, current_column = col)
m = _RE_PROCESSING_COL.search(msg)
if m:
col = m.group("col")
return ParsedUpdate(stage=STAGE_GENERATING, current_column=col)
return ParsedUpdate(stage = STAGE_GENERATING, current_column = col)
m = _RE_PROGRESS.search(msg)
if m:
p = Progress(
done=int(m.group("done")),
total=int(m.group("total")),
percent=float(m.group("pct")),
ok=int(m.group("ok")),
failed=int(m.group("failed")),
rate=float(m.group("rate")),
eta_sec=float(m.group("eta")),
done = int(m.group("done")),
total = int(m.group("total")),
percent = float(m.group("pct")),
ok = int(m.group("ok")),
failed = int(m.group("failed")),
rate = float(m.group("rate")),
eta_sec = float(m.group("eta")),
)
return ParsedUpdate(stage=STAGE_GENERATING, progress=p)
return ParsedUpdate(stage = STAGE_GENERATING, progress = p)
m = _RE_BATCH.search(msg)
if m:
return ParsedUpdate(
stage=STAGE_BATCH,
batch_idx=int(m.group("idx")),
batch_total=int(m.group("total")),
stage = STAGE_BATCH,
batch_idx = int(m.group("idx")),
batch_total = int(m.group("total")),
)
if "Model usage summary" in msg:
return ParsedUpdate(usage_section_start=True)
return ParsedUpdate(usage_section_start = True)
m = _RE_USAGE_MODEL.search(msg)
if m and "|-- model:" in msg:
return ParsedUpdate(usage_model=str(m.group("model")).strip())
return ParsedUpdate(usage_model = str(m.group("model")).strip())
m = _RE_USAGE_TOKENS.search(msg)
if m:
return ParsedUpdate(
usage_input_tokens=int(m.group("input")),
usage_output_tokens=int(m.group("output")),
usage_total_tokens=int(m.group("total")),
usage_tps=float(m.group("tps")),
usage_input_tokens = int(m.group("input")),
usage_output_tokens = int(m.group("output")),
usage_total_tokens = int(m.group("total")),
usage_tps = float(m.group("tps")),
)
m = _RE_USAGE_REQUESTS.search(msg)
if m:
return ParsedUpdate(
usage_requests_success=int(m.group("success")),
usage_requests_failed=int(m.group("failed")),
usage_requests_total=int(m.group("total")),
usage_rpm=float(m.group("rpm")),
usage_requests_success = int(m.group("success")),
usage_requests_failed = int(m.group("failed")),
usage_requests_total = int(m.group("total")),
usage_rpm = float(m.group("rpm")),
)
return None
@ -146,7 +147,10 @@ def apply_update(job: Job, update: ParsedUpdate) -> None:
job.stage = update.stage
if update.current_column is not None:
job.current_column = update.current_column
if update.stage == STAGE_GENERATING and update.current_column not in job._seen_generation_columns:
if (
update.stage == STAGE_GENERATING
and update.current_column not in job._seen_generation_columns
):
job._seen_generation_columns.append(update.current_column)
if update.rows is not None:
job.rows = update.rows
@ -185,7 +189,7 @@ def apply_update(job: Job, update: ParsedUpdate) -> None:
name = update.usage_model.strip().strip("'").strip('"')
job._current_usage_model = name
if name not in job.model_usage:
job.model_usage[name] = ModelUsage(model=name)
job.model_usage[name] = ModelUsage(model = name)
if job._current_usage_model is None:
return
@ -227,7 +231,9 @@ def _compute_overall_progress(job: Job, column_progress: Progress) -> Progress:
if len(job._column_done) == 0:
done = current_done
else:
sum_done = sum(max(0, min(value, total_rows)) for value in job._column_done.values())
sum_done = sum(
max(0, min(value, total_rows)) for value in job._column_done.values()
)
done = int(sum_done / total_columns)
prev_done = int(job.progress.done or 0)
@ -241,13 +247,13 @@ def _compute_overall_progress(job: Job, column_progress: Progress) -> Progress:
percent = prev_percent
return Progress(
done=done,
total=total_rows,
percent=percent,
eta_sec=column_progress.eta_sec,
rate=column_progress.rate,
ok=column_progress.ok,
failed=column_progress.failed,
done = done,
total = total_rows,
percent = percent,
eta_sec = column_progress.eta_sec,
rate = column_progress.rate,
ok = column_progress.ok,
failed = column_progress.failed,
)

14
studio/backend/core/data_recipe/jobs/types.py
View file

@ -54,9 +54,9 @@ class Job:
status: JobStatus = "created"
stage: str | None = None
current_column: str | None = None
progress: Progress = field(default_factory=Progress)
column_progress: Progress = field(default_factory=Progress)
batch: BatchProgress = field(default_factory=BatchProgress)
progress: Progress = field(default_factory = Progress)
column_progress: Progress = field(default_factory = Progress)
batch: BatchProgress = field(default_factory = BatchProgress)
rows: int | None = None
cols: int | None = None
error: str | None = None
@ -67,10 +67,10 @@ class Job:
artifact_path: str | None = None
dataset: list[dict[str, Any]] | None = None
processor_artifacts: dict[str, Any] | None = None
model_usage: dict[str, ModelUsage] = field(default_factory=dict)
model_usage: dict[str, ModelUsage] = field(default_factory = dict)
progress_columns_total: int | None = None
completed_columns: list[str] = field(default_factory=list)
completed_columns: list[str] = field(default_factory = list)
_current_usage_model: str | None = None
_in_usage_summary: bool = False
_seen_generation_columns: list[str] = field(default_factory=list)
_column_done: dict[str, int] = field(default_factory=dict)
_seen_generation_columns: list[str] = field(default_factory = list)
_column_done: dict[str, int] = field(default_factory = dict)

56
studio/backend/core/data_recipe/jobs/worker.py
View file

@ -51,7 +51,9 @@ def _slugify_run_name(value: str) -> str:
return slug[:80].strip("-")
def _build_dataset_name(*, run_name: str | None, job_id: str, artifact_root: Path) -> str:
def _build_dataset_name(
*, run_name: str | None, job_id: str, artifact_root: Path
) -> str:
fallback = f"recipe_{job_id}"
slug = _slugify_run_name(run_name or "")
base_name = f"recipe_{slug}" if slug else fallback
@ -74,16 +76,20 @@ def run_job_process(
Sends events to `event_queue`.
"""
import os
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
os.environ["PYTHONWARNINGS"] = (
"ignore" # Suppress warnings at C-level before imports
)
import warnings
from loggers.config import LogConfig
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
LogConfig.setup_logging(
service_name="unsloth-studio-data-worker",
env=os.getenv("ENVIRONMENT_TYPE", "production"),
service_name = "unsloth-studio-data-worker",
env = os.getenv("ENVIRONMENT_TYPE", "production"),
)
event_queue.put({"type": EVENT_JOB_STARTED, "ts": time.time()})
@ -98,16 +104,16 @@ def run_job_process(
run_name_raw = run.get("run_name")
run_name = run_name_raw if isinstance(run_name_raw, str) else None
dataset_name = _build_dataset_name(
run_name=run_name,
job_id=job_id,
artifact_root=_ARTIFACT_ROOT,
run_name = run_name,
job_id = job_id,
artifact_root = _ARTIFACT_ROOT,
)
merge_batches = bool(run.get("merge_batches"))
ensure_dir(_ARTIFACT_ROOT)
run_config_raw = run.get("run_config") or {}
builder = build_config_builder(recipe)
designer = create_data_designer(recipe, artifact_path=str(_ARTIFACT_ROOT))
designer = create_data_designer(recipe, artifact_path = str(_ARTIFACT_ROOT))
# DataDesigner configures root logging in DataDesigner.__init__.
# Attach queue logger directly to `data_designer` so parser events survive root resets.
@ -123,16 +129,16 @@ def run_job_process(
execution_type = str(run.get("execution_type") or "full").strip().lower()
if execution_type == "preview":
results = designer.preview(builder, num_records=rows)
results = designer.preview(builder, num_records = rows)
analysis = (
None
if results.analysis is None
else to_jsonable(results.analysis.model_dump(mode="json"))
else to_jsonable(results.analysis.model_dump(mode = "json"))
)
dataset = (
[]
if results.dataset is None
else to_preview_jsonable(results.dataset.to_dict(orient="records"))
else to_preview_jsonable(results.dataset.to_dict(orient = "records"))
)
processor_artifacts = (
None
@ -151,10 +157,14 @@ def run_job_process(
}
)
else:
results = designer.create(builder, num_records=rows, dataset_name=dataset_name)
analysis = to_jsonable(results.load_analysis().model_dump(mode="json"))
results = designer.create(
builder, num_records = rows, dataset_name = dataset_name
)
analysis = to_jsonable(results.load_analysis().model_dump(mode = "json"))
if merge_batches:
_merge_batches_to_single_parquet(results.artifact_storage.base_dataset_path)
_merge_batches_to_single_parquet(
results.artifact_storage.base_dataset_path
)
artifact_path = str(results.artifact_storage.base_dataset_path)
event_queue.put(
{
@ -171,7 +181,7 @@ def run_job_process(
"type": EVENT_JOB_ERROR,
"ts": time.time(),
"error": str(exc),
"stack": traceback.format_exc(limit=20),
"stack": traceback.format_exc(limit = 20),
}
)
@ -189,12 +199,12 @@ def _merge_batches_to_single_parquet(base_dataset_path: Path) -> None:
dataframe = read_parquet_dataset(parquet_dir)
shutil.rmtree(parquet_dir)
parquet_dir.mkdir(parents=True, exist_ok=True)
parquet_dir.mkdir(parents = True, exist_ok = True)
merged_file = parquet_dir / "batch_00000.parquet"
dataframe.to_parquet(merged_file, index=False)
dataframe.to_parquet(merged_file, index = False)
_rewrite_merged_metadata(
base_dataset_path=base_dataset_path,
parquet_file=merged_file,
base_dataset_path = base_dataset_path,
parquet_file = merged_file,
)
@ -204,7 +214,7 @@ def _rewrite_merged_metadata(*, base_dataset_path: Path, parquet_file: Path) ->
return
try:
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
metadata = json.loads(metadata_path.read_text(encoding = "utf-8"))
except (OSError, TypeError, ValueError):
return
@ -222,8 +232,8 @@ def _rewrite_merged_metadata(*, base_dataset_path: Path, parquet_file: Path) ->
try:
metadata_path.write_text(
json.dumps(metadata, indent=2, sort_keys=True),
encoding="utf-8",
json.dumps(metadata, indent = 2, sort_keys = True),
encoding = "utf-8",
)
except OSError:
return

2
studio/backend/core/data_recipe/jsonable.py
View file

@ -11,7 +11,7 @@ from typing import Any
def _pil_to_preview_payload(image: Any) -> dict[str, Any]:
buffer = io.BytesIO()
image.convert("RGB").save(buffer, format="JPEG", quality=85)
image.convert("RGB").save(buffer, format = "JPEG", quality = 85)
return {
"type": "image",
"mime": "image/jpeg",

95
studio/backend/core/data_recipe/local_callable_validators.py
View file

@ -33,7 +33,7 @@ _OXC_TOOL_DIR = Path(__file__).resolve().parent / "oxc-validator"
_OXC_RUNNER_PATH = _OXC_TOOL_DIR / "validate.mjs"
@dataclass(frozen=True)
@dataclass(frozen = True)
class OxcLocalCallableValidatorSpec:
name: str
drop: bool
@ -64,7 +64,7 @@ def split_oxc_local_callable_validators(
kept_columns.append(column)
continue
maybe_spec = _parse_oxc_spec(column=column)
maybe_spec = _parse_oxc_spec(column = column)
if maybe_spec is None:
kept_columns.append(column)
continue
@ -96,14 +96,14 @@ def register_oxc_local_callable_validators(
)
builder.add_column(
ValidationColumnConfig(
name=spec.name,
drop=spec.drop,
target_columns=spec.target_columns,
validator_type=ValidatorType.LOCAL_CALLABLE,
validator_params=LocalCallableValidatorParams(
validation_function=validation_function,
name = spec.name,
drop = spec.drop,
target_columns = spec.target_columns,
validator_type = ValidatorType.LOCAL_CALLABLE,
validator_params = LocalCallableValidatorParams(
validation_function = validation_function,
),
batch_size=spec.batch_size,
batch_size = spec.batch_size,
)
)
@ -132,7 +132,11 @@ def _parse_oxc_spec(
target_columns_raw = column.get("target_columns")
target_columns = (
[value.strip() for value in target_columns_raw if isinstance(value, str) and value.strip()]
[
value.strip()
for value in target_columns_raw
if isinstance(value, str) and value.strip()
]
if isinstance(target_columns_raw, list)
else []
)
@ -144,13 +148,13 @@ def _parse_oxc_spec(
drop = bool(column.get("drop") is True)
return OxcLocalCallableValidatorSpec(
name=name,
drop=drop,
target_columns=target_columns,
batch_size=batch_size,
code_lang=code_lang,
validation_mode=validation_mode,
code_shape=code_shape,
name = name,
drop = drop,
target_columns = target_columns,
batch_size = batch_size,
code_lang = code_lang,
validation_mode = validation_mode,
code_shape = code_shape,
)
@ -172,18 +176,16 @@ def _parse_oxc_validation_marker(fn_name: str) -> tuple[str, str, str]:
return "javascript", "syntax", "auto"
code_lang = parts[0] if parts[0] in _OXC_LANG_TO_NODE_LANG else "javascript"
mode = parts[1] if parts[1] in _OXC_VALIDATION_MODES else "syntax"
code_shape = parts[2] if len(parts) >= 3 and parts[2] in _OXC_CODE_SHAPES else "auto"
code_shape = (
parts[2] if len(parts) >= 3 and parts[2] in _OXC_CODE_SHAPES else "auto"
)
return code_lang, mode, code_shape
@lru_cache(maxsize=8)
@lru_cache(maxsize = 8)
def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape: str):
node_lang = _OXC_LANG_TO_NODE_LANG.get(lang, "js")
mode = (
validation_mode
if validation_mode in _OXC_VALIDATION_MODES
else "syntax"
)
mode = validation_mode if validation_mode in _OXC_VALIDATION_MODES else "syntax"
normalized_code_shape = code_shape if code_shape in _OXC_CODE_SHAPES else "auto"
def _validator(df):
@ -197,14 +199,17 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape:
code_values = (
["" for _ in range(row_count)]
if not code_column
else ["" if value is None else str(value) for value in df[code_column].tolist()]
else [
"" if value is None else str(value)
for value in df[code_column].tolist()
]
)
results = _run_oxc_batch(
node_lang=node_lang,
validation_mode=mode,
code_shape=normalized_code_shape,
code_values=code_values,
node_lang = node_lang,
validation_mode = mode,
code_shape = normalized_code_shape,
code_values = code_values,
)
if len(results) != row_count:
results = _fallback_results(
@ -213,9 +218,7 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape:
)
return pd.DataFrame(results)
_validator.__name__ = (
f"{OXC_VALIDATION_FN_MARKER}_{node_lang}_{mode.replace('+', '_')}_{normalized_code_shape}"
)
_validator.__name__ = f"{OXC_VALIDATION_FN_MARKER}_{node_lang}_{mode.replace('+', '_')}_{normalized_code_shape}"
return _validator
@ -247,12 +250,12 @@ def _run_oxc_batch(
env["TEMP"] = tmp_dir_str
proc = subprocess.run(
["node", str(_OXC_RUNNER_PATH)],
cwd=str(_OXC_TOOL_DIR),
input=json.dumps(payload),
text=True,
capture_output=True,
check=False,
env=env,
cwd = str(_OXC_TOOL_DIR),
input = json.dumps(payload),
text = True,
capture_output = True,
check = False,
env = env,
)
except (OSError, ValueError) as exc:
logger.warning("OXC subprocess launch failed: %s", exc)
@ -298,13 +301,21 @@ def _run_oxc_batch(
warning_count_raw = item.get("warning_count")
out.append(
{
"is_valid": bool(is_valid_raw) if isinstance(is_valid_raw, bool) else False,
"error_count": int(error_count_raw) if isinstance(error_count_raw, int) else 0,
"is_valid": bool(is_valid_raw)
if isinstance(is_valid_raw, bool)
else False,
"error_count": int(error_count_raw)
if isinstance(error_count_raw, int)
else 0,
"error_message": str(message_raw or ""),
"severity": str(severity_raw) if isinstance(severity_raw, str) else None,
"severity": str(severity_raw)
if isinstance(severity_raw, str)
else None,
"code": str(code_raw) if isinstance(code_raw, str) else None,
"labels": labels_raw if isinstance(labels_raw, list) else [],
"codeframe": str(codeframe_raw) if isinstance(codeframe_raw, str) else None,
"codeframe": str(codeframe_raw)
if isinstance(codeframe_raw, str)
else None,
"warning_count": int(warning_count_raw)
if isinstance(warning_count_raw, int)
else 0,

3
studio/backend/core/data_recipe/oxc-validator/validate.mjs
View file

@ -1,3 +1,6 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { spawnSync } from "node:child_process";
import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";

61
studio/backend/core/data_recipe/service.py
View file

@ -14,6 +14,7 @@ from .local_callable_validators import (
register_oxc_local_callable_validators,
split_oxc_local_callable_validators,
)
_IMAGE_CONTEXT_PATCHED = False
@ -21,7 +22,9 @@ def _encode_bytes_to_base64(value: bytes | bytearray) -> str:
return base64.b64encode(bytes(value)).decode("utf-8")
def _load_image_file_to_base64(path_value: str, *, base_path: str | None = None) -> str | None:
def _load_image_file_to_base64(
path_value: str, *, base_path: str | None = None
) -> str | None:
try:
path = Path(path_value)
candidates: list[Path] = []
@ -53,7 +56,7 @@ def _pil_image_to_base64(value: Any) -> str | None:
image_format = str(getattr(value, "format", "") or "").upper()
if image_format not in {"PNG", "JPEG", "JPG", "WEBP", "GIF"}:
image_format = "PNG"
value.save(buffer, format=image_format)
value.save(buffer, format = image_format)
return _encode_bytes_to_base64(buffer.getvalue())
@ -93,7 +96,7 @@ def _normalize_image_context_value(value: Any, *, base_path: str | None = None)
path_value = value.get("path")
if isinstance(path_value, str) and path_value.strip():
if as_base64 := _load_image_file_to_base64(path_value, base_path=base_path):
if as_base64 := _load_image_file_to_base64(path_value, base_path = base_path):
return as_base64
return path_value
@ -116,8 +119,10 @@ def _apply_data_designer_image_context_patch() -> None:
original_auto_resolve = ImageContext._auto_resolve_context_value
def _patched_auto_resolve(self: Any, context_value: Any, base_path: str | None) -> Any:
normalized = _normalize_image_context_value(context_value, base_path=base_path)
def _patched_auto_resolve(
self: Any, context_value: Any, base_path: str | None
) -> Any:
normalized = _normalize_image_context_value(context_value, base_path = base_path)
return original_auto_resolve(self, normalized, base_path)
ImageContext._auto_resolve_context_value = _patched_auto_resolve
@ -137,12 +142,12 @@ def build_model_providers(recipe: dict[str, Any]):
api_key = os.getenv(api_key_env)
providers.append(
ModelProvider(
name=provider["name"],
endpoint=provider["endpoint"],
provider_type=provider.get("provider_type", "openai"),
api_key=api_key,
extra_headers=provider.get("extra_headers"),
extra_body=provider.get("extra_body"),
name = provider["name"],
endpoint = provider["endpoint"],
provider_type = provider.get("provider_type", "openai"),
api_key = api_key,
extra_headers = provider.get("extra_headers"),
extra_body = provider.get("extra_body"),
)
)
@ -170,10 +175,10 @@ def build_mcp_providers(
args = []
providers.append(
LocalStdioMCPProvider(
name=str(provider.get("name", "")),
command=str(provider.get("command", "")),
args=[str(value) for value in args],
env={str(key): str(value) for key, value in env.items()},
name = str(provider.get("name", "")),
command = str(provider.get("command", "")),
args = [str(value) for value in args],
env = {str(key): str(value) for key, value in env.items()},
)
)
continue
@ -185,10 +190,10 @@ def build_mcp_providers(
api_key = os.getenv(str(api_key_env))
providers.append(
MCPProvider(
name=str(provider.get("name", "")),
endpoint=str(provider.get("endpoint", "")),
provider_type=str(provider_type),
api_key=str(api_key) if api_key else None,
name = str(provider.get("name", "")),
endpoint = str(provider.get("endpoint", "")),
provider_type = str(provider_type),
api_key = str(api_key) if api_key else None,
)
)
return providers
@ -209,8 +214,8 @@ def build_config_builder(recipe: dict[str, Any]):
)
builder = DataDesignerConfigBuilder.from_config({"data_designer": recipe_core})
register_oxc_local_callable_validators(
builder=builder,
specs=oxc_local_callable_specs,
builder = builder,
specs = oxc_local_callable_specs,
)
# DataDesignerConfigBuilder.from_config currently skips processors.
@ -223,7 +228,7 @@ def build_config_builder(recipe: dict[str, Any]):
continue
kwargs = {k: v for k, v in processor.items() if k != "processor_type"}
builder.add_processor(
processor_type=ProcessorType(processor_type_raw),
processor_type = ProcessorType(processor_type_raw),
**kwargs,
)
@ -239,9 +244,9 @@ def create_data_designer(
from data_designer.interface.data_designer import DataDesigner
return DataDesigner(
artifact_path=artifact_path,
model_providers=build_model_providers(recipe),
mcp_providers=build_mcp_providers(recipe),
artifact_path = artifact_path,
model_providers = build_model_providers(recipe),
mcp_providers = build_mcp_providers(recipe),
)
@ -257,11 +262,11 @@ def preview_recipe(
) -> tuple[list[dict[str, Any]], dict[str, Any] | None, dict[str, Any] | None]:
builder = build_config_builder(recipe)
designer = create_data_designer(recipe)
results = designer.preview(builder, num_records=num_records)
results = designer.preview(builder, num_records = num_records)
dataset: list[dict[str, Any]] = []
if results.dataset is not None:
raw_rows = results.dataset.to_dict(orient="records")
raw_rows = results.dataset.to_dict(orient = "records")
dataset = [to_jsonable(row) for row in raw_rows]
artifacts = (
@ -272,7 +277,7 @@ def preview_recipe(
analysis = (
None
if results.analysis is None
else to_jsonable(results.analysis.model_dump(mode="json"))
else to_jsonable(results.analysis.model_dump(mode = "json"))
)
return dataset, artifacts, analysis

7
studio/backend/core/export/__init__.py
View file

@ -8,13 +8,14 @@ The default get_export_backend() returns an ExportOrchestrator that
delegates to a subprocess. The original ExportBackend runs inside
the subprocess and can be imported directly from .export when needed.
"""
from .orchestrator import ExportOrchestrator, get_export_backend
# Expose ExportOrchestrator as ExportBackend for backward compat
ExportBackend = ExportOrchestrator
__all__ = [
'ExportBackend',
'ExportOrchestrator',
'get_export_backend',
"ExportBackend",
"ExportOrchestrator",
"get_export_backend",
]

287
studio/backend/core/export/export.py
View file

@ -5,6 +5,7 @@
"""
Export backend - handles model exporting in various formats
"""
import glob
import json
import structlog
@ -50,7 +51,7 @@ def _apply_wsl_sudo_patch():
try:
import unsloth_zoo.llama_cpp as llama_cpp_module
def _wsl_do_we_need_sudo(system_type="debian"):
def _wsl_do_we_need_sudo(system_type = "debian"):
logger.info(
"WSL detected — skipping sudo check "
"(build deps pre-installed by setup.sh)"
@ -59,16 +60,14 @@ def _apply_wsl_sudo_patch():
llama_cpp_module.do_we_need_sudo = _wsl_do_we_need_sudo
logger.info(
"Applied WSL sudo patch to "
"unsloth_zoo.llama_cpp.do_we_need_sudo"
"Applied WSL sudo patch to " "unsloth_zoo.llama_cpp.do_we_need_sudo"
)
except Exception as e:
logger.warning(f"Could not apply WSL sudo patch: {e}")
# Model card template
MODEL_CARD = \
"""---
MODEL_CARD = """---
base_model: {base_model}
tags:
- text-generation-inference
@ -92,6 +91,7 @@ This {model_type} model was trained 2x faster with [Unsloth](https://github.com/
[<img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20made%20with%20love.png" width="200"/>](https://github.com/unslothai/unsloth)
"""
class ExportBackend:
"""Handles model export operations"""
@ -130,7 +130,9 @@ class ExportBackend:
logger.error(f"Error during memory cleanup: {e}")
return False
def scan_checkpoints(self, outputs_dir: str = str(outputs_root())) -> List[Tuple[str, List[Tuple[str, str]]]]:
def scan_checkpoints(
self, outputs_dir: str = str(outputs_root())
) -> List[Tuple[str, List[Tuple[str, str]]]]:
"""
Scan outputs folder for training runs and their checkpoints.
@ -138,13 +140,16 @@ class ExportBackend:
List of tuples: [(model_name, [(display_name, checkpoint_path), ...]), ...]
"""
from utils.models.checkpoints import scan_checkpoints
return scan_checkpoints(outputs_dir=outputs_dir)
def load_checkpoint(self,
checkpoint_path: str,
max_seq_length: int = 2048,
load_in_4bit: bool = True,
trust_remote_code: bool = False) -> Tuple[bool, str]:
return scan_checkpoints(outputs_dir = outputs_dir)
def load_checkpoint(
self,
checkpoint_path: str,
max_seq_length: int = 2048,
load_in_4bit: bool = True,
trust_remote_code: bool = False,
) -> Tuple[bool, str]:
"""
Load a checkpoint for export.
@ -174,81 +179,85 @@ class ExportBackend:
self.is_vision = not self._audio_type and is_vision_model(model_id)
# Load model based on type
if self._audio_type == 'csm':
if self._audio_type == "csm":
from unsloth import FastModel
from transformers import CsmForConditionalGeneration
logger.info("Loading as CSM audio model...")
model, tokenizer = FastModel.from_pretrained(
model_name=checkpoint_path,
max_seq_length=max_seq_length,
dtype=None,
auto_model=CsmForConditionalGeneration,
load_in_4bit=False,
trust_remote_code=trust_remote_code,
model_name = checkpoint_path,
max_seq_length = max_seq_length,
dtype = None,
auto_model = CsmForConditionalGeneration,
load_in_4bit = False,
trust_remote_code = trust_remote_code,
)
elif self._audio_type == 'whisper':
elif self._audio_type == "whisper":
from unsloth import FastModel
from transformers import WhisperForConditionalGeneration
logger.info("Loading as Whisper audio model...")
model, tokenizer = FastModel.from_pretrained(
model_name=checkpoint_path,
dtype=None,
load_in_4bit=False,
auto_model=WhisperForConditionalGeneration,
trust_remote_code=trust_remote_code,
model_name = checkpoint_path,
dtype = None,
load_in_4bit = False,
auto_model = WhisperForConditionalGeneration,
trust_remote_code = trust_remote_code,
)
elif self._audio_type == 'snac':
elif self._audio_type == "snac":
logger.info("Loading as SNAC (Orpheus) audio model...")
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=checkpoint_path,
max_seq_length=max_seq_length,
dtype=None,
load_in_4bit=load_in_4bit,
trust_remote_code=trust_remote_code,
model_name = checkpoint_path,
max_seq_length = max_seq_length,
dtype = None,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
)
elif self._audio_type == 'bicodec':
elif self._audio_type == "bicodec":
from unsloth import FastModel
logger.info("Loading as BiCodec (Spark-TTS) audio model...")
model, tokenizer = FastModel.from_pretrained(
model_name=checkpoint_path,
max_seq_length=max_seq_length,
dtype=torch.float32,
load_in_4bit=False,
trust_remote_code=trust_remote_code,
model_name = checkpoint_path,
max_seq_length = max_seq_length,
dtype = torch.float32,
load_in_4bit = False,
trust_remote_code = trust_remote_code,
)
elif self._audio_type == 'dac':
elif self._audio_type == "dac":
from unsloth import FastModel
logger.info("Loading as DAC (OuteTTS) audio model...")
model, tokenizer = FastModel.from_pretrained(
model_name=checkpoint_path,
max_seq_length=max_seq_length,
load_in_4bit=False,
trust_remote_code=trust_remote_code,
model_name = checkpoint_path,
max_seq_length = max_seq_length,
load_in_4bit = False,
trust_remote_code = trust_remote_code,
)
elif self.is_vision:
logger.info("Loading as vision model...")
model, processor = FastVisionModel.from_pretrained(
model_name=checkpoint_path,
max_seq_length=max_seq_length,
dtype=None,
load_in_4bit=load_in_4bit,
trust_remote_code=trust_remote_code,
model_name = checkpoint_path,
max_seq_length = max_seq_length,
dtype = None,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
)
tokenizer = processor # For vision models, processor acts as tokenizer
else:
logger.info("Loading as text model...")
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=checkpoint_path,
max_seq_length=max_seq_length,
dtype=None,
load_in_4bit=load_in_4bit,
trust_remote_code=trust_remote_code,
model_name = checkpoint_path,
max_seq_length = max_seq_length,
dtype = None,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
)
# Check if PEFT model
@ -273,28 +282,35 @@ class ExportBackend:
except Exception as e:
logger.error(f"Error loading checkpoint: {e}")
import traceback
logger.error(traceback.format_exc())
return False, f"Failed to load checkpoint: {str(e)}"
def _write_export_metadata(self, save_directory: str):
"""Write export_metadata.json with base model info for Chat page discovery."""
try:
base_model = get_base_model_from_lora(self.current_checkpoint) if self.current_checkpoint else None
base_model = (
get_base_model_from_lora(self.current_checkpoint)
if self.current_checkpoint
else None
)
metadata = {"base_model": base_model}
metadata_path = os.path.join(save_directory, "export_metadata.json")
with open(metadata_path, "w") as f:
json.dump(metadata, f, indent=2)
json.dump(metadata, f, indent = 2)
logger.info(f"Wrote export metadata to {metadata_path}")
except Exception as e:
logger.warning(f"Could not write export metadata: {e}")
def export_merged_model(self,
save_directory: str,
format_type: str = "16-bit (FP16)",
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False) -> Tuple[bool, str]:
def export_merged_model(
self,
save_directory: str,
format_type: str = "16-bit (FP16)",
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False,
) -> Tuple[bool, str]:
"""
Export merged model (for PEFT models).
@ -319,7 +335,7 @@ class ExportBackend:
# Determine save method
if format_type == "4-bit (FP4)":
save_method = "merged_4bit_forced"
elif self._audio_type == 'whisper':
elif self._audio_type == "whisper":
# Whisper uses save_method=None for local 16-bit merged save
save_method = None
else: # 16-bit (FP16)
@ -332,9 +348,7 @@ class ExportBackend:
ensure_dir(Path(save_directory))
self.current_model.save_pretrained_merged(
save_directory,
self.current_tokenizer,
save_method=save_method
save_directory, self.current_tokenizer, save_method = save_method
)
# Write export metadata so the Chat page can identify the base model
@ -344,18 +358,23 @@ class ExportBackend:
# Push to hub if requested
if push_to_hub:
if not repo_id or not hf_token:
return False, "Repository ID and Hugging Face token required for Hub upload"
return (
False,
"Repository ID and Hugging Face token required for Hub upload",
)
logger.info(f"Pushing merged model to Hub: {repo_id}")
# Whisper uses save_method=None for local but "merged_16bit" for hub push
hub_save_method = save_method if save_method is not None else "merged_16bit"
hub_save_method = (
save_method if save_method is not None else "merged_16bit"
)
self.current_model.push_to_hub_merged(
repo_id,
self.current_tokenizer,
save_method=hub_save_method,
token=hf_token,
private=private
save_method = hub_save_method,
token = hf_token,
private = private,
)
logger.info(f"Model pushed successfully to {repo_id}")
@ -364,16 +383,19 @@ class ExportBackend:
except Exception as e:
logger.error(f"Error exporting merged model: {e}")
import traceback
logger.error(traceback.format_exc())
return False, f"Export failed: {str(e)}"
def export_base_model(self,
save_directory: str,
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False,
base_model_id: Optional[str] = None) -> Tuple[bool, str]:
def export_base_model(
self,
save_directory: str,
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False,
base_model_id: Optional[str] = None,
) -> Tuple[bool, str]:
"""
Export base model (for non-PEFT models).
@ -384,7 +406,10 @@ class ExportBackend:
return False, "No model loaded. Please select a checkpoint first."
if self.is_peft:
return False, "This is a PEFT model. Use 'Merged Model' export type instead."
return (
False,
"This is a PEFT model. Use 'Merged Model' export type instead.",
)
try:
# Save locally if requested
@ -403,40 +428,47 @@ class ExportBackend:
# Push to hub if requested
if push_to_hub:
if not repo_id or not hf_token:
return False, "Repository ID and Hugging Face token required for Hub upload"
return (
False,
"Repository ID and Hugging Face token required for Hub upload",
)
logger.info(f"Pushing base model to Hub: {repo_id}")
# Get base model name from request or model config
base_model = base_model_id or self.current_model.config._name_or_path or "unknown"
base_model = (
base_model_id
or self.current_model.config._name_or_path
or "unknown"
)
# Create repo
hf_api = HfApi(token=hf_token)
hf_api = HfApi(token = hf_token)
repo_id = PushToHubMixin._create_repo(
PushToHubMixin,
repo_id=repo_id,
private=private,
token=hf_token,
repo_id = repo_id,
private = private,
token = hf_token,
)
username = repo_id.split("/")[0]
# Create and push model card
content = MODEL_CARD.format(
username=username,
base_model=base_model,
model_type=self.current_model.config.model_type,
method="",
extra="unsloth",
username = username,
base_model = base_model,
model_type = self.current_model.config.model_type,
method = "",
extra = "unsloth",
)
card = ModelCard(content)
card.push_to_hub(repo_id, token=hf_token, commit_message="Unsloth Model Card")
card.push_to_hub(
repo_id, token = hf_token, commit_message = "Unsloth Model Card"
)
# Upload model files
if save_directory:
hf_api.upload_folder(
folder_path=save_directory,
repo_id=repo_id,
repo_type="model"
folder_path = save_directory, repo_id = repo_id, repo_type = "model"
)
logger.info(f"Model pushed successfully to {repo_id}")
else:
@ -447,16 +479,18 @@ class ExportBackend:
except Exception as e:
logger.error(f"Error exporting base model: {e}")
import traceback
logger.error(traceback.format_exc())
return False, f"Export failed: {str(e)}"
def export_gguf(self,
save_directory: str,
quantization_method: str = "Q4_K_M",
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None) -> Tuple[bool, str]:
def export_gguf(
self,
save_directory: str,
quantization_method: str = "Q4_K_M",
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
) -> Tuple[bool, str]:
"""
Export model in GGUF format.
@ -505,17 +539,21 @@ class ExportBackend:
self.current_model.save_pretrained_gguf(
model_save_path,
self.current_tokenizer,
quantization_method=quant_method
quantization_method = quant_method,
)
# Relocate GGUF artifacts into the export directory.
# convert_to_gguf writes .gguf files to cwd (repo root)
# because --outfile is a relative path like "model.Q4_K_M.gguf".
new_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs
new_ggufs = (
set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs
)
for src in sorted(new_ggufs):
dest = os.path.join(abs_save_dir, os.path.basename(src))
shutil.move(src, dest)
logger.info(f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/")
logger.info(
f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/"
)
# Flatten any .gguf files from subdirectories into abs_save_dir.
# save_pretrained_gguf may create subdirs (e.g. model_gguf/)
@ -528,7 +566,7 @@ class ExportBackend:
shutil.move(str(src), dest)
logger.info(f"Relocated GGUF: {src.name} → {abs_save_dir}/")
# Clean up the subdirectory (intermediate HF files, etc.)
shutil.rmtree(str(sub), ignore_errors=True)
shutil.rmtree(str(sub), ignore_errors = True)
logger.info(f"Cleaned up subdirectory: {sub.name}")
# Write export metadata so the Chat page can identify the base model
@ -546,15 +584,18 @@ class ExportBackend:
# Push to hub if requested
if push_to_hub:
if not repo_id or not hf_token:
return False, "Repository ID and Hugging Face token required for Hub upload"
return (
False,
"Repository ID and Hugging Face token required for Hub upload",
)
logger.info(f"Pushing GGUF model to Hub: {repo_id}")
self.current_model.push_to_hub_gguf(
repo_id,
self.current_tokenizer,
quantization_method=quant_method,
token=hf_token
quantization_method = quant_method,
token = hf_token,
)
logger.info(f"GGUF model pushed successfully to {repo_id}")
@ -563,15 +604,18 @@ class ExportBackend:
except Exception as e:
logger.error(f"Error exporting GGUF model: {e}")
import traceback
logger.error(traceback.format_exc())
return False, f"GGUF export failed: {str(e)}"
def export_lora_adapter(self,
save_directory: str,
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False) -> Tuple[bool, str]:
def export_lora_adapter(
self,
save_directory: str,
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False,
) -> Tuple[bool, str]:
"""
Export LoRA adapter only (not merged).
@ -598,19 +642,16 @@ class ExportBackend:
# Push to hub if requested
if push_to_hub:
if not repo_id or not hf_token:
return False, "Repository ID and Hugging Face token required for Hub upload"
return (
False,
"Repository ID and Hugging Face token required for Hub upload",
)
logger.info(f"Pushing LoRA adapter to Hub: {repo_id}")
self.current_model.push_to_hub(
repo_id,
token=hf_token,
private=private
)
self.current_model.push_to_hub(repo_id, token = hf_token, private = private)
self.current_tokenizer.push_to_hub(
repo_id,
token=hf_token,
private=private
repo_id, token = hf_token, private = private
)
logger.info(f"Adapter pushed successfully to {repo_id}")
@ -619,6 +660,7 @@ class ExportBackend:
except Exception as e:
logger.error(f"Error exporting LoRA adapter: {e}")
import traceback
logger.error(traceback.format_exc())
return False, f"Adapter export failed: {str(e)}"
@ -626,6 +668,7 @@ class ExportBackend:
# Global export backend instance
_export_backend = None
def get_export_backend() -> ExportBackend:
"""Get or create the global export backend instance"""
global _export_backend

111
studio/backend/core/export/orchestrator.py
View file

@ -13,6 +13,7 @@ the old subprocess is killed and a new one is spawned with the correct version.
Pattern follows core/inference/orchestrator.py.
"""
import atexit
import structlog
from loggers import get_logger
@ -65,13 +66,13 @@ class ExportOrchestrator:
self._resp_queue = _CTX.Queue()
self._proc = _CTX.Process(
target=run_export_process,
kwargs={
target = run_export_process,
kwargs = {
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,
"config": config,
},
daemon=True,
daemon = True,
)
self._proc.start()
logger.info("Export subprocess started (pid=%s)", self._proc.pid)
@ -93,7 +94,7 @@ class ExportOrchestrator:
# 3. Wait for graceful shutdown
try:
self._proc.join(timeout=timeout)
self._proc.join(timeout = timeout)
except Exception:
pass
@ -102,14 +103,14 @@ class ExportOrchestrator:
logger.warning("Export subprocess did not exit gracefully, terminating")
try:
self._proc.terminate()
self._proc.join(timeout=5)
self._proc.join(timeout = 5)
except Exception:
pass
if self._proc is not None and self._proc.is_alive():
logger.warning("Subprocess still alive after terminate, killing")
try:
self._proc.kill()
self._proc.join(timeout=3)
self._proc.join(timeout = 3)
except Exception:
pass
@ -120,7 +121,7 @@ class ExportOrchestrator:
def _cleanup(self):
"""atexit handler."""
self._shutdown_subprocess(timeout=5.0)
self._shutdown_subprocess(timeout = 5.0)
def _ensure_subprocess_alive(self) -> bool:
"""Check if subprocess is alive."""
@ -144,15 +145,13 @@ class ExportOrchestrator:
if self._resp_queue is None:
return None
try:
return self._resp_queue.get(timeout=timeout)
return self._resp_queue.get(timeout = timeout)
except queue.Empty:
return None
except (EOFError, OSError, ValueError):
return None
def _wait_response(
self, expected_type: str, timeout: float = 3600.0
) -> dict:
def _wait_response(self, expected_type: str, timeout: float = 3600.0) -> dict:
"""Block until a response of the expected type arrives.
Export operations can take a very long time — GGUF conversion for
@ -163,7 +162,7 @@ class ExportOrchestrator:
while time.monotonic() < deadline:
remaining = max(0.1, deadline - time.monotonic())
resp = self._read_resp(timeout=min(remaining, 2.0))
resp = self._read_resp(timeout = min(remaining, 2.0))
if resp is None:
# Check subprocess health
@ -187,7 +186,8 @@ class ExportOrchestrator:
# Other response types during wait — skip
logger.debug(
"Skipping response type '%s' while waiting for '%s'",
rtype, expected_type,
rtype,
expected_type,
)
raise RuntimeError(
@ -233,15 +233,15 @@ class ExportOrchestrator:
if self._ensure_subprocess_alive():
self._shutdown_subprocess()
elif self._proc is not None:
self._shutdown_subprocess(timeout=2)
self._shutdown_subprocess(timeout = 2)
logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path)
self._spawn_subprocess(sub_config)
try:
resp = self._wait_response("loaded", timeout=300)
resp = self._wait_response("loaded", timeout = 300)
except RuntimeError as exc:
self._shutdown_subprocess(timeout=5)
self._shutdown_subprocess(timeout = 5)
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
@ -271,14 +271,17 @@ class ExportOrchestrator:
private: bool = False,
) -> Tuple[bool, str]:
"""Export merged PEFT model."""
return self._run_export("merged", {
"save_directory": save_directory,
"format_type": format_type,
"push_to_hub": push_to_hub,
"repo_id": repo_id,
"hf_token": hf_token,
"private": private,
})
return self._run_export(
"merged",
{
"save_directory": save_directory,
"format_type": format_type,
"push_to_hub": push_to_hub,
"repo_id": repo_id,
"hf_token": hf_token,
"private": private,
},
)
def export_base_model(
self,
@ -290,14 +293,17 @@ class ExportOrchestrator:
base_model_id: Optional[str] = None,
) -> Tuple[bool, str]:
"""Export base model (non-PEFT)."""
return self._run_export("base", {
"save_directory": save_directory,
"push_to_hub": push_to_hub,
"repo_id": repo_id,
"hf_token": hf_token,
"private": private,
"base_model_id": base_model_id,
})
return self._run_export(
"base",
{
"save_directory": save_directory,
"push_to_hub": push_to_hub,
"repo_id": repo_id,
"hf_token": hf_token,
"private": private,
"base_model_id": base_model_id,
},
)
def export_gguf(
self,
@ -308,13 +314,16 @@ class ExportOrchestrator:
hf_token: Optional[str] = None,
) -> Tuple[bool, str]:
"""Export model in GGUF format."""
return self._run_export("gguf", {
"save_directory": save_directory,
"quantization_method": quantization_method,
"push_to_hub": push_to_hub,
"repo_id": repo_id,
"hf_token": hf_token,
})
return self._run_export(
"gguf",
{
"save_directory": save_directory,
"quantization_method": quantization_method,
"push_to_hub": push_to_hub,
"repo_id": repo_id,
"hf_token": hf_token,
},
)
def export_lora_adapter(
self,
@ -325,13 +334,16 @@ class ExportOrchestrator:
private: bool = False,
) -> Tuple[bool, str]:
"""Export LoRA adapter only."""
return self._run_export("lora", {
"save_directory": save_directory,
"push_to_hub": push_to_hub,
"repo_id": repo_id,
"hf_token": hf_token,
"private": private,
})
return self._run_export(
"lora",
{
"save_directory": save_directory,
"push_to_hub": push_to_hub,
"repo_id": repo_id,
"hf_token": hf_token,
"private": private,
},
)
def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str]:
"""Send an export command to the subprocess and wait for result."""
@ -344,7 +356,7 @@ class ExportOrchestrator:
self._send_cmd(cmd)
resp = self._wait_response(
f"export_{export_type}_done",
timeout=3600, # GGUF for 30B+ models can take 30+ min
timeout = 3600, # GGUF for 30B+ models can take 30+ min
)
return resp.get("success", False), resp.get("message", "")
except RuntimeError as exc:
@ -361,7 +373,7 @@ class ExportOrchestrator:
try:
self._send_cmd({"type": "cleanup"})
resp = self._wait_response("cleanup_done", timeout=30)
resp = self._wait_response("cleanup_done", timeout = 30)
success = resp.get("success", False)
except RuntimeError:
success = False
@ -379,7 +391,8 @@ class ExportOrchestrator:
) -> List[Tuple[str, list]]:
"""Scan for checkpoints — no ML imports needed, runs locally."""
from utils.models.checkpoints import scan_checkpoints
return scan_checkpoints(outputs_dir=outputs_dir)
return scan_checkpoints(outputs_dir = outputs_dir)
# ========== GLOBAL INSTANCE ==========

341
studio/backend/core/export/worker.py
View file

@ -14,6 +14,7 @@ shutdown) via mp.Queue.
Pattern follows core/inference/worker.py and core/training/worker.py.
"""
from __future__ import annotations
import structlog
@ -43,25 +44,45 @@ def _activate_transformers_version(model_name: str) -> None:
resolved = _resolve_base_model(model_name)
if needs_transformers_5(resolved):
venv_t5 = os.path.join(os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5")
venv_t5 = os.path.join(
os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5"
)
if os.path.isdir(venv_t5):
sys.path.insert(0, venv_t5)
logger.info("Activated transformers 5.x from %s", venv_t5)
else:
# Fallback: pip install at runtime (slower, ~10-15s)
logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5)
import subprocess as sp
os.makedirs(venv_t5, exist_ok=True)
os.makedirs(venv_t5, exist_ok = True)
r1 = sp.run(
[sys.executable, "-m", "pip", "install", "--target", venv_t5,
"--no-deps", "transformers==5.2.0"],
stdout=sp.PIPE, stderr=sp.STDOUT,
[
sys.executable,
"-m",
"pip",
"install",
"--target",
venv_t5,
"--no-deps",
"transformers==5.2.0",
],
stdout = sp.PIPE,
stderr = sp.STDOUT,
)
r2 = sp.run(
[sys.executable, "-m", "pip", "install", "--target", venv_t5,
"--no-deps", "huggingface_hub==1.3.0"],
stdout=sp.PIPE, stderr=sp.STDOUT,
[
sys.executable,
"-m",
"pip",
"install",
"--target",
venv_t5,
"--no-deps",
"huggingface_hub==1.3.0",
],
stdout = sp.PIPE,
stderr = sp.STDOUT,
)
if r1.returncode != 0 or r2.returncode != 0:
raise RuntimeError(
@ -92,37 +113,46 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
trust_remote_code = cmd.get("trust_remote_code", False)
try:
_send_response(resp_queue, {
"type": "status",
"message": f"Loading checkpoint: {checkpoint_path}",
"ts": time.time(),
})
success, message = backend.load_checkpoint(
checkpoint_path=checkpoint_path,
max_seq_length=max_seq_length,
load_in_4bit=load_in_4bit,
trust_remote_code=trust_remote_code,
_send_response(
resp_queue,
{
"type": "status",
"message": f"Loading checkpoint: {checkpoint_path}",
"ts": time.time(),
},
)
_send_response(resp_queue, {
"type": "loaded",
"success": success,
"message": message,
"checkpoint": checkpoint_path if success else None,
"is_vision": backend.is_vision if success else False,
"is_peft": backend.is_peft if success else False,
"ts": time.time(),
})
success, message = backend.load_checkpoint(
checkpoint_path = checkpoint_path,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
)
_send_response(
resp_queue,
{
"type": "loaded",
"success": success,
"message": message,
"checkpoint": checkpoint_path if success else None,
"is_vision": backend.is_vision if success else False,
"is_peft": backend.is_peft if success else False,
"ts": time.time(),
},
)
except Exception as exc:
_send_response(resp_queue, {
"type": "loaded",
"success": False,
"message": str(exc),
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "loaded",
"success": False,
"message": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
@ -133,74 +163,86 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
try:
if export_type == "merged":
success, message = backend.export_merged_model(
save_directory=cmd.get("save_directory", ""),
format_type=cmd.get("format_type", "16-bit (FP16)"),
push_to_hub=cmd.get("push_to_hub", False),
repo_id=cmd.get("repo_id"),
hf_token=cmd.get("hf_token"),
private=cmd.get("private", False),
save_directory = cmd.get("save_directory", ""),
format_type = cmd.get("format_type", "16-bit (FP16)"),
push_to_hub = cmd.get("push_to_hub", False),
repo_id = cmd.get("repo_id"),
hf_token = cmd.get("hf_token"),
private = cmd.get("private", False),
)
elif export_type == "base":
success, message = backend.export_base_model(
save_directory=cmd.get("save_directory", ""),
push_to_hub=cmd.get("push_to_hub", False),
repo_id=cmd.get("repo_id"),
hf_token=cmd.get("hf_token"),
private=cmd.get("private", False),
base_model_id=cmd.get("base_model_id"),
save_directory = cmd.get("save_directory", ""),
push_to_hub = cmd.get("push_to_hub", False),
repo_id = cmd.get("repo_id"),
hf_token = cmd.get("hf_token"),
private = cmd.get("private", False),
base_model_id = cmd.get("base_model_id"),
)
elif export_type == "gguf":
success, message = backend.export_gguf(
save_directory=cmd.get("save_directory", ""),
quantization_method=cmd.get("quantization_method", "Q4_K_M"),
push_to_hub=cmd.get("push_to_hub", False),
repo_id=cmd.get("repo_id"),
hf_token=cmd.get("hf_token"),
save_directory = cmd.get("save_directory", ""),
quantization_method = cmd.get("quantization_method", "Q4_K_M"),
push_to_hub = cmd.get("push_to_hub", False),
repo_id = cmd.get("repo_id"),
hf_token = cmd.get("hf_token"),
)
elif export_type == "lora":
success, message = backend.export_lora_adapter(
save_directory=cmd.get("save_directory", ""),
push_to_hub=cmd.get("push_to_hub", False),
repo_id=cmd.get("repo_id"),
hf_token=cmd.get("hf_token"),
private=cmd.get("private", False),
save_directory = cmd.get("save_directory", ""),
push_to_hub = cmd.get("push_to_hub", False),
repo_id = cmd.get("repo_id"),
hf_token = cmd.get("hf_token"),
private = cmd.get("private", False),
)
else:
success, message = False, f"Unknown export type: {export_type}"
_send_response(resp_queue, {
"type": response_type,
"success": success,
"message": message,
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": response_type,
"success": success,
"message": message,
"ts": time.time(),
},
)
except Exception as exc:
_send_response(resp_queue, {
"type": response_type,
"success": False,
"message": str(exc),
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": response_type,
"success": False,
"message": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
def _handle_cleanup(backend, resp_queue: Any) -> None:
"""Handle a cleanup command."""
try:
success = backend.cleanup_memory()
_send_response(resp_queue, {
"type": "cleanup_done",
"success": success,
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "cleanup_done",
"success": success,
"ts": time.time(),
},
)
except Exception as exc:
_send_response(resp_queue, {
"type": "cleanup_done",
"success": False,
"message": str(exc),
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "cleanup_done",
"success": False,
"message": str(exc),
"ts": time.time(),
},
)
def run_export_process(
@ -219,16 +261,19 @@ def run_export_process(
import queue as _queue
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
os.environ["PYTHONWARNINGS"] = (
"ignore" # Suppress warnings at C-level before imports
)
import warnings
from loggers.config import LogConfig
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
LogConfig.setup_logging(
service_name="unsloth-studio-export-worker",
env=os.getenv("ENVIRONMENT_TYPE", "production"),
service_name = "unsloth-studio-export-worker",
env = os.getenv("ENVIRONMENT_TYPE", "production"),
)
checkpoint_path = config["checkpoint_path"]
@ -237,18 +282,22 @@ def run_export_process(
try:
_activate_transformers_version(checkpoint_path)
except Exception as exc:
_send_response(resp_queue, {
"type": "error",
"error": f"Failed to activate transformers version: {exc}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "error",
"error": f"Failed to activate transformers version: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
# ── 1b. On Windows, check Triton availability (must be before import torch) ──
if sys.platform == "win32":
try:
import triton # noqa: F401
logger.info("Triton available — torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
@ -259,11 +308,14 @@ def run_export_process(
# ── 2. Import ML libraries (fresh in this clean process) ──
try:
_send_response(resp_queue, {
"type": "status",
"message": "Importing ML libraries...",
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "status",
"message": "Importing ML libraries...",
"ts": time.time(),
},
)
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
@ -272,15 +324,21 @@ def run_export_process(
from core.export.export import ExportBackend
import transformers
logger.info("Export subprocess loaded transformers %s", transformers.__version__)
logger.info(
"Export subprocess loaded transformers %s", transformers.__version__
)
except Exception as exc:
_send_response(resp_queue, {
"type": "error",
"error": f"Failed to import ML libraries: {exc}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "error",
"error": f"Failed to import ML libraries: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
# ── 3. Create export backend and load initial checkpoint ──
@ -290,12 +348,15 @@ def run_export_process(
_handle_load(backend, config, resp_queue)
except Exception as exc:
_send_response(resp_queue, {
"type": "error",
"error": f"Failed to initialize export backend: {exc}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "error",
"error": f"Failed to initialize export backend: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
# ── 4. Command loop — process commands until shutdown ──
@ -303,7 +364,7 @@ def run_export_process(
while True:
try:
cmd = cmd_queue.get(timeout=1.0)
cmd = cmd_queue.get(timeout = 1.0)
except _queue.Empty:
continue
except (EOFError, OSError):
@ -329,13 +390,16 @@ def run_export_process(
_handle_cleanup(backend, resp_queue)
elif cmd_type == "status":
_send_response(resp_queue, {
"type": "status_response",
"checkpoint": backend.current_checkpoint,
"is_vision": backend.is_vision,
"is_peft": backend.is_peft,
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "status_response",
"checkpoint": backend.current_checkpoint,
"is_vision": backend.is_vision,
"is_peft": backend.is_peft,
"ts": time.time(),
},
)
elif cmd_type == "shutdown":
logger.info("Shutdown command received, cleaning up and exiting")
@ -343,25 +407,36 @@ def run_export_process(
backend.cleanup_memory()
except Exception:
pass
_send_response(resp_queue, {
"type": "shutdown_ack",
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "shutdown_ack",
"ts": time.time(),
},
)
return
else:
logger.warning("Unknown command type: %s", cmd_type)
_send_response(resp_queue, {
"type": "error",
"error": f"Unknown command type: {cmd_type}",
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "error",
"error": f"Unknown command type: {cmd_type}",
"ts": time.time(),
},
)
except Exception as exc:
logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info=True)
_send_response(resp_queue, {
"type": "error",
"error": f"Command '{cmd_type}' failed: {exc}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
logger.error(
"Error handling command '%s': %s", cmd_type, exc, exc_info = True
)
_send_response(
resp_queue,
{
"type": "error",
"error": f"Command '{cmd_type}' failed: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)

9
studio/backend/core/inference/__init__.py
View file

@ -8,6 +8,7 @@ The default get_inference_backend() returns an InferenceOrchestrator that
delegates to a subprocess. The original InferenceBackend runs inside
the subprocess and can be imported directly from .inference when needed.
"""
from .orchestrator import InferenceOrchestrator, get_inference_backend
from .llama_cpp import LlamaCppBackend
@ -15,8 +16,8 @@ from .llama_cpp import LlamaCppBackend
InferenceBackend = InferenceOrchestrator
__all__ = [
'InferenceBackend',
'InferenceOrchestrator',
'get_inference_backend',
'LlamaCppBackend',
"InferenceBackend",
"InferenceOrchestrator",
"get_inference_backend",
"LlamaCppBackend",
]

73
studio/backend/core/inference/audio_codecs.py
View file

@ -5,6 +5,7 @@
Audio codec loading and decoding for TTS inference.
Supports: SNAC (Orpheus), CSM (Sesame), BiCodec (Spark), DAC (OuteTTS)
"""
import io
import re
import wave
@ -45,7 +46,12 @@ class AudioCodecManager:
self._bicodec_repo_path = None
self._dac_audio_codec = None
def load_codec(self, audio_type: str, device: str = "cuda", model_repo_path: Optional[str] = None) -> None:
def load_codec(
self,
audio_type: str,
device: str = "cuda",
model_repo_path: Optional[str] = None,
) -> None:
"""Load the appropriate codec for the given audio type."""
if audio_type == "snac":
self._load_snac(device)
@ -64,7 +70,10 @@ class AudioCodecManager:
if self._snac_model is not None:
return
from snac import SNAC
self._snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
self._snac_model = (
SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
)
logger.info("Loaded SNAC codec (24kHz)")
def _load_bicodec(self, device: str, model_repo_path: Optional[str] = None) -> None:
@ -76,13 +85,22 @@ class AudioCodecManager:
# Clone SparkAudio/Spark-TTS GitHub repo for the sparktts Python package
# (same approach as training — the HF model repos don't contain the package)
spark_code_dir = os.path.join(os.path.dirname(model_repo_path or "."), "Spark-TTS")
spark_code_dir = os.path.join(
os.path.dirname(model_repo_path or "."), "Spark-TTS"
)
sparktts_pkg = os.path.join(spark_code_dir, "sparktts")
if not os.path.isdir(sparktts_pkg):
logger.info(f"Cloning SparkAudio/Spark-TTS to {spark_code_dir}...")
subprocess.run(
["git", "clone", "--depth", "1", "https://github.com/SparkAudio/Spark-TTS", spark_code_dir],
check=True,
[
"git",
"clone",
"--depth",
"1",
"https://github.com/SparkAudio/Spark-TTS",
spark_code_dir,
],
check = True,
)
if spark_code_dir not in sys.path:
@ -112,8 +130,15 @@ class AudioCodecManager:
if not os.path.isdir(outetts_pkg):
logger.info(f"Cloning edwko/OuteTTS to {outetts_code_dir}...")
subprocess.run(
["git", "clone", "--depth", "1", "https://github.com/edwko/OuteTTS", outetts_code_dir],
check=True,
[
"git",
"clone",
"--depth",
"1",
"https://github.com/edwko/OuteTTS",
outetts_code_dir,
],
check = True,
)
# Remove files that pull in heavy / incompatible dependencies
# (matches notebook: gguf_model.py is under models/, others under outetts/)
@ -134,17 +159,19 @@ class AudioCodecManager:
from outetts.models.config import ModelConfig as OuteTTSModelConfig
dummy_config = OuteTTSModelConfig(
tokenizer_path="OuteAI/Llama-OuteTTS-1.0-1B",
device=device,
audio_codec_path=None,
tokenizer_path = "OuteAI/Llama-OuteTTS-1.0-1B",
device = device,
audio_codec_path = None,
)
processor = AudioProcessor(config=dummy_config)
processor = AudioProcessor(config = dummy_config)
self._dac_audio_codec = processor.audio_codec
logger.info("Loaded DAC audio codec")
# ── Decoders ─────────────────────────────────────────────────
def decode_snac(self, generated_ids: torch.Tensor, device: str) -> Tuple[bytes, int]:
def decode_snac(
self, generated_ids: torch.Tensor, device: str
) -> Tuple[bytes, int]:
"""
Decode SNAC tokens (Orpheus) into WAV bytes.
@ -155,12 +182,14 @@ class AudioCodecManager:
Returns (wav_bytes, 24000).
"""
# Find START_OF_SPEECH token (128257)
token_indices = (generated_ids == 128257).nonzero(as_tuple=True)
token_indices = (generated_ids == 128257).nonzero(as_tuple = True)
if len(token_indices[1]) > 0:
cropped = generated_ids[:, token_indices[1][-1] + 1:]
cropped = generated_ids[:, token_indices[1][-1] + 1 :]
else:
# Gracefully fall back to using entire output if marker not found
logger.warning("No START_OF_SPEECH token (128257) found — using full generated output")
logger.warning(
"No START_OF_SPEECH token (128257) found — using full generated output"
)
cropped = generated_ids
row = cropped[0]
@ -213,14 +242,20 @@ class AudioCodecManager:
semantic_matches = re.findall(r"<\|bicodec_semantic_(\d+)\|>", generated_text)
global_matches = re.findall(r"<\|bicodec_global_(\d+)\|>", generated_text)
logger.info(f"BiCodec decode: {len(global_matches)} global tokens, {len(semantic_matches)} semantic tokens")
logger.info(
f"BiCodec decode: {len(global_matches)} global tokens, {len(semantic_matches)} semantic tokens"
)
if len(global_matches) < 10:
logger.info(f"BiCodec generated text (first 500 chars): {generated_text[:500]}")
logger.info(
f"BiCodec generated text (first 500 chars): {generated_text[:500]}"
)
if not semantic_matches:
raise ValueError("No bicodec_semantic tokens found in generated output")
semantic_ids = torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0)
semantic_ids = (
torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0)
)
# Speaker encoder expects exactly 32 global tokens (token_num=32 in BiCodec config).
# Pad with zeros or truncate to 32.
@ -260,7 +295,7 @@ class AudioCodecManager:
c1 = c1[:t]
c2 = c2[:t]
codes = torch.tensor([[c1, c2]], dtype=torch.int64).to(device)
codes = torch.tensor([[c1, c2]], dtype = torch.int64).to(device)
with torch.no_grad():
audio = self._dac_audio_codec.decode(codes)

855
studio/backend/core/inference/inference.py
View file

File diff suppressed because it is too large Load diff

109
studio/backend/core/inference/llama_cpp.py
View file

@ -7,6 +7,7 @@ llama-server inference backend for GGUF models.
Manages a llama-server subprocess and proxies chat completions
through its OpenAI-compatible /v1/chat/completions endpoint.
"""
import atexit
import json
import structlog
@ -146,7 +147,9 @@ class LlamaCppBackend:
if build_path.is_file():
return str(build_path)
if sys.platform == "win32":
win_path = project_root / "llama.cpp" / "build" / "bin" / "Release" / binary_name
win_path = (
project_root / "llama.cpp" / "build" / "bin" / "Release" / binary_name
)
if win_path.is_file():
return str(win_path)
@ -168,7 +171,7 @@ class LlamaCppBackend:
def _find_free_port() -> int:
"""Find an available TCP port."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
# ── Stdout drain (prevents pipe deadlock on Windows) ─────────
@ -258,15 +261,19 @@ class LlamaCppBackend:
try:
import re
from huggingface_hub import list_repo_files
files = list_repo_files(hf_repo, token=hf_token)
files = list_repo_files(hf_repo, token = hf_token)
variant_lower = hf_variant.lower()
# Use word-boundary matching so "Q8_0" doesn't also
# match "IQ8_0" or other superset variant names.
boundary = re.compile(
r'(?<![a-zA-Z0-9])' + re.escape(variant_lower) + r'(?![a-zA-Z0-9])'
r"(?<![a-zA-Z0-9])"
+ re.escape(variant_lower)
+ r"(?![a-zA-Z0-9])"
)
gguf_files = sorted(
f for f in files
f
for f in files
if f.endswith(".gguf") and boundary.search(f.lower())
)
if gguf_files:
@ -274,17 +281,20 @@ class LlamaCppBackend:
# For split GGUFs (e.g. model-Q8_0-00001-of-00003.gguf)
# discover siblings by exact basename + total match
# so "model-Q8_0-v2-*" isn't pulled in as a sibling.
shard_pat = re.compile(r'^(.*)-\d{5}-of-(\d{5})\.gguf$')
shard_pat = re.compile(r"^(.*)-\d{5}-of-(\d{5})\.gguf$")
m = shard_pat.match(gguf_filename)
if m:
prefix = m.group(1)
total = m.group(2)
sibling_pat = re.compile(
r'^' + re.escape(prefix) + r'-\d{5}-of-' + re.escape(total) + r'\.gguf$'
r"^"
+ re.escape(prefix)
+ r"-\d{5}-of-"
+ re.escape(total)
+ r"\.gguf$"
)
gguf_extra_shards = [
f for f in gguf_files[1:]
if sibling_pat.match(f)
f for f in gguf_files[1:] if sibling_pat.match(f)
]
except Exception as e:
logger.warning(f"Could not list repo files: {e}")
@ -295,22 +305,28 @@ class LlamaCppBackend:
repo_name = hf_repo.split("/")[-1].replace("-GGUF", "")
gguf_filename = f"{repo_name}-{hf_variant}.gguf"
logger.info(f"Downloading GGUF: {hf_repo}/{gguf_filename}"
+ (f" (+{len(gguf_extra_shards)} shards)" if gguf_extra_shards else ""))
logger.info(
f"Downloading GGUF: {hf_repo}/{gguf_filename}"
+ (
f" (+{len(gguf_extra_shards)} shards)"
if gguf_extra_shards
else ""
)
)
try:
local_path = hf_hub_download(
repo_id=hf_repo,
filename=gguf_filename,
token=hf_token,
repo_id = hf_repo,
filename = gguf_filename,
token = hf_token,
)
# Download remaining shards for split GGUFs — llama-server
# auto-discovers them when they are in the same directory.
for shard in gguf_extra_shards:
logger.info(f"Downloading GGUF shard: {shard}")
hf_hub_download(
repo_id=hf_repo,
filename=shard,
token=hf_token,
repo_id = hf_repo,
filename = shard,
token = hf_token,
)
except Exception as e:
raise RuntimeError(
@ -320,20 +336,28 @@ class LlamaCppBackend:
logger.info(f"GGUF downloaded to: {local_path}")
cmd = [
binary,
"-m", local_path,
"--port", str(self._port),
"-c", str(n_ctx),
"-ngl", str(n_gpu_layers),
"-m",
local_path,
"--port",
str(self._port),
"-c",
str(n_ctx),
"-ngl",
str(n_gpu_layers),
]
elif gguf_path:
if not Path(gguf_path).is_file():
raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
cmd = [
binary,
"-m", gguf_path,
"--port", str(self._port),
"-c", str(n_ctx),
"-ngl", str(n_gpu_layers),
"-m",
gguf_path,
"--port",
str(self._port),
"-c",
str(n_ctx),
"-ngl",
str(n_gpu_layers),
]
else:
raise ValueError("Either gguf_path or hf_repo must be provided")
@ -354,6 +378,7 @@ class LlamaCppBackend:
# Set library paths so llama-server can find its shared libs and CUDA DLLs
import os
import sys
env = os.environ.copy()
binary_dir = str(Path(binary).parent)
@ -375,20 +400,22 @@ class LlamaCppBackend:
else:
# Linux: set LD_LIBRARY_PATH for shared libs next to the binary
existing_ld = env.get("LD_LIBRARY_PATH", "")
env["LD_LIBRARY_PATH"] = f"{binary_dir}:{existing_ld}" if existing_ld else binary_dir
env["LD_LIBRARY_PATH"] = (
f"{binary_dir}:{existing_ld}" if existing_ld else binary_dir
)
self._stdout_lines = []
self._process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
env=env,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
env = env,
)
# Start background thread to drain stdout and prevent pipe deadlock
self._stdout_thread = threading.Thread(
target=self._drain_stdout, daemon=True, name="llama-stdout"
target = self._drain_stdout, daemon = True, name = "llama-stdout"
)
self._stdout_thread.start()
@ -399,7 +426,7 @@ class LlamaCppBackend:
self._model_identifier = model_identifier
# Wait for llama-server to become healthy
if not self._wait_for_health(timeout=120.0):
if not self._wait_for_health(timeout = 120.0):
self._kill_process()
raise RuntimeError(
"llama-server failed to start. "
@ -434,17 +461,17 @@ class LlamaCppBackend:
return
try:
self._process.terminate()
self._process.wait(timeout=5)
self._process.wait(timeout = 5)
except subprocess.TimeoutExpired:
logger.warning("llama-server did not exit on SIGTERM, sending SIGKILL")
self._process.kill()
self._process.wait(timeout=5)
self._process.wait(timeout = 5)
except Exception as e:
logger.warning(f"Error killing llama-server process: {e}")
finally:
self._process = None
if self._stdout_thread is not None:
self._stdout_thread.join(timeout=2)
self._stdout_thread.join(timeout = 2)
self._stdout_thread = None
def _cleanup(self):
@ -465,7 +492,7 @@ class LlamaCppBackend:
if self._process.poll() is not None:
# Give the drain thread a moment to collect final output
if self._stdout_thread is not None:
self._stdout_thread.join(timeout=2)
self._stdout_thread.join(timeout = 2)
output = "\n".join(self._stdout_lines[-50:])
logger.error(
f"llama-server exited with code {self._process.returncode}. "
@ -474,7 +501,7 @@ class LlamaCppBackend:
return False
try:
resp = httpx.get(url, timeout=2.0)
resp = httpx.get(url, timeout = 2.0)
if resp.status_code == 200:
return True
except (httpx.ConnectError, httpx.TimeoutException):
@ -567,8 +594,8 @@ class LlamaCppBackend:
cumulative = ""
try:
with httpx.Client(timeout=None) as client:
with client.stream("POST", url, json=payload) as response:
with httpx.Client(timeout = None) as client:
with client.stream("POST", url, json = payload) as response:
if response.status_code != 200:
error_body = response.read().decode()
raise RuntimeError(
@ -602,7 +629,9 @@ class LlamaCppBackend:
cumulative += token
yield cumulative
except json.JSONDecodeError:
logger.debug(f"Skipping malformed SSE line: {line[:100]}")
logger.debug(
f"Skipping malformed SSE line: {line[:100]}"
)
except httpx.ConnectError:
raise RuntimeError("Lost connection to llama-server")

185
studio/backend/core/inference/orchestrator.py
View file

@ -14,6 +14,7 @@ is killed and a new one is spawned with the correct version.
Pattern follows core/training/training.py.
"""
import atexit
import base64
import structlog
@ -48,7 +49,9 @@ class InferenceOrchestrator:
self._resp_queue: Any = None
self._cancel_event: Any = None # mp.Event — set to cancel generation instantly
self._lock = threading.Lock()
self._gen_lock = threading.Lock() # Serializes generation — one request at a time
self._gen_lock = (
threading.Lock()
) # Serializes generation — one request at a time
# Local state mirrors (updated from subprocess responses)
self.active_model_name: Optional[str] = None
@ -83,14 +86,14 @@ class InferenceOrchestrator:
self._cancel_event = _CTX.Event()
self._proc = _CTX.Process(
target=run_inference_process,
kwargs={
target = run_inference_process,
kwargs = {
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,
"cancel_event": self._cancel_event,
"config": config,
},
daemon=True,
daemon = True,
)
self._proc.start()
logger.info("Inference subprocess started (pid=%s)", self._proc.pid)
@ -121,7 +124,7 @@ class InferenceOrchestrator:
# 4. Wait for graceful shutdown
try:
self._proc.join(timeout=timeout)
self._proc.join(timeout = timeout)
except Exception:
pass
@ -130,14 +133,14 @@ class InferenceOrchestrator:
logger.warning("Inference subprocess did not exit gracefully, terminating")
try:
self._proc.terminate()
self._proc.join(timeout=5)
self._proc.join(timeout = 5)
except Exception:
pass
if self._proc is not None and self._proc.is_alive():
logger.warning("Subprocess still alive after terminate, killing")
try:
self._proc.kill()
self._proc.join(timeout=3)
self._proc.join(timeout = 3)
except Exception:
pass
@ -149,7 +152,7 @@ class InferenceOrchestrator:
def _cleanup(self):
"""atexit handler."""
self._shutdown_subprocess(timeout=5.0)
self._shutdown_subprocess(timeout = 5.0)
def _ensure_subprocess_alive(self) -> bool:
"""Check if subprocess is alive."""
@ -173,15 +176,13 @@ class InferenceOrchestrator:
if self._resp_queue is None:
return None
try:
return self._resp_queue.get(timeout=timeout)
return self._resp_queue.get(timeout = timeout)
except queue.Empty:
return None
except (EOFError, OSError, ValueError):
return None
def _wait_response(
self, expected_type: str, timeout: float = 120.0
) -> dict:
def _wait_response(self, expected_type: str, timeout: float = 120.0) -> dict:
"""Block until a response of the expected type arrives.
Also handles 'status' and 'error' events during the wait.
@ -192,7 +193,7 @@ class InferenceOrchestrator:
while time.monotonic() < deadline:
remaining = max(0.1, deadline - time.monotonic())
resp = self._read_resp(timeout=min(remaining, 1.0))
resp = self._read_resp(timeout = min(remaining, 1.0))
if resp is None:
# Check subprocess health
@ -214,7 +215,11 @@ class InferenceOrchestrator:
continue
# Other response types during wait — skip
logger.debug("Skipping response type '%s' while waiting for '%s'", rtype, expected_type)
logger.debug(
"Skipping response type '%s' while waiting for '%s'",
rtype,
expected_type,
)
raise RuntimeError(
f"Timeout waiting for '{expected_type}' response after {timeout}s"
@ -241,7 +246,7 @@ class InferenceOrchestrator:
"""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
resp = self._read_resp(timeout=min(0.5, deadline - time.monotonic()))
resp = self._read_resp(timeout = min(0.5, deadline - time.monotonic()))
if resp is None:
if not self._ensure_subprocess_alive():
return
@ -259,7 +264,7 @@ class InferenceOrchestrator:
self,
config, # ModelConfig
max_seq_length: int = 2048,
dtype=None,
dtype = None,
load_in_4bit: bool = True,
hf_token: Optional[str] = None,
trust_remote_code: bool = False,
@ -298,14 +303,15 @@ class InferenceOrchestrator:
elif self._proc is not None:
# Dead subprocess — clean up
self._shutdown_subprocess(timeout=2)
self._shutdown_subprocess(timeout = 2)
logger.info(
"Spawning fresh inference subprocess for '%s' (transformers %s.x)",
model_name, needed_major,
model_name,
needed_major,
)
self._spawn_subprocess(sub_config)
resp = self._wait_response("loaded", timeout=180)
resp = self._wait_response("loaded", timeout = 180)
# Update local state from response
if resp.get("success"):
@ -346,11 +352,13 @@ class InferenceOrchestrator:
return True
try:
self._send_cmd({
"type": "unload",
"model_name": model_name,
})
resp = self._wait_response("unloaded", timeout=30)
self._send_cmd(
{
"type": "unload",
"model_name": model_name,
}
)
resp = self._wait_response("unloaded", timeout = 30)
# Update local state
self.models.pop(model_name, None)
@ -372,40 +380,40 @@ class InferenceOrchestrator:
self,
messages: list,
system_prompt: str = "",
image=None,
image = None,
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_new_tokens: int = 256,
repetition_penalty: float = 1.1,
cancel_event=None,
cancel_event = None,
) -> Generator[str, None, None]:
"""Generate response, streaming tokens from subprocess."""
yield from self._generate_inner(
messages=messages,
system_prompt=system_prompt,
image=image,
temperature=temperature,
top_p=top_p,
top_k=top_k,
min_p=min_p,
max_new_tokens=max_new_tokens,
repetition_penalty=repetition_penalty,
cancel_event=cancel_event,
use_adapter=None,
messages = messages,
system_prompt = system_prompt,
image = image,
temperature = temperature,
top_p = top_p,
top_k = top_k,
min_p = min_p,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
cancel_event = cancel_event,
use_adapter = None,
)
def generate_with_adapter_control(
self,
use_adapter: Optional[Union[bool, str]] = None,
cancel_event=None,
cancel_event = None,
**gen_kwargs,
) -> Generator[str, None, None]:
"""Generate with adapter control, streaming tokens from subprocess."""
yield from self._generate_inner(
use_adapter=use_adapter,
cancel_event=cancel_event,
use_adapter = use_adapter,
cancel_event = cancel_event,
**gen_kwargs,
)
@ -413,15 +421,15 @@ class InferenceOrchestrator:
self,
messages: list = None,
system_prompt: str = "",
image=None,
image = None,
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_new_tokens: int = 256,
repetition_penalty: float = 1.1,
cancel_event=None,
use_adapter=None,
cancel_event = None,
use_adapter = None,
) -> Generator[str, None, None]:
"""Inner generation logic — sends command to subprocess, yields tokens.
@ -442,32 +450,32 @@ class InferenceOrchestrator:
# can consume and drop each other's token events.
with self._gen_lock:
yield from self._generate_locked(
messages=messages,
system_prompt=system_prompt,
image=image,
temperature=temperature,
top_p=top_p,
top_k=top_k,
min_p=min_p,
max_new_tokens=max_new_tokens,
repetition_penalty=repetition_penalty,
cancel_event=cancel_event,
use_adapter=use_adapter,
messages = messages,
system_prompt = system_prompt,
image = image,
temperature = temperature,
top_p = top_p,
top_k = top_k,
min_p = min_p,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
cancel_event = cancel_event,
use_adapter = use_adapter,
)
def _generate_locked(
self,
messages: list = None,
system_prompt: str = "",
image=None,
image = None,
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_new_tokens: int = 256,
repetition_penalty: float = 1.1,
cancel_event=None,
use_adapter=None,
cancel_event = None,
use_adapter = None,
) -> Generator[str, None, None]:
"""Actual generation logic — must be called under _gen_lock."""
request_id = str(uuid.uuid4())
@ -503,7 +511,7 @@ class InferenceOrchestrator:
# Yield tokens from response queue — we are the only reader
# because _gen_lock is held.
while True:
resp = self._read_resp(timeout=30.0)
resp = self._read_resp(timeout = 30.0)
if resp is None:
# Check subprocess health
@ -531,7 +539,7 @@ class InferenceOrchestrator:
# Wait for the subprocess to acknowledge cancellation
# (gen_done/gen_error) so stale events don't leak into
# the next generation request.
self._drain_until_gen_done(timeout=5.0)
self._drain_until_gen_done(timeout = 5.0)
return
yield resp.get("text", "")
@ -577,6 +585,7 @@ class InferenceOrchestrator:
raise RuntimeError("No active model")
import uuid
request_id = str(uuid.uuid4())
cmd = {
@ -599,11 +608,13 @@ class InferenceOrchestrator:
deadline = time.monotonic() + 120.0
while time.monotonic() < deadline:
remaining = max(0.1, deadline - time.monotonic())
resp = self._read_resp(timeout=min(remaining, 1.0))
resp = self._read_resp(timeout = min(remaining, 1.0))
if resp is None:
if not self._ensure_subprocess_alive():
raise RuntimeError("Inference subprocess crashed during audio generation")
raise RuntimeError(
"Inference subprocess crashed during audio generation"
)
continue
rtype = resp.get("type", "")
@ -627,15 +638,15 @@ class InferenceOrchestrator:
def generate_whisper_response(
self,
audio_array,
cancel_event=None,
cancel_event = None,
) -> Generator[str, None, None]:
"""Whisper ASR — sends audio to subprocess, yields text."""
yield from self._generate_audio_input_inner(
audio_array=audio_array,
audio_type="whisper",
messages=[],
system_prompt="",
cancel_event=cancel_event,
audio_array = audio_array,
audio_type = "whisper",
messages = [],
system_prompt = "",
cancel_event = cancel_event,
)
def generate_audio_input_response(
@ -649,21 +660,21 @@ class InferenceOrchestrator:
min_p: float = 0.0,
max_new_tokens: int = 512,
repetition_penalty: float = 1.1,
cancel_event=None,
cancel_event = None,
) -> Generator[str, None, None]:
"""Audio input generation (e.g. Gemma 3n) — streams text tokens."""
yield from self._generate_audio_input_inner(
audio_array=audio_array,
audio_type=None, # worker will use generate_audio_input_response
messages=messages,
system_prompt=system_prompt,
temperature=temperature,
top_p=top_p,
top_k=top_k,
min_p=min_p,
max_new_tokens=max_new_tokens,
repetition_penalty=repetition_penalty,
cancel_event=cancel_event,
audio_array = audio_array,
audio_type = None, # worker will use generate_audio_input_response
messages = messages,
system_prompt = system_prompt,
temperature = temperature,
top_p = top_p,
top_k = top_k,
min_p = min_p,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
cancel_event = cancel_event,
)
def _generate_audio_input_inner(
@ -678,7 +689,7 @@ class InferenceOrchestrator:
min_p: float = 0.0,
max_new_tokens: int = 512,
repetition_penalty: float = 1.1,
cancel_event=None,
cancel_event = None,
) -> Generator[str, None, None]:
"""Shared inner logic for audio input generation (Whisper + ASR)."""
if not self._ensure_subprocess_alive():
@ -690,10 +701,15 @@ class InferenceOrchestrator:
with self._gen_lock:
import uuid
request_id = str(uuid.uuid4())
# Convert numpy array to list for mp.Queue serialization
audio_data = audio_array.tolist() if hasattr(audio_array, 'tolist') else list(audio_array)
audio_data = (
audio_array.tolist()
if hasattr(audio_array, "tolist")
else list(audio_array)
)
cmd = {
"type": "generate_audio_input",
@ -718,7 +734,7 @@ class InferenceOrchestrator:
# Yield tokens — same pattern as _generate_locked
while True:
resp = self._read_resp(timeout=30.0)
resp = self._read_resp(timeout = 30.0)
if resp is None:
if not self._ensure_subprocess_alive():
@ -738,7 +754,7 @@ class InferenceOrchestrator:
if rtype == "token":
if cancel_event is not None and cancel_event.is_set():
self._cancel_generation()
self._drain_until_gen_done(timeout=5.0)
self._drain_until_gen_done(timeout = 5.0)
return
yield resp.get("text", "")
@ -761,6 +777,7 @@ class InferenceOrchestrator:
return None
if img.size[0] > max_size or img.size[1] > max_size:
from PIL import Image
ratio = min(max_size / img.size[0], max_size / img.size[1])
new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
return img.resize(new_size, Image.Resampling.LANCZOS)
@ -770,7 +787,7 @@ class InferenceOrchestrator:
def _pil_to_base64(img) -> str:
"""Convert a PIL Image to base64 string for IPC."""
buf = BytesIO()
img.save(buf, format="PNG")
img.save(buf, format = "PNG")
return base64.b64encode(buf.getvalue()).decode("ascii")
def get_current_model(self) -> Optional[str]:

493
studio/backend/core/inference/worker.py
View file

@ -13,6 +13,7 @@ The subprocess stays alive while a model is loaded, accepting commands
Pattern follows core/training/worker.py.
"""
from __future__ import annotations
import base64
@ -45,7 +46,9 @@ def _activate_transformers_version(model_name: str) -> None:
resolved = _resolve_base_model(model_name)
if needs_transformers_5(resolved):
venv_t5 = os.path.join(os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5")
venv_t5 = os.path.join(
os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5"
)
if os.path.isdir(venv_t5):
sys.path.insert(0, venv_t5)
logger.info("Activated transformers 5.x from %s", venv_t5)
@ -53,16 +56,35 @@ def _activate_transformers_version(model_name: str) -> None:
# Fallback: pip install at runtime (slower, ~10-15s)
logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5)
import subprocess as sp
os.makedirs(venv_t5, exist_ok=True)
os.makedirs(venv_t5, exist_ok = True)
r1 = sp.run(
[sys.executable, "-m", "pip", "install", "--target", venv_t5,
"--no-deps", "transformers==5.2.0"],
stdout=sp.PIPE, stderr=sp.STDOUT,
[
sys.executable,
"-m",
"pip",
"install",
"--target",
venv_t5,
"--no-deps",
"transformers==5.2.0",
],
stdout = sp.PIPE,
stderr = sp.STDOUT,
)
r2 = sp.run(
[sys.executable, "-m", "pip", "install", "--target", venv_t5,
"--no-deps", "huggingface_hub==1.3.0"],
stdout=sp.PIPE, stderr=sp.STDOUT,
[
sys.executable,
"-m",
"pip",
"install",
"--target",
venv_t5,
"--no-deps",
"huggingface_hub==1.3.0",
],
stdout = sp.PIPE,
stderr = sp.STDOUT,
)
if r1.returncode != 0 or r2.returncode != 0:
raise RuntimeError(
@ -80,6 +102,7 @@ def _activate_transformers_version(model_name: str) -> None:
def _decode_image(image_base64: str):
"""Decode base64 string to PIL.Image."""
from PIL import Image
image_data = base64.b64decode(image_base64)
return Image.open(BytesIO(image_data))
@ -90,6 +113,7 @@ def _resize_image(img, max_size: int = 800):
return None
if img.size[0] > max_size or img.size[1] > max_size:
from PIL import Image
ratio = min(max_size / img.size[0], max_size / img.size[1])
new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
return img.resize(new_size, Image.Resampling.LANCZOS)
@ -114,9 +138,9 @@ def _build_model_config(config: dict):
gguf_variant = config.get("gguf_variant")
mc = ModelConfig.from_identifier(
model_id=model_name,
hf_token=hf_token,
gguf_variant=gguf_variant,
model_id = model_name,
hf_token = hf_token,
gguf_variant = gguf_variant,
)
if not mc:
raise ValueError(f"Invalid model identifier: {model_name}")
@ -136,6 +160,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
if mc.is_lora and mc.path:
import json
from pathlib import Path
adapter_cfg_path = Path(mc.path) / "adapter_config.json"
if adapter_cfg_path.exists():
try:
@ -143,24 +168,34 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
adapter_cfg = json.load(f)
training_method = adapter_cfg.get("unsloth_training_method")
if training_method == "lora" and load_in_4bit:
logger.info("adapter_config.json says lora — setting load_in_4bit=False")
logger.info(
"adapter_config.json says lora — setting load_in_4bit=False"
)
load_in_4bit = False
elif training_method == "qlora" and not load_in_4bit:
logger.info("adapter_config.json says qlora — setting load_in_4bit=True")
logger.info(
"adapter_config.json says qlora — setting load_in_4bit=True"
)
load_in_4bit = True
elif not training_method:
if mc.base_model and "-bnb-4bit" not in mc.base_model.lower() and load_in_4bit:
logger.info("No training method, base model has no -bnb-4bit — setting load_in_4bit=False")
if (
mc.base_model
and "-bnb-4bit" not in mc.base_model.lower()
and load_in_4bit
):
logger.info(
"No training method, base model has no -bnb-4bit — setting load_in_4bit=False"
)
load_in_4bit = False
except Exception as e:
logger.warning("Could not read adapter_config.json: %s", e)
success = backend.load_model(
config=mc,
max_seq_length=config.get("max_seq_length", 2048),
load_in_4bit=load_in_4bit,
hf_token=hf_token,
trust_remote_code=config.get("trust_remote_code", False),
config = mc,
max_seq_length = config.get("max_seq_length", 2048),
load_in_4bit = load_in_4bit,
hf_token = hf_token,
trust_remote_code = config.get("trust_remote_code", False),
)
if success:
@ -175,28 +210,37 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"audio_type": getattr(mc, "audio_type", None),
"has_audio_input": getattr(mc, "has_audio_input", False),
}
_send_response(resp_queue, {
"type": "loaded",
"success": True,
"model_info": model_info,
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "loaded",
"success": True,
"model_info": model_info,
"ts": time.time(),
},
)
else:
_send_response(resp_queue, {
"type": "loaded",
"success": False,
"error": "Failed to load model",
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "loaded",
"success": False,
"error": "Failed to load model",
"ts": time.time(),
},
)
except Exception as exc:
_send_response(resp_queue, {
"type": "loaded",
"success": False,
"error": str(exc),
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "loaded",
"success": False,
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
def _handle_generate(
@ -240,7 +284,7 @@ def _handle_generate(
use_adapter = cmd.get("use_adapter")
if use_adapter is not None:
generator = backend.generate_with_adapter_control(
use_adapter=use_adapter,
use_adapter = use_adapter,
**gen_kwargs,
)
else:
@ -254,29 +298,38 @@ def _handle_generate(
logger.info("Generation cancelled for request %s", request_id)
break
_send_response(resp_queue, {
"type": "token",
"request_id": request_id,
"text": cumulative_text,
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "token",
"request_id": request_id,
"text": cumulative_text,
"ts": time.time(),
},
)
_send_response(resp_queue, {
"type": "gen_done",
"request_id": request_id,
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "gen_done",
"request_id": request_id,
"ts": time.time(),
},
)
logger.info("Finished text generation for request_id=%s", request_id)
except Exception as exc:
logger.error("Generation error: %s", exc, exc_info=True)
_send_response(resp_queue, {
"type": "gen_error",
"request_id": request_id,
"error": str(exc),
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
logger.error("Generation error: %s", exc, exc_info = True)
_send_response(
resp_queue,
{
"type": "gen_error",
"request_id": request_id,
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
def _handle_generate_audio(
@ -289,35 +342,41 @@ def _handle_generate_audio(
try:
logger.info("Starting audio generation for request_id=%s", request_id)
wav_bytes, sample_rate = backend.generate_audio_response(
text=cmd["text"],
temperature=cmd.get("temperature", 0.6),
top_p=cmd.get("top_p", 0.95),
top_k=cmd.get("top_k", 50),
min_p=cmd.get("min_p", 0.0),
max_new_tokens=cmd.get("max_new_tokens", 2048),
repetition_penalty=cmd.get("repetition_penalty", 1.1),
use_adapter=cmd.get("use_adapter"),
text = cmd["text"],
temperature = cmd.get("temperature", 0.6),
top_p = cmd.get("top_p", 0.95),
top_k = cmd.get("top_k", 50),
min_p = cmd.get("min_p", 0.0),
max_new_tokens = cmd.get("max_new_tokens", 2048),
repetition_penalty = cmd.get("repetition_penalty", 1.1),
use_adapter = cmd.get("use_adapter"),
)
# Send WAV bytes as base64 (bytes can't go through mp.Queue directly)
_send_response(resp_queue, {
"type": "audio_done",
"request_id": request_id,
"wav_base64": base64.b64encode(wav_bytes).decode("ascii"),
"sample_rate": sample_rate,
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "audio_done",
"request_id": request_id,
"wav_base64": base64.b64encode(wav_bytes).decode("ascii"),
"sample_rate": sample_rate,
"ts": time.time(),
},
)
logger.info("Finished audio generation for request_id=%s", request_id)
except Exception as exc:
logger.error("Audio generation error: %s", exc, exc_info=True)
_send_response(resp_queue, {
"type": "audio_error",
"request_id": request_id,
"error": str(exc),
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
logger.error("Audio generation error: %s", exc, exc_info = True)
_send_response(
resp_queue,
{
"type": "audio_error",
"request_id": request_id,
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
def _handle_generate_audio_input(
@ -333,59 +392,70 @@ def _handle_generate_audio_input(
import numpy as np
# Decode audio array from list (numpy arrays can't go through mp.Queue)
audio_array = np.array(cmd["audio_data"], dtype=np.float32)
audio_array = np.array(cmd["audio_data"], dtype = np.float32)
audio_type = cmd.get("audio_type")
if audio_type == "whisper":
generator = backend.generate_whisper_response(
audio_array=audio_array,
cancel_event=cancel_event,
audio_array = audio_array,
cancel_event = cancel_event,
)
else:
generator = backend.generate_audio_input_response(
messages=cmd.get("messages", []),
system_prompt=cmd.get("system_prompt", ""),
audio_array=audio_array,
temperature=cmd.get("temperature", 0.7),
top_p=cmd.get("top_p", 0.9),
top_k=cmd.get("top_k", 40),
min_p=cmd.get("min_p", 0.0),
max_new_tokens=cmd.get("max_new_tokens", 512),
repetition_penalty=cmd.get("repetition_penalty", 1.1),
cancel_event=cancel_event,
messages = cmd.get("messages", []),
system_prompt = cmd.get("system_prompt", ""),
audio_array = audio_array,
temperature = cmd.get("temperature", 0.7),
top_p = cmd.get("top_p", 0.9),
top_k = cmd.get("top_k", 40),
min_p = cmd.get("min_p", 0.0),
max_new_tokens = cmd.get("max_new_tokens", 512),
repetition_penalty = cmd.get("repetition_penalty", 1.1),
cancel_event = cancel_event,
)
logger.info("Starting audio input generation for request_id=%s", request_id)
for text_chunk in generator:
if cancel_event.is_set():
logger.info("Audio input generation cancelled for request %s", request_id)
logger.info(
"Audio input generation cancelled for request %s", request_id
)
break
_send_response(resp_queue, {
"type": "token",
"request_id": request_id,
"text": text_chunk,
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "token",
"request_id": request_id,
"text": text_chunk,
"ts": time.time(),
},
)
_send_response(resp_queue, {
"type": "gen_done",
"request_id": request_id,
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "gen_done",
"request_id": request_id,
"ts": time.time(),
},
)
logger.info("Finished audio input generation for request_id=%s", request_id)
except Exception as exc:
logger.error("Audio input generation error: %s", exc, exc_info=True)
_send_response(resp_queue, {
"type": "gen_error",
"request_id": request_id,
"error": str(exc),
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
logger.error("Audio input generation error: %s", exc, exc_info = True)
_send_response(
resp_queue,
{
"type": "gen_error",
"request_id": request_id,
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
@ -397,19 +467,25 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
elif backend.active_model_name:
backend.unload_model(backend.active_model_name)
_send_response(resp_queue, {
"type": "unloaded",
"model_name": model_name,
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "unloaded",
"model_name": model_name,
"ts": time.time(),
},
)
except Exception as exc:
logger.error("Unload error: %s", exc)
_send_response(resp_queue, {
"type": "unloaded",
"model_name": model_name,
"error": str(exc),
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "unloaded",
"model_name": model_name,
"error": str(exc),
"ts": time.time(),
},
)
def run_inference_process(
@ -428,16 +504,19 @@ def run_inference_process(
config: Initial configuration dict with model info.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
os.environ["PYTHONWARNINGS"] = (
"ignore" # Suppress warnings at C-level before imports
)
import warnings
from loggers.config import LogConfig
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
LogConfig.setup_logging(
service_name="unsloth-studio-inference-worker",
env=os.getenv("ENVIRONMENT_TYPE", "production"),
service_name = "unsloth-studio-inference-worker",
env = os.getenv("ENVIRONMENT_TYPE", "production"),
)
model_name = config["model_name"]
@ -446,18 +525,22 @@ def run_inference_process(
try:
_activate_transformers_version(model_name)
except Exception as exc:
_send_response(resp_queue, {
"type": "error",
"error": f"Failed to activate transformers version: {exc}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "error",
"error": f"Failed to activate transformers version: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
# ── 1b. On Windows, check Triton availability (must be before import torch) ──
if sys.platform == "win32":
try:
import triton # noqa: F401
logger.info("Triton available — torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
@ -468,11 +551,14 @@ def run_inference_process(
# ── 2. Import ML libraries (fresh in this clean process) ──
try:
_send_response(resp_queue, {
"type": "status",
"message": "Importing ML libraries...",
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "status",
"message": "Importing ML libraries...",
"ts": time.time(),
},
)
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
@ -481,36 +567,46 @@ def run_inference_process(
from core.inference.inference import InferenceBackend
import transformers
logger.info("Subprocess loaded transformers %s", transformers.__version__)
except Exception as exc:
_send_response(resp_queue, {
"type": "error",
"error": f"Failed to import ML libraries: {exc}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "error",
"error": f"Failed to import ML libraries: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
# ── 3. Create inference backend and load initial model ──
try:
backend = InferenceBackend()
_send_response(resp_queue, {
"type": "status",
"message": "Loading model...",
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "status",
"message": "Loading model...",
"ts": time.time(),
},
)
_handle_load(backend, config, resp_queue)
except Exception as exc:
_send_response(resp_queue, {
"type": "error",
"error": f"Failed to initialize inference backend: {exc}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "error",
"error": f"Failed to initialize inference backend: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
# ── 4. Command loop — process commands until shutdown ──
@ -520,7 +616,7 @@ def run_inference_process(
while True:
try:
cmd = cmd_queue.get(timeout=1.0)
cmd = cmd_queue.get(timeout = 1.0)
except _queue.Empty:
continue
except (EOFError, OSError):
@ -564,26 +660,32 @@ def run_inference_process(
elif cmd_type == "reset":
cancel_event.set()
backend.reset_generation_state()
_send_response(resp_queue, {
"type": "reset_ack",
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "reset_ack",
"ts": time.time(),
},
)
elif cmd_type == "status":
# Return current status
_send_response(resp_queue, {
"type": "status_response",
"active_model": backend.active_model_name,
"models": {
name: {
"is_vision": info.get("is_vision", False),
"is_lora": info.get("is_lora", False),
}
for name, info in backend.models.items()
_send_response(
resp_queue,
{
"type": "status_response",
"active_model": backend.active_model_name,
"models": {
name: {
"is_vision": info.get("is_vision", False),
"is_lora": info.get("is_lora", False),
}
for name, info in backend.models.items()
},
"loading": list(backend.loading_models),
"ts": time.time(),
},
"loading": list(backend.loading_models),
"ts": time.time(),
})
)
elif cmd_type == "shutdown":
logger.info("Shutdown command received, exiting")
@ -593,25 +695,36 @@ def run_inference_process(
backend.unload_model(model_name)
except Exception:
pass
_send_response(resp_queue, {
"type": "shutdown_ack",
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "shutdown_ack",
"ts": time.time(),
},
)
return
else:
logger.warning("Unknown command type: %s", cmd_type)
_send_response(resp_queue, {
"type": "error",
"error": f"Unknown command type: {cmd_type}",
"ts": time.time(),
})
_send_response(
resp_queue,
{
"type": "error",
"error": f"Unknown command type: {cmd_type}",
"ts": time.time(),
},
)
except Exception as exc:
logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info=True)
_send_response(resp_queue, {
"type": "error",
"error": f"Command '{cmd_type}' failed: {exc}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
logger.error(
"Error handling command '%s': %s", cmd_type, exc, exc_info = True
)
_send_response(
resp_queue,
{
"type": "error",
"error": f"Command '{cmd_type}' failed: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)

7
studio/backend/core/training/__init__.py
View file

@ -4,10 +4,11 @@
"""
Training submodule - Training backends and trainer classes
"""
from .training import TrainingBackend, TrainingProgress, get_training_backend
__all__ = [
'TrainingProgress',
'TrainingBackend',
'get_training_backend',
"TrainingProgress",
"TrainingBackend",
"get_training_backend",
]

1959
studio/backend/core/training/trainer.py
View file

File diff suppressed because it is too large Load diff

151
studio/backend/core/training/training.py
View file

@ -13,6 +13,7 @@ worker's mp.Queue, and exposes the same API surface to routes/training.py.
Pattern follows core/data_recipe/jobs/manager.py.
"""
import math
import multiprocessing as mp
import queue
@ -39,6 +40,7 @@ PLOT_HEIGHT = 3.5
class TrainingProgress:
"""Mirror of trainer.TrainingProgress — kept here so the parent process
never needs to import the heavy ML modules."""
epoch: float = 0
step: int = 0
total_steps: int = 0
@ -109,7 +111,7 @@ class TrainingBackend:
# Join prior pump thread to prevent it from consuming events
# from the new job's queue (it reads self._event_queue dynamically).
if self._pump_thread is not None and self._pump_thread.is_alive():
self._pump_thread.join(timeout=5.0)
self._pump_thread.join(timeout = 5.0)
if self._pump_thread.is_alive():
logger.warning("Previous pump thread did not exit within 5s")
self._pump_thread = None
@ -117,7 +119,9 @@ class TrainingBackend:
# Reset state
self._should_stop = False
self._cancel_requested = False
self._progress = TrainingProgress(is_training=True, status_message="Initializing training...")
self._progress = TrainingProgress(
is_training = True, status_message = "Initializing training..."
)
self.loss_history.clear()
self.lr_history.clear()
self.step_history.clear()
@ -172,7 +176,9 @@ class TrainingBackend:
"train_on_completions": kwargs.get("train_on_completions", False),
"finetune_vision_layers": kwargs.get("finetune_vision_layers", True),
"finetune_language_layers": kwargs.get("finetune_language_layers", True),
"finetune_attention_modules": kwargs.get("finetune_attention_modules", True),
"finetune_attention_modules": kwargs.get(
"finetune_attention_modules", True
),
"finetune_mlp_modules": kwargs.get("finetune_mlp_modules", True),
"enable_wandb": kwargs.get("enable_wandb", False),
"wandb_token": kwargs.get("wandb_token"),
@ -193,19 +199,19 @@ class TrainingBackend:
self._stop_queue = _CTX.Queue()
self._proc = _CTX.Process(
target=run_training_process,
kwargs={
target = run_training_process,
kwargs = {
"event_queue": self._event_queue,
"stop_queue": self._stop_queue,
"config": config,
},
daemon=True,
daemon = True,
)
self._proc.start()
logger.info("Training subprocess started (pid=%s)", self._proc.pid)
# Start event pump thread
self._pump_thread = threading.Thread(target=self._pump_loop, daemon=True)
self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread.start()
return True
@ -224,7 +230,8 @@ class TrainingBackend:
# Update progress immediately for responsive UI
self._progress.status_message = (
"Stopping training and saving checkpoint..."
if save else "Cancelling training..."
if save
else "Cancelling training..."
)
return True
@ -232,15 +239,17 @@ class TrainingBackend:
"""Force-kill the training subprocess so state can be reset immediately."""
with self._lock:
if self._proc is not None and self._proc.is_alive():
logger.info("Force-terminating training subprocess (pid=%s)", self._proc.pid)
logger.info(
"Force-terminating training subprocess (pid=%s)", self._proc.pid
)
self._proc.terminate()
proc = self._proc
if proc is not None:
proc.join(timeout=5.0)
proc.join(timeout = 5.0)
if proc.is_alive():
proc.kill()
proc.join(timeout=2.0)
proc.join(timeout = 2.0)
def is_training_active(self) -> bool:
"""Check if training is currently active."""
@ -262,9 +271,29 @@ class TrainingBackend:
# Check status message for activity indicators
status_lower = (p.status_message or "").lower()
if any(k in status_lower for k in ["cancelled", "canceled", "stopped", "completed", "ready to train"]):
if any(
k in status_lower
for k in [
"cancelled",
"canceled",
"stopped",
"completed",
"ready to train",
]
):
return False
if any(k in status_lower for k in ["loading", "preparing", "training", "configuring", "tokenizing", "starting", "importing"]):
if any(
k in status_lower
for k in [
"loading",
"preparing",
"training",
"configuring",
"tokenizing",
"starting",
"importing",
]
):
return True
return False
@ -282,7 +311,7 @@ class TrainingBackend:
def refresh_plot_for_theme(self, theme: str) -> Optional[plt.Figure]:
"""Refresh plot with new theme."""
if theme and isinstance(theme, str) and theme in ['light', 'dark']:
if theme and isinstance(theme, str) and theme in ["light", "dark"]:
self.current_theme = theme
if self.loss_history:
with self._lock:
@ -296,6 +325,7 @@ class TrainingBackend:
class _TrainerShim:
"""Minimal shim so routes that access backend.trainer.* still work."""
def __init__(self, backend: "TrainingBackend"):
self._backend = backend
self.should_stop = False
@ -333,7 +363,7 @@ class TrainingBackend:
return
# Try to read an event
event = self._read_queue(self._event_queue, timeout_sec=0.25)
event = self._read_queue(self._event_queue, timeout_sec = 0.25)
if event is not None:
self._handle_event(event)
continue
@ -354,7 +384,10 @@ class TrainingBackend:
self._progress.status_message = "Training stopped."
else:
self._progress.is_training = False
self._progress.error = self._progress.error or "Training process exited unexpectedly"
self._progress.error = (
self._progress.error
or "Training process exited unexpectedly"
)
return
def _handle_event(self, event: dict) -> None:
@ -366,8 +399,12 @@ class TrainingBackend:
self._progress.step = event.get("step", self._progress.step)
self._progress.epoch = event.get("epoch", self._progress.epoch)
self._progress.loss = event.get("loss", self._progress.loss)
self._progress.learning_rate = event.get("learning_rate", self._progress.learning_rate)
self._progress.total_steps = event.get("total_steps", self._progress.total_steps)
self._progress.learning_rate = event.get(
"learning_rate", self._progress.learning_rate
)
self._progress.total_steps = event.get(
"total_steps", self._progress.total_steps
)
self._progress.elapsed_seconds = event.get("elapsed_seconds")
self._progress.eta_seconds = event.get("eta_seconds")
self._progress.grad_norm = event.get("grad_norm")
@ -428,7 +465,7 @@ class TrainingBackend:
@staticmethod
def _read_queue(q: Any, timeout_sec: float) -> Optional[dict]:
try:
return q.get(timeout=timeout_sec)
return q.get(timeout = timeout_sec)
except queue.Empty:
return None
except (EOFError, OSError, ValueError):
@ -449,28 +486,30 @@ class TrainingBackend:
# Plot generation (unchanged from original)
# ------------------------------------------------------------------
def _create_loss_plot(self, progress: TrainingProgress, theme: str = "light") -> plt.Figure:
def _create_loss_plot(
self, progress: TrainingProgress, theme: str = "light"
) -> plt.Figure:
"""Create training loss plot with theme-aware styling."""
plt.close('all')
plt.close("all")
LIGHT_STYLE = {
"facecolor": "#ffffff",
"grid_color": "#d1d5db",
"line": "#16b88a",
"text": "#1f2937",
"empty_text": "#6b7280"
"empty_text": "#6b7280",
}
DARK_STYLE = {
"facecolor": "#292929",
"grid_color": "#404040",
"line": "#4ade80",
"text": "#e5e7eb",
"empty_text": "#9ca3af"
"empty_text": "#9ca3af",
}
style = LIGHT_STYLE if theme == "light" else DARK_STYLE
fig, ax = plt.subplots(figsize=(PLOT_WIDTH, PLOT_HEIGHT))
fig, ax = plt.subplots(figsize = (PLOT_WIDTH, PLOT_HEIGHT))
fig.patch.set_facecolor(style["facecolor"])
ax.set_facecolor(style["facecolor"])
@ -478,8 +517,15 @@ class TrainingBackend:
steps = self.step_history
losses = self.loss_history
scatter_color = "#60a5fa"
ax.scatter(steps, losses, s=16, alpha=0.6, color=scatter_color,
linewidths=0, label="Training Loss (raw)")
ax.scatter(
steps,
losses,
s = 16,
alpha = 0.6,
color = scatter_color,
linewidths = 0,
label = "Training Loss (raw)",
)
MA_WINDOW = 20
window = min(MA_WINDOW, len(losses))
@ -495,15 +541,21 @@ class TrainingBackend:
denom = i - start + 1
ma.append((cumsum[i + 1] - cumsum[start]) / denom)
ax.plot(steps, ma, color=style["line"], linewidth=2.5, alpha=0.95,
label=f"Moving Avg ({ma[-1]:.4f})")
ax.plot(
steps,
ma,
color = style["line"],
linewidth = 2.5,
alpha = 0.95,
label = f"Moving Avg ({ma[-1]:.4f})",
)
leg = ax.legend(frameon=False, fontsize=9)
leg = ax.legend(frameon = False, fontsize = 9)
for t in leg.get_texts():
t.set_color(style["text"])
ax.set_xlabel('Steps', fontsize=10, color=style["text"])
ax.set_ylabel('Loss', fontsize=10, color=style["text"])
ax.set_xlabel("Steps", fontsize = 10, color = style["text"])
ax.set_ylabel("Loss", fontsize = 10, color = style["text"])
if progress.error:
title = f"Error: {progress.error}"
@ -516,17 +568,31 @@ class TrainingBackend:
else:
title = "Training Loss"
ax.set_title(title, fontsize=11, fontweight='bold', pad=10, color=style["text"])
ax.grid(True, alpha=0.4, linestyle='--', color=style["grid_color"])
ax.tick_params(colors=style["text"], which='both')
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_color(style["text"])
ax.spines['left'].set_color(style["text"])
ax.set_title(
title, fontsize = 11, fontweight = "bold", pad = 10, color = style["text"]
)
ax.grid(True, alpha = 0.4, linestyle = "--", color = style["grid_color"])
ax.tick_params(colors = style["text"], which = "both")
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
ax.spines["bottom"].set_color(style["text"])
ax.spines["left"].set_color(style["text"])
else:
display_msg = progress.status_message if progress.status_message else 'Waiting for training data...'
ax.text(0.5, 0.5, display_msg, ha='center', va='center', fontsize=16,
color=style["empty_text"], transform=ax.transAxes)
display_msg = (
progress.status_message
if progress.status_message
else "Waiting for training data..."
)
ax.text(
0.5,
0.5,
display_msg,
ha = "center",
va = "center",
fontsize = 16,
color = style["empty_text"],
transform = ax.transAxes,
)
ax.set_xticks([])
ax.set_yticks([])
for spine in ax.spines.values():
@ -544,7 +610,8 @@ class TrainingBackend:
"""
logger.info(
"_transfer_to_inference_backend: subprocess training — "
"model must be loaded from disk (output_dir=%s)", self._output_dir
"model must be loaded from disk (output_dir=%s)",
self._output_dir,
)
return False

658
studio/backend/core/training/worker.py
View file

@ -10,6 +10,7 @@ solving the transformers version-switching problem completely.
Pattern follows core/data_recipe/jobs/worker.py.
"""
from __future__ import annotations
import structlog
@ -39,7 +40,9 @@ def _activate_transformers_version(model_name: str) -> None:
resolved = _resolve_base_model(model_name)
if needs_transformers_5(resolved):
venv_t5 = os.path.join(os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5")
venv_t5 = os.path.join(
os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5"
)
if os.path.isdir(venv_t5):
sys.path.insert(0, venv_t5)
logger.info("Activated transformers 5.x from %s", venv_t5)
@ -47,16 +50,35 @@ def _activate_transformers_version(model_name: str) -> None:
# Fallback: pip install at runtime (slower, ~10-15s)
logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5)
import subprocess as sp
os.makedirs(venv_t5, exist_ok=True)
os.makedirs(venv_t5, exist_ok = True)
r1 = sp.run(
[sys.executable, "-m", "pip", "install", "--target", venv_t5,
"--no-deps", "transformers==5.2.0"],
stdout=sp.PIPE, stderr=sp.STDOUT,
[
sys.executable,
"-m",
"pip",
"install",
"--target",
venv_t5,
"--no-deps",
"transformers==5.2.0",
],
stdout = sp.PIPE,
stderr = sp.STDOUT,
)
r2 = sp.run(
[sys.executable, "-m", "pip", "install", "--target", venv_t5,
"--no-deps", "huggingface_hub==1.3.0"],
stdout=sp.PIPE, stderr=sp.STDOUT,
[
sys.executable,
"-m",
"pip",
"install",
"--target",
venv_t5,
"--no-deps",
"huggingface_hub==1.3.0",
],
stdout = sp.PIPE,
stderr = sp.STDOUT,
)
if r1.returncode != 0 or r2.returncode != 0:
raise RuntimeError(
@ -85,16 +107,19 @@ def run_training_process(
config: Training configuration dict with all parameters.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
os.environ["PYTHONWARNINGS"] = (
"ignore" # Suppress warnings at C-level before imports
)
import warnings
from loggers.config import LogConfig
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
LogConfig.setup_logging(
service_name="unsloth-studio-training-worker",
env=os.getenv("ENVIRONMENT_TYPE", "production"),
service_name = "unsloth-studio-training-worker",
env = os.getenv("ENVIRONMENT_TYPE", "production"),
)
model_name = config["model_name"]
@ -103,18 +128,21 @@ def run_training_process(
try:
_activate_transformers_version(model_name)
except Exception as exc:
event_queue.put({
"type": "error",
"error": f"Failed to activate transformers version: {exc}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": f"Failed to activate transformers version: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
return
# ── 1b. On Windows, check Triton availability (must be before import torch) ──
if sys.platform == "win32":
try:
import triton # noqa: F401
logger.info("Triton available — torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
@ -132,17 +160,25 @@ def run_training_process(
sys.path.insert(0, backend_path)
from core.training.trainer import UnslothTrainer, TrainingProgress
from utils.paths import ensure_dir, resolve_output_dir, resolve_tensorboard_dir, datasets_root
from utils.paths import (
ensure_dir,
resolve_output_dir,
resolve_tensorboard_dir,
datasets_root,
)
import transformers
logger.info("Subprocess loaded transformers %s", transformers.__version__)
except Exception as exc:
event_queue.put({
"type": "error",
"error": f"Failed to import ML libraries: {exc}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": f"Failed to import ML libraries: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
return
# ── 2b. EMBEDDING MODEL FAST-PATH ──
@ -153,12 +189,14 @@ def run_training_process(
try:
_run_embedding_training(event_queue, stop_queue, config)
except Exception as exc:
event_queue.put({
"type": "error",
"error": str(exc),
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
return
# ── 3. Create a fresh trainer instance ──
@ -169,21 +207,23 @@ def run_training_process(
has_train_loss = progress.step >= 0 and progress.loss > 0
has_eval_loss = progress.eval_loss is not None
if has_train_loss or has_eval_loss:
event_queue.put({
"type": "progress",
"step": progress.step,
"epoch": progress.epoch,
"loss": progress.loss,
"learning_rate": progress.learning_rate,
"total_steps": progress.total_steps,
"elapsed_seconds": progress.elapsed_seconds,
"eta_seconds": progress.eta_seconds,
"grad_norm": progress.grad_norm,
"num_tokens": progress.num_tokens,
"eval_loss": progress.eval_loss,
"status_message": progress.status_message,
"ts": time.time(),
})
event_queue.put(
{
"type": "progress",
"step": progress.step,
"epoch": progress.epoch,
"loss": progress.loss,
"learning_rate": progress.learning_rate,
"total_steps": progress.total_steps,
"elapsed_seconds": progress.elapsed_seconds,
"eta_seconds": progress.eta_seconds,
"grad_norm": progress.grad_norm,
"num_tokens": progress.num_tokens,
"eval_loss": progress.eval_loss,
"status_message": progress.status_message,
"ts": time.time(),
}
)
if progress.status_message:
_send_status(event_queue, progress.status_message)
@ -196,7 +236,7 @@ def run_training_process(
def _poll_stop():
while True:
try:
msg = stop_queue.get(timeout=1.0)
msg = stop_queue.get(timeout = 1.0)
if msg and msg.get("type") == "stop":
save = msg.get("save", True)
trainer.should_stop = True
@ -208,7 +248,7 @@ def run_training_process(
except (EOFError, OSError):
return
stop_thread = threading.Thread(target=_poll_stop, daemon=True)
stop_thread = threading.Thread(target = _poll_stop, daemon = True)
stop_thread.start()
# ── 4. Execute the training pipeline ──
@ -222,12 +262,12 @@ def run_training_process(
# ── 4a. Lightweight detection + tokenizer (no VRAM) ──
_send_status(event_queue, "Detecting model type...")
trainer.pre_detect_and_load_tokenizer(
model_name=model_name,
max_seq_length=config["max_seq_length"],
hf_token=hf_token,
is_dataset_image=config.get("is_dataset_image", False),
is_dataset_audio=config.get("is_dataset_audio", False),
trust_remote_code=config.get("trust_remote_code", False),
model_name = model_name,
max_seq_length = config["max_seq_length"],
hf_token = hf_token,
is_dataset_image = config.get("is_dataset_image", False),
is_dataset_audio = config.get("is_dataset_audio", False),
trust_remote_code = config.get("trust_remote_code", False),
)
if trainer.should_stop:
event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
@ -237,16 +277,16 @@ def run_training_process(
_send_status(event_queue, "Loading and formatting dataset...")
hf_dataset = config.get("hf_dataset", "")
dataset_result = trainer.load_and_format_dataset(
dataset_source=hf_dataset if hf_dataset and hf_dataset.strip() else None,
format_type=config.get("format_type", ""),
local_datasets=config.get("local_datasets") or None,
custom_format_mapping=config.get("custom_format_mapping"),
subset=config.get("subset"),
train_split=config.get("train_split", "train"),
eval_split=config.get("eval_split"),
eval_steps=config.get("eval_steps", 0.00),
dataset_slice_start=config.get("dataset_slice_start"),
dataset_slice_end=config.get("dataset_slice_end"),
dataset_source = hf_dataset if hf_dataset and hf_dataset.strip() else None,
format_type = config.get("format_type", ""),
local_datasets = config.get("local_datasets") or None,
custom_format_mapping = config.get("custom_format_mapping"),
subset = config.get("subset"),
train_split = config.get("train_split", "train"),
eval_split = config.get("eval_split"),
eval_steps = config.get("eval_steps", 0.00),
dataset_slice_start = config.get("dataset_slice_start"),
dataset_slice_end = config.get("dataset_slice_end"),
)
if isinstance(dataset_result, tuple):
@ -260,13 +300,19 @@ def run_training_process(
# or a raw Dataset for audio paths
try:
ds = dataset["dataset"] if isinstance(dataset, dict) else dataset
print(f"\n[DEBUG] Dataset loaded BEFORE model. type={type(ds).__name__}, len={len(ds)}", flush=True)
print(f"[DEBUG] Columns: {ds.column_names}", flush=True)
print(
f"\n[DEBUG] Dataset loaded BEFORE model. type={type(ds).__name__}, len={len(ds)}",
flush = True,
)
print(f"[DEBUG] Columns: {ds.column_names}", flush = True)
sample = ds[0]
preview = {k: str(v)[:300] for k, v in sample.items()}
print(f"[DEBUG] First sample: {preview}\n", flush=True)
print(f"[DEBUG] First sample: {preview}\n", flush = True)
except Exception as e:
print(f"[DEBUG] Could not preview first sample: {type(e).__name__}: {e}", flush=True)
print(
f"[DEBUG] Could not preview first sample: {type(e).__name__}: {e}",
flush = True,
)
# Disable eval if eval_steps <= 0
eval_steps = config.get("eval_steps", 0.00)
@ -276,88 +322,114 @@ def run_training_process(
# Tell the parent process that eval is configured so the frontend
# shows "Waiting for first evaluation step..." instead of "not configured"
if eval_dataset is not None:
event_queue.put({
"type": "eval_configured",
"ts": time.time(),
})
event_queue.put(
{
"type": "eval_configured",
"ts": time.time(),
}
)
if dataset is None or trainer.should_stop:
if trainer.should_stop:
event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
event_queue.put(
{"type": "complete", "output_dir": None, "ts": time.time()}
)
else:
event_queue.put({
"type": "error",
"error": trainer.training_progress.error or "Failed to load dataset",
"stack": "", "ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": trainer.training_progress.error
or "Failed to load dataset",
"stack": "",
"ts": time.time(),
}
)
return
# ── 4c. Load training model (uses VRAM — dataset already formatted) ──
_send_status(event_queue, "Loading model...")
success = trainer.load_model(
model_name=model_name,
max_seq_length=config["max_seq_length"],
load_in_4bit=config["load_in_4bit"],
hf_token=hf_token,
is_dataset_image=config.get("is_dataset_image", False),
is_dataset_audio=config.get("is_dataset_audio", False),
trust_remote_code=config.get("trust_remote_code", False),
model_name = model_name,
max_seq_length = config["max_seq_length"],
load_in_4bit = config["load_in_4bit"],
hf_token = hf_token,
is_dataset_image = config.get("is_dataset_image", False),
is_dataset_audio = config.get("is_dataset_audio", False),
trust_remote_code = config.get("trust_remote_code", False),
)
if not success or trainer.should_stop:
if trainer.should_stop:
event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
event_queue.put(
{"type": "complete", "output_dir": None, "ts": time.time()}
)
else:
error_msg = trainer.training_progress.error or "Failed to load model"
event_queue.put({
"type": "error",
"error": error_msg,
"stack": "", "ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": error_msg,
"stack": "",
"ts": time.time(),
}
)
return
# ── 4d. Prepare model (LoRA or full finetuning) ──
training_type = config.get("training_type", "LoRA/QLoRA")
use_lora = (training_type == "LoRA/QLoRA")
use_lora = training_type == "LoRA/QLoRA"
if use_lora:
_send_status(event_queue, "Configuring LoRA adapters...")
success = trainer.prepare_model_for_training(
use_lora=True,
finetune_vision_layers=config.get("finetune_vision_layers", True),
finetune_language_layers=config.get("finetune_language_layers", True),
finetune_attention_modules=config.get("finetune_attention_modules", True),
finetune_mlp_modules=config.get("finetune_mlp_modules", True),
target_modules=config.get("target_modules"),
lora_r=config.get("lora_r", 16),
lora_alpha=config.get("lora_alpha", 16),
lora_dropout=config.get("lora_dropout", 0.0),
use_gradient_checkpointing=config.get("gradient_checkpointing", "unsloth"),
use_rslora=config.get("use_rslora", False),
use_loftq=config.get("use_loftq", False),
use_lora = True,
finetune_vision_layers = config.get("finetune_vision_layers", True),
finetune_language_layers = config.get("finetune_language_layers", True),
finetune_attention_modules = config.get(
"finetune_attention_modules", True
),
finetune_mlp_modules = config.get("finetune_mlp_modules", True),
target_modules = config.get("target_modules"),
lora_r = config.get("lora_r", 16),
lora_alpha = config.get("lora_alpha", 16),
lora_dropout = config.get("lora_dropout", 0.0),
use_gradient_checkpointing = config.get(
"gradient_checkpointing", "unsloth"
),
use_rslora = config.get("use_rslora", False),
use_loftq = config.get("use_loftq", False),
)
else:
_send_status(event_queue, "Preparing model for full finetuning...")
success = trainer.prepare_model_for_training(use_lora=False)
success = trainer.prepare_model_for_training(use_lora = False)
if not success or trainer.should_stop:
if trainer.should_stop:
event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
event_queue.put(
{"type": "complete", "output_dir": None, "ts": time.time()}
)
else:
event_queue.put({
"type": "error",
"error": trainer.training_progress.error or "Failed to prepare model",
"stack": "", "ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": trainer.training_progress.error
or "Failed to prepare model",
"stack": "",
"ts": time.time(),
}
)
return
# Convert learning rate
try:
lr_value = float(config.get("learning_rate", "2e-4"))
except ValueError:
event_queue.put({
"type": "error",
"error": f"Invalid learning rate: {config.get('learning_rate')}",
"stack": "", "ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": f"Invalid learning rate: {config.get('learning_rate')}",
"stack": "",
"ts": time.time(),
}
)
return
# Generate output dir
@ -379,64 +451,72 @@ def run_training_process(
trainer._train_worker(
dataset,
output_dir=output_dir,
num_epochs=config.get("num_epochs", 3),
learning_rate=lr_value,
batch_size=config.get("batch_size", 2),
gradient_accumulation_steps=config.get("gradient_accumulation_steps", 4),
warmup_steps=config.get("warmup_steps"),
warmup_ratio=config.get("warmup_ratio"),
max_steps=max_steps if max_steps and max_steps > 0 else 0,
save_steps=save_steps if save_steps and save_steps > 0 else 0,
weight_decay=config.get("weight_decay", 0.01),
random_seed=config.get("random_seed", 3407),
packing=config.get("packing", False),
train_on_completions=config.get("train_on_completions", False),
enable_wandb=config.get("enable_wandb", False),
wandb_project=config.get("wandb_project", "unsloth-training"),
wandb_token=config.get("wandb_token"),
enable_tensorboard=config.get("enable_tensorboard", False),
tensorboard_dir=tensorboard_dir,
eval_dataset=eval_dataset,
eval_steps=eval_steps,
max_seq_length=config.get("max_seq_length", 2048),
optim=config.get("optim", "adamw_8bit"),
lr_scheduler_type=config.get("lr_scheduler_type", "linear"),
output_dir = output_dir,
num_epochs = config.get("num_epochs", 3),
learning_rate = lr_value,
batch_size = config.get("batch_size", 2),
gradient_accumulation_steps = config.get("gradient_accumulation_steps", 4),
warmup_steps = config.get("warmup_steps"),
warmup_ratio = config.get("warmup_ratio"),
max_steps = max_steps if max_steps and max_steps > 0 else 0,
save_steps = save_steps if save_steps and save_steps > 0 else 0,
weight_decay = config.get("weight_decay", 0.01),
random_seed = config.get("random_seed", 3407),
packing = config.get("packing", False),
train_on_completions = config.get("train_on_completions", False),
enable_wandb = config.get("enable_wandb", False),
wandb_project = config.get("wandb_project", "unsloth-training"),
wandb_token = config.get("wandb_token"),
enable_tensorboard = config.get("enable_tensorboard", False),
tensorboard_dir = tensorboard_dir,
eval_dataset = eval_dataset,
eval_steps = eval_steps,
max_seq_length = config.get("max_seq_length", 2048),
optim = config.get("optim", "adamw_8bit"),
lr_scheduler_type = config.get("lr_scheduler_type", "linear"),
)
# Check final state
progress = trainer.get_training_progress()
if progress.error:
event_queue.put({
"type": "error",
"error": progress.error,
"stack": "",
"ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": progress.error,
"stack": "",
"ts": time.time(),
}
)
else:
event_queue.put({
"type": "complete",
"output_dir": output_dir,
"status_message": progress.status_message or "Training completed",
"ts": time.time(),
})
event_queue.put(
{
"type": "complete",
"output_dir": output_dir,
"status_message": progress.status_message or "Training completed",
"ts": time.time(),
}
)
except Exception as exc:
event_queue.put({
"type": "error",
"error": str(exc),
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
def _send_status(event_queue: Any, message: str) -> None:
"""Send a status update to the parent process."""
event_queue.put({
"type": "status",
"message": message,
"ts": time.time(),
})
event_queue.put(
{
"type": "status",
"message": message,
"ts": time.time(),
}
)
def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> None:
@ -469,14 +549,17 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
from sentence_transformers.training_args import BatchSamplers
from datasets import load_dataset, Dataset
from transformers import TrainerCallback
from utils.paths import datasets_root, resolve_output_dir
except ImportError as e:
event_queue.put({
"type": "error",
"error": f"Failed to import embedding libraries: {e}. "
"Ensure 'sentence_transformers' and 'unsloth' are installed.",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": f"Failed to import embedding libraries: {e}. "
"Ensure 'sentence_transformers' and 'unsloth' are installed.",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
return
# ── Stop signal handling ──
@ -487,18 +570,21 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
nonlocal _should_stop, _save_on_stop
while True:
try:
msg = stop_queue.get(timeout=1.0)
msg = stop_queue.get(timeout = 1.0)
if msg and msg.get("type") == "stop":
_save_on_stop = msg.get("save", True)
_should_stop = True
logger.info("Embedding training: stop signal received (save=%s)", _save_on_stop)
logger.info(
"Embedding training: stop signal received (save=%s)",
_save_on_stop,
)
return
except _queue.Empty:
continue
except (EOFError, OSError):
return
stop_thread = threading.Thread(target=_poll_stop, daemon=True)
stop_thread = threading.Thread(target = _poll_stop, daemon = True)
stop_thread.start()
# ── 2. Load model ──
@ -508,21 +594,23 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
hf_token = hf_token if hf_token and hf_token.strip() else None
max_seq_length = config.get("max_seq_length", 512)
training_type = config.get("training_type", "LoRA/QLoRA")
use_lora = (training_type == "LoRA/QLoRA")
use_lora = training_type == "LoRA/QLoRA"
model = FastSentenceTransformer.from_pretrained(
model_name=model_name,
max_seq_length=max_seq_length,
full_finetuning=not use_lora,
token=hf_token,
model_name = model_name,
max_seq_length = max_seq_length,
full_finetuning = not use_lora,
token = hf_token,
)
except Exception as e:
event_queue.put({
"type": "error",
"error": f"Failed to load embedding model '{model_name}': {e}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": f"Failed to load embedding model '{model_name}': {e}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
return
if _should_stop:
@ -540,24 +628,29 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
model = FastSentenceTransformer.get_peft_model(
model,
r=config.get("lora_r", 32),
target_modules=config.get("target_modules") or ["q_proj", "k_proj", "v_proj", "o_proj"],
lora_alpha=config.get("lora_alpha", 64),
lora_dropout=config.get("lora_dropout", 0.0),
bias="none",
use_gradient_checkpointing=gradient_checkpointing,
random_state=config.get("random_seed", 3407),
use_rslora=config.get("use_rslora", False),
loftq_config={"loftq_bits": 4, "loftq_iter": 1} if config.get("use_loftq") else None,
task_type="FEATURE_EXTRACTION",
r = config.get("lora_r", 32),
target_modules = config.get("target_modules")
or ["q_proj", "k_proj", "v_proj", "o_proj"],
lora_alpha = config.get("lora_alpha", 64),
lora_dropout = config.get("lora_dropout", 0.0),
bias = "none",
use_gradient_checkpointing = gradient_checkpointing,
random_state = config.get("random_seed", 3407),
use_rslora = config.get("use_rslora", False),
loftq_config = {"loftq_bits": 4, "loftq_iter": 1}
if config.get("use_loftq")
else None,
task_type = "FEATURE_EXTRACTION",
)
except Exception as e:
event_queue.put({
"type": "error",
"error": f"Failed to configure LoRA for embedding model: {e}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": f"Failed to configure LoRA for embedding model: {e}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
return
if _should_stop:
@ -578,16 +671,21 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
dataset = load_dataset(
hf_dataset.strip(),
subset,
split=train_split,
token=hf_token,
split = train_split,
token = hf_token,
)
elif local_datasets:
# Load from local file(s) — mirrors the non-embedding pipeline's
# directory handling so recipe outputs (parquet-files/) work.
all_files: list[str] = []
for dataset_file in local_datasets:
file_path = dataset_file if os.path.isabs(dataset_file) else os.path.join(
str(datasets_root()), dataset_file,
file_path = (
dataset_file
if os.path.isabs(dataset_file)
else os.path.join(
str(datasets_root()),
dataset_file,
)
)
if os.path.isdir(file_path):
file_path_obj = Path(file_path)
@ -606,7 +704,9 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
if candidates:
all_files.extend(str(c) for c in candidates)
continue
raise ValueError(f"No supported data files in directory: {file_path_obj}")
raise ValueError(
f"No supported data files in directory: {file_path_obj}"
)
else:
all_files.append(file_path)
@ -619,14 +719,19 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
elif first_ext == ".parquet":
loader = "parquet"
else:
raise ValueError(f"Unsupported local dataset format: {all_files[0]}")
dataset = load_dataset(loader, data_files=all_files, split="train")
raise ValueError(
f"Unsupported local dataset format: {all_files[0]}"
)
dataset = load_dataset(loader, data_files = all_files, split = "train")
else:
event_queue.put({
"type": "error",
"error": "No dataset specified for embedding training.",
"stack": "", "ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": "No dataset specified for embedding training.",
"stack": "",
"ts": time.time(),
}
)
return
# Apply dataset slicing if specified
@ -639,12 +744,14 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
logger.info(f"Embedding dataset loaded: {len(dataset)} samples")
except Exception as e:
event_queue.put({
"type": "error",
"error": f"Failed to load dataset: {e}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": f"Failed to load dataset: {e}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
return
if _should_stop:
@ -659,16 +766,21 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
try:
lr_value = float(config.get("learning_rate", "2e-4"))
except ValueError:
event_queue.put({
"type": "error",
"error": f"Invalid learning rate: {config.get('learning_rate')}",
"stack": "", "ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": f"Invalid learning rate: {config.get('learning_rate')}",
"stack": "",
"ts": time.time(),
}
)
return
output_dir = config.get("output_dir")
if not output_dir:
output_dir = f"./outputs/{model_name.replace('/', '_')}_{int(time.time())}"
output_dir = str(
resolve_output_dir(f"{model_name.replace('/', '_')}_{int(time.time())}")
)
num_epochs = config.get("num_epochs", 2)
batch_size = config.get("batch_size", 256)
@ -728,7 +840,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
class _EmbeddingProgressCallback(TrainerCallback):
"""Sends training progress events to the parent process via event_queue."""
def on_log(self, args, state, control, logs=None, **kwargs):
def on_log(self, args, state, control, logs = None, **kwargs):
if not logs:
return
loss_value = logs.get("loss", logs.get("train_loss", 0.0))
@ -741,21 +853,23 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
if remaining > 0:
eta = (elapsed / current_step) * remaining
event_queue.put({
"type": "progress",
"step": current_step,
"epoch": round(state.epoch, 2) if state.epoch else 0,
"loss": loss_value,
"learning_rate": logs.get("learning_rate", 0.0),
"total_steps": total_steps,
"elapsed_seconds": elapsed,
"eta_seconds": eta,
"grad_norm": logs.get("grad_norm"),
"num_tokens": getattr(state, "num_input_tokens_seen", None),
"eval_loss": logs.get("eval_loss"),
"status_message": "",
"ts": time.time(),
})
event_queue.put(
{
"type": "progress",
"step": current_step,
"epoch": round(state.epoch, 2) if state.epoch else 0,
"loss": loss_value,
"learning_rate": logs.get("learning_rate", 0.0),
"total_steps": total_steps,
"elapsed_seconds": elapsed,
"eta_seconds": eta,
"grad_norm": logs.get("grad_norm"),
"num_tokens": getattr(state, "num_input_tokens_seen", None),
"eval_loss": logs.get("eval_loss"),
"status_message": "",
"ts": time.time(),
}
)
def on_step_end(self, args, state, control, **kwargs):
if _should_stop:
@ -767,31 +881,35 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
_send_status(event_queue, "Starting embedding training...")
try:
trainer = SentenceTransformerTrainer(
model=model,
train_dataset=dataset,
loss=loss,
args=args,
callbacks=[_EmbeddingProgressCallback()],
model = model,
train_dataset = dataset,
loss = loss,
args = args,
callbacks = [_EmbeddingProgressCallback()],
)
trainer.train()
except Exception as e:
event_queue.put({
"type": "error",
"error": f"Embedding training failed: {e}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": f"Embedding training failed: {e}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
return
# ── 10. Save model ──
if _should_stop and not _save_on_stop:
event_queue.put({
"type": "complete",
"output_dir": None,
"status_message": "Training cancelled",
"ts": time.time(),
})
event_queue.put(
{
"type": "complete",
"output_dir": None,
"status_message": "Training cancelled",
"ts": time.time(),
}
)
return
_send_status(event_queue, "Saving model...")
@ -801,18 +919,22 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
logger.info("Embedding model saved to %s", output_dir)
except Exception as e:
logger.error("Failed to save embedding model: %s", e)
event_queue.put({
"type": "error",
"error": f"Training completed but failed to save: {e}",
"stack": traceback.format_exc(limit=20),
"ts": time.time(),
})
event_queue.put(
{
"type": "error",
"error": f"Training completed but failed to save: {e}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
return
# ── 11. Done ──
event_queue.put({
"type": "complete",
"output_dir": output_dir,
"status_message": "Embedding training completed",
"ts": time.time(),
})
event_queue.put(
{
"type": "complete",
"output_dir": output_dir,
"status_message": "Embedding training completed",
"ts": time.time(),
}
)

3
studio/backend/loggers/__init__.py
View file

@ -1,3 +1,6 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from .handlers import get_logger
__all__ = ["get_logger"]

16
studio/backend/loggers/config.py
View file

@ -1,3 +1,6 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Logging configuration for structured logging with structlog.
This module provides centralized logging configuration with environment-specific
@ -19,6 +22,7 @@ from typing import Optional
import structlog
class LogConfig:
"""Structured logging configuration for the application.
@ -41,9 +45,9 @@ class LogConfig:
log_level = getattr(logging, log_level_name, logging.INFO)
structlog.configure(
processors=[
processors = [
# Reorder processors to control field order
structlog.processors.TimeStamper(fmt="iso"), # timestamp first
structlog.processors.TimeStamper(fmt = "iso"), # timestamp first
structlog.processors.add_log_level, # level second
structlog.contextvars.merge_contextvars,
# Custom processor to flatten the extra field
@ -59,14 +63,14 @@ class LogConfig:
},
},
(
structlog.processors.JSONRenderer(sort_keys=False) # Preserve order
structlog.processors.JSONRenderer(sort_keys = False) # Preserve order
if env == "production"
else structlog.dev.ConsoleRenderer()
),
],
wrapper_class=structlog.make_filtering_bound_logger(log_level),
logger_factory=structlog.PrintLoggerFactory(file=sys.stdout),
cache_logger_on_first_use=True,
wrapper_class = structlog.make_filtering_bound_logger(log_level),
logger_factory = structlog.PrintLoggerFactory(file = sys.stdout),
cache_logger_on_first_use = True,
)
return structlog.get_logger(service_name)

25
studio/backend/loggers/handlers.py
View file

@ -1,3 +1,6 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Logging handlers and middleware for structured logging.
This module provides FastAPI middleware and structlog processors for:
@ -36,21 +39,23 @@ class LoggingMiddleware(BaseHTTPMiddleware):
"/api/train/status",
"/api/train/metrics",
"/api/train/hardware",
"/api/system"
"/api/system",
}
is_excluded = (
request.url.path in EXCLUDED_PATHS
or request.url.path.startswith("/assets/")
or request.url.path.endswith((".png", ".jpg", ".jpeg", ".ico", ".woff", ".woff2", ".ttf"))
or request.url.path.endswith(
(".png", ".jpg", ".jpeg", ".ico", ".woff", ".woff2", ".ttf")
)
)
if not is_excluded:
logger.info(
"request_completed",
method=request.method,
path=request.url.path,
status_code=response.status_code,
process_time_ms=round(process_time, 2),
method = request.method,
path = request.url.path,
status_code = response.status_code,
process_time_ms = round(process_time, 2),
)
return response
@ -58,10 +63,10 @@ class LoggingMiddleware(BaseHTTPMiddleware):
except Exception as e:
logger.error(
"request_failed",
path=request.url.path,
method=request.method,
error=str(e),
exc_info=True,
path = request.url.path,
method = request.method,
error = str(e),
exc_info = True,
)
raise

84
studio/backend/main.py
View file

@ -4,9 +4,11 @@
"""
Main FastAPI application for Unsloth UI Backend
"""
import os
# Suppress annoying C-level dependency warnings globally
os.environ["PYTHONWARNINGS"] = "ignore"
os.environ["PYTHONWARNINGS"] = "ignore"
import secrets
import shutil
@ -54,7 +56,7 @@ async def lifespan(app: FastAPI):
# Version switching now uses .venv_t5/ (pre-installed by setup.sh).
overlay_dir = Path(__file__).resolve().parent.parent.parent / ".venv_overlay"
if overlay_dir.is_dir():
shutil.rmtree(overlay_dir, ignore_errors=True)
shutil.rmtree(overlay_dir, ignore_errors = True)
# Detect hardware first — sets DEVICE global used everywhere
detect_hardware()
@ -62,12 +64,14 @@ async def lifespan(app: FastAPI):
# Disable flex attention on Blackwell+ GPUs (sm_120 and above)
if get_device() == DeviceType.CUDA:
import torch
props = torch.cuda.get_device_properties(0)
sm_version = props.major * 10 + props.minor
if sm_version >= 120:
os.environ["UNSLOTH_ENABLE_FLEX_ATTENTION"] = "0"
import structlog
from loggers import get_logger
get_logger(__name__).info(
f"GPU sm_{sm_version} detected — setting UNSLOTH_FLEX_ATTENTION=0"
)
@ -75,13 +79,16 @@ async def lifespan(app: FastAPI):
# Pre-cache the helper GGUF model for LLM-assisted dataset detection.
# Runs in a background thread so it doesn't block server startup.
import threading
def _precache():
try:
from utils.datasets.llm_assist import precache_helper_gguf
precache_helper_gguf()
except Exception:
pass # non-critical
threading.Thread(target=_precache, daemon=True).start()
threading.Thread(target = _precache, daemon = True).start()
if not storage.is_initialized():
setup_token = secrets.token_urlsafe(32)
@ -100,10 +107,10 @@ async def lifespan(app: FastAPI):
# Create FastAPI app
app = FastAPI(
title="Unsloth UI Backend",
version="1.0.0",
description="Backend API for Unsloth UI - Training and Model Management",
lifespan=lifespan,
title = "Unsloth UI Backend",
version = "1.0.0",
description = "Backend API for Unsloth UI - Training and Model Management",
lifespan = lifespan,
)
# Initialize structured logging
@ -111,8 +118,8 @@ from loggers.config import LogConfig
from loggers.handlers import LoggingMiddleware
logger = LogConfig.setup_logging(
service_name="unsloth-studio-backend",
env=os.getenv("ENVIRONMENT_TYPE", "production")
service_name = "unsloth-studio-backend",
env = os.getenv("ENVIRONMENT_TYPE", "production"),
)
app.add_middleware(LoggingMiddleware)
@ -120,38 +127,39 @@ app.add_middleware(LoggingMiddleware)
# CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # In production, specify allowed origins
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
allow_origins = ["*"], # In production, specify allowed origins
allow_credentials = True,
allow_methods = ["*"],
allow_headers = ["*"],
)
# ============ Register API Routes ============
# Register routers
app.include_router(auth_router, prefix="/api/auth", tags=["auth"])
app.include_router(training_router, prefix="/api/train", tags=["training"])
app.include_router(models_router, prefix="/api/models", tags=["models"])
app.include_router(inference_router, prefix="/api/inference", tags=["inference"])
app.include_router(auth_router, prefix = "/api/auth", tags = ["auth"])
app.include_router(training_router, prefix = "/api/train", tags = ["training"])
app.include_router(models_router, prefix = "/api/models", tags = ["models"])
app.include_router(inference_router, prefix = "/api/inference", tags = ["inference"])
# OpenAI-compatible endpoints: mount the same inference router at /v1
# so external tools (Open WebUI, SillyTavern, etc.) can use the
# standard /v1/chat/completions path.
app.include_router(inference_router, prefix="/v1", tags=["openai-compat"])
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"])
app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"])
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 ============
@app.get("/api/health")
async def health_check():
"""Health check endpoint"""
return {
"status": "healthy",
"timestamp": datetime.now().isoformat(),
"service": "Unsloth UI Backend"
"service": "Unsloth UI Backend",
}
@ -167,11 +175,13 @@ async def get_system_info():
gpu_info = {"available": mem_info.get("available", False), "devices": []}
if mem_info.get("available"):
gpu_info["devices"].append({
"index": mem_info.get("device", 0),
"name": mem_info.get("device_name", "Unknown"),
"memory_total_gb": round(mem_info.get("total_gb", 0), 2),
})
gpu_info["devices"].append(
{
"index": mem_info.get("device", 0),
"name": mem_info.get("device_name", "Unknown"),
"memory_total_gb": round(mem_info.get("total_gb", 0), 2),
}
)
# CPU & Memory
memory = psutil.virtual_memory()
@ -203,6 +213,7 @@ async def get_hardware_info():
# ============ Serve Frontend (Optional) ============
def setup_frontend(app: FastAPI, build_path: Path):
"""Mount frontend static files (optional)"""
if not build_path.exists():
@ -211,15 +222,15 @@ def setup_frontend(app: FastAPI, build_path: Path):
# Mount assets
assets_dir = build_path / "assets"
if assets_dir.exists():
app.mount("/assets", StaticFiles(directory=assets_dir), name="assets")
app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
@app.get("/")
async def serve_root():
content = (build_path / "index.html").read_bytes()
return Response(
content=content,
media_type="text/html",
headers={"Cache-Control": "no-cache, no-store, must-revalidate"},
content = content,
media_type = "text/html",
headers = {"Cache-Control": "no-cache, no-store, must-revalidate"},
)
@app.get("/{full_path:path}")
@ -230,8 +241,8 @@ def setup_frontend(app: FastAPI, build_path: Path):
file_path = (build_path / full_path).resolve()
# Block path traversal — ensure resolved path stays inside build_path
if not str(file_path).startswith(str(build_path.resolve())):
return Response(status_code=403)
if not file_path.is_relative_to(build_path.resolve()):
return Response(status_code = 403)
if file_path.is_file():
return FileResponse(file_path)
@ -239,10 +250,9 @@ def setup_frontend(app: FastAPI, build_path: Path):
# Serve index.html as bytes — avoids Content-Length mismatch
content = (build_path / "index.html").read_bytes()
return Response(
content=content,
media_type="text/html",
headers={"Cache-Control": "no-cache, no-store, must-revalidate"},
content = content,
media_type = "text/html",
headers = {"Cache-Control": "no-cache, no-store, must-revalidate"},
)
return True

1
studio/backend/models/__init__.py
View file

@ -4,6 +4,7 @@
"""
Pydantic models for API request/response schemas
"""
from .training import (
TrainingStartRequest,
TrainingJobResponse,

24
studio/backend/models/auth.py
View file

@ -4,28 +4,38 @@
"""
Pydantic schemas for Authentication API
"""
from pydantic import BaseModel, Field
class AuthSetupRequest(BaseModel):
"""First-time setup: create the initial admin user + password."""
setup_token: str = Field(..., description="One-time setup token printed to the server console")
username: str = Field(..., description="Admin username")
password: str = Field(..., min_length=8, description="Admin password (minimum 8 characters)")
setup_token: str = Field(
..., description = "One-time setup token printed to the server console"
)
username: str = Field(..., description = "Admin username")
password: str = Field(
..., min_length = 8, description = "Admin password (minimum 8 characters)"
)
class AuthLoginRequest(BaseModel):
"""Login payload: username/password to obtain a JWT."""
username: str = Field(..., description="Username")
password: str = Field(..., description="Password")
username: str = Field(..., description = "Username")
password: str = Field(..., description = "Password")
class RefreshTokenRequest(BaseModel):
"""Refresh token payload to obtain new access + refresh tokens."""
refresh_token: str = Field(..., description="Refresh token from a previous login or refresh")
refresh_token: str = Field(
..., description = "Refresh token from a previous login or refresh"
)
class AuthStatusResponse(BaseModel):
"""Indicate whether auth has been initialized."""
initialized: bool = Field(..., description="True if auth setup has been completed")
initialized: bool = Field(..., description = "True if auth setup has been completed")

34
studio/backend/models/data_recipe.py
View file

@ -13,13 +13,13 @@ from pydantic import BaseModel, Field
class RecipePayload(BaseModel):
recipe: dict[str, Any] = Field(default_factory=dict)
recipe: dict[str, Any] = Field(default_factory = dict)
run: dict[str, Any] | None = None
ui: dict[str, Any] | None = None
class PreviewResponse(BaseModel):
dataset: list[dict[str, Any]] = Field(default_factory=list)
dataset: list[dict[str, Any]] = Field(default_factory = list)
processor_artifacts: dict[str, Any] | None = None
analysis: dict[str, Any] | None = None
@ -32,7 +32,7 @@ class ValidateError(BaseModel):
class ValidateResponse(BaseModel):
valid: bool
errors: list[ValidateError] = Field(default_factory=list)
errors: list[ValidateError] = Field(default_factory = list)
raw_detail: str | None = None
@ -41,42 +41,42 @@ class JobCreateResponse(BaseModel):
class SeedInspectRequest(BaseModel):
dataset_name: str = Field(min_length=1)
dataset_name: str = Field(min_length = 1)
hf_token: str | None = None
subset: str | None = None
split: str | None = "train"
preview_size: int = Field(default=10, ge=1, le=50)
preview_size: int = Field(default = 10, ge = 1, le = 50)
class SeedInspectUploadRequest(BaseModel):
filename: str = Field(min_length=1)
content_base64: str = Field(min_length=1)
preview_size: int = Field(default=10, ge=1, le=50)
filename: str = Field(min_length = 1)
content_base64: str = Field(min_length = 1)
preview_size: int = Field(default = 10, ge = 1, le = 50)
seed_source_type: str | None = None
unstructured_chunk_size: int | None = Field(default=None, ge=1, le=20000)
unstructured_chunk_overlap: int | None = Field(default=None, ge=0, le=20000)
unstructured_chunk_size: int | None = Field(default = None, ge = 1, le = 20000)
unstructured_chunk_overlap: int | None = Field(default = None, ge = 0, le = 20000)
class SeedInspectResponse(BaseModel):
dataset_name: str
resolved_path: str
columns: list[str] = Field(default_factory=list)
preview_rows: list[dict[str, Any]] = Field(default_factory=list)
columns: list[str] = Field(default_factory = list)
preview_rows: list[dict[str, Any]] = Field(default_factory = list)
split: str | None = None
subset: str | None = None
class McpToolsListRequest(BaseModel):
mcp_providers: list[dict[str, Any]] = Field(default_factory=list)
timeout_sec: float | None = Field(default=None, gt=0)
mcp_providers: list[dict[str, Any]] = Field(default_factory = list)
timeout_sec: float | None = Field(default = None, gt = 0)
class McpToolsProviderResult(BaseModel):
name: str
tools: list[str] = Field(default_factory=list)
tools: list[str] = Field(default_factory = list)
error: str | None = None
class McpToolsListResponse(BaseModel):
providers: list[McpToolsProviderResult] = Field(default_factory=list)
duplicate_tools: dict[str, list[str]] = Field(default_factory=dict)
providers: list[McpToolsProviderResult] = Field(default_factory = list)
duplicate_tools: dict[str, list[str]] = Field(default_factory = dict)

14
studio/backend/models/datasets.py
View file

@ -4,6 +4,7 @@
"""
Dataset-related Pydantic models for API requests and responses.
"""
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field, model_validator
@ -11,13 +12,14 @@ from pydantic import BaseModel, Field, model_validator
class CheckFormatRequest(BaseModel):
"""Request for dataset format check"""
dataset_name: str # HuggingFace dataset name or local path
is_vlm: bool = False
hf_token: Optional[str] = None
subset: Optional[str] = None
train_split: Optional[str] = "train"
@model_validator(mode="before")
@model_validator(mode = "before")
@classmethod
def _compat_split(cls, values: Any) -> Any:
"""Accept legacy 'split' field as alias for 'train_split'."""
@ -28,6 +30,7 @@ class CheckFormatRequest(BaseModel):
class CheckFormatResponse(BaseModel):
"""Response for dataset format check"""
requires_manual_mapping: bool
detected_format: str
columns: List[str]
@ -46,6 +49,7 @@ class CheckFormatResponse(BaseModel):
class AiAssistMappingRequest(BaseModel):
"""Request for LLM-assisted column classification (user-triggered)."""
columns: List[str]
samples: List[Dict[str, Any]] # Preview rows already loaded in the dialog
dataset_name: Optional[str] = None # For LLM context
@ -56,6 +60,7 @@ class AiAssistMappingRequest(BaseModel):
class AiAssistMappingResponse(BaseModel):
"""Response from LLM-assisted column classification and conversion advice."""
success: bool
suggested_mapping: Optional[Dict[str, str]] = None
warning: Optional[str] = None
@ -69,8 +74,9 @@ class AiAssistMappingResponse(BaseModel):
class UploadDatasetResponse(BaseModel):
"""Response with stored dataset path for training."""
filename: str = Field(..., description="Original filename")
stored_path: str = Field(..., description="Absolute path stored on backend")
filename: str = Field(..., description = "Original filename")
stored_path: str = Field(..., description = "Absolute path stored on backend")
class LocalDatasetItem(BaseModel):
@ -90,4 +96,4 @@ class LocalDatasetItem(BaseModel):
class LocalDatasetsResponse(BaseModel):
datasets: List[LocalDatasetItem] = Field(default_factory=list)
datasets: List[LocalDatasetItem] = Field(default_factory = list)

55
studio/backend/models/export.py
View file

@ -4,6 +4,7 @@
"""
Pydantic schemas for Export API.
"""
from pydantic import BaseModel, Field
from typing import List, Optional, Literal, Dict, Any
@ -11,20 +12,20 @@ from typing import List, Optional, Literal, Dict, Any
class LoadCheckpointRequest(BaseModel):
"""Request for loading a checkpoint into the export backend."""
checkpoint_path: str = Field(..., description="Path to the checkpoint directory")
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",
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",
description = "Whether to load the model in 4-bit quantization",
)
trust_remote_code: bool = Field(
False,
description="Allow loading models with custom code. Only enable for checkpoints/base models you trust.",
description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.",
)
@ -33,26 +34,26 @@ class ExportStatusResponse(BaseModel):
current_checkpoint: Optional[str] = Field(
None,
description="Path to the currently loaded checkpoint, if any",
description = "Path to the currently loaded checkpoint, if any",
)
is_vision: bool = Field(
False,
description="True if the loaded checkpoint is a vision model",
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",
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")
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",
default = None,
description = "Optional extra details about the operation",
)
@ -61,27 +62,27 @@ class ExportCommonOptions(BaseModel):
save_directory: str = Field(
...,
description="Local directory where the exported artifacts will be written",
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",
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)",
description = "Hugging Face Hub repository ID (username/model-name)",
)
hf_token: Optional[str] = Field(
None,
description="Hugging Face access token used for Hub operations",
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)",
description = "If True, create a private repository on the Hub (where applicable)",
)
base_model_id: Optional[str] = Field(
None,
description="HuggingFace model ID of the base model (for model card metadata)",
description = "HuggingFace model ID of the base model (for model card metadata)",
)
@ -90,7 +91,7 @@ class ExportMergedModelRequest(ExportCommonOptions):
format_type: Literal["16-bit (FP16)", "4-bit (FP4)"] = Field(
"16-bit (FP16)",
description="Export precision / format for the merged model",
description = "Export precision / format for the merged model",
)
@ -98,7 +99,6 @@ class ExportBaseModelRequest(ExportCommonOptions):
"""Request for exporting a non-PEFT (base) model."""
# Uses fields from ExportCommonOptions only
pass
class ExportGGUFRequest(BaseModel):
@ -106,23 +106,23 @@ class ExportGGUFRequest(BaseModel):
save_directory: str = Field(
...,
description="Directory where GGUF files will be saved",
description = "Directory where GGUF files will be saved",
)
quantization_method: str = Field(
"Q4_K_M",
description='GGUF quantization method (e.g. "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",
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",
description = "Hugging Face Hub repository ID for GGUF upload",
)
hf_token: Optional[str] = Field(
None,
description="Hugging Face token for GGUF upload",
description = "Hugging Face token for GGUF upload",
)
@ -130,6 +130,3 @@ class ExportLoRAAdapterRequest(ExportCommonOptions):
"""Request for exporting only the LoRA adapter (not merged)."""
# Uses fields from ExportCommonOptions only
pass

220
studio/backend/models/inference.py
View file

@ -4,6 +4,7 @@
"""
Pydantic schemas for Inference API
"""
from __future__ import annotations
import time
@ -15,21 +16,29 @@ from pydantic import BaseModel, Discriminator, Field, Tag
class LoadRequest(BaseModel):
"""Request to load a model for inference"""
model_path: str = Field(..., description="Model identifier or local path")
hf_token: Optional[str] = Field(None, description="HuggingFace token for gated models")
max_seq_length: int = Field(4096, ge=128, le=32768, description="Maximum sequence length")
load_in_4bit: bool = Field(True, description="Load model in 4-bit quantization")
is_lora: bool = Field(False, description="Whether this is a LoRA adapter")
gguf_variant: Optional[str] = Field(None, description="GGUF quantization variant (e.g. 'Q4_K_M')")
model_path: str = Field(..., description = "Model identifier or local path")
hf_token: Optional[str] = Field(
None, description = "HuggingFace token for gated models"
)
max_seq_length: int = Field(
4096, ge = 128, le = 32768, description = "Maximum sequence length"
)
load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization")
is_lora: bool = Field(False, description = "Whether this is a LoRA adapter")
gguf_variant: Optional[str] = Field(
None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
)
trust_remote_code: bool = Field(
False,
description="Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
)
class UnloadRequest(BaseModel):
"""Request to unload a model"""
model_path: str = Field(..., description="Model identifier to unload")
model_path: str = Field(..., description = "Model identifier to unload")
class ValidateModelRequest(BaseModel):
@ -39,9 +48,14 @@ class ValidateModelRequest(BaseModel):
This does NOT actually load weights into GPU memory.
"""
model_path: str = Field(..., description="Model identifier or local path")
hf_token: Optional[str] = Field(None, description="HuggingFace token for gated models")
gguf_variant: Optional[str] = Field(None, description="GGUF quantization variant (e.g. 'Q4_K_M')")
model_path: str = Field(..., description = "Model identifier or local path")
hf_token: Optional[str] = Field(
None, description = "HuggingFace token for gated models"
)
gguf_variant: Optional[str] = Field(
None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
)
class ValidateModelResponse(BaseModel):
@ -51,58 +65,99 @@ class ValidateModelResponse(BaseModel):
valid == True means ModelConfig.from_identifier() succeeded and basic
introspection (GGUF / LoRA / vision flags) is available.
"""
valid: bool = Field(..., description="Whether the model identifier looks valid")
message: str = Field(..., description="Human-readable validation message")
identifier: Optional[str] = Field(None, description="Resolved model identifier")
display_name: Optional[str] = Field(None, description="Display name derived from identifier")
is_gguf: bool = Field(False, description="Whether this is a GGUF model (llama.cpp)")
is_lora: bool = Field(False, description="Whether this is a LoRA adapter")
is_vision: bool = Field(False, description="Whether this is a vision-capable model")
valid: bool = Field(..., description = "Whether the model identifier looks valid")
message: str = Field(..., description = "Human-readable validation message")
identifier: Optional[str] = Field(None, description = "Resolved model identifier")
display_name: Optional[str] = Field(
None, description = "Display name derived from identifier"
)
is_gguf: bool = Field(False, description = "Whether this is a GGUF model (llama.cpp)")
is_lora: bool = Field(False, description = "Whether this is a LoRA adapter")
is_vision: bool = Field(False, description = "Whether this is a vision-capable model")
class GenerateRequest(BaseModel):
"""Request for text generation (legacy /generate/stream endpoint)"""
messages: List[dict] = Field(..., description="Chat messages in OpenAI format")
system_prompt: str = Field("You are a helpful AI assistant.", description="System prompt")
temperature: float = Field(0.7, ge=0.0, le=2.0, description="Sampling temperature")
top_p: float = Field(0.9, ge=0.0, le=1.0, description="Top-p sampling")
top_k: int = Field(40, ge=-1, le=100, description="Top-k sampling")
max_new_tokens: int = Field(2048, ge=1, le=4096, description="Maximum tokens to generate")
repetition_penalty: float = Field(1.1, ge=1.0, le=2.0, description="Repetition penalty")
image_base64: Optional[str] = Field(None, description="Base64 encoded image for vision models")
messages: List[dict] = Field(..., description = "Chat messages in OpenAI format")
system_prompt: str = Field(
"You are a helpful AI assistant.", description = "System prompt"
)
temperature: float = Field(0.7, ge = 0.0, le = 2.0, description = "Sampling temperature")
top_p: float = Field(0.9, ge = 0.0, le = 1.0, description = "Top-p sampling")
top_k: int = Field(40, ge = -1, le = 100, description = "Top-k sampling")
max_new_tokens: int = Field(
2048, ge = 1, le = 4096, description = "Maximum tokens to generate"
)
repetition_penalty: float = Field(
1.1, ge = 1.0, le = 2.0, description = "Repetition penalty"
)
image_base64: Optional[str] = Field(
None, description = "Base64 encoded image for vision models"
)
class LoadResponse(BaseModel):
"""Response after loading a model"""
status: str = Field(..., description="Load status")
model: str = Field(..., description="Model identifier")
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")
is_gguf: bool = Field(False, description="Whether model is a GGUF model (llama.cpp)")
is_audio: bool = Field(False, description="Whether model is a TTS audio model")
audio_type: Optional[str] = Field(None, description="Audio codec type: snac, csm, bicodec, dac")
has_audio_input: bool = Field(False, description="Whether model accepts audio input (ASR)")
inference: dict = Field(..., description="Inference parameters (temperature, top_p, top_k, min_p)")
status: str = Field(..., description = "Load status")
model: str = Field(..., description = "Model identifier")
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")
is_gguf: bool = Field(
False, description = "Whether model is a GGUF model (llama.cpp)"
)
is_audio: bool = Field(False, description = "Whether model is a TTS audio model")
audio_type: Optional[str] = Field(
None, description = "Audio codec type: snac, csm, bicodec, dac"
)
has_audio_input: bool = Field(
False, description = "Whether model accepts audio input (ASR)"
)
inference: dict = Field(
..., description = "Inference parameters (temperature, top_p, top_k, min_p)"
)
class UnloadResponse(BaseModel):
"""Response after unloading a model"""
status: str = Field(..., description="Unload status")
model: str = Field(..., description="Model identifier that was unloaded")
status: str = Field(..., description = "Unload status")
model: str = Field(..., description = "Model identifier that was unloaded")
class InferenceStatusResponse(BaseModel):
"""Current inference backend status"""
active_model: Optional[str] = Field(None, description="Currently active model identifier")
is_vision: bool = Field(False, description="Whether the active model is a vision model")
is_gguf: bool = Field(False, description="Whether the active model is a GGUF model (llama.cpp)")
gguf_variant: Optional[str] = Field(None, description="GGUF quantization variant (e.g. Q4_K_M)")
is_audio: bool = Field(False, description="Whether the active model is a TTS audio model")
audio_type: Optional[str] = Field(None, description="Audio codec type: snac, csm, bicodec, dac")
has_audio_input: bool = Field(False, description="Whether model accepts audio input (ASR)")
loading: List[str] = Field(default_factory=list, description="Models currently being loaded")
loaded: List[str] = Field(default_factory=list, description="Models currently loaded")
active_model: Optional[str] = Field(
None, description = "Currently active model identifier"
)
is_vision: bool = Field(
False, description = "Whether the active model is a vision model"
)
is_gguf: bool = Field(
False, description = "Whether the active model is a GGUF model (llama.cpp)"
)
gguf_variant: Optional[str] = Field(
None, description = "GGUF quantization variant (e.g. Q4_K_M)"
)
is_audio: bool = Field(
False, description = "Whether the active model is a TTS audio model"
)
audio_type: Optional[str] = Field(
None, description = "Audio codec type: snac, csm, bicodec, dac"
)
has_audio_input: bool = Field(
False, description = "Whether model accepts audio input (ASR)"
)
loading: List[str] = Field(
default_factory = list, description = "Models currently being loaded"
)
loaded: List[str] = Field(
default_factory = list, description = "Models currently loaded"
)
# =====================================================================
@ -112,20 +167,24 @@ class InferenceStatusResponse(BaseModel):
# ── Multimodal content parts (OpenAI vision format) ──────────────
class TextContentPart(BaseModel):
"""Text content part in a multimodal message."""
type: Literal["text"]
text: str
class ImageUrl(BaseModel):
"""Image URL object — supports data URIs and remote URLs."""
url: str = Field(..., description="data:image/png;base64,... or https://...")
url: str = Field(..., description = "data:image/png;base64,... or https://...")
detail: Optional[Literal["auto", "low", "high"]] = "auto"
class ImageContentPart(BaseModel):
"""Image content part in a multimodal message."""
type: Literal["image_url"]
image_url: ImageUrl
@ -148,6 +207,7 @@ ContentPart = Annotated[
# ── Messages ─────────────────────────────────────────────────────
class ChatMessage(BaseModel):
"""
A single message in the conversation.
@ -155,8 +215,13 @@ class ChatMessage(BaseModel):
``content`` may be a plain string (text-only) or a list of
content parts for multimodal messages (OpenAI vision format).
"""
role: Literal["system", "user", "assistant"] = Field(..., description="Message role")
content: Union[str, list[ContentPart]] = Field(..., description="Message content (string or multimodal parts)")
role: Literal["system", "user", "assistant"] = Field(
..., description = "Message role"
)
content: Union[str, list[ContentPart]] = Field(
..., description = "Message content (string or multimodal parts)"
)
class ChatCompletionRequest(BaseModel):
@ -165,22 +230,36 @@ class ChatCompletionRequest(BaseModel):
Extensions (non-OpenAI fields) are marked with 'x-unsloth'.
"""
model: str = Field("default", description="Model identifier (informational; the active model is used)")
messages: list[ChatMessage] = Field(..., description="Conversation messages")
stream: bool = Field(True, description="Whether to stream the response via SSE")
temperature: float = Field(0.7, ge=0.0, le=2.0)
top_p: float = Field(0.9, ge=0.0, le=1.0)
max_tokens: Optional[int] = Field(2048, ge=1, le=4096, description="Maximum tokens to generate")
model: str = Field(
"default",
description = "Model identifier (informational; the active model is used)",
)
messages: list[ChatMessage] = Field(..., description = "Conversation messages")
stream: bool = Field(True, description = "Whether to stream the response via SSE")
temperature: float = Field(0.7, ge = 0.0, le = 2.0)
top_p: float = Field(0.9, ge = 0.0, le = 1.0)
max_tokens: Optional[int] = Field(
2048, ge = 1, le = 4096, description = "Maximum tokens to generate"
)
# ── Unsloth extensions (ignored by standard OpenAI clients) ──
top_k: int = Field(40, ge=-1, le=100, description="[x-unsloth] Top-k sampling")
min_p: float = Field(0.0, ge=0.0, le=1.0, description="[x-unsloth] Min-p sampling threshold")
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")
audio_base64: Optional[str] = Field(None, description="[x-unsloth] Base64-encoded WAV for audio-input models (ASR)")
top_k: int = Field(40, ge = -1, le = 100, description = "[x-unsloth] Top-k sampling")
min_p: float = Field(
0.0, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold"
)
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"
)
audio_base64: Optional[str] = Field(
None, description = "[x-unsloth] Base64-encoded WAV for audio-input models (ASR)"
)
use_adapter: Optional[Union[bool, str]] = Field(
None,
description=(
description = (
"[x-unsloth] Adapter control for compare mode. "
"null = no change (default), "
"false = disable adapters (base model), "
@ -195,12 +274,14 @@ class ChatCompletionRequest(BaseModel):
class ChoiceDelta(BaseModel):
"""Delta content for a streaming chunk."""
role: Optional[str] = None
content: Optional[str] = None
class ChunkChoice(BaseModel):
"""A single choice in a streaming chunk."""
index: int = 0
delta: ChoiceDelta
finish_reason: Optional[Literal["stop", "length"]] = None
@ -208,9 +289,10 @@ class ChunkChoice(BaseModel):
class ChatCompletionChunk(BaseModel):
"""A single SSE chunk in OpenAI streaming format."""
id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
id: str = Field(default_factory = lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
created: int = Field(default_factory=lambda: int(time.time()))
created: int = Field(default_factory = lambda: int(time.time()))
model: str = "default"
choices: list[ChunkChoice]
@ -220,12 +302,14 @@ class ChatCompletionChunk(BaseModel):
class CompletionMessage(BaseModel):
"""The assistant's complete response message."""
role: Literal["assistant"] = "assistant"
content: str
class CompletionChoice(BaseModel):
"""A single choice in a non-streaming response."""
index: int = 0
message: CompletionMessage
finish_reason: Literal["stop", "length"] = "stop"
@ -233,6 +317,7 @@ class CompletionChoice(BaseModel):
class CompletionUsage(BaseModel):
"""Token usage statistics (approximate)."""
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
@ -240,9 +325,10 @@ class CompletionUsage(BaseModel):
class ChatCompletion(BaseModel):
"""Non-streaming chat completion response."""
id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
id: str = Field(default_factory = lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
object: Literal["chat.completion"] = "chat.completion"
created: int = Field(default_factory=lambda: int(time.time()))
created: int = Field(default_factory = lambda: int(time.time()))
model: str = "default"
choices: list[CompletionChoice]
usage: CompletionUsage = Field(default_factory=CompletionUsage)
usage: CompletionUsage = Field(default_factory = CompletionUsage)

149
studio/backend/models/models.py
View file

@ -4,6 +4,7 @@
"""
Pydantic schemas for Model Management API
"""
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any, Literal
@ -13,123 +14,169 @@ ModelType = Literal["text", "vision", "audio", "embeddings"]
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")
loss: Optional[float] = Field(None, description="Training loss at this checkpoint")
display_name: str = Field(
..., description = "User-friendly checkpoint name (folder name)"
)
path: str = Field(..., description = "Full path to the checkpoint directory")
loss: Optional[float] = Field(None, description = "Training loss at this checkpoint")
class ModelCheckpoints(BaseModel):
"""A training run and its associated checkpoints."""
name: str = Field(..., description="Training run folder name")
name: str = Field(..., description = "Training run folder name")
checkpoints: List[CheckpointInfo] = Field(
default_factory=list,
description="List of checkpoints for this training run (final + intermediate)",
default_factory = list,
description = "List of checkpoints for this training run (final + intermediate)",
)
base_model: Optional[str] = Field(
None,
description="Base model name from adapter_config.json or config.json",
description = "Base model name from adapter_config.json or config.json",
)
peft_type: Optional[str] = Field(
None,
description="PEFT type (e.g. LORA) if adapter training, None for full fine-tune",
description = "PEFT type (e.g. LORA) if adapter training, None for full fine-tune",
)
lora_rank: Optional[int] = Field(
None,
description="LoRA rank (r) if applicable",
description = "LoRA rank (r) if applicable",
)
class CheckpointListResponse(BaseModel):
"""Response for listing available checkpoints in an outputs directory."""
outputs_dir: str = Field(..., description="Directory that was scanned")
outputs_dir: str = Field(..., description = "Directory that was scanned")
models: List[ModelCheckpoints] = Field(
default_factory=list,
description="List of training runs with their checkpoints",
default_factory = list,
description = "List of training runs with their checkpoints",
)
class ModelDetails(BaseModel):
"""Detailed model configuration and metadata - can be used for both list and detail views"""
id: str = Field(..., description="Model identifier")
model_name: Optional[str] = Field(None, description="Model identifier (alias for id, for backward compatibility)")
name: Optional[str] = Field(None, description="Display name for the model")
config: Optional[Dict[str, Any]] = Field(None, description="Model configuration dictionary")
is_vision: bool = Field(False, description="Whether model is a vision model")
is_embedding: bool = Field(False, description="Whether model is an embedding/sentence-transformer model")
is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
is_gguf: bool = Field(False, description="Whether model is a GGUF model (llama.cpp format)")
is_audio: bool = Field(False, description="Whether model is a TTS audio model")
audio_type: Optional[str] = Field(None, description="Audio codec type: snac, csm, bicodec, dac")
has_audio_input: bool = Field(False, description="Whether model accepts audio input (ASR)")
model_type: Optional[ModelType] = Field(None, description="Collapsed model modality: text, vision, audio, or embeddings")
base_model: Optional[str] = Field(None, description="Base model if this is a LoRA adapter")
id: str = Field(..., description = "Model identifier")
model_name: Optional[str] = Field(
None, description = "Model identifier (alias for id, for backward compatibility)"
)
name: Optional[str] = Field(None, description = "Display name for the model")
config: Optional[Dict[str, Any]] = Field(
None, description = "Model configuration dictionary"
)
is_vision: bool = Field(False, description = "Whether model is a vision model")
is_embedding: bool = Field(
False, description = "Whether model is an embedding/sentence-transformer model"
)
is_lora: bool = Field(False, description = "Whether model is a LoRA adapter")
is_gguf: bool = Field(
False, description = "Whether model is a GGUF model (llama.cpp format)"
)
is_audio: bool = Field(False, description = "Whether model is a TTS audio model")
audio_type: Optional[str] = Field(
None, description = "Audio codec type: snac, csm, bicodec, dac"
)
has_audio_input: bool = Field(
False, description = "Whether model accepts audio input (ASR)"
)
model_type: Optional[ModelType] = Field(
None, description = "Collapsed model modality: text, vision, audio, or embeddings"
)
base_model: Optional[str] = Field(
None, description = "Base model if this is a LoRA adapter"
)
class LoRAInfo(BaseModel):
"""LoRA adapter or exported model information"""
display_name: str = Field(..., description="Display name for the LoRA")
adapter_path: str = Field(..., description="Path to the LoRA adapter or exported model")
base_model: Optional[str] = Field(None, description="Base model identifier")
source: Optional[str] = Field(None, description="'training' or 'exported'")
export_type: Optional[str] = Field(None, description="'lora', 'merged', or 'gguf' (for exports)")
display_name: str = Field(..., description = "Display name for the LoRA")
adapter_path: str = Field(
..., description = "Path to the LoRA adapter or exported model"
)
base_model: Optional[str] = Field(None, description = "Base model identifier")
source: Optional[str] = Field(None, description = "'training' or 'exported'")
export_type: Optional[str] = Field(
None, description = "'lora', 'merged', or 'gguf' (for exports)"
)
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")
loras: List[LoRAInfo] = Field(
default_factory = list, description = "List of found LoRA adapters"
)
outputs_dir: str = Field(..., description = "Directory that was scanned")
class ModelListResponse(BaseModel):
"""Response schema for listing models"""
models: List[ModelDetails] = Field(default_factory=list, description="List of models")
default_models: List[str] = Field(default_factory=list, description="List of default model IDs")
models: List[ModelDetails] = Field(
default_factory = list, description = "List of models"
)
default_models: List[str] = Field(
default_factory = list, description = "List of default model IDs"
)
class GgufVariantDetail(BaseModel):
"""A single GGUF quantization variant in a HuggingFace repo."""
filename: str = Field(..., description="GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')")
quant: str = Field(..., description="Quantization label (e.g., 'Q4_K_M')")
size_bytes: int = Field(0, description="File size in bytes")
filename: str = Field(
..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')"
)
quant: str = Field(..., description = "Quantization label (e.g., 'Q4_K_M')")
size_bytes: int = Field(0, description = "File size in bytes")
class GgufVariantsResponse(BaseModel):
"""Response for listing GGUF quantization variants in a HuggingFace repo."""
repo_id: str = Field(..., description="HuggingFace repo ID")
variants: List[GgufVariantDetail] = Field(default_factory=list, description="Available GGUF variants")
has_vision: bool = Field(False, description="Whether the model has vision support (mmproj files)")
default_variant: Optional[str] = Field(None, description="Recommended default quantization variant")
repo_id: str = Field(..., description = "HuggingFace repo ID")
variants: List[GgufVariantDetail] = Field(
default_factory = list, description = "Available GGUF variants"
)
has_vision: bool = Field(
False, description = "Whether the model has vision support (mmproj files)"
)
default_variant: Optional[str] = Field(
None, description = "Recommended default quantization variant"
)
class LocalModelInfo(BaseModel):
"""Discovered local model candidate."""
id: str = Field(..., description="Identifier to use for loading/training")
display_name: str = Field(..., description="Display label")
path: str = Field(..., description="Local path where model data was discovered")
id: str = Field(..., description = "Identifier to use for loading/training")
display_name: str = Field(..., description = "Display label")
path: str = Field(..., description = "Local path where model data was discovered")
source: Literal["models_dir", "hf_cache"] = Field(
...,
description="Discovery source",
description = "Discovery source",
)
model_id: Optional[str] = Field(
None,
description="HF repo id for cached models, e.g. org/model",
description = "HF repo id for cached models, e.g. org/model",
)
updated_at: Optional[float] = Field(
None,
description="Unix timestamp of latest observed update",
description = "Unix timestamp of latest observed update",
)
class LocalModelListResponse(BaseModel):
"""Response schema for listing local/cached models."""
models_dir: str = Field(..., description="Directory scanned for custom local models")
models_dir: str = Field(
..., description = "Directory scanned for custom local models"
)
hf_cache_dir: Optional[str] = Field(
None,
description="HF cache root that was scanned",
description = "HF cache root that was scanned",
)
models: List[LocalModelInfo] = Field(
default_factory=list,
description="Discovered local/cached models",
default_factory = list,
description = "Discovered local/cached models",
)

50
studio/backend/models/responses.py
View file

@ -5,45 +5,63 @@
Pydantic response schemas for endpoints that previously returned raw dicts.
These are small response models for training and model management routes.
"""
from pydantic import BaseModel, Field
from typing import Optional, List
# --- Training route response models ---
class TrainingStopResponse(BaseModel):
"""Response for stopping a training job"""
status: str = Field(..., description="Current status: 'stopped' or 'idle'")
message: str = Field(..., description="Human-readable status message")
status: str = Field(..., description = "Current status: 'stopped' or 'idle'")
message: str = Field(..., description = "Human-readable status message")
class TrainingMetricsResponse(BaseModel):
"""Response for training metrics history"""
loss_history: List[float] = Field(default_factory=list, description="Loss values per step")
lr_history: List[float] = Field(default_factory=list, description="Learning rate per step")
step_history: List[int] = Field(default_factory=list, description="Step numbers")
grad_norm_history: List[float] = Field(default_factory=list, description="Gradient norm values")
grad_norm_step_history: List[int] = Field(default_factory=list, description="Step numbers for gradient norm values")
current_loss: Optional[float] = Field(None, description="Most recent loss value")
current_lr: Optional[float] = Field(None, description="Most recent learning rate")
current_step: Optional[int] = Field(None, description="Most recent step number")
loss_history: List[float] = Field(
default_factory = list, description = "Loss values per step"
)
lr_history: List[float] = Field(
default_factory = list, description = "Learning rate per step"
)
step_history: List[int] = Field(default_factory = list, description = "Step numbers")
grad_norm_history: List[float] = Field(
default_factory = list, description = "Gradient norm values"
)
grad_norm_step_history: List[int] = Field(
default_factory = list, description = "Step numbers for gradient norm values"
)
current_loss: Optional[float] = Field(None, description = "Most recent loss value")
current_lr: Optional[float] = Field(None, description = "Most recent learning rate")
current_step: Optional[int] = Field(None, description = "Most recent step number")
# --- Model management route response models ---
class LoRABaseModelResponse(BaseModel):
"""Response for getting a LoRA's base model"""
lora_path: str = Field(..., description="Path to the LoRA adapter")
base_model: str = Field(..., description="Base model identifier")
lora_path: str = Field(..., description = "Path to the LoRA adapter")
base_model: str = Field(..., description = "Base model identifier")
class VisionCheckResponse(BaseModel):
"""Response for checking if a model is a vision model"""
model_name: str = Field(..., description="Model identifier")
is_vision: bool = Field(..., description="Whether the model is a vision model")
model_name: str = Field(..., description = "Model identifier")
is_vision: bool = Field(..., description = "Whether the model is a vision model")
class EmbeddingCheckResponse(BaseModel):
"""Response for checking if a model is an embedding model"""
model_name: str = Field(..., description="Model identifier")
is_embedding: bool = Field(..., description="Whether the model is an embedding/sentence-transformer model")
model_name: str = Field(..., description = "Model identifier")
is_embedding: bool = Field(
..., description = "Whether the model is an embedding/sentence-transformer model"
)

207
studio/backend/models/training.py
View file

@ -4,44 +4,63 @@
"""
Pydantic schemas for Training API
"""
from pydantic import BaseModel, Field, model_validator
from typing import Any, Optional, List, Dict, 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")
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")
trust_remote_code: bool = Field(
False,
description="Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
)
# 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")
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")
subset: Optional[str] = None
train_split: Optional[str] = Field("train", description="Training split name")
eval_split: Optional[str] = Field(None, description="Eval split name. None = auto-detect")
eval_steps: float = Field(0.00, description="Fraction of total steps between evals (0-1)")
dataset_slice_start: Optional[int] = Field(None, description="Inclusive start row index for dataset slicing")
dataset_slice_end: Optional[int] = Field(None, description="Inclusive end row index for dataset slicing")
train_split: Optional[str] = Field("train", description = "Training split name")
eval_split: Optional[str] = Field(
None, description = "Eval split name. None = auto-detect"
)
eval_steps: float = Field(
0.00, description = "Fraction of total steps between evals (0-1)"
)
dataset_slice_start: Optional[int] = Field(
None, description = "Inclusive start row index for dataset slicing"
)
dataset_slice_end: Optional[int] = Field(
None, description = "Inclusive end row index for dataset slicing"
)
@model_validator(mode="before")
@model_validator(mode = "before")
@classmethod
def _compat_split(cls, values: Any) -> Any:
"""Accept legacy 'split' field as alias for 'train_split'."""
if isinstance(values, dict) and "split" in values:
values.setdefault("train_split", values.pop("split"))
return values
custom_format_mapping: Optional[Dict[str, Any]] = Field(
None,
description=(
description = (
"User-provided column-to-role mapping, e.g. {'image': 'image', 'caption': 'text'} "
"for VLM or {'instruction': 'user', 'output': 'assistant'} for LLM. "
"Enhanced format includes __system_prompt, __user_template, "
@ -49,59 +68,77 @@ class TrainingStartRequest(BaseModel):
),
)
# 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")
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")
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")
is_dataset_image: bool = Field(False, description="Whether the dataset contains image data")
is_dataset_audio: bool = Field(False, description="Whether the dataset contains audio data")
is_embedding: bool = Field(False, description="Whether model is an embedding/sentence-transformer model")
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")
is_dataset_image: bool = Field(
False, description = "Whether the dataset contains image data"
)
is_dataset_audio: bool = Field(
False, description = "Whether the dataset contains audio data"
)
is_embedding: bool = Field(
False, description = "Whether model is an embedding/sentence-transformer model"
)
# 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")
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'")
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")
job_id: str = Field(..., description = "Training job identifier")
phase: Literal[
"idle",
"loading_model",
@ -110,31 +147,49 @@ class TrainingStatus(BaseModel):
"training",
"completed",
"error",
"stopped"
] = Field(..., description="Current phase of training pipeline")
is_training_running: bool = Field(..., description="True if training loop is actively running")
eval_enabled: bool = Field(False, description="True if evaluation dataset is configured for this training run")
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'}")
"stopped",
] = Field(..., description = "Current phase of training pipeline")
is_training_running: bool = Field(
..., description = "True if training loop is actively running"
)
eval_enabled: bool = Field(
False,
description = "True if evaluation dataset is configured for this training run",
)
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'}"
)
metric_history: Optional[dict] = Field(
None,
description="Full metric history arrays for chart recovery after SSE reconnection. "
"Keys: 'steps', 'loss', 'lr', 'grad_norm', 'grad_norm_steps' — each a list of numeric values.",
description = "Full metric history arrays for chart recovery after SSE reconnection. "
"Keys: 'steps', 'loss', 'lr', 'grad_norm', 'grad_norm_steps' — each a list of numeric values.",
)
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[float] = 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")
eval_loss: Optional[float] = Field(None, description="Eval loss from the most recent evaluation step")
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[float] = 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"
)
eval_loss: Optional[float] = Field(
None, description = "Eval loss from the most recent evaluation step"
)

7
studio/backend/models/users.py
View file

@ -12,7 +12,6 @@ from pydantic import BaseModel, Field
class Token(BaseModel):
"""Authentication token model with access and refresh tokens."""
access_token: str = Field(..., description="JWT access token (60 min expiry)")
refresh_token: str = Field(..., description="Opaque refresh token (7 day expiry)")
token_type: str = Field(..., description="Token type, always 'bearer'")
access_token: str = Field(..., description = "JWT access token (60 min expiry)")
refresh_token: str = Field(..., description = "Opaque refresh token (7 day expiry)")
token_type: str = Field(..., description = "Token type, always 'bearer'")

0
__init__.py → studio/backend/plugins/__init__.py
View file

2
studio/backend/plugins/data-designer-unstructured-seed/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

3
studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml
View file

@ -1,3 +1,6 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
[build-system]
requires = ["setuptools>=68", "wheel"]
build-backend = "setuptools.build_meta"

32
studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py
View file

@ -36,9 +36,9 @@ def build_unstructured_preview_rows(
chunk_overlap: Any,
) -> list[dict[str, str]]:
parquet_path, rows = materialize_unstructured_seed_dataset(
source_path=source_path,
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
source_path = source_path,
chunk_size = chunk_size,
chunk_overlap = chunk_overlap,
)
count = max(0, int(preview_size))
if rows:
@ -47,12 +47,14 @@ def build_unstructured_preview_rows(
try:
import pandas as pd
except ImportError as exc: # pragma: no cover
raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc
raise RuntimeError(
f"pandas is required for unstructured seed processing: {exc}"
) from exc
dataframe = pd.read_parquet(parquet_path).head(count)
return [
{"chunk_text": str(value.get("chunk_text", "")).strip()}
for value in dataframe.to_dict(orient="records")
for value in dataframe.to_dict(orient = "records")
if str(value.get("chunk_text", "")).strip()
]
@ -69,9 +71,9 @@ def materialize_unstructured_seed_dataset(
size, overlap = resolve_chunking(chunk_size, chunk_overlap)
key = _compute_cache_key(
source_path=resolved,
chunk_size=size,
chunk_overlap=overlap,
source_path = resolved,
chunk_size = size,
chunk_overlap = overlap,
)
parquet_path = _CACHE_DIR / f"{key}.parquet"
if parquet_path.exists():
@ -79,9 +81,9 @@ def materialize_unstructured_seed_dataset(
text = load_unstructured_text_file(resolved)
chunks = split_text_into_chunks(
text=text,
chunk_size=size,
chunk_overlap=overlap,
text = text,
chunk_size = size,
chunk_overlap = overlap,
)
if not chunks:
raise ValueError("No text found in unstructured seed source.")
@ -91,10 +93,12 @@ def materialize_unstructured_seed_dataset(
try:
import pandas as pd
except ImportError as exc: # pragma: no cover
raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc
raise RuntimeError(
f"pandas is required for unstructured seed processing: {exc}"
) from exc
tmp_path = _CACHE_DIR / f"{key}.tmp.parquet"
pd.DataFrame(rows).to_parquet(tmp_path, index=False)
pd.DataFrame(rows).to_parquet(tmp_path, index = False)
tmp_path.replace(parquet_path)
return parquet_path, rows
@ -104,7 +108,7 @@ def load_unstructured_text_file(path: Path) -> str:
if ext not in {".txt", ".md"}:
raise ValueError(f"Unsupported unstructured seed file type: {ext}")
raw = path.read_text(encoding="utf-8", errors="ignore")
raw = path.read_text(encoding = "utf-8", errors = "ignore")
return normalize_unstructured_text(raw)

8
studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py
View file

@ -15,11 +15,11 @@ from .chunking import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE, resolve_chunkin
class UnstructuredSeedSource(SeedSource):
seed_type: Literal["unstructured"] = "unstructured"
path: str = Field(..., min_length=1)
path: str = Field(..., min_length = 1)
chunk_size: int = DEFAULT_CHUNK_SIZE
chunk_overlap: int = DEFAULT_CHUNK_OVERLAP
@field_validator("path", mode="after")
@field_validator("path", mode = "after")
@classmethod
def _validate_path(cls, value: str) -> str:
path = Path(value).expanduser()
@ -27,13 +27,13 @@ class UnstructuredSeedSource(SeedSource):
raise ValueError(f"Unstructured seed path is not a file: {path}")
return value
@field_validator("chunk_size", mode="after")
@field_validator("chunk_size", mode = "after")
@classmethod
def _validate_chunk_size(cls, value: int) -> int:
size, _ = resolve_chunking(value, 0)
return size
@field_validator("chunk_overlap", mode="after")
@field_validator("chunk_overlap", mode = "after")
@classmethod
def _validate_chunk_overlap(cls, value: int, info) -> int:
size = info.data.get("chunk_size", cls.model_fields["chunk_size"].default)

6
studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py
View file

@ -18,8 +18,8 @@ class UnstructuredSeedReader(SeedReader[UnstructuredSeedSource]):
def get_dataset_uri(self) -> str:
path, _ = materialize_unstructured_seed_dataset(
source_path=Path(self.source.path),
chunk_size=self.source.chunk_size,
chunk_overlap=self.source.chunk_overlap,
source_path = Path(self.source.path),
chunk_size = self.source.chunk_size,
chunk_overlap = self.source.chunk_overlap,
)
return str(path)

6
studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py
View file

@ -4,7 +4,7 @@
from data_designer.plugins.plugin import Plugin, PluginType
unstructured_seed_plugin = Plugin(
impl_qualified_name="data_designer_unstructured_seed.impl.UnstructuredSeedReader",
config_qualified_name="data_designer_unstructured_seed.config.UnstructuredSeedSource",
plugin_type=PluginType.SEED_READER,
impl_qualified_name = "data_designer_unstructured_seed.impl.UnstructuredSeedReader",
config_qualified_name = "data_designer_unstructured_seed.config.UnstructuredSeedSource",
plugin_type = PluginType.SEED_READER,
)

2
studio/backend/requirements/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

5
studio/backend/requirements/single-env/patch_metadata.py
View file

@ -1,6 +1,7 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Relax strict metadata pins so pip check matches known working single-env stack.
Why:
@ -48,13 +49,13 @@ def metadata_path(dist_name: str) -> Path | None:
def patch_file(path: Path) -> bool:
original = path.read_text(encoding="utf-8")
original = path.read_text(encoding = "utf-8")
updated = original
for pattern, repl in PATCHES:
updated = pattern.sub(repl, updated)
if updated == original:
return False
path.write_text(updated, encoding="utf-8")
path.write_text(updated, encoding = "utf-8")
return True

68
studio/backend/routes/auth.py
View file

@ -4,6 +4,7 @@
"""
Authentication API routes
"""
from fastapi import APIRouter, HTTPException, status
import secrets
@ -25,7 +26,7 @@ from auth.authentication import (
router = APIRouter()
@router.get("/status", response_model=AuthStatusResponse)
@router.get("/status", response_model = AuthStatusResponse)
async def auth_status() -> AuthStatusResponse:
"""
Check whether auth has already been initialized.
@ -33,10 +34,10 @@ async def auth_status() -> AuthStatusResponse:
- initialized = False -> frontend should show "Set admin password" screen.
- initialized = True -> frontend should show normal login.
"""
return AuthStatusResponse(initialized=storage.is_initialized())
return AuthStatusResponse(initialized = storage.is_initialized())
@router.post("/setup", response_model=Token, status_code=status.HTTP_201_CREATED)
@router.post("/setup", response_model = Token, status_code = status.HTTP_201_CREATED)
async def setup_auth(payload: AuthSetupRequest) -> Token:
"""
First-time setup: create the admin user and a JWT secret.
@ -46,15 +47,15 @@ async def setup_auth(payload: AuthSetupRequest) -> Token:
"""
if storage.is_initialized():
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Auth is already initialized.",
status_code = status.HTTP_400_BAD_REQUEST,
detail = "Auth is already initialized.",
)
# Validate the one-time setup token
if not storage.consume_setup_token(payload.setup_token):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Invalid or expired setup token.",
status_code = status.HTTP_403_FORBIDDEN,
detail = "Invalid or expired setup token.",
)
# Generate a strong random JWT secret for this installation
@ -63,34 +64,34 @@ async def setup_auth(payload: AuthSetupRequest) -> Token:
# Create user + generate tokens atomically — rollback if anything fails
try:
storage.create_initial_user(
username=payload.username,
password=payload.password,
jwt_secret=jwt_secret,
username = payload.username,
password = payload.password,
jwt_secret = jwt_secret,
)
# Reload JWT secret from DB (so authentication.py picks it up)
reload_secret()
# Issue access + refresh tokens for the new user
access_token = create_access_token(subject=payload.username)
refresh_token = create_refresh_token(subject=payload.username)
access_token = create_access_token(subject = payload.username)
refresh_token = create_refresh_token(subject = payload.username)
except Exception as e:
# Rollback: remove the user row so setup can be retried
storage.delete_user(payload.username)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Setup failed (rolled back): {str(e)}",
status_code = status.HTTP_500_INTERNAL_SERVER_ERROR,
detail = f"Setup failed (rolled back): {str(e)}",
)
return Token(
access_token=access_token,
refresh_token=refresh_token,
token_type="bearer",
access_token = access_token,
refresh_token = refresh_token,
token_type = "bearer",
)
@router.post("/login", response_model=Token)
@router.post("/login", response_model = Token)
async def login(payload: AuthLoginRequest) -> Token:
"""
Login with username/password and receive access + refresh tokens.
@ -98,27 +99,27 @@ async def login(payload: AuthLoginRequest) -> Token:
record = storage.get_user_and_secret(payload.username)
if record is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Incorrect username or password",
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Incorrect username or password",
)
salt, pwd_hash, _jwt_secret = record
if not hashing.verify_password(payload.password, salt, pwd_hash):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Incorrect username or password",
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Incorrect username or password",
)
access_token = create_access_token(subject=payload.username)
refresh_token = create_refresh_token(subject=payload.username)
access_token = create_access_token(subject = payload.username)
refresh_token = create_refresh_token(subject = payload.username)
return Token(
access_token=access_token,
refresh_token=refresh_token,
token_type="bearer",
access_token = access_token,
refresh_token = refresh_token,
token_type = "bearer",
)
@router.post("/refresh", response_model=Token)
@router.post("/refresh", response_model = Token)
async def refresh(payload: RefreshTokenRequest) -> Token:
"""
Exchange a valid refresh token for a new access token.
@ -128,13 +129,12 @@ async def refresh(payload: RefreshTokenRequest) -> Token:
new_access_token = refresh_access_token(payload.refresh_token)
if new_access_token is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid or expired refresh token",
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Invalid or expired refresh token",
)
return Token(
access_token=new_access_token,
refresh_token=payload.refresh_token,
token_type="bearer",
access_token = new_access_token,
refresh_token = payload.refresh_token,
token_type = "bearer",
)

2
studio/backend/routes/data_recipe/__init__.py
View file

@ -21,7 +21,7 @@ from .mcp import router as mcp_router
from .seed import router as seed_router
from .validate import router as validate_router
router = APIRouter(dependencies=[Depends(get_current_subject)])
router = APIRouter(dependencies = [Depends(get_current_subject)])
router.include_router(seed_router)
router.include_router(validate_router)
router.include_router(jobs_router)

49
studio/backend/routes/data_recipe/jobs.py
View file

@ -21,25 +21,30 @@ def _normalize_run_name(value: Any) -> str | None:
if value is None:
return None
if not isinstance(value, str):
raise HTTPException(status_code=400, detail="invalid run_name: must be a string")
raise HTTPException(
status_code = 400, detail = "invalid run_name: must be a string"
)
trimmed = value.strip()
if not trimmed:
return None
return trimmed[:120]
@router.post("/jobs", response_class=JSONResponse, response_model=JobCreateResponse)
@router.post("/jobs", response_class = JSONResponse, response_model = JobCreateResponse)
def create_job(payload: RecipePayload):
recipe = payload.recipe
if not recipe.get("columns"):
raise HTTPException(status_code=400, detail="Recipe must include columns.")
raise HTTPException(status_code = 400, detail = "Recipe must include columns.")
run: dict[str, Any] = payload.run or {}
run.pop("artifact_path", None)
run.pop("dataset_name", None)
execution_type = str(run.get("execution_type") or "full").strip().lower()
if execution_type not in {"preview", "full"}:
raise HTTPException(status_code=400, detail="invalid execution_type: must be 'preview' or 'full'")
raise HTTPException(
status_code = 400,
detail = "invalid execution_type: must be 'preview' or 'full'",
)
run["execution_type"] = execution_type
run["run_name"] = _normalize_run_name(run.get("run_name"))
run_config_raw = run.get("run_config")
@ -49,15 +54,17 @@ def create_job(payload: RecipePayload):
RunConfig.model_validate(run_config_raw)
except (ImportError, ValidationError, TypeError, ValueError) as exc:
raise HTTPException(status_code=400, detail=f"invalid run_config: {exc}") from exc
raise HTTPException(
status_code = 400, detail = f"invalid run_config: {exc}"
) from exc
mgr = get_job_manager()
try:
job_id = mgr.start(recipe=recipe, run=run)
job_id = mgr.start(recipe = recipe, run = run)
except RuntimeError as exc:
raise HTTPException(status_code=409, detail=str(exc)) from exc
raise HTTPException(status_code = 409, detail = str(exc)) from exc
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
raise HTTPException(status_code = 400, detail = str(exc)) from exc
return {"job_id": job_id}
@ -67,7 +74,7 @@ def job_status(job_id: str):
mgr = get_job_manager()
state = mgr.get_status(job_id)
if state is None:
raise HTTPException(status_code=404, detail="job not found")
raise HTTPException(status_code = 404, detail = "job not found")
return state
@ -76,7 +83,7 @@ def current_job():
mgr = get_job_manager()
state = mgr.get_current_status()
if state is None:
raise HTTPException(status_code=404, detail="no job")
raise HTTPException(status_code = 404, detail = "no job")
return state
@ -85,7 +92,7 @@ def cancel_job(job_id: str):
mgr = get_job_manager()
ok = mgr.cancel(job_id)
if not ok:
raise HTTPException(status_code=404, detail="job not found")
raise HTTPException(status_code = 404, detail = "job not found")
return mgr.get_status(job_id)
@ -94,22 +101,22 @@ def job_analysis(job_id: str):
mgr = get_job_manager()
analysis = mgr.get_analysis(job_id)
if analysis is None:
raise HTTPException(status_code=404, detail="analysis not ready")
raise HTTPException(status_code = 404, detail = "analysis not ready")
return analysis
@router.get("/jobs/{job_id}/dataset")
def job_dataset(
job_id: str,
limit: int = Query(default=20, ge=1, le=500),
offset: int = Query(default=0, ge=0),
limit: int = Query(default = 20, ge = 1, le = 500),
offset: int = Query(default = 0, ge = 0),
):
mgr = get_job_manager()
result = mgr.get_dataset(job_id, limit=limit, offset=offset)
result = mgr.get_dataset(job_id, limit = limit, offset = offset)
if result is None:
raise HTTPException(status_code=404, detail="dataset not ready")
raise HTTPException(status_code = 404, detail = "dataset not ready")
if "error" in result:
raise HTTPException(status_code=422, detail=result["error"])
raise HTTPException(status_code = 422, detail = result["error"])
return {
"dataset": result["dataset"],
"total": result["total"],
@ -136,9 +143,9 @@ async def job_events(request: Request, job_id: str):
except (TypeError, ValueError):
pass
sub = mgr.subscribe(job_id, after_seq=after_seq)
sub = mgr.subscribe(job_id, after_seq = after_seq)
if sub is None:
raise HTTPException(status_code=404, detail="job not found")
raise HTTPException(status_code = 404, detail = "job not found")
async def gen():
try:
@ -148,11 +155,11 @@ async def job_events(request: Request, job_id: str):
while True:
if await request.is_disconnected():
break
event = await sub.next_event(timeout_sec=1.0)
event = await sub.next_event(timeout_sec = 1.0)
if event is None:
continue
yield sub.format_sse(event)
finally:
mgr.unsubscribe(sub)
return StreamingResponse(gen(), media_type="text/event-stream")
return StreamingResponse(gen(), media_type = "text/event-stream")

30
studio/backend/routes/data_recipe/mcp.py
View file

@ -19,16 +19,16 @@ from models.data_recipe import (
router = APIRouter()
@router.post("/mcp/tools", response_model=McpToolsListResponse)
@router.post("/mcp/tools", response_model = McpToolsListResponse)
def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse:
try:
from data_designer.engine.mcp import io as mcp_io
except ImportError as exc:
return McpToolsListResponse(
providers=[
providers = [
McpToolsProviderResult(
name="",
error=f"MCP dependencies unavailable: {exc}",
name = "",
error = f"MCP dependencies unavailable: {exc}",
)
]
)
@ -42,29 +42,31 @@ def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse:
if len(built) != 1:
providers.append(
McpToolsProviderResult(
name=provider_name,
error="Unsupported MCP provider config.",
name = provider_name,
error = "Unsupported MCP provider config.",
)
)
continue
provider = built[0]
try:
tools = mcp_io.list_tools(provider, timeout_sec=payload.timeout_sec)
tool_names = sorted({tool.name for tool in tools if getattr(tool, "name", "")})
tools = mcp_io.list_tools(provider, timeout_sec = payload.timeout_sec)
tool_names = sorted(
{tool.name for tool in tools if getattr(tool, "name", "")}
)
for tool_name in tool_names:
tool_to_providers[tool_name].append(provider.name)
providers.append(
McpToolsProviderResult(
name=provider.name,
tools=tool_names,
name = provider.name,
tools = tool_names,
)
)
except Exception as exc:
providers.append(
McpToolsProviderResult(
name=provider.name or provider_name,
error=str(exc).strip() or "Failed to load tools.",
name = provider.name or provider_name,
error = str(exc).strip() or "Failed to load tools.",
)
)
@ -75,6 +77,6 @@ def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse:
}
return McpToolsListResponse(
providers=providers,
duplicate_tools=duplicate_tools,
providers = providers,
duplicate_tools = duplicate_tools,
)

153
studio/backend/routes/data_recipe/seed.py
View file

@ -61,7 +61,7 @@ def _list_hf_data_files(*, dataset_name: str, token: str | None) -> list[str]:
return []
try:
api = HfApi()
repo_files = api.list_repo_files(dataset_name, repo_type="dataset", token=token)
repo_files = api.list_repo_files(dataset_name, repo_type = "dataset", token = token)
return [file for file in repo_files if file.lower().endswith(DATA_EXTS)]
except (HfHubHTTPError, OSError, ValueError):
return []
@ -86,10 +86,12 @@ def _select_best_file(data_files: list[str], split: str = DEFAULT_SPLIT) -> str
return (1, len(path))
return (2, len(path))
return sorted(data_files, key=score)[0]
return sorted(data_files, key = score)[0]
def _resolve_seed_hf_path(dataset_name: str, data_files: list[str], split: str = DEFAULT_SPLIT) -> str | None:
def _resolve_seed_hf_path(
dataset_name: str, data_files: list[str], split: str = DEFAULT_SPLIT
) -> str | None:
selected = _select_best_file(data_files, split)
if not selected:
return None
@ -156,36 +158,42 @@ def _decode_base64_payload(content_base64: str) -> bytes:
if "," in raw and raw.lower().startswith("data:"):
raw = raw.split(",", 1)[1]
try:
return base64.b64decode(raw, validate=True)
return base64.b64decode(raw, validate = True)
except binascii.Error as exc:
raise HTTPException(status_code=400, detail="invalid base64 payload") from exc
raise HTTPException(status_code = 400, detail = "invalid base64 payload") from exc
def _read_preview_rows_from_local_file(path: Path, preview_size: int) -> list[dict[str, Any]]:
def _read_preview_rows_from_local_file(
path: Path, preview_size: int
) -> list[dict[str, Any]]:
try:
import pandas as pd
except ImportError as exc:
raise HTTPException(status_code=500, detail=f"seed inspect dependencies unavailable: {exc}") from exc
raise HTTPException(
status_code = 500, detail = f"seed inspect dependencies unavailable: {exc}"
) from exc
ext = path.suffix.lower()
try:
if ext == ".csv":
df = pd.read_csv(path, nrows=preview_size)
df = pd.read_csv(path, nrows = preview_size)
elif ext == ".jsonl":
df = pd.read_json(path, lines=True).head(preview_size)
df = pd.read_json(path, lines = True).head(preview_size)
elif ext == ".json":
try:
df = pd.read_json(path).head(preview_size)
except ValueError:
df = pd.read_json(path, lines=True).head(preview_size)
df = pd.read_json(path, lines = True).head(preview_size)
else:
raise HTTPException(status_code=422, detail=f"unsupported file type: {ext}")
raise HTTPException(status_code = 422, detail = f"unsupported file type: {ext}")
except HTTPException:
raise
except (ValueError, OSError) as exc:
raise HTTPException(status_code=422, detail=f"seed inspect failed: {exc}") from exc
raise HTTPException(
status_code = 422, detail = f"seed inspect failed: {exc}"
) from exc
rows = df.to_dict(orient="records")
rows = df.to_dict(orient = "records")
return _serialize_preview_rows(rows)
@ -199,26 +207,33 @@ def _read_preview_rows_from_unstructured_file(
size, overlap = resolve_chunking(chunk_size, chunk_overlap)
try:
rows = build_unstructured_preview_rows(
source_path=path,
preview_size=preview_size,
chunk_size=size,
chunk_overlap=overlap,
source_path = path,
preview_size = preview_size,
chunk_size = size,
chunk_overlap = overlap,
)
except (FileNotFoundError, RuntimeError, ValueError, OSError) as exc:
raise HTTPException(status_code=422, detail=f"seed inspect failed: {exc}") from exc
raise HTTPException(
status_code = 422, detail = f"seed inspect failed: {exc}"
) from exc
return _serialize_preview_rows(rows)
@router.post("/seed/inspect", response_model=SeedInspectResponse)
@router.post("/seed/inspect", response_model = SeedInspectResponse)
def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
dataset_name = payload.dataset_name.strip()
if not dataset_name or dataset_name.count("/") < 1:
raise HTTPException(status_code=400, detail="dataset_name must be a Hugging Face repo id like org/repo")
raise HTTPException(
status_code = 400,
detail = "dataset_name must be a Hugging Face repo id like org/repo",
)
try:
from datasets import load_dataset
except ImportError as exc:
raise HTTPException(status_code=500, detail=f"seed inspect dependencies unavailable: {exc}") from exc
raise HTTPException(
status_code = 500, detail = f"seed inspect dependencies unavailable: {exc}"
) from exc
split = _normalize_optional_text(payload.split) or DEFAULT_SPLIT
subset = _normalize_optional_text(payload.subset)
@ -226,22 +241,22 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
preview_size = int(payload.preview_size)
preview_rows: list[dict[str, Any]] = []
data_files = _list_hf_data_files(dataset_name=dataset_name, token=token)
data_files = _list_hf_data_files(dataset_name = dataset_name, token = token)
selected_file = _select_best_file(data_files, split)
if selected_file:
try:
single_file_kwargs = _build_stream_load_kwargs(
dataset_name=dataset_name,
split=split,
subset=subset,
token=token,
data_file=selected_file,
dataset_name = dataset_name,
split = split,
subset = subset,
token = token,
data_file = selected_file,
)
preview_rows = _load_preview_rows(
load_dataset_fn=load_dataset,
load_kwargs=single_file_kwargs,
preview_size=preview_size,
load_dataset_fn = load_dataset,
load_kwargs = single_file_kwargs,
preview_size = preview_size,
)
except (ValueError, OSError, RuntimeError):
preview_rows = []
@ -249,21 +264,25 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
if not preview_rows:
try:
split_kwargs = _build_stream_load_kwargs(
dataset_name=dataset_name,
split=split,
subset=subset,
token=token,
dataset_name = dataset_name,
split = split,
subset = subset,
token = token,
)
preview_rows = _load_preview_rows(
load_dataset_fn=load_dataset,
load_kwargs=split_kwargs,
preview_size=preview_size,
load_dataset_fn = load_dataset,
load_kwargs = split_kwargs,
preview_size = preview_size,
)
except (ValueError, OSError, RuntimeError) as exc:
raise HTTPException(status_code=422, detail=f"seed inspect failed: {exc}") from exc
raise HTTPException(
status_code = 422, detail = f"seed inspect failed: {exc}"
) from exc
if not preview_rows:
raise HTTPException(status_code=422, detail="dataset appears empty or unreadable")
raise HTTPException(
status_code = 422, detail = "dataset appears empty or unreadable"
)
preview_rows = _serialize_preview_rows(preview_rows)
columns = _extract_columns(preview_rows)
@ -272,19 +291,21 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
else:
resolved_path = _resolve_seed_hf_path(dataset_name, data_files, split)
if not resolved_path:
raise HTTPException(status_code=422, detail="unable to resolve seed dataset path")
raise HTTPException(
status_code = 422, detail = "unable to resolve seed dataset path"
)
return SeedInspectResponse(
dataset_name=dataset_name,
resolved_path=resolved_path,
columns=columns,
preview_rows=preview_rows,
split=split,
subset=subset,
dataset_name = dataset_name,
resolved_path = resolved_path,
columns = columns,
preview_rows = preview_rows,
split = split,
subset = subset,
)
@router.post("/seed/inspect-upload", response_model=SeedInspectResponse)
@router.post("/seed/inspect-upload", response_model = SeedInspectResponse)
def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectResponse:
seed_source_type = _normalize_optional_text(payload.seed_source_type) or "local"
filename = _sanitize_filename(payload.filename)
@ -292,18 +313,24 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons
if seed_source_type == "unstructured":
if ext not in UNSTRUCTURED_UPLOAD_EXTS:
allowed = ", ".join(sorted(UNSTRUCTURED_UPLOAD_EXTS))
raise HTTPException(status_code=400, detail=f"unsupported file type: {ext}. allowed: {allowed}")
raise HTTPException(
status_code = 400,
detail = f"unsupported file type: {ext}. allowed: {allowed}",
)
else:
if ext not in LOCAL_UPLOAD_EXTS:
allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS))
raise HTTPException(status_code=400, detail=f"unsupported file type: {ext}. allowed: {allowed}")
raise HTTPException(
status_code = 400,
detail = f"unsupported file type: {ext}. allowed: {allowed}",
)
file_bytes = _decode_base64_payload(payload.content_base64)
if not file_bytes:
raise HTTPException(status_code=400, detail="empty upload payload")
raise HTTPException(status_code = 400, detail = "empty upload payload")
max_size_bytes = 50 * 1024 * 1024
if len(file_bytes) > max_size_bytes:
raise HTTPException(status_code=413, detail="file too large (max 50MB)")
raise HTTPException(status_code = 413, detail = "file too large (max 50MB)")
ensure_dir(SEED_UPLOAD_DIR)
stored_name = f"{uuid4().hex}_{filename}"
@ -312,10 +339,10 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons
if seed_source_type == "unstructured":
preview_rows = _read_preview_rows_from_unstructured_file(
path=stored_path,
preview_size=int(payload.preview_size),
chunk_size=payload.unstructured_chunk_size,
chunk_overlap=payload.unstructured_chunk_overlap,
path = stored_path,
preview_size = int(payload.preview_size),
chunk_size = payload.unstructured_chunk_size,
chunk_overlap = payload.unstructured_chunk_overlap,
)
else:
preview_rows = _read_preview_rows_from_local_file(
@ -323,14 +350,16 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons
int(payload.preview_size),
)
if not preview_rows:
raise HTTPException(status_code=422, detail="dataset appears empty or unreadable")
raise HTTPException(
status_code = 422, detail = "dataset appears empty or unreadable"
)
columns = _extract_columns(preview_rows)
return SeedInspectResponse(
dataset_name=filename,
resolved_path=str(stored_path),
columns=columns,
preview_rows=preview_rows,
split=None,
subset=None,
dataset_name = filename,
resolved_path = str(stored_path),
columns = columns,
preview_rows = preview_rows,
split = None,
subset = None,
)

28
studio/backend/routes/data_recipe/validate.py
View file

@ -44,9 +44,9 @@ def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]:
_resolve_and_add_seed_columns(config, resource_provider.seed_reader)
_add_internal_row_id_column_if_needed(config)
violations = validate_data_designer_config(
columns=config.columns,
processor_configs=config.processors or [],
allowed_references=_get_allowed_references(config),
columns = config.columns,
processor_configs = config.processors or [],
allowed_references = _get_allowed_references(config),
)
except (TypeError, ValueError, AttributeError):
return []
@ -60,34 +60,34 @@ def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]:
message = str(violation.message).strip() or "Validation failed."
errors.append(
ValidateError(
message=message,
path=path,
code=code,
message = message,
path = path,
code = code,
)
)
return errors
@router.post("/validate", response_model=ValidateResponse)
@router.post("/validate", response_model = ValidateResponse)
def validate(payload: RecipePayload) -> ValidateResponse:
recipe = payload.recipe
if not recipe.get("columns"):
return ValidateResponse(
valid=False,
errors=[ValidateError(message="Recipe must include columns.")],
valid = False,
errors = [ValidateError(message = "Recipe must include columns.")],
)
try:
validate_recipe(recipe)
except RuntimeError as exc:
raise HTTPException(status_code=503, detail=str(exc)) from exc
raise HTTPException(status_code = 503, detail = str(exc)) from exc
except Exception as exc:
detail = str(exc).strip() or "Validation failed."
parsed_errors = _collect_validation_errors(recipe)
return ValidateResponse(
valid=False,
errors=parsed_errors or [ValidateError(message=detail)],
raw_detail=detail,
valid = False,
errors = parsed_errors or [ValidateError(message = detail)],
raw_detail = detail,
)
return ValidateResponse(valid=True)
return ValidateResponse(valid = True)

216
studio/backend/routes/datasets.py
View file

@ -4,6 +4,7 @@
"""
Datasets API routes
"""
import base64
import io
import json
@ -27,8 +28,6 @@ router = APIRouter()
logger = get_logger(__name__)
from models.datasets import (
AiAssistMappingRequest,
AiAssistMappingResponse,
@ -53,9 +52,10 @@ def _serialize_preview_value(value):
try:
from PIL.Image import Image as PILImage
if isinstance(value, PILImage):
buffer = io.BytesIO()
value.convert("RGB").save(buffer, format="JPEG", quality=85)
value.convert("RGB").save(buffer, format = "JPEG", quality = 85)
return {
"type": "image",
"mime": "image/jpeg",
@ -88,10 +88,10 @@ def _serialize_preview_rows(rows):
# Tabular formats are preferred over archives for Tier 1 preview because
# archives (e.g. images.zip) may be loaded as ImageFolder datasets with
# synthetic columns (image/label) that don't match the real dataset schema.
_TABULAR_EXTS = ('.parquet', '.json', '.jsonl', '.csv', '.tsv', '.arrow')
_ARCHIVE_EXTS = ('.tar', '.tar.gz', '.tgz', '.gz', '.zst', '.zip', '.txt')
_TABULAR_EXTS = (".parquet", ".json", ".jsonl", ".csv", ".tsv", ".arrow")
_ARCHIVE_EXTS = (".tar", ".tar.gz", ".tgz", ".gz", ".zst", ".zip", ".txt")
DATA_EXTS = _TABULAR_EXTS + _ARCHIVE_EXTS
LOCAL_FILE_EXTS = ('.json', '.jsonl', '.csv', '.parquet')
LOCAL_FILE_EXTS = (".json", ".jsonl", ".csv", ".parquet")
LOCAL_UPLOAD_EXTS = {".csv", ".json", ".jsonl", ".parquet"}
LOCAL_DATASETS_ROOT = recipe_datasets_root()
DATASET_UPLOAD_DIR = dataset_uploads_root()
@ -99,7 +99,7 @@ DATASET_UPLOAD_DIR = dataset_uploads_root()
def _safe_read_metadata(path: Path) -> dict | None:
try:
payload = json.loads(path.read_text(encoding="utf-8"))
payload = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, ValueError, TypeError):
return None
if not isinstance(payload, dict):
@ -200,30 +200,36 @@ def _build_local_dataset_items() -> list[LocalDatasetItem]:
items.append(
LocalDatasetItem(
id=entry.name,
label=entry.name,
path=str(parquet_dir.resolve()),
rows=rows,
updated_at=updated_at,
metadata=metadata_summary,
id = entry.name,
label = entry.name,
path = str(parquet_dir.resolve()),
rows = rows,
updated_at = updated_at,
metadata = metadata_summary,
)
)
items.sort(key=lambda item: item.updated_at or 0, reverse=True)
items.sort(key = lambda item: item.updated_at or 0, reverse = True)
return items
def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int):
def _load_local_preview_slice(
*, dataset_path: Path, train_split: str, preview_size: int
):
from datasets import load_dataset
if dataset_path.is_dir():
parquet_dir = dataset_path / "parquet-files" if (dataset_path / "parquet-files").exists() else dataset_path
parquet_dir = (
dataset_path / "parquet-files"
if (dataset_path / "parquet-files").exists()
else dataset_path
)
parquet_files = sorted(parquet_dir.glob("*.parquet"))
if parquet_files:
dataset = load_dataset(
"parquet",
data_files=[str(path) for path in parquet_files],
split=train_split,
data_files = [str(path) for path in parquet_files],
split = train_split,
)
total_rows = len(dataset)
preview_slice = dataset.select(range(min(preview_size, total_rows)))
@ -234,21 +240,22 @@ def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_s
candidate_files.extend(sorted(dataset_path.glob(f"*{ext}")))
if not candidate_files:
raise HTTPException(
status_code=400,
detail="Unsupported local dataset directory (expected parquet/json/jsonl/csv files)",
status_code = 400,
detail = "Unsupported local dataset directory (expected parquet/json/jsonl/csv files)",
)
dataset_path = candidate_files[0]
if dataset_path.suffix in ['.json', '.jsonl']:
dataset = load_dataset('json', data_files=str(dataset_path), split=train_split)
elif dataset_path.suffix == '.csv':
dataset = load_dataset('csv', data_files=str(dataset_path), split=train_split)
elif dataset_path.suffix == '.parquet':
dataset = load_dataset('parquet', data_files=str(dataset_path), split=train_split)
if dataset_path.suffix in [".json", ".jsonl"]:
dataset = load_dataset("json", data_files = str(dataset_path), split = train_split)
elif dataset_path.suffix == ".csv":
dataset = load_dataset("csv", data_files = str(dataset_path), split = train_split)
elif dataset_path.suffix == ".parquet":
dataset = load_dataset(
"parquet", data_files = str(dataset_path), split = train_split
)
else:
raise HTTPException(
status_code=400,
detail=f"Unsupported file format: {dataset_path.suffix}"
status_code = 400, detail = f"Unsupported file format: {dataset_path.suffix}"
)
total_rows = len(dataset)
@ -263,7 +270,7 @@ def _sanitize_filename(filename: str) -> str:
return name
@router.post("/upload", response_model=UploadDatasetResponse)
@router.post("/upload", response_model = UploadDatasetResponse)
async def upload_dataset(
file: UploadFile,
current_subject: str = Depends(get_current_subject),
@ -273,8 +280,8 @@ async def upload_dataset(
if ext not in LOCAL_UPLOAD_EXTS:
allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS))
raise HTTPException(
status_code=400,
detail=f"Unsupported file type: {ext}. Allowed: {allowed}",
status_code = 400,
detail = f"Unsupported file type: {ext}. Allowed: {allowed}",
)
max_size_bytes = 512 * 1024 * 1024
@ -289,25 +296,27 @@ async def upload_dataset(
while chunk := await file.read(1024 * 1024):
size += len(chunk)
if size > max_size_bytes:
stored_path.unlink(missing_ok=True)
raise HTTPException(status_code=413, detail="File too large (max 512MB)")
stored_path.unlink(missing_ok = True)
raise HTTPException(
status_code = 413, detail = "File too large (max 512MB)"
)
f.write(chunk)
if size == 0:
stored_path.unlink(missing_ok=True)
raise HTTPException(status_code=400, detail="Empty upload payload")
stored_path.unlink(missing_ok = True)
raise HTTPException(status_code = 400, detail = "Empty upload payload")
return UploadDatasetResponse(filename=filename, stored_path=str(stored_path))
return UploadDatasetResponse(filename = filename, stored_path = str(stored_path))
@router.get("/local", response_model=LocalDatasetsResponse)
@router.get("/local", response_model = LocalDatasetsResponse)
def list_local_datasets(
current_subject: str = Depends(get_current_subject),
) -> LocalDatasetsResponse:
return LocalDatasetsResponse(datasets=_build_local_dataset_items())
return LocalDatasetsResponse(datasets = _build_local_dataset_items())
@router.post("/check-format", response_model=CheckFormatResponse)
@router.post("/check-format", response_model = CheckFormatResponse)
def check_format(
request: CheckFormatRequest,
current_subject: str = Depends(get_current_subject),
@ -341,9 +350,9 @@ def check_format(
# ── Local file ──────────────────────────────────────────
train_split = request.train_split or "train"
preview_slice, total_rows = _load_local_preview_slice(
dataset_path=dataset_path,
train_split=train_split,
preview_size=PREVIEW_SIZE,
dataset_path = dataset_path,
train_split = train_split,
preview_size = PREVIEW_SIZE,
)
else:
# ── HuggingFace dataset ─────────────────────────────────
@ -352,23 +361,32 @@ def check_format(
try:
from huggingface_hub import HfApi
api = HfApi()
repo_files = api.list_repo_files(
request.dataset_name,
repo_type="dataset",
token=request.hf_token or None,
repo_type = "dataset",
token = request.hf_token or None,
)
data_files = [f for f in repo_files if any(f.endswith(ext) for ext in DATA_EXTS)]
data_files = [
f for f in repo_files if any(f.endswith(ext) for ext in DATA_EXTS)
]
# Prefer tabular formats over archives (e.g. images.zip → ImageFolder
# with synthetic image/label columns that don't match the real schema).
tabular_files = [f for f in data_files if any(f.endswith(ext) for ext in _TABULAR_EXTS)]
tabular_files = [
f
for f in data_files
if any(f.endswith(ext) for ext in _TABULAR_EXTS)
]
candidates = tabular_files or data_files
# When a subset is specified, narrow to files whose name matches
# (e.g. subset="testmini" → prefer "testmini.parquet").
if request.subset and candidates:
subset_matches = [f for f in candidates if request.subset in Path(f).stem]
subset_matches = [
f for f in candidates if request.subset in Path(f).stem
]
if subset_matches:
candidates = subset_matches
@ -394,7 +412,11 @@ def check_format(
if preview_slice is None:
# Tier 2: full streaming (resolves all files — slow for large repos)
logger.info("Tier 2: falling back to full streaming load_dataset")
load_kwargs = {"path": request.dataset_name, "split": request.train_split, "streaming": True}
load_kwargs = {
"path": request.dataset_name,
"split": request.train_split,
"streaming": True,
}
if request.subset:
load_kwargs["name"] = request.subset
if request.hf_token:
@ -405,17 +427,19 @@ def check_format(
rows = list(islice(streamed_ds, PREVIEW_SIZE))
if not rows:
raise HTTPException(
status_code=400,
detail="Dataset appears to be empty or could not be streamed"
status_code = 400,
detail = "Dataset appears to be empty or could not be streamed",
)
preview_slice = Dataset.from_list(rows)
total_rows = None
# Run lightweight format check on the preview slice
result = check_dataset_format(preview_slice, is_vlm=request.is_vlm)
result = check_dataset_format(preview_slice, is_vlm = request.is_vlm)
logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_image={result.get('is_image', False)}")
logger.info(
f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_image={result.get('is_image', False)}"
)
# Generate preview samples
preview_samples = None
@ -428,13 +452,15 @@ def check_format(
try:
format_result = format_dataset(
preview_slice,
format_type="auto",
num_proc=1, # Only 10 preview rows — no need for multiprocessing
format_type = "auto",
num_proc = 1, # Only 10 preview rows — no need for multiprocessing
)
processed = format_result["dataset"]
preview_samples = _serialize_preview_rows(processed)
except Exception as e:
logger.warning(f"Processed preview generation failed (non-fatal): {e}")
logger.warning(
f"Processed preview generation failed (non-fatal): {e}"
)
preview_samples = _serialize_preview_rows(preview_slice)
else:
preview_samples = _serialize_preview_rows(preview_slice)
@ -445,7 +471,9 @@ def check_format(
if image_col and image_col in (result.get("columns") or []):
try:
sample_val = preview_slice[0][image_col]
if isinstance(sample_val, str) and sample_val.startswith(("http://", "https://")):
if isinstance(sample_val, str) and sample_val.startswith(
("http://", "https://")
):
url_warning = (
"This dataset contains image URLs instead of embedded images. "
"Images will be downloaded during training, which may be slow for large datasets."
@ -456,33 +484,32 @@ def check_format(
pass
return CheckFormatResponse(
requires_manual_mapping=result["requires_manual_mapping"],
detected_format=result["detected_format"],
columns=result["columns"],
is_image=result.get("is_image", False),
is_audio=result.get("is_audio", False),
multimodal_columns=result.get("multimodal_columns"),
suggested_mapping=result.get("suggested_mapping"),
detected_image_column=result.get("detected_image_column"),
detected_audio_column=result.get("detected_audio_column"),
detected_text_column=result.get("detected_text_column"),
detected_speaker_column=result.get("detected_speaker_column"),
preview_samples=preview_samples,
total_rows=total_rows,
warning=warning,
requires_manual_mapping = result["requires_manual_mapping"],
detected_format = result["detected_format"],
columns = result["columns"],
is_image = result.get("is_image", False),
is_audio = result.get("is_audio", False),
multimodal_columns = result.get("multimodal_columns"),
suggested_mapping = result.get("suggested_mapping"),
detected_image_column = result.get("detected_image_column"),
detected_audio_column = result.get("detected_audio_column"),
detected_text_column = result.get("detected_text_column"),
detected_speaker_column = result.get("detected_speaker_column"),
preview_samples = preview_samples,
total_rows = total_rows,
warning = warning,
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error checking dataset format: {e}", exc_info=True)
logger.error(f"Error checking dataset format: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to check dataset format: {str(e)}"
status_code = 500, detail = f"Failed to check dataset format: {str(e)}"
)
@router.post("/ai-assist-mapping", response_model=AiAssistMappingResponse)
@router.post("/ai-assist-mapping", response_model = AiAssistMappingResponse)
def ai_assist_mapping(
request: AiAssistMappingRequest,
current_subject: str = Depends(get_current_subject),
@ -507,35 +534,32 @@ def ai_assist_mapping(
]
result = llm_conversion_advisor(
column_names=request.columns,
samples=truncated,
dataset_name=request.dataset_name,
hf_token=request.hf_token,
model_name=request.model_name,
model_type=request.model_type,
column_names = request.columns,
samples = truncated,
dataset_name = request.dataset_name,
hf_token = request.hf_token,
model_name = request.model_name,
model_type = request.model_type,
)
if result and result.get("success"):
return AiAssistMappingResponse(
success=True,
suggested_mapping=result.get("suggested_mapping"),
system_prompt=result.get("system_prompt"),
user_template=result.get("user_template"),
assistant_template=result.get("assistant_template"),
label_mapping=result.get("label_mapping"),
dataset_type=result.get("dataset_type"),
is_conversational=result.get("is_conversational"),
user_notification=result.get("user_notification"),
success = True,
suggested_mapping = result.get("suggested_mapping"),
system_prompt = result.get("system_prompt"),
user_template = result.get("user_template"),
assistant_template = result.get("assistant_template"),
label_mapping = result.get("label_mapping"),
dataset_type = result.get("dataset_type"),
is_conversational = result.get("is_conversational"),
user_notification = result.get("user_notification"),
)
return AiAssistMappingResponse(
success=False,
warning="AI could not determine column roles. Please assign them manually.",
success = False,
warning = "AI could not determine column roles. Please assign them manually.",
)
except Exception as e:
logger.error(f"AI assist mapping failed: {e}", exc_info=True)
raise HTTPException(
status_code=500,
detail=f"AI assist failed: {str(e)}"
)
logger.error(f"AI assist mapping failed: {e}", exc_info = True)
raise HTTPException(status_code = 500, detail = f"AI assist failed: {str(e)}")

158
studio/backend/routes/export.py
View file

@ -43,11 +43,7 @@ router = APIRouter()
logger = get_logger(__name__)
@router.post("/load-checkpoint", response_model=ExportOperationResponse)
@router.post("/load-checkpoint", response_model = ExportOperationResponse)
async def load_checkpoint(
request: LoadCheckpointRequest,
current_subject: str = Depends(get_current_subject),
@ -65,6 +61,7 @@ async def load_checkpoint(
# before loading the export checkpoint (they'd compete for VRAM).
try:
from core.inference import get_inference_backend
inf = get_inference_backend()
if inf.active_model_name:
logger.info(
@ -79,6 +76,7 @@ async def load_checkpoint(
try:
from core.training import get_training_backend
trn = get_training_backend()
if trn.is_training_active():
logger.info("Stopping active training to free GPU memory for export")
@ -88,35 +86,39 @@ async def load_checkpoint(
for _ in range(60): # up to 30s
if not trn.is_training_active():
break
import time; time.sleep(0.5)
import time
time.sleep(0.5)
else:
logger.warning("Training subprocess did not exit within 30s, proceeding anyway")
logger.warning(
"Training subprocess did not exit within 30s, proceeding anyway"
)
except Exception as e:
logger.warning("Could not stop training: %s", e)
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,
trust_remote_code=request.trust_remote_code,
checkpoint_path = request.checkpoint_path,
max_seq_length = request.max_seq_length,
load_in_4bit = request.load_in_4bit,
trust_remote_code = request.trust_remote_code,
)
if not success:
raise HTTPException(status_code=400, detail=message)
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(success=True, message=message)
return ExportOperationResponse(success = True, message = message)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error loading checkpoint: {e}", exc_info=True)
logger.error(f"Error loading checkpoint: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to load checkpoint: {str(e)}",
status_code = 500,
detail = f"Failed to load checkpoint: {str(e)}",
)
@router.post("/cleanup", response_model=ExportOperationResponse)
@router.post("/cleanup", response_model = ExportOperationResponse)
async def cleanup_export_memory(
current_subject: str = Depends(get_current_subject),
):
@ -131,25 +133,25 @@ async def cleanup_export_memory(
if not success:
raise HTTPException(
status_code=500,
detail="Memory cleanup failed. See server logs for details.",
status_code = 500,
detail = "Memory cleanup failed. See server logs for details.",
)
return ExportOperationResponse(
success=True,
message="Memory cleanup completed successfully",
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)
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)}",
status_code = 500,
detail = f"Failed to cleanup export memory: {str(e)}",
)
@router.get("/status", response_model=ExportStatusResponse)
@router.get("/status", response_model = ExportStatusResponse)
async def get_export_status(
current_subject: str = Depends(get_current_subject),
):
@ -159,19 +161,19 @@ async def get_export_status(
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)),
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)
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)}",
status_code = 500,
detail = f"Failed to get export status: {str(e)}",
)
@router.post("/export/merged", response_model=ExportOperationResponse)
@router.post("/export/merged", response_model = ExportOperationResponse)
async def export_merged_model(
request: ExportMergedModelRequest,
current_subject: str = Depends(get_current_subject),
@ -184,29 +186,29 @@ async def 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,
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)
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(success=True, message=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)
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)}",
status_code = 500,
detail = f"Failed to export merged model: {str(e)}",
)
@router.post("/export/base", response_model=ExportOperationResponse)
@router.post("/export/base", response_model = ExportOperationResponse)
async def export_base_model(
request: ExportBaseModelRequest,
current_subject: str = Depends(get_current_subject),
@ -219,29 +221,29 @@ async def 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,
base_model_id=request.base_model_id,
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,
base_model_id = request.base_model_id,
)
if not success:
raise HTTPException(status_code=400, detail=message)
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(success=True, message=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)
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)}",
status_code = 500,
detail = f"Failed to export base model: {str(e)}",
)
@router.post("/export/gguf", response_model=ExportOperationResponse)
@router.post("/export/gguf", response_model = ExportOperationResponse)
async def export_gguf(
request: ExportGGUFRequest,
current_subject: str = Depends(get_current_subject),
@ -254,28 +256,28 @@ async def 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,
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)
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(success=True, message=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)
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)}",
status_code = 500,
detail = f"Failed to export GGUF model: {str(e)}",
)
@router.post("/export/lora", response_model=ExportOperationResponse)
@router.post("/export/lora", response_model = ExportOperationResponse)
async def export_lora_adapter(
request: ExportLoRAAdapterRequest,
current_subject: str = Depends(get_current_subject),
@ -288,24 +290,22 @@ async def 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,
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)
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(success=True, message=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)
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)}",
status_code = 500,
detail = f"Failed to export LoRA adapter: {str(e)}",
)

713
studio/backend/routes/inference.py
View file

File diff suppressed because it is too large Load diff

395
studio/backend/routes/models.py
View file

@ -4,6 +4,8 @@
"""
Model Management API routes
"""
import os
import sys
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Query
@ -31,9 +33,18 @@ try:
list_gguf_variants,
ModelConfig,
)
from utils.models.model_config import _pick_best_gguf, _extract_quant_label, is_audio_input_type
from utils.models.model_config import (
_pick_best_gguf,
_extract_quant_label,
is_audio_input_type,
)
from core.inference import get_inference_backend
from utils.paths import outputs_root, exports_root, resolve_output_dir, resolve_export_dir
from utils.paths import (
outputs_root,
exports_root,
resolve_output_dir,
resolve_export_dir,
)
except ImportError:
# Fallback: try to import from parent directory
parent_backend = backend_path.parent / "backend"
@ -50,9 +61,18 @@ except ImportError:
list_gguf_variants,
ModelConfig,
)
from utils.models.model_config import _pick_best_gguf, _extract_quant_label, is_audio_input_type
from utils.models.model_config import (
_pick_best_gguf,
_extract_quant_label,
is_audio_input_type,
)
from core.inference import get_inference_backend
from utils.paths import outputs_root, exports_root, resolve_output_dir, resolve_export_dir
from utils.paths import (
outputs_root,
exports_root,
resolve_output_dir,
resolve_export_dir,
)
from models import (
CheckpointInfo,
@ -66,13 +86,19 @@ from models import (
ModelListResponse,
)
from models.models import GgufVariantDetail, GgufVariantsResponse, ModelType
from models.responses import LoRABaseModelResponse, VisionCheckResponse, EmbeddingCheckResponse
from models.responses import (
LoRABaseModelResponse,
VisionCheckResponse,
EmbeddingCheckResponse,
)
router = APIRouter()
logger = get_logger(__name__)
def derive_model_type(is_vision: bool, audio_type: Optional[str], is_embedding: bool = False) -> ModelType:
def derive_model_type(
is_vision: bool, audio_type: Optional[str], is_embedding: bool = False
) -> ModelType:
"""Collapse individual capability flags into a single model modality string."""
if is_embedding:
return "embeddings"
@ -83,12 +109,11 @@ def derive_model_type(is_vision: bool, audio_type: Optional[str], is_embedding:
return "text"
def _resolve_hf_cache_dir() -> Path:
"""Resolve local HF cache root used by hub downloads."""
try:
from huggingface_hub.constants import HF_HUB_CACHE
return Path(HF_HUB_CACHE)
except Exception:
return Path.home() / ".cache" / "huggingface" / "hub"
@ -117,11 +142,11 @@ def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]:
updated_at = None
found.append(
LocalModelInfo(
id=str(child),
display_name=child.name,
path=str(child),
source="models_dir",
updated_at=updated_at,
id = str(child),
display_name = child.name,
path = str(child),
source = "models_dir",
updated_at = updated_at,
),
)
# Also scan for standalone .gguf files directly in the models directory
@ -133,11 +158,11 @@ def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]:
updated_at = None
found.append(
LocalModelInfo(
id=str(gguf_file),
display_name=gguf_file.stem,
path=str(gguf_file),
source="models_dir",
updated_at=updated_at,
id = str(gguf_file),
display_name = gguf_file.stem,
path = str(gguf_file),
source = "models_dir",
updated_at = updated_at,
),
)
@ -153,7 +178,7 @@ def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]:
if not repo_dir.is_dir():
continue
repo_name = repo_dir.name[len("models--"):]
repo_name = repo_dir.name[len("models--") :]
if not repo_name:
continue
model_id = repo_name.replace("--", "/")
@ -165,28 +190,53 @@ def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]:
found.append(
LocalModelInfo(
id=model_id,
model_id=model_id,
display_name=model_id.split("/")[-1],
path=str(repo_dir),
source="hf_cache",
updated_at=updated_at,
id = model_id,
model_id = model_id,
display_name = model_id.split("/")[-1],
path = str(repo_dir),
source = "hf_cache",
updated_at = updated_at,
),
)
return found
@router.get("/local", response_model=LocalModelListResponse)
@router.get("/local", response_model = LocalModelListResponse)
async def list_local_models(
models_dir: str = Query(default="./models", description="Directory to scan for local model folders"),
models_dir: str = Query(
default = "./models", description = "Directory to scan for local model folders"
),
current_subject: str = Depends(get_current_subject),
):
"""
List local model candidates from custom models dir and HF cache.
"""
# Validate models_dir against an allowlist of trusted directories.
# Only the trusted Path objects are used for filesystem access -- the
# user-supplied string is only used for matching, never for path construction.
hf_cache_dir = _resolve_hf_cache_dir()
allowed_roots = [Path("./models").resolve(), hf_cache_dir]
try:
from utils.paths import studio_root, outputs_root
allowed_roots.extend([studio_root(), outputs_root()])
except Exception:
pass
requested = os.path.realpath(os.path.expanduser(models_dir))
models_root = None
for root in allowed_roots:
root_str = os.path.realpath(str(root))
if requested == root_str or requested.startswith(root_str + os.sep):
models_root = root # Use the trusted root, not the user-supplied path
break
if models_root is None:
raise HTTPException(
status_code = 403,
detail = "Directory not allowed",
)
try:
models_root = Path(models_dir).expanduser().resolve()
hf_cache_dir = _resolve_hf_cache_dir()
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
deduped: dict[str, LocalModelInfo] = {}
@ -196,88 +246,80 @@ async def list_local_models(
models = sorted(
deduped.values(),
key=lambda item: (item.updated_at or 0),
reverse=True,
key = lambda item: (item.updated_at or 0),
reverse = True,
)
return LocalModelListResponse(
models_dir=str(models_root),
hf_cache_dir=str(hf_cache_dir),
models=models,
models_dir = str(models_root),
hf_cache_dir = str(hf_cache_dir),
models = models,
)
except Exception as e:
logger.error(f"Error listing local models: {e}", exc_info=True)
logger.error(f"Error listing local models: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to list local models: {str(e)}",
status_code = 500,
detail = f"Failed to list local models: {str(e)}",
)
@router.get("/list")
async def list_models(
current_subject: str = Depends(get_current_subject),
):
"""
List available models (default models and loaded models).
This endpoint returns the default models and any currently loaded models.
"""
try:
inference_backend = get_inference_backend()
# Get default models
default_models = inference_backend.default_models
# Get loaded models
loaded_models = []
for model_name, model_data in inference_backend.models.items():
_is_vision = model_data.get("is_vision", False)
_audio_type = model_data.get("audio_type")
model_info = ModelDetails(
id=model_name,
name=model_name.split("/")[-1] if "/" in model_name else model_name,
is_vision=_is_vision,
is_lora=model_data.get("is_lora", False),
is_audio=model_data.get("is_audio", False),
audio_type=_audio_type,
has_audio_input=model_data.get("has_audio_input", False),
model_type=derive_model_type(_is_vision, _audio_type),
id = model_name,
name = model_name.split("/")[-1] if "/" in model_name else model_name,
is_vision = _is_vision,
is_lora = model_data.get("is_lora", False),
is_audio = model_data.get("is_audio", False),
audio_type = _audio_type,
has_audio_input = model_data.get("has_audio_input", False),
model_type = derive_model_type(_is_vision, _audio_type),
)
loaded_models.append(model_info)
# Combine default and loaded models
all_models = []
seen_ids = set()
# Add default models
for model_id in default_models:
if model_id not in seen_ids:
model_info = ModelDetails(
id=model_id,
name=model_id.split("/")[-1] if "/" in model_id else model_id
id = model_id,
name = model_id.split("/")[-1] if "/" in model_id else model_id,
)
all_models.append(model_info)
seen_ids.add(model_id)
# Add loaded models
for model_info in loaded_models:
if model_info.id not in seen_ids:
all_models.append(model_info)
seen_ids.add(model_info.id)
return ModelListResponse(
models=all_models,
default_models=default_models
)
return ModelListResponse(models = all_models, default_models = default_models)
except Exception as e:
logger.error(f"Error listing models: {e}", exc_info=True)
raise HTTPException(
status_code=500,
detail=f"Failed to list models: {str(e)}"
)
logger.error(f"Error listing models: {e}", exc_info = True)
raise HTTPException(status_code = 500, detail = f"Failed to list models: {str(e)}")
@router.get("/config/{model_name:path}")
@ -293,18 +335,20 @@ async def get_model_config(
"""
try:
from utils.models.model_config import is_local_path
if not is_local_path(model_name):
model_name = model_name.lower()
logger.info(f"Getting model config for: {model_name}")
from utils.models.model_config import detect_audio_type
# Load model defaults from backend
config_dict = load_model_defaults(model_name)
# Detect model capabilities (pass HF token for gated models)
is_vision = is_vision_model(model_name)
is_embedding = is_embedding_model(model_name, hf_token=hf_token)
audio_type = detect_audio_type(model_name, hf_token=hf_token)
is_embedding = is_embedding_model(model_name, hf_token = hf_token)
audio_type = detect_audio_type(model_name, hf_token = hf_token)
# Check if it's a LoRA adapter
is_lora = False
@ -316,33 +360,38 @@ async def get_model_config(
except Exception:
pass
logger.info(f"Model config result for {model_name}: is_vision={is_vision}, is_embedding={is_embedding}, audio_type={audio_type}, is_lora={is_lora}")
return ModelDetails(
id=model_name,
model_name=model_name,
config=config_dict,
is_vision=is_vision,
is_embedding=is_embedding,
is_lora=is_lora,
is_audio=audio_type is not None,
audio_type=audio_type,
has_audio_input=is_audio_input_type(audio_type),
model_type=derive_model_type(is_vision, audio_type, is_embedding),
base_model=base_model,
logger.info(
f"Model config result for {model_name}: is_vision={is_vision}, is_embedding={is_embedding}, audio_type={audio_type}, is_lora={is_lora}"
)
return ModelDetails(
id = model_name,
model_name = model_name,
config = config_dict,
is_vision = is_vision,
is_embedding = is_embedding,
is_lora = is_lora,
is_audio = audio_type is not None,
audio_type = audio_type,
has_audio_input = is_audio_input_type(audio_type),
model_type = derive_model_type(is_vision, audio_type, is_embedding),
base_model = base_model,
)
except Exception as e:
logger.error(f"Error getting model config: {e}", exc_info=True)
logger.error(f"Error getting model config: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to get model config: {str(e)}"
status_code = 500, detail = f"Failed to get model config: {str(e)}"
)
@router.get("/loras")
async def scan_loras(
outputs_dir: str = Query(default=str(outputs_root()), description="Directory to scan for LoRA adapters"),
exports_dir: str = Query(default=str(exports_root()), description="Directory to scan for exported models"),
outputs_dir: str = Query(
default = str(outputs_root()), description = "Directory to scan for LoRA adapters"
),
exports_dir: str = Query(
default = str(exports_root()), description = "Directory to scan for exported models"
),
current_subject: str = Depends(get_current_subject),
):
"""
@ -357,102 +406,101 @@ async def scan_loras(
lora_list = []
# Scan training outputs
trained_loras = scan_trained_loras(outputs_dir=resolved_outputs_dir)
trained_loras = scan_trained_loras(outputs_dir = resolved_outputs_dir)
for display_name, adapter_path in trained_loras:
base_model = get_base_model_from_lora(adapter_path)
lora_list.append(LoRAInfo(
display_name=display_name,
adapter_path=adapter_path,
base_model=base_model,
source="training",
))
lora_list.append(
LoRAInfo(
display_name = display_name,
adapter_path = adapter_path,
base_model = base_model,
source = "training",
)
)
# Scan exported models (merged, LoRA, base — skips GGUF)
exported = scan_exported_models(exports_dir=resolved_exports_dir)
exported = scan_exported_models(exports_dir = resolved_exports_dir)
for display_name, model_path, export_type, base_model in exported:
lora_list.append(LoRAInfo(
display_name=display_name,
adapter_path=model_path,
base_model=base_model,
source="exported",
export_type=export_type,
))
lora_list.append(
LoRAInfo(
display_name = display_name,
adapter_path = model_path,
base_model = base_model,
source = "exported",
export_type = export_type,
)
)
return LoRAScanResponse(
loras=lora_list,
outputs_dir=resolved_outputs_dir
)
return LoRAScanResponse(loras = lora_list, outputs_dir = resolved_outputs_dir)
except Exception as e:
logger.error(f"Error scanning LoRAs: {e}", exc_info=True)
logger.error(f"Error scanning LoRAs: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to scan LoRA adapters: {str(e)}"
status_code = 500, detail = f"Failed to scan LoRA adapters: {str(e)}"
)
@router.get("/loras/{lora_path:path}/base-model", response_model=LoRABaseModelResponse)
@router.get("/loras/{lora_path:path}/base-model", response_model = LoRABaseModelResponse)
async def get_lora_base_model(
lora_path: str,
current_subject: str = Depends(get_current_subject),
):
"""
Get the base model for a LoRA adapter.
This endpoint wraps the backend get_base_model_from_lora function.
"""
try:
base_model = get_base_model_from_lora(lora_path)
if base_model is None:
raise HTTPException(
status_code=404,
detail=f"Could not determine base model for LoRA: {lora_path}"
status_code = 404,
detail = f"Could not determine base model for LoRA: {lora_path}",
)
return LoRABaseModelResponse(
lora_path=lora_path,
base_model=base_model,
lora_path = lora_path,
base_model = base_model,
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error getting LoRA base model: {e}", exc_info=True)
logger.error(f"Error getting LoRA base model: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to get base model: {str(e)}"
status_code = 500, detail = f"Failed to get base model: {str(e)}"
)
@router.get("/check-vision/{model_name:path}", response_model=VisionCheckResponse)
@router.get("/check-vision/{model_name:path}", response_model = VisionCheckResponse)
async def check_vision_model(
model_name: str,
current_subject: str = Depends(get_current_subject),
):
"""
Check if a model is a vision model.
This endpoint wraps the backend is_vision_model function.
"""
try:
logger.info(f"Checking if vision model: {model_name}")
is_vision = is_vision_model(model_name)
logger.info(f"Vision check result for {model_name}: is_vision={is_vision}")
return VisionCheckResponse(
model_name=model_name,
is_vision=is_vision,
)
except Exception as e:
logger.error(f"Error checking vision model: {e}", exc_info=True)
raise HTTPException(
status_code=500,
detail=f"Failed to check vision model: {str(e)}"
model_name = model_name,
is_vision = is_vision,
)
@router.get("/check-embedding/{model_name:path}", response_model=EmbeddingCheckResponse)
except Exception as e:
logger.error(f"Error checking vision model: {e}", exc_info = True)
raise HTTPException(
status_code = 500, detail = f"Failed to check vision model: {str(e)}"
)
@router.get("/check-embedding/{model_name:path}", response_model = EmbeddingCheckResponse)
async def check_embedding_model(
model_name: str,
hf_token: Optional[str] = Query(None),
@ -465,26 +513,31 @@ async def check_embedding_model(
"""
try:
logger.info(f"Checking if embedding model: {model_name}")
is_embedding = is_embedding_model(model_name, hf_token=hf_token)
is_embedding = is_embedding_model(model_name, hf_token = hf_token)
logger.info(f"Embedding check result for {model_name}: is_embedding={is_embedding}")
logger.info(
f"Embedding check result for {model_name}: is_embedding={is_embedding}"
)
return EmbeddingCheckResponse(
model_name=model_name,
is_embedding=is_embedding,
model_name = model_name,
is_embedding = is_embedding,
)
except Exception as e:
logger.error(f"Error checking embedding model: {e}", exc_info=True)
logger.error(f"Error checking embedding model: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to check embedding model: {str(e)}"
status_code = 500, detail = f"Failed to check embedding model: {str(e)}"
)
@router.get("/gguf-variants", response_model=GgufVariantsResponse)
@router.get("/gguf-variants", response_model = GgufVariantsResponse)
async def get_gguf_variants(
repo_id: str = Query(..., description="HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"),
hf_token: Optional[str] = Query(None, description="HuggingFace token for private repos"),
repo_id: str = Query(
..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"
),
hf_token: Optional[str] = Query(
None, description = "HuggingFace token for private repos"
),
current_subject: str = Depends(get_current_subject),
):
"""
@ -495,7 +548,7 @@ async def get_gguf_variants(
default variant.
"""
try:
variants, has_vision = list_gguf_variants(repo_id, hf_token=hf_token)
variants, has_vision = list_gguf_variants(repo_id, hf_token = hf_token)
# Determine default variant
filenames = [v.filename for v in variants]
@ -503,32 +556,32 @@ async def get_gguf_variants(
default_variant = _extract_quant_label(best) if best else None
return GgufVariantsResponse(
repo_id=repo_id,
variants=[
repo_id = repo_id,
variants = [
GgufVariantDetail(
filename=v.filename,
quant=v.quant,
size_bytes=v.size_bytes,
filename = v.filename,
quant = v.quant,
size_bytes = v.size_bytes,
)
for v in variants
],
has_vision=has_vision,
default_variant=default_variant,
has_vision = has_vision,
default_variant = default_variant,
)
except Exception as e:
logger.error(f"Error listing GGUF variants for '{repo_id}': {e}", exc_info=True)
logger.error(f"Error listing GGUF variants for '{repo_id}': {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to list GGUF variants: {str(e)}",
status_code = 500,
detail = f"Failed to list GGUF variants: {str(e)}",
)
@router.get("/checkpoints", response_model=CheckpointListResponse)
@router.get("/checkpoints", response_model = CheckpointListResponse)
async def list_checkpoints(
outputs_dir: str = Query(
default=str(outputs_root()),
description="Directory to scan for checkpoints",
default = str(outputs_root()),
description = "Directory to scan for checkpoints",
),
current_subject: str = Depends(get_current_subject),
):
@ -539,29 +592,29 @@ async def list_checkpoints(
"""
try:
resolved_outputs_dir = str(resolve_output_dir(outputs_dir))
raw_models = scan_checkpoints(outputs_dir=resolved_outputs_dir)
raw_models = scan_checkpoints(outputs_dir = resolved_outputs_dir)
models = [
ModelCheckpoints(
name=model_name,
checkpoints=[
CheckpointInfo(display_name=display_name, path=path, loss=loss)
name = model_name,
checkpoints = [
CheckpointInfo(display_name = display_name, path = path, loss = loss)
for display_name, path, loss in checkpoints
],
base_model=metadata.get("base_model"),
peft_type=metadata.get("peft_type"),
lora_rank=metadata.get("lora_rank"),
base_model = metadata.get("base_model"),
peft_type = metadata.get("peft_type"),
lora_rank = metadata.get("lora_rank"),
)
for model_name, checkpoints, metadata in raw_models
]
return CheckpointListResponse(
outputs_dir=resolved_outputs_dir,
models=models,
outputs_dir = resolved_outputs_dir,
models = models,
)
except Exception as e:
logger.error(f"Error listing checkpoints: {e}", exc_info=True)
logger.error(f"Error listing checkpoints: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to list checkpoints: {str(e)}",
status_code = 500,
detail = f"Failed to list checkpoints: {str(e)}",
)

334
studio/backend/routes/training.py
View file

@ -4,6 +4,7 @@
"""
Training API routes
"""
import sys
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Request
@ -50,12 +51,11 @@ from pydantic import BaseModel as PydanticBaseModel
class TrainingStopRequest(PydanticBaseModel):
save: bool = True
router = APIRouter()
logger = get_logger(__name__)
@router.get("/hardware")
async def get_hardware_utilization(
current_subject: str = Depends(get_current_subject),
@ -100,13 +100,13 @@ async def start_training(
if backend.is_training_active():
existing_job_id: Optional[str] = getattr(backend, "current_job_id", "")
return TrainingJobResponse(
job_id=existing_job_id or job_id,
status="error",
message=(
job_id = existing_job_id or job_id,
status = "error",
message = (
"Training is already in progress. "
"Stop current training before starting a new one."
),
error="Training already active",
error = "Training already active",
)
# Validate dataset paths if provided
@ -128,8 +128,8 @@ async def start_training(
if missing_datasets:
missing_detail = "; ".join(missing_datasets[:3])
raise HTTPException(
status_code=400,
detail=f"Local dataset not found: {missing_detail}",
status_code = 400,
detail = f"Local dataset not found: {missing_detail}",
)
request.local_datasets = validated_datasets
@ -167,7 +167,9 @@ async def start_training(
"lora_r": request.lora_r,
"lora_alpha": request.lora_alpha,
"lora_dropout": request.lora_dropout,
"target_modules": request.target_modules if request.target_modules else None,
"target_modules": request.target_modules
if request.target_modules
else None,
"gradient_checkpointing": request.gradient_checkpointing.strip()
if request.gradient_checkpointing and request.gradient_checkpointing.strip()
else "unsloth",
@ -194,15 +196,20 @@ async def start_training(
# net, consult the YAML directly so models that need it always get it.
if not training_kwargs["trust_remote_code"]:
model_defaults = load_model_defaults(request.model_name)
yaml_trust = model_defaults.get("training", {}).get("trust_remote_code", False)
yaml_trust = model_defaults.get("training", {}).get(
"trust_remote_code", False
)
if yaml_trust:
logger.info(f"YAML config sets trust_remote_code=True for {request.model_name}")
logger.info(
f"YAML config sets trust_remote_code=True for {request.model_name}"
)
training_kwargs["trust_remote_code"] = True
# Free GPU memory: shut down any running inference/export subprocesses
# before training starts (they'd compete for VRAM otherwise)
try:
from core.inference import get_inference_backend
inf_backend = get_inference_backend()
if inf_backend.active_model_name:
logger.info(
@ -217,9 +224,12 @@ async def start_training(
try:
from core.export import get_export_backend
exp_backend = get_export_backend()
if exp_backend.current_checkpoint:
logger.info("Shutting down export subprocess to free GPU memory for training")
logger.info(
"Shutting down export subprocess to free GPU memory for training"
)
exp_backend._shutdown_subprocess()
exp_backend.current_checkpoint = None
exp_backend.is_vision = False
@ -233,28 +243,28 @@ async def start_training(
if not success:
progress_error = backend.trainer.training_progress.error
return TrainingJobResponse(
job_id=job_id,
status="error",
message=progress_error or "Failed to start training subprocess",
error=progress_error or "subprocess_start_failed",
job_id = job_id,
status = "error",
message = progress_error or "Failed to start training subprocess",
error = progress_error or "subprocess_start_failed",
)
return TrainingJobResponse(
job_id=job_id,
status="queued",
message="Training job queued and starting in subprocess",
error=None,
job_id = job_id,
status = "queued",
message = "Training job queued and starting in subprocess",
error = None,
)
except Exception as e:
logger.error(f"Error starting training: {e}", exc_info=True)
logger.error(f"Error starting training: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to start training: {str(e)}",
status_code = 500,
detail = f"Failed to start training: {str(e)}",
)
@router.post("/stop", response_model=TrainingStopResponse)
@router.post("/stop", response_model = TrainingStopResponse)
async def stop_training(
body: TrainingStopRequest = TrainingStopRequest(),
current_subject: str = Depends(get_current_subject),
@ -272,23 +282,21 @@ async def stop_training(
if not is_active:
return TrainingStopResponse(
status="idle",
message="No training job is currently running"
status = "idle", message = "No training job is currently running"
)
# Call backend stop method
backend.stop_training(save=body.save)
backend.stop_training(save = body.save)
return TrainingStopResponse(
status="stopped",
message="Stop requested. Training will stop at the next safe step."
status = "stopped",
message = "Stop requested. Training will stop at the next safe step.",
)
except Exception as e:
logger.error(f"Error stopping training: {e}", exc_info=True)
logger.error(f"Error stopping training: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to stop training: {str(e)}"
status_code = 500, detail = f"Failed to stop training: {str(e)}"
)
@ -306,21 +314,30 @@ async def reset_training(
if is_active:
if backend._cancel_requested:
# Cancel (save=False) was requested — force-terminate so we can reset immediately
logger.info("Force-terminating subprocess for immediate reset (cancel path)")
logger.info(
"Force-terminating subprocess for immediate reset (cancel path)"
)
backend.force_terminate()
else:
logger.warning("Rejected reset while training active: is_active=%s", is_active)
logger.warning(
"Rejected reset while training active: is_active=%s", is_active
)
raise HTTPException(
status_code=409,
detail="Training is still running. Stop training and wait for it to finish before resetting.",
status_code = 409,
detail = "Training is still running. Stop training and wait for it to finish before resetting.",
)
logger.info("Reset training state: clearing runtime + metric history")
backend._should_stop = False # Clear stop flag so status returns to idle
backend.trainer._update_progress(
is_training=False, is_completed=False, error=None,
status_message="Ready to train", step=0, loss=0.0, epoch=0,
total_steps=0,
is_training = False,
is_completed = False,
error = None,
status_message = "Ready to train",
step = 0,
loss = 0.0,
epoch = 0,
total_steps = 0,
)
backend.loss_history = []
backend.lr_history = []
@ -331,10 +348,10 @@ async def reset_training(
except HTTPException:
raise
except Exception as e:
logger.error(f"Error resetting training: {e}", exc_info=True)
logger.error(f"Error resetting training: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to reset training: {str(e)}",
status_code = 500,
detail = f"Failed to reset training: {str(e)}",
)
@ -410,25 +427,24 @@ async def get_training_status(
}
return TrainingStatus(
job_id=job_id,
phase=phase,
is_training_running=is_active,
eval_enabled=backend.eval_enabled,
message=status_message,
error=error_message,
details=details,
metric_history=metric_history,
job_id = job_id,
phase = phase,
is_training_running = is_active,
eval_enabled = backend.eval_enabled,
message = status_message,
error = error_message,
details = details,
metric_history = metric_history,
)
except Exception as e:
logger.error(f"Error getting training status: {e}", exc_info=True)
logger.error(f"Error getting training status: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to get training status: {str(e)}"
status_code = 500, detail = f"Failed to get training status: {str(e)}"
)
@router.get("/metrics", response_model=TrainingMetricsResponse)
@router.get("/metrics", response_model = TrainingMetricsResponse)
async def get_training_metrics(
current_subject: str = Depends(get_current_subject),
):
@ -437,7 +453,7 @@ async def get_training_metrics(
"""
try:
backend = get_training_backend()
# Get metrics from backend
loss_history = backend.loss_history
lr_history = backend.lr_history
@ -451,21 +467,20 @@ async def get_training_metrics(
current_step = step_history[-1] if step_history else None
return TrainingMetricsResponse(
loss_history=loss_history,
lr_history=lr_history,
step_history=step_history,
grad_norm_history=grad_norm_history,
grad_norm_step_history=grad_norm_step_history,
current_loss=current_loss,
current_lr=current_lr,
current_step=current_step,
loss_history = loss_history,
lr_history = lr_history,
step_history = step_history,
grad_norm_history = grad_norm_history,
grad_norm_step_history = grad_norm_step_history,
current_loss = current_loss,
current_lr = current_lr,
current_step = current_step,
)
except Exception as e:
logger.error(f"Error getting training metrics: {e}", exc_info=True)
logger.error(f"Error getting training metrics: {e}", exc_info = True)
raise HTTPException(
status_code=500,
detail=f"Failed to get training metrics: {str(e)}"
status_code = 500, detail = f"Failed to get training metrics: {str(e)}"
)
@ -518,29 +533,31 @@ async def stream_training_progress(
)
# Get actual values from progress object if available
elapsed_seconds = getattr(progress, 'elapsed_seconds', None) if progress else None
eta_seconds = getattr(progress, 'eta_seconds', None) if progress else None
elapsed_seconds = (
getattr(progress, "elapsed_seconds", None) if progress else None
)
eta_seconds = getattr(progress, "eta_seconds", None) if progress else None
grad_norm = grad_norm_override
if grad_norm is None and progress:
grad_norm = getattr(progress, 'grad_norm', None)
num_tokens = getattr(progress, 'num_tokens', None) if progress else None
grad_norm = getattr(progress, "grad_norm", None)
num_tokens = getattr(progress, "num_tokens", None) if progress else None
eval_loss = eval_loss_override
if eval_loss is None and progress:
eval_loss = getattr(progress, 'eval_loss', None)
eval_loss = getattr(progress, "eval_loss", None)
return TrainingProgress(
job_id=job_id,
step=step,
total_steps=total,
loss=loss,
learning_rate=learning_rate,
progress_percent=progress_percent,
epoch=epoch,
elapsed_seconds=elapsed_seconds,
eta_seconds=eta_seconds,
grad_norm=grad_norm,
num_tokens=num_tokens,
eval_loss=eval_loss,
job_id = job_id,
step = step,
total_steps = total,
loss = loss,
learning_rate = learning_rate,
progress_percent = progress_percent,
epoch = epoch,
elapsed_seconds = elapsed_seconds,
eta_seconds = eta_seconds,
grad_norm = grad_norm,
num_tokens = num_tokens,
eval_loss = eval_loss,
)
def format_sse(
@ -574,23 +591,37 @@ async def stream_training_progress(
}
for i, step_val in enumerate(backend.step_history):
if step_val > resume_from_step:
loss_val = backend.loss_history[i] if i < len(backend.loss_history) else 0.0
lr_val = backend.lr_history[i] if i < len(backend.lr_history) else 0.0
loss_val = (
backend.loss_history[i]
if i < len(backend.loss_history)
else 0.0
)
lr_val = (
backend.lr_history[i] if i < len(backend.lr_history) else 0.0
)
tp_replay = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
total_replay = getattr(tp_replay, "total_steps", step_val) if tp_replay else step_val
epoch_replay = getattr(tp_replay, "epoch", None) if tp_replay else None
total_replay = (
getattr(tp_replay, "total_steps", step_val)
if tp_replay
else step_val
)
epoch_replay = (
getattr(tp_replay, "epoch", None) if tp_replay else None
)
payload = build_progress(
step_val,
loss_val,
lr_val,
total_replay,
epoch_replay,
progress=tp_replay,
grad_norm_override=grad_norm_by_step.get(step_val),
progress = tp_replay,
grad_norm_override = grad_norm_by_step.get(step_val),
)
yield format_sse(
payload.model_dump_json(), event = "progress", event_id = step_val
)
yield format_sse(payload.model_dump_json(), event="progress", event_id=step_val)
replayed += 1
if replayed:
logger.info(f"SSE reconnect: replayed {replayed} missed steps")
@ -603,45 +634,62 @@ async def stream_training_progress(
initial_epoch = getattr(tp, "epoch", None) if tp else None
initial_progress = build_progress(
step=0,
loss=0.0,
learning_rate=0.0,
total_steps=initial_total_steps,
epoch=initial_epoch,
progress=tp,
step = 0,
loss = 0.0,
learning_rate = 0.0,
total_steps = initial_total_steps,
epoch = initial_epoch,
progress = tp,
)
yield format_sse(
initial_progress.model_dump_json(), event = "progress", event_id = 0
)
yield format_sse(initial_progress.model_dump_json(), event="progress", event_id=0)
# If not active, send final state and exit
if not is_active:
if backend.step_history:
final_step = backend.step_history[-1]
final_loss = backend.loss_history[-1] if backend.loss_history else 0.0
final_loss = (
backend.loss_history[-1] if backend.loss_history else 0.0
)
final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
final_total_steps = (
getattr(tp, "total_steps", final_step) if tp else final_step
)
final_epoch = getattr(tp, "epoch", None) if tp else None
payload = build_progress(final_step, final_loss, final_lr, final_total_steps, final_epoch, progress=tp)
yield format_sse(payload.model_dump_json(), event="complete", event_id=final_step)
payload = build_progress(
final_step,
final_loss,
final_lr,
final_total_steps,
final_epoch,
progress = tp,
)
yield format_sse(
payload.model_dump_json(), event = "complete", event_id = final_step
)
else:
yield format_sse(
build_progress(-1, 0.0, 0.0, 0, progress=tp).model_dump_json(),
event="complete",
event_id=0,
build_progress(-1, 0.0, 0.0, 0, progress = tp).model_dump_json(),
event = "complete",
event_id = 0,
)
return
# ── Live polling loop ────────────────────────────────────
last_step = resume_from_step if resume_from_step is not None else -1
no_update_count = 0
max_no_updates = 1800 # Timeout after 30 minutes (large models need time for compilation)
max_no_updates = (
1800 # Timeout after 30 minutes (large models need time for compilation)
)
while backend.is_training_active():
try:
if backend.step_history:
current_step = backend.step_history[-1]
current_loss = backend.loss_history[-1] if backend.loss_history else 0.0
current_loss = (
backend.loss_history[-1] if backend.loss_history else 0.0
)
current_lr = backend.lr_history[-1] if backend.lr_history else 0.0
tp_inner = getattr(
getattr(backend, "trainer", None), "training_progress", None
@ -651,7 +699,9 @@ async def stream_training_progress(
if tp_inner
else current_step
)
current_epoch = getattr(tp_inner, "epoch", None) if tp_inner else None
current_epoch = (
getattr(tp_inner, "epoch", None) if tp_inner else None
)
# Only send if step changed
if current_step != last_step:
@ -661,12 +711,12 @@ async def stream_training_progress(
current_lr,
current_total_steps,
current_epoch,
progress=tp_inner,
progress = tp_inner,
)
yield format_sse(
progress_payload.model_dump_json(),
event="progress",
event_id=current_step,
event = "progress",
event_id = current_step,
)
last_step = current_step
no_update_count = 0
@ -680,12 +730,12 @@ async def stream_training_progress(
current_lr,
current_total_steps,
current_epoch,
progress=tp_inner,
progress = tp_inner,
)
yield format_sse(
heartbeat_payload.model_dump_json(),
event="heartbeat",
event_id=current_step,
event = "heartbeat",
event_id = current_step,
)
else:
# No steps yet, but training is active (model loading, etc.)
@ -695,43 +745,53 @@ async def stream_training_progress(
# the frontend can show "Tokenizing…" etc.
tp_prep = getattr(
getattr(backend, "trainer", None),
"training_progress", None,
"training_progress",
None,
)
prep_total = (
getattr(tp_prep, "total_steps", 0)
if tp_prep else 0
getattr(tp_prep, "total_steps", 0) if tp_prep else 0
)
preparing_payload = build_progress(
0, 0.0, 0.0, prep_total, progress=tp_prep,
0,
0.0,
0.0,
prep_total,
progress = tp_prep,
)
yield format_sse(
preparing_payload.model_dump_json(),
event="heartbeat",
event_id=0,
event = "heartbeat",
event_id = 0,
)
# Timeout check
if no_update_count > max_no_updates:
logger.warning("Progress stream timeout - no updates received")
tp_timeout = getattr(getattr(backend, "trainer", None), "training_progress", None)
timeout_payload = build_progress(last_step, 0.0, 0.0, 0, progress=tp_timeout)
tp_timeout = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
timeout_payload = build_progress(
last_step, 0.0, 0.0, 0, progress = tp_timeout
)
yield format_sse(
timeout_payload.model_dump_json(),
event="error",
event_id=last_step if last_step >= 0 else 0,
event = "error",
event_id = last_step if last_step >= 0 else 0,
)
break
await asyncio.sleep(1) # Poll every second
except Exception as e:
logger.error(f"Error in progress stream: {e}", exc_info=True)
tp_error = getattr(getattr(backend, "trainer", None), "training_progress", None)
error_payload = build_progress(0, 0.0, 0.0, 0, progress=tp_error)
logger.error(f"Error in progress stream: {e}", exc_info = True)
tp_error = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
error_payload = build_progress(0, 0.0, 0.0, 0, progress = tp_error)
yield format_sse(
error_payload.model_dump_json(),
event="error",
event_id=last_step if last_step >= 0 else 0,
event = "error",
event_id = last_step if last_step >= 0 else 0,
)
break
@ -739,9 +799,7 @@ async def stream_training_progress(
final_step = backend.step_history[-1] if backend.step_history else last_step
final_loss = backend.loss_history[-1] if backend.loss_history else 0.0
final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
final_tp = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
final_tp = getattr(getattr(backend, "trainer", None), "training_progress", None)
final_total_steps = (
getattr(final_tp, "total_steps", final_step) if final_tp else final_step
)
@ -752,20 +810,20 @@ async def stream_training_progress(
final_lr,
final_total_steps,
final_epoch,
progress=final_tp,
progress = final_tp,
)
yield format_sse(
final_payload.model_dump_json(),
event="complete",
event_id=final_step if final_step >= 0 else 0,
event = "complete",
event_id = final_step if final_step >= 0 else 0,
)
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
}
},
)

40
studio/backend/run.py
View file

@ -5,11 +5,12 @@
Run script for Unsloth UI Backend.
Works independently and can be moved to any directory.
"""
import os
import sys
# Suppress annoying C-level dependency warnings globally (e.g. SwigPyPacked)
os.environ["PYTHONWARNINGS"] = "ignore"
os.environ["PYTHONWARNINGS"] = "ignore"
from pathlib import Path
@ -19,6 +20,7 @@ if str(backend_dir) not in sys.path:
sys.path.insert(0, str(backend_dir))
from loggers import get_logger
logger = get_logger(__name__)
@ -38,9 +40,9 @@ def _resolve_external_ip() -> str:
try:
req = urllib.request.Request(
"http://metadata.google.internal/computeMetadata/v1/instance/network-interfaces/0/access-configs/0/external-ip",
headers={"Metadata-Flavor": "Google"},
headers = {"Metadata-Flavor": "Google"},
)
with urllib.request.urlopen(req, timeout=1) as resp:
with urllib.request.urlopen(req, timeout = 1) as resp:
ip = resp.read().decode().strip()
if ip:
return ip
@ -49,7 +51,7 @@ def _resolve_external_ip() -> str:
# 2. Try public IP service
try:
with urllib.request.urlopen("https://ifconfig.me", timeout=3) as resp:
with urllib.request.urlopen("https://ifconfig.me", timeout = 3) as resp:
ip = resp.read().decode().strip()
if ip:
return ip
@ -70,7 +72,7 @@ def _resolve_external_ip() -> str:
def run_server(
host: str = "0.0.0.0",
port: int = 8000,
frontend_path: Path = "studio/frontend/dist",
frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist",
silent: bool = False,
):
"""
@ -108,11 +110,13 @@ def run_server(
# Run server
def _run():
config = uvicorn.Config(app, host=host, port=port, log_level="info", access_log=False)
config = uvicorn.Config(
app, host = host, port = port, log_level = "info", access_log = False
)
server = uvicorn.Server(config)
asyncio.run(server.serve())
thread = Thread(target=_run, daemon=True)
thread = Thread(target = _run, daemon = True)
thread.start()
time.sleep(3)
@ -135,24 +139,26 @@ def run_server(
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Run Unsloth UI Backend server")
parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
parser = argparse.ArgumentParser(description = "Run Unsloth UI Backend server")
parser.add_argument("--host", default = "0.0.0.0", help = "Host to bind to")
parser.add_argument("--port", type = int, default = 8000, help = "Port to bind to")
parser.add_argument(
"--frontend", type=str, default="studio/frontend/dist", help="Path to frontend build"
"--frontend",
type = str,
default = Path(__file__).resolve().parent.parent / "frontend" / "dist",
help = "Path to frontend build",
)
parser.add_argument("--silent", action="store_true", help="Suppress output")
parser.add_argument("--silent", action = "store_true", help = "Suppress output")
args = parser.parse_args()
frontend_path = Path(args.frontend) if args.frontend else None
run_server(
host=args.host, port=args.port, frontend_path=frontend_path, silent=args.silent
)
kwargs = dict(host = args.host, port = args.port, silent = args.silent)
if args.frontend is not None:
kwargs["frontend_path"] = Path(args.frontend)
run_server(**kwargs)
# Keep running
import time
while True:
time.sleep(1)

2
studio/backend/state/__init__.py
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

2
studio/backend/tests/__init__.py
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

1
studio/backend/tests/conftest.py
View file

@ -6,6 +6,7 @@ Shared pytest configuration for the backend test suite.
Ensures that the backend root is on sys.path so that
`import utils.utils` (and similar flat imports) resolve correctly.
"""
import sys
from pathlib import Path

112
studio/backend/tests/test_utils.py
View file

@ -16,6 +16,7 @@ Run with:
cd studio/backend
python -m pytest tests/test_utils.py -v
"""
import platform
from unittest.mock import patch, MagicMock
@ -24,18 +25,20 @@ import pytest
# --- Conditional framework imports ---
try:
import torch
HAS_TORCH = True
except ImportError:
HAS_TORCH = False
try:
import mlx.core as mx
HAS_MLX = True
except ImportError:
HAS_MLX = False
needs_torch = pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not installed")
needs_mlx = pytest.mark.skipif(not HAS_MLX, reason="MLX not installed")
needs_torch = pytest.mark.skipif(not HAS_TORCH, reason = "PyTorch not installed")
needs_mlx = pytest.mark.skipif(not HAS_MLX, reason = "MLX not installed")
from utils.hardware import (
get_device,
@ -52,6 +55,7 @@ from utils.utils import format_error_message
# ========== Helpers ==========
def _actual_device() -> str:
"""Return the real device string for the current machine."""
if HAS_TORCH and torch.cuda.is_available():
@ -69,6 +73,7 @@ def _reset_and_detect():
# ========== get_device() ==========
class TestGetDevice:
"""Tests for get_device() — should agree with the real hardware."""
@ -89,28 +94,34 @@ class TestGetDevice:
@needs_torch
def test_returns_cuda_when_cuda_available(self):
with patch("utils.hardware.hardware._has_torch", return_value=True), \
patch("torch.cuda.is_available", return_value=True):
with (
patch("utils.hardware.hardware._has_torch", return_value = True),
patch("torch.cuda.is_available", return_value = True),
):
assert _reset_and_detect() == DeviceType.CUDA
@needs_mlx
def test_returns_mlx_when_on_apple_silicon_with_mlx(self):
with patch("utils.hardware.hardware._has_torch", return_value=False), \
patch("utils.hardware.hardware.is_apple_silicon", return_value=True), \
patch("utils.hardware.hardware._has_mlx", return_value=True):
with (
patch("utils.hardware.hardware._has_torch", return_value = False),
patch("utils.hardware.hardware.is_apple_silicon", return_value = True),
patch("utils.hardware.hardware._has_mlx", return_value = True),
):
assert _reset_and_detect() == DeviceType.MLX
def test_returns_cpu_when_nothing_available(self):
with patch("utils.hardware.hardware._has_torch", return_value=False), \
patch("utils.hardware.hardware.is_apple_silicon", return_value=False), \
patch("utils.hardware.hardware._has_mlx", return_value=False):
with (
patch("utils.hardware.hardware._has_torch", return_value = False),
patch("utils.hardware.hardware.is_apple_silicon", return_value = False),
patch("utils.hardware.hardware._has_mlx", return_value = False),
):
assert _reset_and_detect() == DeviceType.CPU
# ========== is_apple_silicon() ==========
class TestIsAppleSilicon:
class TestIsAppleSilicon:
def test_returns_bool(self):
assert isinstance(is_apple_silicon(), bool)
@ -136,6 +147,7 @@ class TestIsAppleSilicon:
# ========== clear_gpu_cache() ==========
class TestClearGpuCache:
"""clear_gpu_cache() must never raise, regardless of platform."""
@ -144,9 +156,11 @@ class TestClearGpuCache:
@needs_torch
def test_calls_cuda_cache_when_cuda(self):
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \
patch("torch.cuda.empty_cache") as mock_empty, \
patch("torch.cuda.ipc_collect") as mock_ipc:
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
patch("torch.cuda.empty_cache") as mock_empty,
patch("torch.cuda.ipc_collect") as mock_ipc,
):
clear_gpu_cache()
mock_empty.assert_called_once()
mock_ipc.assert_called_once()
@ -154,18 +168,18 @@ class TestClearGpuCache:
@needs_mlx
def test_mlx_does_not_raise(self):
"""MLX cache clear is a no-op — should just succeed."""
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MLX):
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX):
clear_gpu_cache()
def test_noop_on_cpu(self):
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU):
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU):
clear_gpu_cache()
# ========== get_gpu_memory_info() ==========
class TestGetGpuMemoryInfo:
class TestGetGpuMemoryInfo:
def test_returns_dict(self):
result = get_gpu_memory_info()
assert isinstance(result, dict)
@ -183,8 +197,7 @@ class TestGetGpuMemoryInfo:
# --- When a GPU IS available ---
@pytest.mark.skipif(
_actual_device() == "cpu",
reason="No GPU available on this machine"
_actual_device() == "cpu", reason = "No GPU available on this machine"
)
def test_gpu_available_fields(self):
result = get_gpu_memory_info()
@ -200,14 +213,16 @@ class TestGetGpuMemoryInfo:
@needs_torch
def test_cuda_path_returns_correct_fields(self):
mock_props = MagicMock()
mock_props.total_memory = 16 * (1024 ** 3)
mock_props.total_memory = 16 * (1024**3)
mock_props.name = "NVIDIA Test GPU"
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \
patch("torch.cuda.current_device", return_value=0), \
patch("torch.cuda.get_device_properties", return_value=mock_props), \
patch("torch.cuda.memory_allocated", return_value=4 * (1024 ** 3)), \
patch("torch.cuda.memory_reserved", return_value=6 * (1024 ** 3)):
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
patch("torch.cuda.current_device", return_value = 0),
patch("torch.cuda.get_device_properties", return_value = mock_props),
patch("torch.cuda.memory_allocated", return_value = 4 * (1024**3)),
patch("torch.cuda.memory_reserved", return_value = 6 * (1024**3)),
):
result = get_gpu_memory_info()
assert result["available"] is True
@ -223,13 +238,15 @@ class TestGetGpuMemoryInfo:
@needs_mlx
def test_mlx_path_returns_correct_fields(self):
mock_psutil_mem = MagicMock()
mock_psutil_mem.total = 32 * (1024 ** 3) # 32 GB unified
mock_psutil_mem.total = 32 * (1024**3) # 32 GB unified
mock_psutil = MagicMock()
mock_psutil.virtual_memory.return_value = mock_psutil_mem
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MLX), \
patch.dict("sys.modules", {"psutil": mock_psutil}):
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX),
patch.dict("sys.modules", {"psutil": mock_psutil}),
):
result = get_gpu_memory_info()
assert result["available"] is True
@ -240,7 +257,7 @@ class TestGetGpuMemoryInfo:
# --- CPU-only path ---
def test_cpu_path_returns_unavailable(self):
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU):
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU):
result = get_gpu_memory_info()
assert result["available"] is False
assert result["backend"] == "cpu"
@ -249,8 +266,13 @@ class TestGetGpuMemoryInfo:
@needs_torch
def test_cuda_error_returns_unavailable(self):
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \
patch("torch.cuda.current_device", side_effect=RuntimeError("CUDA init failed")):
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
patch(
"torch.cuda.current_device",
side_effect = RuntimeError("CUDA init failed"),
),
):
result = get_gpu_memory_info()
assert result["available"] is False
assert "error" in result
@ -258,8 +280,8 @@ class TestGetGpuMemoryInfo:
# ========== log_gpu_memory() ==========
class TestLogGpuMemory:
class TestLogGpuMemory:
def test_does_not_raise(self):
log_gpu_memory("test")
@ -275,8 +297,13 @@ class TestLogGpuMemory:
}
import structlog
from loggers import get_logger
with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info), \
caplog.at_level(logging.INFO, logger="utils.hardware.hardware"):
with (
patch(
"utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
),
caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"),
):
log_gpu_memory("unit-test")
assert "unit-test" in caplog.text
@ -287,8 +314,13 @@ class TestLogGpuMemory:
fake_info = {"available": False, "backend": "cpu"}
import structlog
from loggers import get_logger
with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info), \
caplog.at_level(logging.INFO, logger="utils.hardware.hardware"):
with (
patch(
"utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
),
caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"),
):
log_gpu_memory("cpu-test")
assert "No GPU available" in caplog.text
@ -296,8 +328,8 @@ class TestLogGpuMemory:
# ========== format_error_message() ==========
class TestFormatErrorMessage:
class TestFormatErrorMessage:
def test_not_found(self):
err = Exception("Repository not found for unsloth/test")
msg = format_error_message(err, "unsloth/test")
@ -324,7 +356,7 @@ class TestFormatErrorMessage:
@needs_torch
def test_cuda_oom(self):
err = Exception("CUDA out of memory")
with patch("utils.hardware.get_device", return_value=DeviceType.CUDA):
with patch("utils.hardware.get_device", return_value = DeviceType.CUDA):
msg = format_error_message(err, "big/model")
assert "GPU" in msg
assert "big/model" not in msg
@ -335,7 +367,7 @@ class TestFormatErrorMessage:
@needs_mlx
def test_mlx_oom(self):
err = Exception("MLX backend out of memory")
with patch("utils.hardware.get_device", return_value=DeviceType.MLX):
with patch("utils.hardware.get_device", return_value = DeviceType.MLX):
msg = format_error_message(err, "unsloth/huge-model")
assert "Apple Silicon" in msg
@ -343,7 +375,7 @@ class TestFormatErrorMessage:
def test_cpu_oom(self):
err = Exception("not enough memory to allocate")
with patch("utils.hardware.get_device", return_value=DeviceType.CPU):
with patch("utils.hardware.get_device", return_value = DeviceType.CPU):
msg = format_error_message(err, "any/model")
assert "system" in msg.lower()

2
studio/backend/utils/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

7
studio/backend/utils/cache_cleanup.py
View file

@ -8,6 +8,7 @@ The unsloth_compiled_cache is created by unsloth_zoo/compiler.py during
FastModel.from_pretrained() and contains model-type-specific compiled Python
files. It should be cleared between model loads to avoid stale artefacts.
"""
import shutil
import structlog
from loggers import get_logger
@ -16,8 +17,8 @@ from pathlib import Path
logger = get_logger(__name__)
# Possible locations where unsloth_compiled_cache may appear
_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend
_PROJECT_ROOT = _BACKEND_DIR.parent.parent # repo root
_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend
_PROJECT_ROOT = _BACKEND_DIR.parent.parent # repo root
_CACHE_DIRS = [
_BACKEND_DIR / "unsloth_compiled_cache",
@ -31,4 +32,4 @@ def clear_unsloth_compiled_cache() -> None:
for cache_dir in _CACHE_DIRS:
if cache_dir.exists():
logger.info(f"Removing unsloth compiled cache: {cache_dir}")
shutil.rmtree(cache_dir, ignore_errors=True)
shutil.rmtree(cache_dir, ignore_errors = True)

28
studio/backend/utils/datasets/chat_templates.py
View file

@ -61,7 +61,7 @@ def get_tokenizer_chat_template(tokenizer, model_name):
try:
tokenizer = get_chat_template(
tokenizer,
chat_template=matched_template,
chat_template = matched_template,
)
except Exception as e:
logger.info(f"⚠️ Failed to apply Unsloth template '{matched_template}': {e}")
@ -80,7 +80,7 @@ def get_tokenizer_chat_template(tokenizer, model_name):
try:
tokenizer = get_chat_template(
tokenizer,
chat_template="chatml",
chat_template = "chatml",
)
except Exception as e:
logger.info(f"⚠️ Failed to apply default ChatML template: {e}")
@ -119,14 +119,14 @@ def get_dataset_info_summary(dataset_info):
def apply_chat_template_to_dataset(
dataset_info,
tokenizer,
model_name=None,
custom_prompt_template=None,
add_eos_token=False,
remove_bos_prefix=False,
custom_format_mapping=None,
auto_detect_mapping=True,
batch_size=1000,
num_proc=None,
model_name = None,
custom_prompt_template = None,
add_eos_token = False,
remove_bos_prefix = False,
custom_format_mapping = None,
auto_detect_mapping = True,
batch_size = 1000,
num_proc = None,
):
"""
Applies chat template to dataset based on its format.
@ -233,7 +233,7 @@ def apply_chat_template_to_dataset(
return result
try:
dataset = dataset.map(_apply_custom_mapping, batched=True, batch_size=batch_size)
dataset = dataset.map(_apply_custom_mapping, batched = True, batch_size = batch_size)
# Update to use conversations format
final_format = "chatml_conversations"
chat_column = "conversations"
@ -256,7 +256,7 @@ def apply_chat_template_to_dataset(
if not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template):
try:
from unsloth.chat_templates import get_chat_template
tokenizer = get_chat_template(tokenizer, chat_template="alpaca")
tokenizer = get_chat_template(tokenizer, chat_template = "alpaca")
logger.info(f"📝 Set alpaca chat template on tokenizer for model saving")
except Exception as e:
logger.info(f"⚠️ Could not set alpaca template on tokenizer: {e}")
@ -333,8 +333,8 @@ def apply_chat_template_to_dataset(
try:
text = tokenizer.apply_chat_template(
convo,
tokenize=False,
add_generation_prompt=False
tokenize = False,
add_generation_prompt = False
)
if remove_bos_prefix:

52
studio/backend/utils/datasets/data_collators.py
View file

@ -12,11 +12,10 @@ import torch
from dataclasses import dataclass
from typing import Any, List, Optional, Union
from loggers import get_logger
logger = get_logger(__name__)
@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
"""
@ -26,16 +25,23 @@ class DataCollatorSpeechSeq2SeqWithPadding:
masks padding in labels with -100, and strips leading BOS token.
Mirrors the collator from the Whisper.ipynb notebook.
"""
processor: Any
def __call__(self, features: List[dict]) -> dict:
input_features = [{"input_features": feature["input_features"]} for feature in features]
batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
input_features = [
{"input_features": feature["input_features"]} for feature in features
]
batch = self.processor.feature_extractor.pad(
input_features, return_tensors = "pt"
)
label_features = [{"input_ids": feature["labels"]} for feature in features]
labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
labels_batch = self.processor.tokenizer.pad(label_features, return_tensors = "pt")
labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
labels = labels_batch["input_ids"].masked_fill(
labels_batch.attention_mask.ne(1), -100
)
if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
labels = labels[:, 1:]
@ -54,6 +60,7 @@ class DeepSeekOCRDataCollator:
- Text tokenization
- Proper label masking for instruction fine-tuning
"""
processor: Any # Qwen2VLProcessor or similar
max_length: int = 2048
ignore_index: int = -100
@ -86,7 +93,7 @@ class DeepSeekOCRDataCollator:
for item in content:
if isinstance(item, dict) and item.get("type") == "image":
img = item.get("image")
if img is not None and hasattr(img, 'size'): # PIL Image
if img is not None and hasattr(img, "size"): # PIL Image
all_images.append(img)
# Process with the VL processor
@ -94,19 +101,19 @@ class DeepSeekOCRDataCollator:
# Qwen2VL style processing
texts = [
self.processor.apply_chat_template(
msgs, tokenize=False, add_generation_prompt=False
msgs, tokenize = False, add_generation_prompt = False
)
for msgs in all_messages
]
# Process with images
inputs = self.processor(
text=texts,
images=all_images if all_images else None,
return_tensors="pt",
padding=True,
truncation=True,
max_length=self.max_length,
text = texts,
images = all_images if all_images else None,
return_tensors = "pt",
padding = True,
truncation = True,
max_length = self.max_length,
)
# Create labels (mask input, keep output)
@ -134,6 +141,7 @@ class VLMDataCollator:
- LLaVA
- Other VL models with compatible processors
"""
processor: Any
max_length: int = 2048
ignore_index: int = -100
@ -163,26 +171,26 @@ class VLMDataCollator:
# Apply chat template
texts = [
self.processor.apply_chat_template(
msgs, tokenize=False, add_generation_prompt=False
msgs, tokenize = False, add_generation_prompt = False
)
for msgs in all_messages
]
# Process inputs
inputs = self.processor(
text=texts,
images=all_images if all_images else None,
return_tensors="pt",
padding=True,
truncation=True,
max_length=self.max_length,
text = texts,
images = all_images if all_images else None,
return_tensors = "pt",
padding = True,
truncation = True,
max_length = self.max_length,
)
# Create labels
labels = inputs["input_ids"].clone()
# Mask padding
if hasattr(self.processor, 'tokenizer'):
if hasattr(self.processor, "tokenizer"):
pad_token_id = self.processor.tokenizer.pad_token_id
else:
pad_token_id = self.processor.pad_token_id

393
studio/backend/utils/datasets/dataset_utils.py
View file

@ -45,22 +45,21 @@ from .vlm_processing import generate_smart_vlm_instruction
from .data_collators import DeepSeekOCRDataCollator, VLMDataCollator
from .model_mappings import TEMPLATE_TO_MODEL_MAPPER
from loggers import get_logger
logger = get_logger(__name__)
def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
"""
Lightweight format check without processing - for frontend validation.
Use this to quickly determine if user needs to manually map columns
before calling the full format_and_template_dataset().
Args:
dataset: HuggingFace dataset
is_vlm: Whether this is a Vision-Language Model dataset
Returns:
dict: {
"requires_manual_mapping": bool - True if user must map columns,
@ -71,8 +70,12 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
"detected_text_column": str or None - For VLM only,
}
"""
columns = list(dataset.column_names) if hasattr(dataset, 'column_names') else list(next(iter(dataset)).keys())
columns = (
list(dataset.column_names)
if hasattr(dataset, "column_names")
else list(next(iter(dataset)).keys())
)
# Auto-detect multimodal data regardless of is_vlm flag
multimodal_info = detect_multimodal_dataset(dataset)
is_audio = multimodal_info.get("is_audio", False)
@ -185,11 +188,17 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
**audio_fields,
}
# Normalise any format-specific role to canonical chatml (user/assistant/system)
_TO_CHATML = {
"user": "user", "human": "user", "instruction": "user",
"assistant": "assistant", "gpt": "assistant", "output": "assistant",
"system": "system", "input": "system",
"user": "user",
"human": "user",
"instruction": "user",
"assistant": "assistant",
"gpt": "assistant",
"output": "assistant",
"system": "system",
"input": "system",
}
_CHATML_ROLE_ORDER = ("system", "user", "assistant")
_CHATML_TO_ALPACA = {"user": "instruction", "system": "input", "assistant": "output"}
@ -232,11 +241,21 @@ def _apply_user_mapping(dataset, mapping: dict, batch_size: int = 1000):
for col in role_groups[chatml_role]:
if col in examples:
content = examples[col][i]
convo.append({"role": chatml_role, "content": str(content) if content else ""})
convo.append(
{
"role": chatml_role,
"content": str(content) if content else "",
}
)
conversations.append(convo)
return {"conversations": conversations}
return dataset.map(_convert, batched=True, batch_size=batch_size, remove_columns=dataset.column_names)
return dataset.map(
_convert,
batched = True,
batch_size = batch_size,
remove_columns = dataset.column_names,
)
def _extract_column_value(val, col: str, label_mapping: dict) -> str:
@ -248,7 +267,7 @@ def _extract_column_value(val, col: str, label_mapping: dict) -> str:
inner = val["text"]
str_val = inner[0] if isinstance(inner, list) and inner else str(inner)
else:
str_val = json.dumps(val, ensure_ascii=False)
str_val = json.dumps(val, ensure_ascii = False)
elif isinstance(val, list):
str_val = val[0] if len(val) == 1 else ", ".join(str(v) for v in val)
else:
@ -286,6 +305,7 @@ def _apply_template_mapping(
role_groups[canonical].append(col)
import logging as _log
_log.getLogger(__name__).info(
f"Applying role mapping: sys={bool(system_prompt)}, "
f"user_cols={role_groups['user']}, asst_cols={role_groups['assistant']}, "
@ -326,8 +346,10 @@ def _apply_template_mapping(
return {"conversations": conversations}
return dataset.map(
_convert, batched=True, batch_size=batch_size,
remove_columns=dataset.column_names,
_convert,
batched = True,
batch_size = batch_size,
remove_columns = dataset.column_names,
)
@ -341,7 +363,11 @@ def _apply_user_mapping_alpaca(dataset, mapping: dict, batch_size: int = 1000):
Returns:
Dataset with instruction/input/output columns
"""
col_for: dict[str, str | None] = {"instruction": None, "input": None, "output": None}
col_for: dict[str, str | None] = {
"instruction": None,
"input": None,
"output": None,
}
for col_name, role in mapping.items():
canonical = _TO_CHATML.get(role)
alpaca_field = _CHATML_TO_ALPACA.get(canonical) if canonical else None
@ -352,25 +378,48 @@ def _apply_user_mapping_alpaca(dataset, mapping: dict, batch_size: int = 1000):
num = len(next(iter(examples.values())))
instructions, inputs, outputs = [], [], []
for i in range(num):
for field, dest in (("instruction", instructions), ("input", inputs), ("output", outputs)):
for field, dest in (
("instruction", instructions),
("input", inputs),
("output", outputs),
):
col = col_for[field]
val = str(examples[col][i]) if col and col in examples and examples[col][i] else ""
val = (
str(examples[col][i])
if col and col in examples and examples[col][i]
else ""
)
dest.append(val)
return {"instruction": instructions, "input": inputs, "output": outputs}
return dataset.map(_convert, batched=True, batch_size=batch_size, remove_columns=dataset.column_names)
return dataset.map(
_convert,
batched = True,
batch_size = batch_size,
remove_columns = dataset.column_names,
)
def format_dataset(
dataset,
format_type = "auto",
tokenizer = None,
aliases_for_system = ["system",],
aliases_for_user = ["user", "human", "input",],
aliases_for_assistant = ["gpt", "assistant", "output",],
batch_size = 1000,
num_proc = None,
auto_detect_custom = True,
format_type = "auto",
tokenizer = None,
aliases_for_system = [
"system",
],
aliases_for_user = [
"user",
"human",
"input",
],
aliases_for_assistant = [
"gpt",
"assistant",
"output",
],
batch_size = 1000,
num_proc = None,
auto_detect_custom = True,
custom_format_mapping = None,
):
"""
@ -395,13 +444,17 @@ def format_dataset(
if custom_format_mapping:
try:
if format_type == "alpaca":
mapped_dataset = _apply_user_mapping_alpaca(dataset, custom_format_mapping, batch_size)
mapped_dataset = _apply_user_mapping_alpaca(
dataset, custom_format_mapping, batch_size
)
final_format = "alpaca"
chat_column = None
else:
# auto / chatml / sharegpt / conversational — all produce chatml conversations
# (sharegpt is always standardized to role/content internally)
mapped_dataset = _apply_user_mapping(dataset, custom_format_mapping, batch_size)
mapped_dataset = _apply_user_mapping(
dataset, custom_format_mapping, batch_size
)
final_format = "chatml_conversations"
chat_column = "conversations"
@ -414,7 +467,9 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": [f"Applied user-provided column mapping ({format_type}): {custom_format_mapping}"]
"warnings": [
f"Applied user-provided column mapping ({format_type}): {custom_format_mapping}"
],
}
except Exception as e:
return {
@ -426,10 +481,9 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": [f"Failed to apply user mapping: {e}"]
"warnings": [f"Failed to apply user mapping: {e}"],
}
# Detect current format
detected = detect_dataset_format(dataset)
warnings = []
@ -442,7 +496,6 @@ def format_dataset(
# AUTO MODE: Keep format but standardize if needed
if format_type == "auto":
# Alpaca - keep as is
if detected["format"] == "alpaca":
return {
@ -454,16 +507,20 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
# ShareGPT - needs standardization
elif detected["format"] == "sharegpt":
try:
standardized = standardize_chat_format(
dataset, tokenizer, aliases_for_system,
aliases_for_user, aliases_for_assistant,
batch_size, num_proc
dataset,
tokenizer,
aliases_for_system,
aliases_for_user,
aliases_for_assistant,
batch_size,
num_proc,
)
return {
"dataset": standardized,
@ -474,7 +531,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
except Exception as e:
warnings.append(f"Failed to standardize ShareGPT format: {e}")
@ -487,10 +544,14 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
elif detected["format"] == "chatml" and detected["chat_column"] in ["conversations", "messages", "texts"]:
elif detected["format"] == "chatml" and detected["chat_column"] in [
"conversations",
"messages",
"texts",
]:
return {
"dataset": dataset,
"detected_format": f"chatml_{detected['chat_column']}",
@ -500,13 +561,14 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
# Unknown - try standardization, if fails pass as is
else:
warnings.append(f"Unknown format detected. Keys found: {detected['sample_keys']}")
warnings.append(
f"Unknown format detected. Keys found: {detected['sample_keys']}"
)
# NEW: Try heuristic detection
if auto_detect_custom:
@ -514,7 +576,6 @@ def format_dataset(
if custom_mapping:
warnings.append(f"Auto-detected column mapping: {custom_mapping}")
def _apply_auto_mapping(examples):
conversations = []
num_examples = len(examples[list(examples.keys())[0]])
@ -523,25 +584,27 @@ def format_dataset(
all_columns = set(examples.keys())
mapped_columns = set(custom_mapping.keys())
preserved_columns = {
col: examples[col]
for col in all_columns - mapped_columns
col: examples[col] for col in all_columns - mapped_columns
}
for i in range(num_examples):
convo = []
for target_role in ['system', 'user', 'assistant']:
for target_role in ["system", "user", "assistant"]:
for col_name, role in custom_mapping.items():
if role == target_role and col_name in examples:
content = examples[col_name][i]
if content and str(content).strip():
convo.append({"role": role, "content": str(content)})
convo.append(
{"role": role, "content": str(content)}
)
conversations.append(convo)
return {"conversations": conversations, **preserved_columns}
try:
dataset = dataset.map(_apply_auto_mapping, batched=True, batch_size=batch_size)
dataset = dataset.map(
_apply_auto_mapping, batched = True, batch_size = batch_size
)
return {
"dataset": dataset,
"detected_format": "unknown",
@ -551,7 +614,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
except Exception as e:
warnings.append(f"Auto-detection failed: {e}")
@ -560,9 +623,13 @@ def format_dataset(
if detected["chat_column"]:
try:
standardized = standardize_chat_format(
dataset, tokenizer, aliases_for_system,
aliases_for_user, aliases_for_assistant,
batch_size, num_proc
dataset,
tokenizer,
aliases_for_system,
aliases_for_user,
aliases_for_assistant,
batch_size,
num_proc,
)
warnings.append("Successfully standardized unknown format")
return {
@ -574,10 +641,12 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
except Exception as e:
warnings.append(f"Could not standardize: {e}. Passing dataset as-is.")
warnings.append(
f"Could not standardize: {e}. Passing dataset as-is."
)
# Return as-is with warnings
return {
@ -589,12 +658,11 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
# ALPACA MODE: Convert to Alpaca
elif format_type == "alpaca":
if detected["format"] == "alpaca":
return {
"dataset": dataset,
@ -605,16 +673,20 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
elif detected["format"] in ["sharegpt", "chatml"]:
# First standardize if ShareGPT
if detected["format"] == "sharegpt":
dataset = standardize_chat_format(
dataset, tokenizer, aliases_for_system,
aliases_for_user, aliases_for_assistant,
batch_size, num_proc
dataset,
tokenizer,
aliases_for_system,
aliases_for_user,
aliases_for_assistant,
batch_size,
num_proc,
)
# Then convert to Alpaca
@ -628,7 +700,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
else:
@ -642,12 +714,11 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
# CHATML MODE: Convert to ChatML
elif format_type in ["chatml", "conversational", "sharegpt"]:
if detected["format"] == "alpaca":
converted = convert_alpaca_to_chatml(dataset, batch_size, num_proc)
return {
@ -659,14 +730,18 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
elif detected["format"] == "sharegpt":
standardized = standardize_chat_format(
dataset, tokenizer, aliases_for_system,
aliases_for_user, aliases_for_assistant,
batch_size, num_proc
dataset,
tokenizer,
aliases_for_system,
aliases_for_user,
aliases_for_assistant,
batch_size,
num_proc,
)
return {
"dataset": standardized,
@ -677,7 +752,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
elif detected["format"] == "chatml":
@ -690,7 +765,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
else:
@ -698,9 +773,13 @@ def format_dataset(
if detected["chat_column"]:
try:
standardized = standardize_chat_format(
dataset, tokenizer, aliases_for_system,
aliases_for_user, aliases_for_assistant,
batch_size, num_proc
dataset,
tokenizer,
aliases_for_system,
aliases_for_user,
aliases_for_assistant,
batch_size,
num_proc,
)
return {
"dataset": standardized,
@ -711,7 +790,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
except Exception as e:
warnings.append(f"Standardization failed: {e}")
@ -725,7 +804,7 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
else:
@ -737,25 +816,34 @@ def format_and_template_dataset(
model_name,
tokenizer,
is_vlm = False,
format_type="auto",
format_type = "auto",
# VLM-specific parameters
vlm_instruction=None, # Now optional - will auto-generate
vlm_text_column=None,
vlm_image_column=None,
dataset_name=None,
custom_prompt_template=None,
add_eos_token=False,
remove_bos_prefix=False,
custom_format_mapping=None,
auto_detect_custom=True,
auto_detect_mapping=True,
aliases_for_system=["system",],
aliases_for_user=["user", "human", "input",],
aliases_for_assistant=["gpt", "assistant", "output",],
batch_size=1000,
num_proc=None,
progress_callback=None,
vlm_instruction = None, # Now optional - will auto-generate
vlm_text_column = None,
vlm_image_column = None,
dataset_name = None,
custom_prompt_template = None,
add_eos_token = False,
remove_bos_prefix = False,
custom_format_mapping = None,
auto_detect_custom = True,
auto_detect_mapping = True,
aliases_for_system = [
"system",
],
aliases_for_user = [
"user",
"human",
"input",
],
aliases_for_assistant = [
"gpt",
"assistant",
"output",
],
batch_size = 1000,
num_proc = None,
progress_callback = None,
):
"""
Convenience function that combines format_dataset and apply_chat_template_to_dataset.
@ -786,25 +874,27 @@ def format_and_template_dataset(
# Expect mapping like: {"image_col": "image", "caption_col": "text"}
user_vlm_image_column = None
user_vlm_text_column = None
for col, role in custom_format_mapping.items():
if role == "image":
user_vlm_image_column = col
elif role in ["text", "user", "caption", "assistant"]:
user_vlm_text_column = col
if user_vlm_image_column and user_vlm_text_column:
try:
dataset = convert_to_vlm_format(
dataset,
instruction=vlm_instruction,
text_column=user_vlm_text_column,
image_column=user_vlm_image_column,
dataset_name=dataset_name,
progress_callback=progress_callback,
instruction = vlm_instruction,
text_column = user_vlm_text_column,
image_column = user_vlm_image_column,
dataset_name = dataset_name,
progress_callback = progress_callback,
)
warnings.append(f"Applied user VLM mapping: image='{user_vlm_image_column}', text='{user_vlm_text_column}'")
warnings.append(
f"Applied user VLM mapping: image='{user_vlm_image_column}', text='{user_vlm_text_column}'"
)
return {
"dataset": dataset,
"detected_format": "user_mapped",
@ -826,7 +916,9 @@ def format_and_template_dataset(
f"text='{user_vlm_text_column}') failed: {e} — "
f"falling back to auto-detection"
)
logger.info(f"⚠️ User VLM mapping failed, falling back to auto-detection...")
logger.info(
f"⚠️ User VLM mapping failed, falling back to auto-detection..."
)
custom_format_mapping = None # clear so auto-detection runs below
else:
errors.append(
@ -850,10 +942,13 @@ def format_and_template_dataset(
if vlm_structure["format"] == "vlm_messages_llava":
try:
dataset = convert_llava_to_vlm_format(dataset)
warnings.append("Converted from Llava format (image indices) to standard VLM format")
warnings.append(
"Converted from Llava format (image indices) to standard VLM format"
)
except Exception as e:
errors.append(f"Failed to convert Llava format: {e}")
import traceback
traceback.print_exc()
return {
@ -872,15 +967,18 @@ def format_and_template_dataset(
try:
dataset = convert_sharegpt_with_images_to_vlm_format(
dataset,
image_column=vlm_structure["image_column"],
messages_column=vlm_structure["messages_column"],
dataset_name=dataset_name,
progress_callback=progress_callback,
image_column = vlm_structure["image_column"],
messages_column = vlm_structure["messages_column"],
dataset_name = dataset_name,
progress_callback = progress_callback,
)
warnings.append(
"Converted from ShareGPT+image format to standard VLM format"
)
warnings.append("Converted from ShareGPT+image format to standard VLM format")
except Exception as e:
errors.append(f"Failed to convert ShareGPT+image format: {e}")
import traceback
traceback.print_exc()
return {
@ -910,14 +1008,18 @@ def format_and_template_dataset(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
friendly = llm_generate_dataset_warning(
issues, dataset_name=dataset_name, modality="vision",
column_names=columns,
issues,
dataset_name = dataset_name,
modality = "vision",
column_names = columns,
)
except Exception:
pass
errors.append(
friendly or f"Could not auto-detect image/text columns. Found: {vlm_structure}. "
friendly
or f"Could not auto-detect image/text columns. Found: {vlm_structure}. "
)
return {
"dataset": dataset,
@ -933,21 +1035,26 @@ def format_and_template_dataset(
try:
dataset = convert_to_vlm_format(
dataset,
instruction=vlm_instruction,
text_column=vlm_text_column,
image_column=vlm_image_column,
dataset_name=dataset_name,
progress_callback=progress_callback,
instruction = vlm_instruction,
text_column = vlm_text_column,
image_column = vlm_image_column,
dataset_name = dataset_name,
progress_callback = progress_callback,
)
if vlm_instruction:
warnings.append(f"Using user-provided instruction: '{vlm_instruction}'")
warnings.append(
f"Using user-provided instruction: '{vlm_instruction}'"
)
else:
warnings.append("Auto-generated instruction based on dataset analysis")
warnings.append(
"Auto-generated instruction based on dataset analysis"
)
except Exception as e:
errors.append(f"Failed to convert to VLM format: {e}")
import traceback
traceback.print_exc()
return {
@ -987,41 +1094,45 @@ def format_and_template_dataset(
# Step 1: Format the dataset
dataset_info = format_dataset(
dataset,
format_type=format_type,
tokenizer=tokenizer,
auto_detect_custom=auto_detect_custom,
custom_format_mapping=custom_format_mapping,
aliases_for_system=aliases_for_system,
aliases_for_user=aliases_for_user,
aliases_for_assistant=aliases_for_assistant,
batch_size=batch_size,
num_proc=num_proc,
format_type = format_type,
tokenizer = tokenizer,
auto_detect_custom = auto_detect_custom,
custom_format_mapping = custom_format_mapping,
aliases_for_system = aliases_for_system,
aliases_for_user = aliases_for_user,
aliases_for_assistant = aliases_for_assistant,
batch_size = batch_size,
num_proc = num_proc,
)
# Step 2: Apply chat template
# Gemma emits a leading <bos> that must be stripped for text-only chatml/sharegpt.
is_alpaca = format_type == "alpaca" or (format_type == "auto" and dataset_info["detected_format"] == "alpaca")
is_alpaca = format_type == "alpaca" or (
format_type == "auto" and dataset_info["detected_format"] == "alpaca"
)
is_gemma = "gemma" in model_name.lower()
if is_gemma and not dataset_info["is_image"] and not is_alpaca:
remove_bos_prefix = True
template_result = apply_chat_template_to_dataset(
dataset_info=dataset_info,
tokenizer=tokenizer,
model_name=model_name,
custom_prompt_template=custom_prompt_template,
add_eos_token=add_eos_token,
remove_bos_prefix=remove_bos_prefix,
custom_format_mapping=custom_format_mapping,
auto_detect_mapping=auto_detect_mapping,
batch_size=batch_size,
num_proc=num_proc,
dataset_info = dataset_info,
tokenizer = tokenizer,
model_name = model_name,
custom_prompt_template = custom_prompt_template,
add_eos_token = add_eos_token,
remove_bos_prefix = remove_bos_prefix,
custom_format_mapping = custom_format_mapping,
auto_detect_mapping = auto_detect_mapping,
batch_size = batch_size,
num_proc = num_proc,
)
# Step 3: Generate summary
summary = get_dataset_info_summary(dataset_info)
# Combine results
all_warnings = dataset_info.get("warnings", []) + template_result.get("warnings", [])
all_warnings = dataset_info.get("warnings", []) + template_result.get(
"warnings", []
)
all_errors = template_result.get("errors", [])
# If format_dataset returned "unknown" but apply_chat_template rescued

308
studio/backend/utils/datasets/format_conversion.py
View file

@ -12,19 +12,28 @@ import os
from datasets import IterableDataset
from loggers import get_logger
logger = get_logger(__name__)
def standardize_chat_format(
dataset,
tokenizer=None,
aliases_for_system=["system",],
aliases_for_user=["user", "human", "input",],
aliases_for_assistant=["gpt", "assistant", "output",],
batch_size=1000,
num_proc=None,
tokenizer = None,
aliases_for_system = [
"system",
],
aliases_for_user = [
"user",
"human",
"input",
],
aliases_for_assistant = [
"gpt",
"assistant",
"output",
],
batch_size = 1000,
num_proc = None,
):
"""
Our own standardization function that handles BOTH messages and conversations.
@ -67,22 +76,25 @@ def standardize_chat_format(
return dataset # Unexpected structure
keys = list(uniques.keys())
length_first = len(set(uniques[keys[0]]))
length_first = len(set(uniques[keys[0]]))
length_second = len(set(uniques[keys[1]]))
# Determine which is role and which is content
if length_first < length_second:
role_key = keys[0]
role_key = keys[0]
content_key = keys[1]
else:
role_key = keys[1]
role_key = keys[1]
content_key = keys[0]
# Mapping for aliases
aliases_mapping = {}
for x in aliases_for_system: aliases_mapping[x] = "system"
for x in aliases_for_user: aliases_mapping[x] = "user"
for x in aliases_for_assistant: aliases_mapping[x] = "assistant"
for x in aliases_for_system:
aliases_mapping[x] = "system"
for x in aliases_for_user:
aliases_mapping[x] = "user"
for x in aliases_for_assistant:
aliases_mapping[x] = "assistant"
def _standardize_dataset(examples):
convos = examples[chat_column]
@ -109,10 +121,9 @@ def standardize_chat_format(
return {chat_column: all_convos}
dataset_map_kwargs = {
'batched': True,
'batch_size': batch_size,
"batched": True,
"batch_size": batch_size,
}
if not isinstance(dataset, IterableDataset):
@ -123,13 +134,13 @@ def standardize_chat_format(
else:
num_proc = safe_num_proc(num_proc)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Standardizing chat format"
dataset_map_kwargs["num_proc"] = num_proc
dataset_map_kwargs["desc"] = "Standardizing chat format"
return dataset.map(_standardize_dataset, **dataset_map_kwargs)
def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None):
"""
Converts ChatML format (messages OR conversations) to Alpaca format.
Handles both standardized and ShareGPT formats.
@ -142,10 +153,16 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
def _convert(examples):
# Auto-detect which column name is used
chatml_data = examples.get("messages") or examples.get("conversations") or examples.get("texts")
chatml_data = (
examples.get("messages")
or examples.get("conversations")
or examples.get("texts")
)
if chatml_data is None:
raise ValueError("No 'messages' or 'conversations' or 'texts' column found.")
raise ValueError(
"No 'messages' or 'conversations' or 'texts' column found."
)
instructions = []
outputs = []
@ -172,15 +189,11 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
inputs.append("") # Alpaca typically has empty input
outputs.append(output)
return {
"instruction": instructions,
"input": inputs,
"output": outputs
}
return {"instruction": instructions, "input": inputs, "output": outputs}
dataset_map_kwargs = {
'batched': True,
'batch_size': batch_size,
"batched": True,
"batch_size": batch_size,
}
if not isinstance(dataset, IterableDataset):
@ -191,13 +204,13 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
else:
num_proc = safe_num_proc(num_proc)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Converting ChatML to Alpaca format"
dataset_map_kwargs["num_proc"] = num_proc
dataset_map_kwargs["desc"] = "Converting ChatML to Alpaca format"
return dataset.map(_convert, **dataset_map_kwargs)
def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None):
def convert_alpaca_to_chatml(dataset, batch_size = 1000, num_proc = None):
"""
Converts Alpaca format to ChatML format.
@ -222,15 +235,15 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None):
# Build conversation in standard ChatML format
convo = [
{"role": "user", "content": user_content},
{"role": "assistant", "content": output}
{"role": "assistant", "content": output},
]
conversations.append(convo)
return {"conversations": conversations}
dataset_map_kwargs = {
'batched': True,
'batch_size': batch_size,
"batched": True,
"batch_size": batch_size,
}
if not isinstance(dataset, IterableDataset):
@ -241,8 +254,8 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None):
else:
num_proc = safe_num_proc(num_proc)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Converting Alpaca to ChatML format"
dataset_map_kwargs["num_proc"] = num_proc
dataset_map_kwargs["desc"] = "Converting Alpaca to ChatML format"
return dataset.map(_convert, **dataset_map_kwargs)
@ -262,11 +275,11 @@ def _format_eta(seconds):
def convert_to_vlm_format(
dataset,
instruction=None,
text_column="text",
image_column="image",
dataset_name=None,
progress_callback=None,
instruction = None,
text_column = "text",
image_column = "image",
dataset_name = None,
progress_callback = None,
):
"""
Converts simple {image, text} format to VLM messages format.
@ -290,27 +303,31 @@ def convert_to_vlm_format(
def _notify(msg):
"""Send status update to the training overlay if callback is available."""
if progress_callback:
progress_callback(status_message=msg)
progress_callback(status_message = msg)
# Generate smart instruction if not provided
if instruction is None:
instruction_info = generate_smart_vlm_instruction(
dataset,
text_column=text_column,
image_column=image_column,
dataset_name=dataset_name,
text_column = text_column,
image_column = image_column,
dataset_name = dataset_name,
)
instruction = instruction_info["instruction"]
instruction_column = instruction_info.get("instruction_column")
uses_dynamic = instruction_info["uses_dynamic_instruction"]
logger.info(f"📝 Auto-detected instruction type: {instruction_info['instruction_type']}")
logger.info(
f"📝 Auto-detected instruction type: {instruction_info['instruction_type']}"
)
logger.info(f"📝 Confidence: {instruction_info['confidence']:.2f}")
if not uses_dynamic:
logger.info(f"📝 Using instruction: '{instruction}'")
else:
logger.info(f"📝 Using dynamic instructions from column: '{instruction_column}'")
logger.info(
f"📝 Using dynamic instructions from column: '{instruction_column}'"
)
else:
instruction_column = None
uses_dynamic = False
@ -324,13 +341,17 @@ def convert_to_vlm_format(
if image_data.startswith(("http://", "https://")):
import fsspec
from io import BytesIO
with fsspec.open(image_data, "rb", expand=True) as f:
with fsspec.open(image_data, "rb", expand = True) as f:
image_data = Image.open(BytesIO(f.read())).convert("RGB")
elif _image_lookup is not None and image_data in _image_lookup:
# Bare filename → resolve via HF repo lookup
from huggingface_hub import hf_hub_download
local_path = hf_hub_download(
dataset_name, _image_lookup[image_data], repo_type="dataset",
dataset_name,
_image_lookup[image_data],
repo_type = "dataset",
)
image_data = Image.open(local_path).convert("RGB")
else:
@ -340,6 +361,7 @@ def convert_to_vlm_format(
text_data = sample[text_column]
if isinstance(text_data, list) and len(text_data) > 0:
import random
text_data = random.choice(text_data)
# Get instruction (static or dynamic)
@ -354,15 +376,10 @@ def convert_to_vlm_format(
"role": "user",
"content": [
{"type": "text", "text": current_instruction},
{"type": "image", "image": image_data} # PIL object
]
{"type": "image", "image": image_data}, # PIL object
],
},
{
"role": "assistant",
"content": [
{"type": "text", "text": text_data}
]
}
{"role": "assistant", "content": [{"type": "text", "text": text_data}]},
]
# Return dict with messages
@ -370,13 +387,15 @@ def convert_to_vlm_format(
total = len(dataset)
first_image = next(iter(dataset))[image_column]
has_urls = isinstance(first_image, str) and first_image.startswith(("http://", "https://"))
has_urls = isinstance(first_image, str) and first_image.startswith(
("http://", "https://")
)
# ── Bare-filename detection: images stored as filenames (e.g. "img_001.png")
# that don't exist locally. Build a basename→repo_path lookup so we can
# resolve them via hf_hub_download during conversion.
_image_lookup = None
_IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp', '.tiff')
_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff")
if (
not has_urls
and isinstance(first_image, str)
@ -385,18 +404,25 @@ def convert_to_vlm_format(
):
try:
from huggingface_hub import HfApi
_notify("Resolving image filenames from HF repo...")
logger.info(f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup...")
repo_files = HfApi().list_repo_files(dataset_name, repo_type="dataset")
logger.info(
f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup..."
)
repo_files = HfApi().list_repo_files(dataset_name, repo_type = "dataset")
_image_lookup = {
os.path.basename(f): f
for f in repo_files
if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS)
}
if first_image in _image_lookup:
logger.info(f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')")
logger.info(
f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')"
)
else:
logger.info(f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open")
logger.info(
f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open"
)
_image_lookup = None
except Exception as e:
logger.info(f"⚠️ Failed to build HF repo image lookup: {e}")
@ -413,15 +439,19 @@ def convert_to_vlm_format(
num_workers = safe_num_proc()
_notify(f"Probing {PROBE_SIZE} image URLs with {num_workers} workers...")
logger.info(f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers...")
logger.info(
f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers..."
)
probe_samples = [dataset[i] for i in range(PROBE_SIZE)]
probe_ok = 0
probe_fail = 0
probe_start = time.time()
with ThreadPoolExecutor(max_workers=num_workers) as executor:
futures = {executor.submit(_convert_single_sample, s): s for s in probe_samples}
with ThreadPoolExecutor(max_workers = num_workers) as executor:
futures = {
executor.submit(_convert_single_sample, s): s for s in probe_samples
}
for future in as_completed(futures):
try:
future.result()
@ -443,9 +473,12 @@ def convert_to_vlm_format(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
friendly = llm_generate_dataset_warning(
issues, dataset_name=dataset_name, modality="vision",
column_names=[image_column, text_column],
issues,
dataset_name = dataset_name,
modality = "vision",
column_names = [image_column, text_column],
)
except Exception:
pass
@ -471,7 +504,9 @@ def convert_to_vlm_format(
if probe_fail > 0:
info_msg += f" | {fail_rate:.0%} broken URLs will be skipped"
logger.info(f"✅ Probe passed: {probe_ok}/{probe_total} ok, {probe_fail} failed ({fail_rate:.0%}), {throughput:.1f} img/s")
logger.info(
f"✅ Probe passed: {probe_ok}/{probe_total} ok, {probe_fail} failed ({fail_rate:.0%}), {throughput:.1f} img/s"
)
logger.info(f"⏱️ Estimated time for {total:,} samples: ~{eta_str}")
_notify(info_msg)
@ -496,8 +531,11 @@ def convert_to_vlm_format(
batch_end = min(batch_start + batch_size, total)
batch_samples = [dataset[i] for i in range(batch_start, batch_end)]
with ThreadPoolExecutor(max_workers=num_workers) as executor:
futures = {executor.submit(_convert_single_sample, s): i for i, s in enumerate(batch_samples)}
with ThreadPoolExecutor(max_workers = num_workers) as executor:
futures = {
executor.submit(_convert_single_sample, s): i
for i, s in enumerate(batch_samples)
}
batch_results = [None] * len(batch_samples)
for future in as_completed(futures):
idx = futures[future]
@ -506,9 +544,13 @@ def convert_to_vlm_format(
except Exception as e:
failed_count += 1
if failed_count == 1:
print(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
print(
f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}"
)
if failed_count == 1:
logger.info(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
logger.info(
f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}"
)
converted_list.extend(r for r in batch_results if r is not None)
@ -519,11 +561,13 @@ def convert_to_vlm_format(
remaining_time = (total - done) / rate if rate > 0 else 0
eta_str = _format_eta(remaining_time)
progress_msg = f"Downloading images: {done:,}/{total:,} ({done*100//total}%) | ~{eta_str} remaining | {failed_count} skipped"
logger.info(f" [{done}/{total}] {rate:.1f} img/s, {failed_count} failed, ETA {eta_str}")
logger.info(
f" [{done}/{total}] {rate:.1f} img/s, {failed_count} failed, ETA {eta_str}"
)
_notify(progress_msg)
else:
# Sequential conversion for local/embedded images (fast, no I/O bottleneck)
pbar = tqdm(dataset, total=total, desc="Converting VLM samples", unit="sample")
pbar = tqdm(dataset, total = total, desc = "Converting VLM samples", unit = "sample")
for sample in pbar:
try:
converted_list.append(_convert_single_sample(sample))
@ -534,13 +578,17 @@ def convert_to_vlm_format(
print(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
if failed_count == 1:
# Log the first failure to aid debugging
logger.info(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
pbar.set_postfix(ok=len(converted_list), failed=failed_count, refresh=False)
logger.info(
f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}"
)
pbar.set_postfix(ok = len(converted_list), failed = failed_count, refresh = False)
pbar.close()
if failed_count > 0:
fail_rate = failed_count / total
logger.info(f"⚠️ Skipped {failed_count}/{total} ({fail_rate:.0%}) samples with broken/unreachable images")
logger.info(
f"⚠️ Skipped {failed_count}/{total} ({fail_rate:.0%}) samples with broken/unreachable images"
)
# For datasets that skipped the probe (small URL datasets), check fail rate now
if has_urls and fail_rate >= MAX_FAIL_RATE:
issues = [
@ -550,9 +598,12 @@ def convert_to_vlm_format(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
friendly = llm_generate_dataset_warning(
issues, dataset_name=dataset_name, modality="vision",
column_names=[image_column, text_column],
issues,
dataset_name = dataset_name,
modality = "vision",
column_names = [image_column, text_column],
)
except Exception:
pass
@ -573,14 +624,18 @@ def convert_to_vlm_format(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
friendly = llm_generate_dataset_warning(
issues, dataset_name=dataset_name, modality="vision",
column_names=[image_column, text_column],
issues,
dataset_name = dataset_name,
modality = "vision",
column_names = [image_column, text_column],
)
except Exception:
pass
raise ValueError(
friendly or (
friendly
or (
f"All {total} samples failed during VLM conversion — no usable images found. "
"This dataset may contain only image URLs that are no longer accessible."
)
@ -595,10 +650,10 @@ def convert_to_vlm_format(
def convert_sharegpt_with_images_to_vlm_format(
dataset,
image_column="image",
messages_column="conversations",
dataset_name=None,
progress_callback=None,
image_column = "image",
messages_column = "conversations",
dataset_name = None,
progress_callback = None,
):
"""
Converts ShareGPT/ChatML datasets that have a separate image column and
@ -619,16 +674,18 @@ def convert_sharegpt_with_images_to_vlm_format(
from PIL import Image
from tqdm import tqdm
_IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp', '.tiff')
_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff")
_ROLE_MAP = {
"human": "user", "user": "user",
"gpt": "assistant", "assistant": "assistant",
"human": "user",
"user": "user",
"gpt": "assistant",
"assistant": "assistant",
"system": "system",
}
def _notify(msg):
if progress_callback:
progress_callback(status_message=msg)
progress_callback(status_message = msg)
# ── Resolve image loading strategy (same 3-tier as convert_to_vlm_format) ──
total = len(dataset)
@ -643,9 +700,12 @@ def convert_sharegpt_with_images_to_vlm_format(
):
try:
from huggingface_hub import HfApi
_notify("Resolving image filenames from HF repo...")
logger.info(f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup...")
repo_files = HfApi().list_repo_files(dataset_name, repo_type="dataset")
logger.info(
f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup..."
)
repo_files = HfApi().list_repo_files(dataset_name, repo_type = "dataset")
_image_lookup = {
os.path.basename(f): f
for f in repo_files
@ -656,9 +716,13 @@ def convert_sharegpt_with_images_to_vlm_format(
if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS):
_image_lookup[f] = f
if first_image in _image_lookup:
logger.info(f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')")
logger.info(
f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')"
)
else:
logger.info(f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open")
logger.info(
f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open"
)
_image_lookup = None
except Exception as e:
logger.info(f"⚠️ Failed to build HF repo image lookup: {e}")
@ -666,25 +730,32 @@ def convert_sharegpt_with_images_to_vlm_format(
def _resolve_image(image_data):
"""Resolve image data to a PIL Image object."""
if hasattr(image_data, 'size') and hasattr(image_data, 'mode'):
if hasattr(image_data, "size") and hasattr(image_data, "mode"):
return image_data # Already PIL
if isinstance(image_data, str):
if image_data.startswith(("http://", "https://")):
import fsspec
from io import BytesIO
with fsspec.open(image_data, "rb", expand=True) as f:
with fsspec.open(image_data, "rb", expand = True) as f:
return Image.open(BytesIO(f.read())).convert("RGB")
elif _image_lookup is not None and image_data in _image_lookup:
from huggingface_hub import hf_hub_download
local_path = hf_hub_download(
dataset_name, _image_lookup[image_data], repo_type="dataset",
dataset_name,
_image_lookup[image_data],
repo_type = "dataset",
)
return Image.open(local_path).convert("RGB")
else:
return Image.open(image_data).convert("RGB")
if isinstance(image_data, dict) and ("bytes" in image_data or "path" in image_data):
if isinstance(image_data, dict) and (
"bytes" in image_data or "path" in image_data
):
if image_data.get("bytes"):
from io import BytesIO
return Image.open(BytesIO(image_data["bytes"])).convert("RGB")
if image_data.get("path"):
return Image.open(image_data["path"]).convert("RGB")
@ -726,7 +797,7 @@ def convert_sharegpt_with_images_to_vlm_format(
converted_list = []
failed_count = 0
pbar = tqdm(dataset, total=total, desc="Converting ShareGPT+image", unit="sample")
pbar = tqdm(dataset, total = total, desc = "Converting ShareGPT+image", unit = "sample")
for sample in pbar:
try:
converted_list.append(_convert_single_sample(sample))
@ -734,11 +805,13 @@ def convert_sharegpt_with_images_to_vlm_format(
failed_count += 1
if failed_count == 1:
logger.info(f"⚠️ First conversion failure: {type(e).__name__}: {e}")
pbar.set_postfix(ok=len(converted_list), failed=failed_count, refresh=False)
pbar.set_postfix(ok = len(converted_list), failed = failed_count, refresh = False)
pbar.close()
if failed_count > 0:
logger.info(f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples")
logger.info(
f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples"
)
if len(converted_list) == 0:
raise ValueError(
@ -764,7 +837,9 @@ def convert_llava_to_vlm_format(dataset):
"""
from PIL import Image
logger.info(f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format...")
logger.info(
f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format..."
)
def _convert_single_sample(sample):
"""Convert a single llava sample to standard VLM format."""
@ -787,10 +862,12 @@ def convert_llava_to_vlm_format(dataset):
if isinstance(pil_image, str):
pil_image = Image.open(pil_image).convert("RGB")
new_content.append({
"type": "image",
"image": pil_image # Actual PIL object
})
new_content.append(
{
"type": "image",
"image": pil_image, # Actual PIL object
}
)
else:
# No index, try to use first image
if len(images) > 0:
@ -798,22 +875,13 @@ def convert_llava_to_vlm_format(dataset):
if isinstance(pil_image, str):
pil_image = Image.open(pil_image).convert("RGB")
new_content.append({
"type": "image",
"image": pil_image
})
new_content.append({"type": "image", "image": pil_image})
elif item["type"] == "text":
# Keep text as-is (only type + text)
new_content.append({
"type": "text",
"text": item.get("text", "")
})
new_content.append({"type": "text", "text": item.get("text", "")})
new_messages.append({
"role": msg["role"],
"content": new_content
})
new_messages.append({"role": msg["role"], "content": new_content})
return {"messages": new_messages}

300
studio/backend/utils/datasets/format_detection.py
View file

@ -13,7 +13,10 @@ import re
def _keyword_in_column(keyword: str, col_name: str) -> bool:
"""Word-boundary keyword match to avoid false positives like 'pic' in 'topic'."""
return re.search(r'\b' + re.escape(keyword) + r'\b', col_name, re.IGNORECASE) is not None
return (
re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE)
is not None
)
def detect_dataset_format(dataset):
@ -37,7 +40,7 @@ def detect_dataset_format(dataset):
"format": "alpaca",
"chat_column": None,
"needs_standardization": False,
"sample_keys": []
"sample_keys": [],
}
# Check for chat-based formats (messages or conversations)
@ -65,7 +68,7 @@ def detect_dataset_format(dataset):
"format": "sharegpt",
"chat_column": chat_column,
"needs_standardization": True,
"sample_keys": list(msg_keys)
"sample_keys": list(msg_keys),
}
# ChatML uses "role" and "content"
@ -74,7 +77,7 @@ def detect_dataset_format(dataset):
"format": "chatml",
"chat_column": chat_column,
"needs_standardization": False,
"sample_keys": list(msg_keys)
"sample_keys": list(msg_keys),
}
# Unknown structure but has chat column
@ -83,7 +86,7 @@ def detect_dataset_format(dataset):
"format": "unknown",
"chat_column": chat_column,
"needs_standardization": None,
"sample_keys": list(msg_keys)
"sample_keys": list(msg_keys),
}
except Exception as e:
return {
@ -91,7 +94,7 @@ def detect_dataset_format(dataset):
"chat_column": chat_column,
"needs_standardization": None,
"sample_keys": [],
"error": str(e)
"error": str(e),
}
# No recognized format
@ -99,7 +102,7 @@ def detect_dataset_format(dataset):
"format": "unknown",
"chat_column": None,
"needs_standardization": None,
"sample_keys": []
"sample_keys": [],
}
@ -120,49 +123,86 @@ def detect_custom_format_heuristic(dataset):
# Keywords
assistant_words = [
'output', 'answer', 'response', 'assistant', 'completion',
'expected', 'recommendation', 'reply', 'result', 'target',
'solution', 'explanation', 'solve'
"output",
"answer",
"response",
"assistant",
"completion",
"expected",
"recommendation",
"reply",
"result",
"target",
"solution",
"explanation",
"solve",
]
# Split into high/low priority
user_words_high_priority = [
'input', 'question', 'query', 'prompt', 'instruction',
'request', 'snippet', 'user', 'text',
'problem', 'exercise'
"input",
"question",
"query",
"prompt",
"instruction",
"request",
"snippet",
"user",
"text",
"problem",
"exercise",
]
user_words_low_priority = ['task'] # Ambiguous - can be user OR system
user_words_low_priority = ["task"] # Ambiguous - can be user OR system
user_words = user_words_high_priority + user_words_low_priority
system_words = [
'system', 'context', 'description', 'persona', 'role',
'template', 'task' # Also in system
"system",
"context",
"description",
"persona",
"role",
"template",
"task", # Also in system
]
# Metadata columns to ignore
metadata_exact_match = {
'id', 'idx', 'index', 'key', 'timestamp', 'date',
'metadata', 'source', 'kind', 'type', 'category',
'score', 'label', 'tag', 'inference_mode'
"id",
"idx",
"index",
"key",
"timestamp",
"date",
"metadata",
"source",
"kind",
"type",
"category",
"score",
"label",
"tag",
"inference_mode",
}
metadata_prefix_patterns = [
'problem_type', 'problem_source',
'generation_model', 'pass_rate',
"problem_type",
"problem_source",
"generation_model",
"pass_rate",
]
priority_patterns = {
'generated': 100,
'gen_': 90,
'model_': 80,
'predicted': 70,
'completion': 60,
"generated": 100,
"gen_": 90,
"model_": 80,
"predicted": 70,
"completion": 60,
}
def has_keyword(col_name, keywords):
"""Check if any keyword appears in column name."""
col_lower = col_name.lower()
col_normalized = col_lower.replace('_', '').replace('-', '').replace(' ', '')
col_normalized = col_lower.replace("_", "").replace("-", "").replace(" ", "")
for keyword in keywords:
if keyword in col_lower or keyword in col_normalized:
@ -180,13 +220,16 @@ def detect_custom_format_heuristic(dataset):
return True
for pattern in metadata_prefix_patterns:
if col_lower.startswith(pattern.split('_')[0] + '_') and col_lower != pattern:
if '_' in col_lower:
prefix = col_lower.split('_')[0]
if prefix in ['generation', 'pass', 'inference']:
if (
col_lower.startswith(pattern.split("_")[0] + "_")
and col_lower != pattern
):
if "_" in col_lower:
prefix = col_lower.split("_")[0]
if prefix in ["generation", "pass", "inference"]:
return True
if len(col_lower) <= 2 and not col_lower in ['qa', 'q', 'a']:
if len(col_lower) <= 2 and not col_lower in ["qa", "q", "a"]:
return True
return False
@ -221,16 +264,18 @@ def detect_custom_format_heuristic(dataset):
score += 10
# Penalize ambiguous keywords when scoring for user
if role_type == 'user':
if role_type == "user":
col_lower = col_name.lower()
# If column is ONLY "task" (or task_xxx), give it lower priority for user role
if 'task' in col_lower and not any(kw in col_lower for kw in user_words_high_priority):
if "task" in col_lower and not any(
kw in col_lower for kw in user_words_high_priority
):
score -= 15 # Significant penalty so other user columns win
priority_bonus = get_priority_score(col_name)
score += priority_bonus
if role_type in ['assistant', 'user']:
if role_type in ["assistant", "user"]:
avg_length = get_content_length(col_name)
if num_candidates > 1:
@ -256,20 +301,24 @@ def detect_custom_format_heuristic(dataset):
content_columns = [col for col in all_columns if not is_metadata(col)]
# Count candidates first
assistant_potential = [col for col in content_columns if has_keyword(col, assistant_words)]
assistant_potential = [
col for col in content_columns if has_keyword(col, assistant_words)
]
user_potential = [col for col in content_columns if has_keyword(col, user_words)]
# STEP 1: Find best ASSISTANT column
assistant_candidates = []
for col in assistant_potential:
score = score_column(col, assistant_words, 'assistant', len(assistant_potential))
score = score_column(
col, assistant_words, "assistant", len(assistant_potential)
)
if score > 0:
assistant_candidates.append((col, score))
if assistant_candidates:
assistant_candidates.sort(key=lambda x: x[1], reverse=True)
assistant_candidates.sort(key = lambda x: x[1], reverse = True)
assistant_col = assistant_candidates[0][0]
mapping[assistant_col] = 'assistant'
mapping[assistant_col] = "assistant"
else:
assistant_col = None
@ -278,14 +327,14 @@ def detect_custom_format_heuristic(dataset):
for col in user_potential:
if col == assistant_col:
continue
score = score_column(col, user_words, 'user', len(user_potential))
score = score_column(col, user_words, "user", len(user_potential))
if score > 0:
user_candidates.append((col, score))
if user_candidates:
user_candidates.sort(key=lambda x: x[1], reverse=True)
user_candidates.sort(key = lambda x: x[1], reverse = True)
user_col = user_candidates[0][0]
mapping[user_col] = 'user'
mapping[user_col] = "user"
else:
user_col = None
@ -296,7 +345,7 @@ def detect_custom_format_heuristic(dataset):
for col in remaining_columns:
if has_keyword(col, system_words):
# Found a system match in remaining columns
mapping[col] = 'system'
mapping[col] = "system"
system_col = col
break
@ -309,22 +358,22 @@ def detect_custom_format_heuristic(dataset):
# If no strong keyword match, decide based on what's missing
if not has_keyword(remaining_col, user_words + assistant_words):
mapping[remaining_col] = 'system'
mapping[remaining_col] = "system"
elif user_col is None:
# No user column yet, assign this as user
mapping[remaining_col] = 'user'
mapping[remaining_col] = "user"
else:
# Already have user + assistant, treat as system context
mapping[remaining_col] = 'system'
mapping[remaining_col] = "system"
# VALIDATION: Ensure we have at least user + assistant
has_user = any(role == 'user' for role in mapping.values())
has_assistant = any(role == 'assistant' for role in mapping.values())
has_user = any(role == "user" for role in mapping.values())
has_assistant = any(role == "assistant" for role in mapping.values())
if not has_user and len(remaining_columns) > 0:
for col in remaining_columns:
if col not in mapping:
mapping[col] = 'user'
mapping[col] = "user"
has_user = True
break
@ -358,14 +407,27 @@ def detect_multimodal_dataset(dataset):
# Keywords that indicate image data
image_keywords = [
'image', 'img', 'pixel',
'jpg', 'jpeg', 'png', 'webp', 'bmp', 'gif', 'tiff', 'svg',
'photo', 'pic', 'picture', 'visual',
'file_name', 'filename',
"image",
"img",
"pixel",
"jpg",
"jpeg",
"png",
"webp",
"bmp",
"gif",
"tiff",
"svg",
"photo",
"pic",
"picture",
"visual",
"file_name",
"filename",
]
# Keywords that indicate audio data
audio_keywords = ['audio', 'speech', 'wav', 'waveform', 'sound']
audio_keywords = ["audio", "speech", "wav", "waveform", "sound"]
multimodal_columns = []
audio_columns = []
@ -419,7 +481,7 @@ def detect_multimodal_dataset(dataset):
# Detect text column for audio datasets
detected_text_col = None
if audio_columns:
text_keywords = ['text', 'sentence', 'transcript', 'transcription', 'label']
text_keywords = ["text", "sentence", "transcript", "transcription", "label"]
for col_name in column_names:
if col_name.lower() in text_keywords:
detected_text_col = col_name
@ -430,7 +492,7 @@ def detect_multimodal_dataset(dataset):
# Detect speaker_id column for TTS datasets (CSM, Orpheus, Spark)
detected_speaker_col = None
if audio_columns:
speaker_keywords = ['source', 'speaker', 'speaker_id']
speaker_keywords = ["source", "speaker", "speaker_id"]
for col_name in column_names:
if col_name.lower() in speaker_keywords:
detected_speaker_col = col_name
@ -456,6 +518,7 @@ def _is_image_value(value) -> bool:
# PIL Image instance
try:
from PIL.Image import Image as PILImage
if isinstance(value, PILImage):
return True
except ImportError:
@ -470,7 +533,9 @@ def _is_image_value(value) -> bool:
if "bytes" in value and "path" in value:
# Check path extension to exclude audio files
path = value.get("path") or ""
if isinstance(path, str) and any(path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS):
if isinstance(path, str) and any(
path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS
):
return False
return True
@ -479,11 +544,13 @@ def _is_image_value(value) -> bool:
return _has_image_header(value)
# String that looks like an image file path or URL
_IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp', '.tiff', '.svg')
_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".svg")
if isinstance(value, str) and len(value) < 1000:
lower = value.strip().lower()
# Image URL (http://... ending in image extension)
if lower.startswith(("http://", "https://")) and any(lower.split("?")[0].endswith(ext) for ext in _IMAGE_EXTS):
if lower.startswith(("http://", "https://")) and any(
lower.split("?")[0].endswith(ext) for ext in _IMAGE_EXTS
):
return True
# Image file path (relative or absolute path ending in image extension)
if any(lower.endswith(ext) for ext in _IMAGE_EXTS):
@ -493,7 +560,15 @@ def _is_image_value(value) -> bool:
_AUDIO_EXTENSIONS = (
".wav", ".mp3", ".flac", ".ogg", ".opus", ".m4a", ".aac", ".wma", ".webm",
".wav",
".mp3",
".flac",
".ogg",
".opus",
".m4a",
".aac",
".wma",
".webm",
)
@ -509,7 +584,9 @@ def _is_audio_value(value) -> bool:
# Undecoded/streaming → {"bytes": b"...", "path": "some.wav"}
if "bytes" in value or "path" in value:
path = value.get("path") or ""
if isinstance(path, str) and any(path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS):
if isinstance(path, str) and any(
path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS
):
return True
return False
@ -520,19 +597,19 @@ def _has_image_header(data: bytes) -> bool:
if len(data) < 4:
return False
# JPEG
if data[:2] == b'\xff\xd8':
if data[:2] == b"\xff\xd8":
return True
# PNG
if data[:4] == b'\x89PNG':
if data[:4] == b"\x89PNG":
return True
# GIF
if data[:3] == b'GIF':
if data[:3] == b"GIF":
return True
# WebP
if data[:4] == b'RIFF' and len(data) >= 12 and data[8:12] == b'WEBP':
if data[:4] == b"RIFF" and len(data) >= 12 and data[8:12] == b"WEBP":
return True
# BMP
if data[:2] == b'BM':
if data[:2] == b"BM":
return True
return False
@ -568,10 +645,13 @@ def detect_vlm_dataset_structure(dataset):
if isinstance(content, list) and len(content) > 0:
if isinstance(content[0], dict) and "type" in content[0]:
# Check for llava format
has_index = any('index' in item for item in content if isinstance(item, dict))
has_images_column = 'images' in column_names
has_index = any(
"index" in item
for item in content
if isinstance(item, dict)
)
has_images_column = "images" in column_names
if has_index and has_images_column:
return {
@ -583,7 +663,11 @@ def detect_vlm_dataset_structure(dataset):
}
# Standard VLM format
has_image = any('image' in item for item in content if isinstance(item, dict))
has_image = any(
"image" in item
for item in content
if isinstance(item, dict)
)
if has_image:
return {
"format": "vlm_messages",
@ -637,26 +721,65 @@ def detect_vlm_dataset_structure(dataset):
# Define metadata patterns to EXCLUDE
metadata_patterns = {
'suffixes': ['_id', '_url', '_name', '_filename', '_uri', '_link', '_key', '_index'],
'prefixes': ['id_', 'url_', 'name_', 'filename_', 'uri_', 'link_', 'key_', 'index_'],
"suffixes": [
"_id",
"_url",
"_name",
"_filename",
"_uri",
"_link",
"_key",
"_index",
],
"prefixes": [
"id_",
"url_",
"name_",
"filename_",
"uri_",
"link_",
"key_",
"index_",
],
}
# Image-related keywords
image_keywords = ['image', 'img', 'photo', 'picture', 'pic', 'visual', 'scan', 'file_name', 'filename']
image_keywords = [
"image",
"img",
"photo",
"picture",
"pic",
"visual",
"scan",
"file_name",
"filename",
]
# Text-related keywords
text_keywords = ['text', 'caption', 'captions', 'description', 'answer', 'output', 'response', 'label']
text_keywords = [
"text",
"caption",
"captions",
"description",
"answer",
"output",
"response",
"label",
]
def is_metadata_column(col_name):
"""Check if column name looks like metadata."""
col_lower = col_name.lower()
# Check suffixes
if any(col_lower.endswith(suffix) for suffix in metadata_patterns['suffixes']):
if any(col_lower.endswith(suffix) for suffix in metadata_patterns["suffixes"]):
return True
# Check prefixes
if any(col_lower.startswith(prefix) for prefix in metadata_patterns['prefixes']):
if any(
col_lower.startswith(prefix) for prefix in metadata_patterns["prefixes"]
):
return True
return False
@ -664,11 +787,13 @@ def detect_vlm_dataset_structure(dataset):
def _score_image_candidate(col, sample_value):
"""Score a candidate image column by how resolvable its value is."""
# PIL Image object (highest priority - already loaded)
if hasattr(sample_value, 'size') and hasattr(sample_value, 'mode'):
if hasattr(sample_value, "size") and hasattr(sample_value, "mode"):
return 100
# Dict with image data (bytes/path from HF Image feature)
if isinstance(sample_value, dict) and ('bytes' in sample_value or 'path' in sample_value):
if isinstance(sample_value, dict) and (
"bytes" in sample_value or "path" in sample_value
):
return 75
if isinstance(sample_value, str):
@ -693,13 +818,16 @@ def detect_vlm_dataset_structure(dataset):
# Local file — check it exists
if not sample_value.startswith(("http://", "https://")):
return os.path.exists(sample_value) # bare filenames return False here, that's OK
return os.path.exists(
sample_value
) # bare filenames return False here, that's OK
# URL — quick HEAD request with short timeout
try:
import urllib.request
req = urllib.request.Request(sample_value, method="HEAD")
resp = urllib.request.urlopen(req, timeout=3)
req = urllib.request.Request(sample_value, method = "HEAD")
resp = urllib.request.urlopen(req, timeout = 3)
return resp.status < 400
except Exception:
return False
@ -732,7 +860,7 @@ def detect_vlm_dataset_structure(dataset):
if not candidates:
return None
candidates.sort(key=lambda x: x[1], reverse=True)
candidates.sort(key = lambda x: x[1], reverse = True)
# Single candidate or top candidate is PIL/dict — no probing needed
if len(candidates) == 1 or candidates[0][1] >= 75:
@ -766,14 +894,18 @@ def detect_vlm_dataset_structure(dataset):
# Longer text = higher priority (likely content, not just a label)
priority = min(len(sample_value), 1000) # Cap at 1000
candidates.append((col, priority))
elif isinstance(sample_value, list) and len(sample_value) > 0 and isinstance(sample_value[0], str):
elif (
isinstance(sample_value, list)
and len(sample_value) > 0
and isinstance(sample_value[0], str)
):
# List of strings (e.g. captions list) — lower priority than plain strings
priority = min(len(sample_value[0]), 1000) // 2
candidates.append((col, priority))
# Return highest priority candidate
if candidates:
candidates.sort(key=lambda x: x[1], reverse=True)
candidates.sort(key = lambda x: x[1], reverse = True)
return candidates[0][0]
return None

207
studio/backend/utils/datasets/llm_assist.py
View file

@ -42,33 +42,45 @@ def precache_helper_gguf():
return
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
variant = os.environ.get(
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
)
try:
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import disable_progress_bars, enable_progress_bars
disable_progress_bars()
logging.getLogger("huggingface_hub").setLevel(logging.WARNING)
# Find the GGUF file matching the variant
api = HfApi()
files = api.list_repo_files(repo, repo_type="model")
files = api.list_repo_files(repo, repo_type = "model")
gguf_files = [f for f in files if f.endswith(".gguf")]
# Find all GGUF files matching the variant (may be split into shards)
variant_lower = variant.lower().replace("-", "_")
matching = sorted(
f for f in gguf_files
if variant_lower in f.lower().replace("-", "_")
f for f in gguf_files if variant_lower in f.lower().replace("-", "_")
)
if matching:
logger.info(f"Pre-caching helper GGUF: {repo}/{matching[0]}"
+ (f" (+{len(matching) - 1} shards)" if len(matching) > 1 else ""))
logger.info(
f"Pre-caching helper GGUF: {repo}/{matching[0]}"
+ (f" (+{len(matching) - 1} shards)" if len(matching) > 1 else "")
)
for target in matching:
hf_hub_download(repo_id=repo, filename=target)
hf_hub_download(repo_id = repo, filename = target)
logger.info(f"Helper GGUF cached: {len(matching)} file(s)")
else:
logger.warning(f"No GGUF matching variant '{variant}' in {repo}")
except Exception as e:
logger.warning(f"Failed to pre-cache helper GGUF: {e}")
finally:
try:
enable_progress_bars()
except Exception as e:
pass
def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
@ -81,7 +93,9 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
return None
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
variant = os.environ.get(
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
)
backend = None
try:
@ -89,15 +103,14 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
backend = LlamaCppBackend()
logger.info(f"Loading helper model: {repo} ({variant})")
print(f"🤖 Loading helper model: {repo} ({variant})...")
ok = backend.load_model(
hf_repo=repo,
hf_variant=variant,
model_identifier=f"helper:{repo}:{variant}",
is_vision=False,
n_ctx=2048,
n_gpu_layers=-1,
hf_repo = repo,
hf_variant = variant,
model_identifier = f"helper:{repo}:{variant}",
is_vision = False,
n_ctx = 2048,
n_gpu_layers = -1,
)
if not ok:
logger.warning("Helper model failed to start")
@ -106,12 +119,12 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
messages = [{"role": "user", "content": prompt}]
cumulative = ""
for text in backend.generate_chat_completion(
messages=messages,
temperature=0.1,
top_p=0.9,
top_k=20,
max_tokens=max_tokens,
repetition_penalty=1.0,
messages = messages,
temperature = 0.1,
top_p = 0.9,
top_k = 20,
max_tokens = max_tokens,
repetition_penalty = 1.0,
):
cumulative = text # cumulative — last value is full text
@ -127,7 +140,7 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
if backend is not None:
try:
backend.unload_model()
print("🤖 Helper model unloaded")
logger.info("Helper model unloaded")
except Exception:
pass
@ -176,7 +189,7 @@ def llm_generate_vlm_instruction(
"Respond with ONLY the instruction sentence, nothing else."
)
result = _run_with_helper(prompt, max_tokens=100)
result = _run_with_helper(prompt, max_tokens = 100)
if not result:
return None
@ -187,7 +200,7 @@ def llm_generate_vlm_instruction(
logger.warning(f"Helper model returned unusable instruction: {instruction!r}")
return None
print(f"🤖 LLM-generated instruction: {instruction}")
logger.info(f"LLM-generated instruction: {instruction}")
return {
"instruction": instruction,
"confidence": 0.85,
@ -231,7 +244,7 @@ def llm_classify_columns(
'Example: {"question": "user", "answer": "assistant", "id": "metadata"}'
)
result = _run_with_helper(prompt, max_tokens=200)
result = _run_with_helper(prompt, max_tokens = 200)
if not result:
return None
@ -248,6 +261,7 @@ def llm_classify_columns(
except json.JSONDecodeError:
# Try to find JSON object in the response
import re
match = re.search(r"\{[^}]+\}", text)
if match:
try:
@ -266,7 +280,11 @@ def llm_classify_columns(
valid_roles = {"user", "assistant", "system", "metadata"}
cleaned = {}
for col, role in mapping.items():
if col in column_names and isinstance(role, str) and role.lower() in valid_roles:
if (
col in column_names
and isinstance(role, str)
and role.lower() in valid_roles
):
cleaned[col] = role.lower()
if not cleaned:
@ -278,7 +296,7 @@ def llm_classify_columns(
logger.warning(f"Helper model mapping missing user/assistant: {cleaned}")
return None
print(f"🤖 LLM-classified columns: {cleaned}")
logger.info(f"LLM-classified columns: {cleaned}")
return cleaned
@ -319,7 +337,7 @@ def llm_generate_dataset_warning(
"Keep it under 3 sentences. Be specific about the dataset."
)
result = _run_with_helper(prompt, max_tokens=200)
result = _run_with_helper(prompt, max_tokens = 200)
if not result:
return None
@ -328,7 +346,7 @@ def llm_generate_dataset_warning(
if len(warning) < 10 or len(warning) > 500:
return None
print(f"🤖 LLM-generated warning: {warning}")
logger.info(f"LLM-generated warning: {warning}")
return warning
@ -369,18 +387,16 @@ def _parse_json_response(text: str) -> Optional[dict]:
return None
def _generate_with_backend(
backend, messages: list[dict], max_tokens: int = 512
) -> str:
def _generate_with_backend(backend, messages: list[dict], max_tokens: int = 512) -> str:
"""Run one chat completion on an already-loaded backend. Returns raw text."""
cumulative = ""
for text in backend.generate_chat_completion(
messages=messages,
temperature=0.1,
top_p=0.9,
top_k=20,
max_tokens=max_tokens,
repetition_penalty=1.0,
messages = messages,
temperature = 0.1,
top_p = 0.9,
top_k = 20,
max_tokens = max_tokens,
repetition_penalty = 1.0,
):
cumulative = text
return cumulative.strip()
@ -398,7 +414,7 @@ def fetch_hf_dataset_card(
try:
from huggingface_hub import DatasetCard
card = DatasetCard.load(dataset_name, token=hf_token)
card = DatasetCard.load(dataset_name, token = hf_token)
readme = card.text or ""
# Truncate at sentence boundary
@ -413,14 +429,21 @@ def fetch_hf_dataset_card(
metadata = {}
if card.data:
for key in (
"task_categories", "task_ids", "language",
"size_categories", "tags", "license", "pretty_name",
"task_categories",
"task_ids",
"language",
"size_categories",
"tags",
"license",
"pretty_name",
):
val = getattr(card.data, key, None)
if val is not None:
metadata[key] = val
logger.info(f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields")
logger.info(
f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields"
)
return readme, metadata
except Exception as e:
@ -447,30 +470,31 @@ def _run_multi_pass_advisor(
return None
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
variant = os.environ.get(
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
)
backend = None
try:
from core.inference.llama_cpp import LlamaCppBackend
backend = LlamaCppBackend()
print(f"🤖 Loading advisor model: {repo} ({variant})...")
logger.info(f"Loading advisor model: {repo} ({variant})")
t0 = time.monotonic()
ok = backend.load_model(
hf_repo=repo,
hf_variant=variant,
model_identifier=f"advisor:{repo}:{variant}",
is_vision=False,
n_ctx=2048,
n_gpu_layers=-1,
hf_repo = repo,
hf_variant = variant,
model_identifier = f"advisor:{repo}:{variant}",
is_vision = False,
n_ctx = 2048,
n_gpu_layers = -1,
)
if not ok:
logger.warning("Advisor model failed to start")
return None
print(f"🤖 Advisor model loaded in {time.monotonic() - t0:.1f}s")
logger.info(f"Advisor model loaded in {time.monotonic() - t0:.1f}s")
# ── Format samples ──
samples_text = ""
for i, row in enumerate(samples[:5], 1):
@ -478,8 +502,9 @@ def _run_multi_pass_advisor(
samples_text += f"Row {i}:\n" + "\n".join(parts) + "\n"
metadata_str = (
json.dumps(dataset_metadata, indent=2, default=str)[:500]
if dataset_metadata else "N/A"
json.dumps(dataset_metadata, indent = 2, default = str)[:500]
if dataset_metadata
else "N/A"
)
card_excerpt = (dataset_card or "")[:1200] or "N/A"
@ -489,13 +514,14 @@ def _run_multi_pass_advisor(
if model_name:
try:
from utils.models.model_config import load_model_config
config = load_model_config(model_name, use_auth=True, token=hf_token)
config = load_model_config(model_name, use_auth = True, token = hf_token)
archs = getattr(config, "architectures", [])
if archs and "Gemma3nForConditionalGeneration" in archs:
is_gemma_3n = True
except Exception:
is_gemma_3n = "gemma-3n" in model_name.lower()
if model_type == "audio" and not is_gemma_3n:
target_hints = (
"\n\nHINT: The user is training an AUDIO model. The dataset MUST contain "
@ -514,7 +540,7 @@ def _run_multi_pass_advisor(
)
# ── Pass 1: Classify ──
print("🤖 Pass 1: Classifying dataset...", flush=True)
logger.info("Pass 1: Classifying dataset...")
t1 = time.monotonic()
messages1 = [
{
@ -559,9 +585,9 @@ def _run_multi_pass_advisor(
Respond with ONLY the JSON object. No markdown, no explanation."""),
},
]
raw1 = _generate_with_backend(backend, messages1, max_tokens=256)
raw1 = _generate_with_backend(backend, messages1, max_tokens = 256)
pass1 = _parse_json_response(raw1)
print(f"🤖 Pass 1 done ({time.monotonic() - t1:.1f}s): {pass1}", flush=True)
logger.info(f"Pass 1 done ({time.monotonic() - t1:.1f}s): {pass1}")
if not pass1:
logger.warning(f"Advisor Pass 1 failed to produce JSON: {raw1[:200]}")
@ -580,7 +606,9 @@ def _run_multi_pass_advisor(
}
# ── Pass 2: Map columns to roles ──
print("🤖 Pass 2: Mapping columns to roles...", flush=True)
logger.info("Pass 2: Mapping columns to roles...")
t2 = time.monotonic()
messages2 = [
{
@ -592,13 +620,13 @@ def _run_multi_pass_advisor(
'- "user" = This column contains INPUT that the model will receive as a prompt.\n'
'- "assistant" = This column contains OUTPUT that the model should learn to generate.\n\n'
"CRITICAL RULES:\n"
"1. There MUST be at least one column assigned to \"user\" AND at least one "
"column assigned to \"assistant\". Never assign all columns to the same role.\n"
'1. There MUST be at least one column assigned to "user" AND at least one '
'column assigned to "assistant". Never assign all columns to the same role.\n'
"2. The column that contains the TARGET or OUTPUT or ANSWER or LABEL must "
"ALWAYS be assigned to \"assistant\". This is the thing the model should learn "
'ALWAYS be assigned to "assistant". This is the thing the model should learn '
"to produce.\n"
"3. The columns that contain the SOURCE or INPUT or CONTEXT or QUESTION must "
"be assigned to \"user\". This is what the model receives.\n"
'be assigned to "user". This is what the model receives.\n'
'4. Metadata columns like "id", "index", "source", "url", "date" should be '
'set to "skip".\n\n'
"You must respond with ONLY a valid JSON object."
@ -611,7 +639,7 @@ def _run_multi_pass_advisor(
Here is a dataset that has been classified:
CLASSIFICATION:
{json.dumps(pass1, indent=2)}
{json.dumps(pass1, indent = 2)}
COLUMNS AVAILABLE: {columns}
@ -659,9 +687,9 @@ def _run_multi_pass_advisor(
Respond with ONLY the JSON object."""),
},
]
raw2 = _generate_with_backend(backend, messages2, max_tokens=512)
raw2 = _generate_with_backend(backend, messages2, max_tokens = 512)
pass2 = _parse_json_response(raw2)
print(f"🤖 Pass 2 done ({time.monotonic() - t2:.1f}s): {pass2}", flush=True)
logger.info(f"Pass 2 done ({time.monotonic() - t2:.1f}s): {pass2}")
if not pass2:
logger.warning(f"Advisor Pass 2 failed to produce JSON: {raw2[:200]}")
@ -674,10 +702,7 @@ def _run_multi_pass_advisor(
# Validate: must have at least one user AND one assistant
roles_present = set(column_roles.values())
if "user" not in roles_present or "assistant" not in roles_present:
print(
f"🤖 Pass 2 sanity fail: missing user or assistant role: {column_roles}",
flush=True,
)
logger.warning(f"Pass 2 sanity fail: missing user or assistant role: {column_roles}")
return None # triggers fallback to simple classification
# ── Pass 3: System prompt (non-conversational datasets only) ──
@ -686,7 +711,7 @@ def _run_multi_pass_advisor(
is_conv = pass1.get("is_conversational", False)
if not is_conv:
print("🤖 Pass 3: Generating system prompt...", flush=True)
logger.info("Pass 3: Generating system prompt...")
t3 = time.monotonic()
# Format label mapping info for the prompt
@ -726,8 +751,8 @@ def _run_multi_pass_advisor(
Write ONLY the system prompt text. No quotes, no labels, no explanation around it."""),
},
]
raw3 = _generate_with_backend(backend, messages3, max_tokens=256)
print(f"🤖 Pass 3 done ({time.monotonic() - t3:.1f}s): {raw3[:200] if raw3 else None}", flush=True)
raw3 = _generate_with_backend(backend, messages3, max_tokens = 256)
logger.info(f"Pass 3 done ({time.monotonic() - t3:.1f}s): {raw3[:200] if raw3 else None}")
if raw3:
# Pass 3 returns raw text, not JSON — clean it up
@ -746,15 +771,14 @@ def _run_multi_pass_advisor(
note_parts = [f"This is a {dtype} dataset (not conversational)."]
if desc:
note_parts.append(desc)
note_parts.append("Columns have been mapped to conversation roles. You can adjust the mapping if needed.")
note_parts.append(
"Columns have been mapped to conversation roles. You can adjust the mapping if needed."
)
user_notification = " ".join(note_parts)
total_time = time.monotonic() - t0
print(
f"🤖 Advisor complete ({total_time:.1f}s): type={dtype}, "
f"mapping={suggested_mapping}, sys_prompt={bool(sys_prompt)}, label_map={bool(label_map)}",
flush=True,
)
logger.info(f"Advisor complete ({total_time:.1f}s): type={dtype}, mapping={suggested_mapping}, sys_prompt={bool(sys_prompt)}, label_map={bool(label_map)}")
return {
"success": True,
@ -774,7 +798,7 @@ def _run_multi_pass_advisor(
if backend is not None:
try:
backend.unload_model()
print("🤖 Advisor model unloaded")
logger.info("Advisor model unloaded")
except Exception:
pass
@ -805,18 +829,18 @@ def llm_conversion_advisor(
# Try multi-pass advisor
result = _run_multi_pass_advisor(
columns=column_names,
samples=samples,
dataset_name=dataset_name,
dataset_card=dataset_card,
dataset_metadata=dataset_metadata,
model_name=model_name,
model_type=model_type,
hf_token=hf_token,
columns = column_names,
samples = samples,
dataset_name = dataset_name,
dataset_card = dataset_card,
dataset_metadata = dataset_metadata,
model_name = model_name,
model_type = model_type,
hf_token = hf_token,
)
if result and result.get("success"):
print(f"🤖 Conversion advisor succeeded: type={result.get('dataset_type')}")
logger.info(f"Conversion advisor succeeded: type={result.get('dataset_type')}")
return result
# Fallback: simple column classification
@ -826,7 +850,8 @@ def llm_conversion_advisor(
return {
"success": True,
"suggested_mapping": {
col: role for col, role in simple_mapping.items()
col: role
for col, role in simple_mapping.items()
if role in ("user", "assistant", "system")
},
"dataset_type": None,

5
studio/backend/utils/datasets/model_mappings.py
View file

@ -8,7 +8,6 @@ This module contains the mapping dictionaries that associate model names
with their corresponding chat templates and response markers.
"""
TEMPLATE_TO_MODEL_MAPPER = {
"phi-3.5": (
"unsloth/Phi-3.5-mini-instruct-bnb-4bit",
@ -407,14 +406,11 @@ MODEL_TO_TEMPLATE_MAPPER = {}
for key, values in TEMPLATE_TO_MODEL_MAPPER.items():
for value in values:
MODEL_TO_TEMPLATE_MAPPER[value] = key
pass
# Get lowercased
lowered_key = key.lower()
for value in values:
MODEL_TO_TEMPLATE_MAPPER[value.lower()] = lowered_key
pass
pass
TEMPLATE_TO_RESPONSES_MAPPER = {
@ -531,4 +527,3 @@ TEMPLATE_TO_RESPONSES_MAPPER = {
"response": "<|assistant|><think>",
},
}

41
studio/backend/utils/datasets/vlm_processing.py
View file

@ -14,9 +14,9 @@ from itertools import islice
def generate_smart_vlm_instruction(
dataset,
text_column="text",
image_column="image",
dataset_name=None,
text_column = "text",
image_column = "image",
dataset_name = None,
):
"""
Generate smart, context-aware instruction for VLM datasets using heuristics.
@ -66,11 +66,12 @@ def generate_smart_vlm_instruction(
# OCR / Transcription
"ocr": {
"keywords": ["ocr", "transcribe", "transcript"],
"content_hints": [r"[A-Za-z\u0600-\u06FF]{10,}"], # Long text passages (Latin/Arabic)
"content_hints": [
r"[A-Za-z\u0600-\u06FF]{10,}"
], # Long text passages (Latin/Arabic)
"instruction": "Transcribe all the text shown in this image.",
"confidence": 0.9,
},
# LaTeX / Math
"latex": {
"keywords": ["latex", "math", "formula", "equation"],
@ -78,7 +79,6 @@ def generate_smart_vlm_instruction(
"instruction": "Convert this image to LaTeX notation.",
"confidence": 0.95,
},
# Caption / Description
"caption": {
"keywords": ["caption", "description", "describe"],
@ -86,15 +86,21 @@ def generate_smart_vlm_instruction(
"instruction": "Provide a detailed description of this image.",
"confidence": 0.85,
},
# Medical / Radiology
"medical": {
"keywords": ["medical", "radiology", "xray", "ct", "mri", "scan", "diagnosis"],
"keywords": [
"medical",
"radiology",
"xray",
"ct",
"mri",
"scan",
"diagnosis",
],
"content_hints": [r"\b(lesion|radiograph|patient|diagnosis|findings)\b"],
"instruction": "Analyze this medical image and describe the key findings.",
"confidence": 0.9,
},
# Code / Programming
"code": {
"keywords": ["code", "program", "function", "algorithm"],
@ -102,7 +108,6 @@ def generate_smart_vlm_instruction(
"instruction": "Explain what this code visualization shows.",
"confidence": 0.85,
},
# Chart / Graph
"chart": {
"keywords": ["chart", "graph", "plot", "visualization", "diagram"],
@ -110,7 +115,6 @@ def generate_smart_vlm_instruction(
"instruction": "Describe this chart or graph, including key data points and trends.",
"confidence": 0.85,
},
# Document / Text Recognition
"document": {
"keywords": ["document", "page", "paragraph", "article"],
@ -132,7 +136,9 @@ def generate_smart_vlm_instruction(
score += 0.5
# Check dataset name if provided
if dataset_name and any(keyword in dataset_name.lower() for keyword in task_info["keywords"]):
if dataset_name and any(
keyword in dataset_name.lower() for keyword in task_info["keywords"]
):
score += 0.3
# Check content patterns
@ -186,7 +192,7 @@ def generate_smart_vlm_instruction(
row = {}
for col in s:
val = s[col]
if hasattr(val, 'size') and hasattr(val, 'mode'): # PIL Image
if hasattr(val, "size") and hasattr(val, "mode"): # PIL Image
row[col] = "<image>"
elif isinstance(val, list):
row[col] = str(val)[:300]
@ -195,15 +201,15 @@ def generate_smart_vlm_instruction(
sample_rows.append(row)
llm_result = llm_generate_vlm_instruction(
column_names=list(column_names),
samples=sample_rows,
dataset_name=dataset_name,
column_names = list(column_names),
samples = sample_rows,
dataset_name = dataset_name,
)
if llm_result and llm_result.get("instruction"):
print(
f"\n[DEBUG] LLM-assisted VLM instruction generated: "
f"'{llm_result['instruction']}' (confidence={llm_result.get('confidence', 'N/A')})\n",
flush=True,
flush = True,
)
return {
"instruction": llm_result["instruction"],
@ -214,6 +220,7 @@ def generate_smart_vlm_instruction(
}
except Exception as e:
import logging
logging.getLogger(__name__).debug(f"LLM-assisted instruction skipped: {e}")
# ===== LEVEL 5: Generic Fallback =====

27
studio/backend/utils/hardware/__init__.py
View file

@ -4,6 +4,7 @@
"""
Hardware detection and GPU utilities
"""
from .hardware import (
DeviceType,
DEVICE,
@ -21,17 +22,17 @@ from .hardware import (
)
__all__ = [
'DeviceType',
'DEVICE',
'detect_hardware',
'get_device',
'is_apple_silicon',
'clear_gpu_cache',
'get_gpu_memory_info',
'log_gpu_memory',
'get_gpu_summary',
'get_package_versions',
'get_gpu_utilization',
'get_physical_gpu_count',
'safe_num_proc',
"DeviceType",
"DEVICE",
"detect_hardware",
"get_device",
"is_apple_silicon",
"clear_gpu_cache",
"get_gpu_memory_info",
"log_gpu_memory",
"get_gpu_summary",
"get_package_versions",
"get_gpu_utilization",
"get_physical_gpu_count",
"safe_num_proc",
]

34
studio/backend/utils/hardware/hardware.py
View file

@ -15,6 +15,7 @@ Usage:
import torch
...
"""
import platform
import structlog
from loggers import get_logger
@ -26,11 +27,13 @@ logger = get_logger(__name__)
# ========== Device Enum ==========
class DeviceType(str, Enum):
"""Supported compute backends. Inherits from str so it serializes cleanly in JSON."""
CUDA = "cuda"
MLX = "mlx"
CPU = "cpu"
MLX = "mlx"
CPU = "cpu"
# ========== Global State (set once by detect_hardware) ==========
@ -40,6 +43,7 @@ DEVICE: Optional[DeviceType] = None
# ========== Detection ==========
def is_apple_silicon() -> bool:
"""Check if running on Apple Silicon hardware (pure platform check, no ML imports)."""
return platform.system() == "Darwin" and platform.machine() == "arm64"
@ -49,6 +53,7 @@ def _has_torch() -> bool:
"""Check if PyTorch is importable."""
try:
import torch
return True
except ImportError:
return False
@ -58,6 +63,7 @@ def _has_mlx() -> bool:
"""Check if MLX is importable."""
try:
import mlx.core
return True
except ImportError:
return False
@ -80,6 +86,7 @@ def detect_hardware() -> DeviceType:
# --- CUDA: try PyTorch ---
if _has_torch():
import torch
if torch.cuda.is_available():
DEVICE = DeviceType.CUDA
device_name = torch.cuda.get_device_properties(0).name
@ -101,6 +108,7 @@ def detect_hardware() -> DeviceType:
# ========== Convenience helpers ==========
def get_device() -> DeviceType:
"""
Return the detected device. Auto-detects if detect_hardware() hasn't been called yet.
@ -118,12 +126,14 @@ def clear_gpu_cache():
Safe to call on any platform — no-ops gracefully.
"""
import gc
gc.collect()
device = get_device()
if device == DeviceType.CUDA:
import torch
torch.cuda.synchronize()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
@ -144,6 +154,7 @@ def get_gpu_memory_info() -> Dict[str, Any]:
if device == DeviceType.CUDA:
try:
import torch
idx = torch.cuda.current_device()
props = torch.cuda.get_device_properties(idx)
@ -215,6 +226,7 @@ def log_gpu_memory(context: str):
# ========== GPU Summary & Package Versions ==========
def get_gpu_summary() -> Dict[str, Any]:
"""
Return a compact summary of the primary GPU.
@ -256,6 +268,7 @@ def get_package_versions() -> Dict[str, Optional[str]]:
# CUDA toolkit version bundled with torch
try:
import torch
versions["cuda"] = getattr(torch.version, "cuda", None)
except Exception:
versions["cuda"] = None
@ -265,6 +278,7 @@ def get_package_versions() -> Dict[str, Optional[str]]:
# ========== Live GPU Utilization (nvidia-smi) ==========
def get_gpu_utilization() -> Dict[str, Any]:
"""
Return a live snapshot of GPU utilization via ``nvidia-smi``.
@ -312,9 +326,9 @@ def get_gpu_utilization() -> Dict[str, Any]:
"memory.used,memory.total,power.draw,power.limit",
"--format=csv,noheader,nounits",
],
capture_output=True,
text=True,
timeout=5,
capture_output = True,
text = True,
timeout = 5,
)
if result.returncode == 0 and result.stdout.strip():
@ -360,7 +374,9 @@ def get_gpu_utilization() -> Dict[str, Any]:
power_limit = smi_data.get("power_limit")
vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None
vram_total_gb = round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
vram_total_gb = (
round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
)
vram_pct = (
round((vram_used_mb / vram_total_mb) * 100, 1)
if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0
@ -395,6 +411,7 @@ def get_gpu_utilization() -> Dict[str, Any]:
_physical_gpu_count: Optional[int] = None
def get_physical_gpu_count() -> int:
"""
Return the number of physical NVIDIA GPUs on the machine.
@ -409,9 +426,12 @@ def get_physical_gpu_count() -> int:
try:
import subprocess
result = subprocess.run(
["nvidia-smi", "-L"],
capture_output=True, text=True, timeout=5,
capture_output = True,
text = True,
timeout = 5,
)
if result.returncode == 0 and result.stdout.strip():
_physical_gpu_count = len(result.stdout.strip().splitlines())

2
studio/backend/utils/inference/__init__.py
View file

@ -4,7 +4,7 @@
"""
Inference utility functions
"""
from utils.inference.inference_config import load_inference_config
__all__ = ["load_inference_config"]

26
studio/backend/utils/inference/inference_config.py
View file

@ -7,6 +7,7 @@ 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
@ -21,15 +22,15 @@ logger = get_logger(__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:
{
@ -41,30 +42,33 @@ def load_inference_config(model_identifier: str) -> Dict[str, Any]:
"""
# 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:
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)),
"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)),
"trust_remote_code": model_inference.get("trust_remote_code", default_inference.get("trust_remote_code", False)),
"trust_remote_code": model_inference.get(
"trust_remote_code", default_inference.get("trust_remote_code", False)
),
}
return inference_config
return inference_config

33
studio/backend/utils/models/__init__.py
View file

@ -4,6 +4,7 @@
"""
Model and LoRA configuration handling
"""
from .model_config import (
ModelConfig,
GgufVariantInfo,
@ -24,20 +25,20 @@ from .model_config import (
from .checkpoints import scan_checkpoints
__all__ = [
'ModelConfig',
'GgufVariantInfo',
'is_vision_model',
'is_embedding_model',
'detect_audio_type',
'is_audio_input_type',
'VALID_AUDIO_TYPES',
'scan_trained_loras',
'scan_exported_models',
'load_model_defaults',
'get_base_model_from_lora',
'load_model_config',
'list_gguf_variants',
'MODEL_NAME_MAPPING',
'UI_STATUS_INDICATORS',
'scan_checkpoints',
"ModelConfig",
"GgufVariantInfo",
"is_vision_model",
"is_embedding_model",
"detect_audio_type",
"is_audio_input_type",
"VALID_AUDIO_TYPES",
"scan_trained_loras",
"scan_exported_models",
"load_model_defaults",
"get_base_model_from_lora",
"load_model_config",
"list_gguf_variants",
"MODEL_NAME_MAPPING",
"UI_STATUS_INDICATORS",
"scan_checkpoints",
]

17
studio/backend/utils/models/checkpoints.py
View file

@ -4,6 +4,7 @@
"""
Checkpoint scanning utilities for discovering training runs and their checkpoints.
"""
import json
import structlog
from loggers import get_logger
@ -86,7 +87,9 @@ def scan_checkpoints(
name_part = parts[0]
idx = name_part.find("_")
if idx > 0:
metadata["base_model"] = name_part[:idx] + "/" + name_part[idx + 1:]
metadata["base_model"] = (
name_part[:idx] + "/" + name_part[idx + 1 :]
)
else:
metadata["base_model"] = name_part
@ -109,13 +112,19 @@ def scan_checkpoints(
# Assign the last checkpoint's loss to the main adapter entry
if len(checkpoints) > 1:
last_checkpoint_loss = checkpoints[-1][2]
checkpoints[0] = (checkpoints[0][0], checkpoints[0][1], last_checkpoint_loss)
checkpoints[0] = (
checkpoints[0][0],
checkpoints[0][1],
last_checkpoint_loss,
)
models.append((item.name, checkpoints, metadata))
logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)")
logger.debug(
f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)"
)
# Sort by modification time (newest first)
models.sort(key=lambda x: Path(x[1][0][1]).stat().st_mtime, reverse=True)
models.sort(key = lambda x: Path(x[1][0][1]).stat().st_mtime, reverse = True)
logger.info(f"Found {len(models)} training runs in {outputs_dir}")
return models

512
studio/backend/utils/models/model_config.py
View file

@ -4,6 +4,7 @@
"""
Model and LoRA configuration handling
"""
from transformers import AutoConfig
from dataclasses import dataclass
from typing import Optional, Dict, Any
@ -77,7 +78,6 @@ MODEL_NAME_MAPPING = {
"unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml": [
"unsloth/ERNIE-4.5-VL-28B-A3B-PT",
],
"tiiuae_Falcon-H1-0.5B-Instruct.yaml": [
"tiiuae/Falcon-H1-0.5B-Instruct",
"unsloth/Falcon-H1-0.5B-Instruct",
@ -132,7 +132,6 @@ MODEL_NAME_MAPPING = {
"unsloth/gpt-oss-20b-unsloth-bnb-4bit",
"unsloth/gpt-oss-20b-BF16",
],
"unsloth_gpt-oss-120b.yaml": [
"openai/gpt-oss-120b",
"unsloth/gpt-oss-120b-unsloth-bnb-4bit",
@ -169,7 +168,6 @@ MODEL_NAME_MAPPING = {
"unsloth/Meta-Llama-3.1-405B-bnb-4bit",
"meta-llama/Meta-Llama-3.1-405B",
],
"unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml": [
"unsloth/Meta-Llama-3.1-8B-Instruct-unsloth-bnb-4bit",
"unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
@ -229,10 +227,9 @@ MODEL_NAME_MAPPING = {
"unsloth/Mistral-Nemo-Base-2407-bnb-4bit",
"unsloth/Mistral-Nemo-Base-2407",
"mistralai/Mistral-Nemo-Base-2407",
"unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit",
"unsloth/Mistral-Nemo-Instruct-2407",
"unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit",
"unsloth/Mistral-Nemo-Instruct-2407",
"mistralai/Mistral-Nemo-Instruct-2407",
],
"unsloth_Mistral-Small-Instruct-2409.yaml": [
"unsloth/Mistral-Small-Instruct-2409-bnb-4bit",
@ -384,7 +381,10 @@ for canonical_file, model_names in MODEL_NAME_MAPPING.items():
for model_name in model_names:
_REVERSE_MODEL_MAPPING[model_name.lower()] = canonical_file
def load_model_config(model_name: str, use_auth: bool = False, token: Optional[str] = None):
def load_model_config(
model_name: str, use_auth: bool = False, token: Optional[str] = None
):
"""
Load model config with optional authentication control.
"""
@ -392,32 +392,30 @@ def load_model_config(model_name: str, use_auth: bool = False, token: Optional[s
if token:
# Explicit token provided - use it
return AutoConfig.from_pretrained(
model_name,
trust_remote_code=True,
token=token
model_name, trust_remote_code = True, token = token
)
if not use_auth:
# Load without any authentication (for public model checks)
with without_hf_auth():
return AutoConfig.from_pretrained(
model_name,
trust_remote_code=True,
token=None
model_name, trust_remote_code = True, token = None
)
# Use default authentication (cached tokens)
return AutoConfig.from_pretrained(
model_name,
trust_remote_code=True
)
return AutoConfig.from_pretrained(model_name, trust_remote_code = True)
# VLM architecture suffixes and known VLM model_type values.
_VLM_ARCH_SUFFIXES = ("ForConditionalGeneration", "ForVisionText2Text")
_VLM_MODEL_TYPES = {
'phi3_v', 'llava', 'llava_next', 'llava_onevision',
'internvl_chat', 'cogvlm2', 'minicpmv',
"phi3_v",
"llava",
"llava_next",
"llava_onevision",
"internvl_chat",
"cogvlm2",
"minicpmv",
}
# Pre-computed .venv_t5 path and backend dir for subprocess version switching.
@ -426,7 +424,7 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent)
# Inline script executed in a subprocess with transformers 5.x activated.
# Receives model_name and token via argv, prints JSON result to stdout.
_VISION_CHECK_SCRIPT = r'''
_VISION_CHECK_SCRIPT = r"""
import sys, os, json
os.environ["TOKENIZERS_PARALLELISM"] = "false"
@ -472,10 +470,12 @@ try:
except Exception as exc:
logger.info(json.dumps({"error": str(exc)}))
sys.exit(1)
'''
"""
def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None) -> bool:
def _is_vision_model_subprocess(
model_name: str, hf_token: Optional[str] = None
) -> bool:
"""Run is_vision_model check in a subprocess with transformers 5.x.
Same pattern as training/inference workers: spawn a clean subprocess
@ -486,16 +486,26 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
try:
result = subprocess.run(
[sys.executable, "-c", _VISION_CHECK_SCRIPT,
_VENV_T5_DIR, _BACKEND_DIR, model_name, token_arg],
capture_output=True, text=True, timeout=60,
[
sys.executable,
"-c",
_VISION_CHECK_SCRIPT,
_VENV_T5_DIR,
_BACKEND_DIR,
model_name,
token_arg,
],
capture_output = True,
text = True,
timeout = 60,
)
if result.returncode != 0:
stderr = result.stderr.strip()
logger.warning(
"Vision check subprocess failed for '%s': %s",
model_name, stderr or result.stdout.strip(),
model_name,
stderr or result.stdout.strip(),
)
return False
@ -503,7 +513,8 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
if "error" in data:
logger.warning(
"Vision check subprocess error for '%s': %s",
model_name, data["error"],
model_name,
data["error"],
)
return False
@ -511,7 +522,10 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
logger.info(
"Vision check (subprocess, transformers 5.x) for '%s': "
"model_type=%s, architectures=%s, is_vision=%s",
model_name, data.get("model_type"), data.get("architectures"), is_vlm,
model_name,
data.get("model_type"),
data.get("architectures"),
is_vlm,
)
return is_vlm
@ -540,52 +554,54 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
# because AutoConfig in the main process (transformers 4.57.x) doesn't
# recognize their architectures.
from utils.transformers_version import needs_transformers_5
if needs_transformers_5(model_name):
logger.info(
"Model '%s' needs transformers 5.x — checking vision via subprocess",
model_name,
)
return _is_vision_model_subprocess(model_name, hf_token=hf_token)
return _is_vision_model_subprocess(model_name, hf_token = hf_token)
try:
config = load_model_config(model_name, use_auth=True, token=hf_token)
config = load_model_config(model_name, use_auth = True, token = hf_token)
# Exclude audio-only models that share ForConditionalGeneration suffix
# (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration)
_audio_only_model_types = {'csm', 'whisper'}
model_type = getattr(config, 'model_type', None)
_audio_only_model_types = {"csm", "whisper"}
model_type = getattr(config, "model_type", None)
if model_type in _audio_only_model_types:
return False
# Check 1: Architecture class name patterns
if hasattr(config, 'architectures'):
is_vlm = any(
x.endswith(_VLM_ARCH_SUFFIXES)
for x in config.architectures
)
if hasattr(config, "architectures"):
is_vlm = any(x.endswith(_VLM_ARCH_SUFFIXES) for x in config.architectures)
if is_vlm:
logger.info(f"Model {model_name} detected as VLM: architecture {config.architectures}")
logger.info(
f"Model {model_name} detected as VLM: architecture {config.architectures}"
)
return True
# Check 2: Has vision_config (most VLMs: LLaVA, Gemma-3, Qwen2-VL, etc.)
if hasattr(config, 'vision_config'):
if hasattr(config, "vision_config"):
logger.info(f"Model {model_name} detected as VLM: has vision_config")
return True
# Check 3: Has img_processor (Phi-3.5 Vision uses this instead of vision_config)
if hasattr(config, 'img_processor'):
if hasattr(config, "img_processor"):
logger.info(f"Model {model_name} detected as VLM: has img_processor")
return True
# Check 4: Has image_token_index (common in VLMs for image placeholder tokens)
if hasattr(config, 'image_token_index'):
if hasattr(config, "image_token_index"):
logger.info(f"Model {model_name} detected as VLM: has image_token_index")
return True
# Check 5: Known VLM model_type values that may not match above checks
if hasattr(config, 'model_type'):
if hasattr(config, "model_type"):
if config.model_type in _VLM_MODEL_TYPES:
logger.info(f"Model {model_name} detected as VLM: model_type={config.model_type}")
logger.info(
f"Model {model_name} detected as VLM: model_type={config.model_type}"
)
return True
return False
@ -595,19 +611,20 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
return False
VALID_AUDIO_TYPES = ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm')
VALID_AUDIO_TYPES = ("snac", "csm", "bicodec", "dac", "whisper", "audio_vlm")
# Cache detection results per session to avoid repeated API calls
_audio_detection_cache: Dict[str, Optional[str]] = {}
# Tokenizer token patterns → audio_type (all 6 types detected from tokenizer_config.json)
_AUDIO_TOKEN_PATTERNS = {
'csm': lambda tokens: '<|AUDIO|>' in tokens and '<|audio_eos|>' in tokens,
'whisper': lambda tokens: '<|startoftranscript|>' in tokens,
'audio_vlm': lambda tokens: '<audio_soft_token>' in tokens,
'bicodec': lambda tokens: any(t.startswith('<|bicodec_') for t in tokens),
'dac': lambda tokens: '<|audio_start|>' in tokens and '<|audio_end|>' in tokens,
'snac': lambda tokens: sum(1 for t in tokens if t.startswith('<custom_token_')) > 10000,
"csm": lambda tokens: "<|AUDIO|>" in tokens and "<|audio_eos|>" in tokens,
"whisper": lambda tokens: "<|startoftranscript|>" in tokens,
"audio_vlm": lambda tokens: "<audio_soft_token>" in tokens,
"bicodec": lambda tokens: any(t.startswith("<|bicodec_") for t in tokens),
"dac": lambda tokens: "<|audio_start|>" in tokens and "<|audio_end|>" in tokens,
"snac": lambda tokens: sum(1 for t in tokens if t.startswith("<custom_token_"))
> 10000,
}
@ -631,17 +648,20 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option
return result
def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]:
def _detect_audio_from_tokenizer(
model_name: str, hf_token: Optional[str] = None
) -> Optional[str]:
"""Detect audio type from tokenizer special tokens (for LLM-based audio models).
First checks local HF cache, then fetches tokenizer_config.json from HuggingFace.
Checks added_tokens_decoder for distinctive patterns.
"""
def _check_token_patterns(tok_config: dict) -> Optional[str]:
added = tok_config.get('added_tokens_decoder', {})
added = tok_config.get("added_tokens_decoder", {})
if not added:
return None
token_contents = [v.get('content', '') for v in added.values()]
token_contents = [v.get("content", "") for v in added.values()]
for audio_type, check_fn in _AUDIO_TOKEN_PATTERNS.items():
if check_fn(token_contents):
return audio_type
@ -650,6 +670,7 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
# 1) Check local HF cache first (works for gated/offline models)
try:
from huggingface_hub.constants import HF_HUB_CACHE
cache_dir = Path(HF_HUB_CACHE)
repo_dir_name = f"models--{model_name.replace('/', '--')}"
repo_dir = cache_dir / repo_dir_name
@ -657,7 +678,10 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
snapshots_dir = repo_dir / "snapshots"
if snapshots_dir.exists():
for snapshot in snapshots_dir.iterdir():
for tok_path in ['tokenizer_config.json', 'LLM/tokenizer_config.json']:
for tok_path in [
"tokenizer_config.json",
"LLM/tokenizer_config.json",
]:
tok_file = snapshot / tok_path
if tok_file.exists():
tok_config = json.loads(tok_file.read_text())
@ -672,16 +696,16 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
import requests
import os
paths_to_try = ['tokenizer_config.json', 'LLM/tokenizer_config.json']
paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"]
# Use provided token, or fall back to env
token = hf_token or os.environ.get('HF_TOKEN')
token = hf_token or os.environ.get("HF_TOKEN")
headers = {}
if token:
headers['Authorization'] = f'Bearer {token}'
headers["Authorization"] = f"Bearer {token}"
for tok_path in paths_to_try:
url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}"
resp = requests.get(url, headers=headers, timeout=15)
resp = requests.get(url, headers = headers, timeout = 15)
if not resp.ok:
continue
@ -692,7 +716,9 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
return None
except Exception as e:
logger.debug(f"Could not detect audio type from tokenizer for {model_name}: {e}")
logger.debug(
f"Could not detect audio type from tokenizer for {model_name}: {e}"
)
return None
@ -701,7 +727,7 @@ def is_audio_input_type(audio_type: Optional[str]) -> bool:
Whisper (ASR) and audio_vlm (Gemma3n) accept audio input.
"""
return audio_type in ('whisper', 'audio_vlm')
return audio_type in ("whisper", "audio_vlm")
def _is_mmproj(filename: str) -> bool:
@ -756,7 +782,8 @@ def detect_gguf_model(path: str) -> Optional[str]:
if p.is_dir():
gguf_files = sorted(
(f for f in p.glob("*.gguf") if not _is_mmproj(f.name)),
key=lambda f: f.stat().st_size, reverse=True,
key = lambda f: f.stat().st_size,
reverse = True,
)
if gguf_files:
return str(gguf_files[0].resolve())
@ -767,9 +794,18 @@ def detect_gguf_model(path: str) -> Optional[str]:
# Preferred GGUF quantization levels, in descending priority.
# Q4_K_M is a good default: small, fast, acceptable quality.
_GGUF_QUANT_PREFERENCE = [
"Q4_K_M", "Q4_K_S", "Q5_K_M", "Q5_K_S",
"Q6_K", "Q8_0", "Q3_K_M", "Q3_K_L", "Q2_K",
"F16", "BF16", "F32",
"Q4_K_M",
"Q4_K_S",
"Q5_K_M",
"Q5_K_S",
"Q6_K",
"Q8_0",
"Q3_K_M",
"Q3_K_L",
"Q2_K",
"F16",
"BF16",
"F32",
]
@ -797,9 +833,10 @@ def _pick_best_gguf(filenames: list[str]) -> Optional[str]:
@dataclass
class GgufVariantInfo:
"""A single GGUF quantization variant from a HuggingFace repo."""
filename: str # e.g., "gemma-3-4b-it-Q4_K_M.gguf"
quant: str # e.g., "Q4_K_M" (extracted from filename)
size_bytes: int # file size
filename: str # e.g., "gemma-3-4b-it-Q4_K_M.gguf"
quant: str # e.g., "Q4_K_M" (extracted from filename)
size_bytes: int # file size
def _extract_quant_label(filename: str) -> str:
@ -815,21 +852,23 @@ def _extract_quant_label(filename: str) -> str:
"MXFP4_MOE/model-MXFP4_MOE-0001.gguf"→ "MXFP4_MOE"
"""
import re
# Use only the basename (rfilename may include directory)
basename = filename.rsplit("/", 1)[-1]
# Strip .gguf and any shard suffix (-00001-of-00010)
stem = re.sub(r'-\d{3,}-of-\d{3,}', '', basename.rsplit(".", 1)[0])
stem = re.sub(r"-\d{3,}-of-\d{3,}", "", basename.rsplit(".", 1)[0])
# Match known quantization patterns
match = re.search(
r'(UD-)?' # Optional UD- prefix (Ultra Discrete)
r'(MXFP[0-9]+(?:_[A-Z0-9]+)*' # MXFP variants: MXFP4, MXFP4_MOE
r'|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?' # IQ variants: IQ4_XS, IQ4_NL, IQ1_S
r'|TQ[0-9]+_[0-9]+' # Ternary quant: TQ1_0, TQ2_0
r'|Q[0-9]+_K_[A-Z]+' # K-quant: Q4_K_M, Q3_K_S
r'|Q[0-9]+_[0-9]+' # Standard: Q8_0, Q5_1
r'|Q[0-9]+_K' # Short K-quant: Q6_K
r'|BF16|F16|F32)', # Full precision
stem, re.IGNORECASE,
r"(UD-)?" # Optional UD- prefix (Ultra Discrete)
r"(MXFP[0-9]+(?:_[A-Z0-9]+)*" # MXFP variants: MXFP4, MXFP4_MOE
r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?" # IQ variants: IQ4_XS, IQ4_NL, IQ1_S
r"|TQ[0-9]+_[0-9]+" # Ternary quant: TQ1_0, TQ2_0
r"|Q[0-9]+_K_[A-Z]+" # K-quant: Q4_K_M, Q3_K_S
r"|Q[0-9]+_[0-9]+" # Standard: Q8_0, Q5_1
r"|Q[0-9]+_K" # Short K-quant: Q6_K
r"|BF16|F16|F32)", # Full precision
stem,
re.IGNORECASE,
)
if match:
prefix = match.group(1) or ""
@ -853,11 +892,11 @@ def list_gguf_variants(
"""
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(repo_id, token=hf_token, files_metadata=True)
info = hf_model_info(repo_id, token = hf_token, files_metadata = True)
variants: list[GgufVariantInfo] = []
has_vision = False
quant_totals: dict[str, int] = {} # quant -> total bytes
quant_totals: dict[str, int] = {} # quant -> total bytes
quant_first_file: dict[str, str] = {} # quant -> first filename (for display)
for sibling in info.siblings:
@ -877,11 +916,13 @@ def list_gguf_variants(
quant_first_file[quant] = fname
for quant, total_size in quant_totals.items():
variants.append(GgufVariantInfo(
filename=quant_first_file[quant],
quant=quant,
size_bytes=total_size,
))
variants.append(
GgufVariantInfo(
filename = quant_first_file[quant],
quant = quant,
size_bytes = total_size,
)
)
return variants, has_vision
@ -898,7 +939,7 @@ def detect_gguf_model_remote(
try:
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(repo_id, token=hf_token)
info = hf_model_info(repo_id, token = hf_token)
repo_files = [s.rfilename for s in info.siblings]
return _pick_best_gguf(repo_files)
except Exception as e:
@ -919,9 +960,9 @@ def download_gguf_file(
from huggingface_hub import hf_hub_download
local_path = hf_hub_download(
repo_id=repo_id,
filename=filename,
token=hf_token,
repo_id = repo_id,
filename = filename,
token = hf_token,
)
return local_path
@ -964,7 +1005,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
try:
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(model_name, token=hf_token)
info = hf_model_info(model_name, token = hf_token)
tags = set(info.tags or [])
pipeline_tag = info.pipeline_tag or ""
@ -1024,16 +1065,21 @@ def scan_trained_loras(outputs_dir: str = str(outputs_root())) -> List[Tuple[str
logger.debug(f"Found trained LoRA: {display_name}")
# Sort by modification time (newest first)
trained_loras.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True)
trained_loras.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True)
logger.info(f"Found {len(trained_loras)} trained LoRA adapters in {outputs_dir}")
logger.info(
f"Found {len(trained_loras)} trained LoRA adapters in {outputs_dir}"
)
return trained_loras
except Exception as e:
logger.error(f"Error scanning outputs folder: {e}")
return []
def scan_exported_models(exports_dir: str = str(exports_root())) -> List[Tuple[str, str, str, Optional[str]]]:
def scan_exported_models(
exports_dir: str = str(exports_root()),
) -> List[Tuple[str, str, str, Optional[str]]]:
"""
Scan exports folder for exported models (merged, LoRA, GGUF).
@ -1082,9 +1128,8 @@ def scan_exported_models(exports_dir: str = str(exports_root())) -> List[Tuple[s
adapter_config = checkpoint_dir / "adapter_config.json"
config_file = checkpoint_dir / "config.json"
has_weights = (
any(checkpoint_dir.glob("*.safetensors"))
or any(checkpoint_dir.glob("*.bin"))
has_weights = any(checkpoint_dir.glob("*.safetensors")) or any(
checkpoint_dir.glob("*.bin")
)
has_gguf = any(checkpoint_dir.glob("*.gguf"))
@ -1134,7 +1179,9 @@ def scan_exported_models(exports_dir: str = str(exports_root())) -> List[Tuple[s
# Fallback: read base model from the original training run's
# adapter_config.json in ./outputs/{run_name}/
if not base_model:
outputs_adapter_cfg = resolve_output_dir(run_dir.name) / "adapter_config.json"
outputs_adapter_cfg = (
resolve_output_dir(run_dir.name) / "adapter_config.json"
)
try:
if outputs_adapter_cfg.exists():
cfg = json.loads(outputs_adapter_cfg.read_text())
@ -1147,7 +1194,7 @@ def scan_exported_models(exports_dir: str = str(exports_root())) -> List[Tuple[s
results.append((display_name, model_path, export_type, base_model))
logger.debug(f"Found exported model: {display_name} ({export_type})")
results.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True)
results.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True)
logger.info(f"Found {len(results)} exported models in {exports_dir}")
return results
@ -1177,11 +1224,13 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
# Try adapter_config.json first
adapter_config_path = lora_path_obj / "adapter_config.json"
if adapter_config_path.exists():
with open(adapter_config_path, 'r') as f:
with open(adapter_config_path, "r") as f:
config = json.load(f)
base_model = config.get("base_model_name_or_path")
if base_model:
logger.info(f"Detected base model from adapter_config.json: {base_model}")
logger.info(
f"Detected base model from adapter_config.json: {base_model}"
)
return base_model
# Fallback: try training_args.bin (requires torch)
@ -1189,10 +1238,13 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
if training_args_path.exists():
try:
import torch
training_args = torch.load(training_args_path)
if hasattr(training_args, 'model_name_or_path'):
if hasattr(training_args, "model_name_or_path"):
base_model = training_args.model_name_or_path
logger.info(f"Detected base model from training_args.bin: {base_model}")
logger.info(
f"Detected base model from training_args.bin: {base_model}"
)
return base_model
except Exception as e:
logger.warning(f"Could not load training_args.bin: {e}")
@ -1216,22 +1268,23 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
except Exception as e:
logger.error(f"Error reading base model from LoRA config: {e}")
return None
pass
# Status indicators that appear in UI dropdowns
UI_STATUS_INDICATORS = [" (Ready)", " (Loading...)", " (Active)", "↓ "]
def load_model_defaults(model_name: str) -> Dict[str, Any]:
"""
Load default training parameters for a model from YAML file.
Args:
model_name: Model identifier (e.g., "unsloth/Meta-Llama-3.1-8B-bnb-4bit")
Returns:
Dictionary with default parameters from YAML file, or empty dict if not found
The function looks for a YAML file in configs/model_defaults/ (including subfolders)
The function looks for a YAML file in configs/model_defaults/ (including subfolders)
based on the model name or its aliases from MODEL_NAME_MAPPING.
If no specific file exists, it falls back to default.yaml.
"""
@ -1239,22 +1292,26 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
# Get the script directory to locate configs
script_dir = Path(__file__).parent.parent.parent
defaults_dir = script_dir / "assets" / "configs" / "model_defaults"
# First, check if model is in the mapping
if model_name.lower() in _REVERSE_MODEL_MAPPING:
canonical_file = _REVERSE_MODEL_MAPPING[model_name.lower()]
# Search in subfolders and root
for config_path in defaults_dir.rglob(canonical_file):
if config_path.is_file():
with open(config_path, 'r', encoding='utf-8') as f:
with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded model defaults from {config_path} (via mapping)")
logger.info(
f"Loaded model defaults from {config_path} (via mapping)"
)
return config
# If model_name is a local path (e.g. /home/.../Spark-TTS-0.5B/LLM from
# adapter_config.json), try matching the last 1-2 path components against
# the registry (e.g. "Spark-TTS-0.5B/LLM").
if model_name not in _REVERSE_MODEL_MAPPING and (model_name.startswith("/") or model_name.startswith(".")):
if model_name not in _REVERSE_MODEL_MAPPING and (
model_name.startswith("/") or model_name.startswith(".")
):
parts = Path(model_name).parts
for depth in [2, 1]:
if len(parts) >= depth:
@ -1263,9 +1320,11 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
canonical_file = _REVERSE_MODEL_MAPPING[suffix]
for config_path in defaults_dir.rglob(canonical_file):
if config_path.is_file():
with open(config_path, 'r', encoding='utf-8') as f:
with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded model defaults from {config_path} (via path suffix '{suffix}')")
logger.info(
f"Loaded model defaults from {config_path} (via path suffix '{suffix}')"
)
return config
# Try exact model name match (for backward compatibility)
@ -1273,48 +1332,58 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
# Search in subfolders and root
for config_path in defaults_dir.rglob(model_filename):
if config_path.is_file():
with open(config_path, 'r', encoding='utf-8') as f:
with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded model defaults from {config_path}")
return config
# Fall back to default.yaml
default_config_path = defaults_dir / "default.yaml"
if default_config_path.exists():
with open(default_config_path, 'r', encoding='utf-8') as f:
with open(default_config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded default model defaults from {default_config_path}")
return config
logger.warning(f"No default config found for model {model_name}")
return {}
except Exception as e:
logger.error(f"Error loading model defaults for {model_name}: {e}")
return {}
@dataclass
class ModelConfig:
"""Configuration for a model to load"""
identifier: str # Clean model identifier (org/name or path)
display_name: str # Original UI display name
path: str # Normalized filesystem path
is_local: bool # Is this a local file vs HF model?
is_cached: bool # Is this already in HF cache?
is_vision: bool # Is this a vision model?
is_lora: bool # Is this a lora adapter?
is_gguf: bool = False # Is this a GGUF model?
is_audio: bool = False # Is this a TTS audio model?
audio_type: Optional[str] = None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac'
identifier: str # Clean model identifier (org/name or path)
display_name: str # Original UI display name
path: str # Normalized filesystem path
is_local: bool # Is this a local file vs HF model?
is_cached: bool # Is this already in HF cache?
is_vision: bool # Is this a vision model?
is_lora: bool # Is this a lora adapter?
is_gguf: bool = False # Is this a GGUF model?
is_audio: bool = False # Is this a TTS audio model?
audio_type: Optional[str] = (
None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac'
)
has_audio_input: bool = False # Accepts audio input (ASR/speech understanding)
gguf_file: Optional[str] = None # Full path to the .gguf file (local mode)
gguf_mmproj_file: Optional[str] = None # Full path to the mmproj .gguf file (vision projection)
gguf_hf_repo: Optional[str] = None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
gguf_mmproj_file: Optional[str] = (
None # Full path to the mmproj .gguf file (vision projection)
)
gguf_hf_repo: Optional[str] = (
None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
)
gguf_variant: Optional[str] = None # Quantization variant (e.g. "Q4_K_M")
base_model: Optional[str] = None # Base model (for LoRAs)
@classmethod
def from_lora_path(cls, lora_path: str, hf_token: Optional[str] = None) -> Optional['ModelConfig']:
def from_lora_path(
cls, lora_path: str, hf_token: Optional[str] = None
) -> Optional["ModelConfig"]:
"""
Create ModelConfig from a local LoRA adapter path.
@ -1341,26 +1410,26 @@ class ModelConfig:
return None
# Check if base model is vision
is_vision = is_vision_model(base_model, hf_token=hf_token)
is_vision = is_vision_model(base_model, hf_token = hf_token)
# Check if base model is audio
audio_type = detect_audio_type(base_model, hf_token=hf_token)
audio_type = detect_audio_type(base_model, hf_token = hf_token)
display_name = lora_path_obj.name
identifier = lora_path # Use path as identifier for local LoRAs
return cls(
identifier=identifier,
display_name=display_name,
path=lora_path,
is_local=True,
is_cached=True, # Local LoRAs are always "cached"
is_vision=is_vision,
is_lora=True,
is_audio=audio_type is not None and audio_type != 'audio_vlm',
audio_type=audio_type,
has_audio_input=is_audio_input_type(audio_type),
base_model=base_model,
identifier = identifier,
display_name = display_name,
path = lora_path,
is_local = True,
is_cached = True, # Local LoRAs are always "cached"
is_vision = is_vision,
is_lora = True,
is_audio = audio_type is not None and audio_type != "audio_vlm",
audio_type = audio_type,
has_audio_input = is_audio_input_type(audio_type),
base_model = base_model,
)
except Exception as e:
@ -1374,7 +1443,7 @@ class ModelConfig:
hf_token: Optional[str] = None,
is_lora: bool = False,
gguf_variant: Optional[str] = None,
) -> Optional['ModelConfig']:
) -> Optional["ModelConfig"]:
"""
Create ModelConfig from a clean model identifier.
@ -1432,7 +1501,7 @@ class ModelConfig:
try:
meta = json.loads(meta_path.read_text())
base = meta.get("base_model")
if base and is_vision_model(base, hf_token=hf_token):
if base and is_vision_model(base, hf_token = hf_token):
base_is_vision = True
logger.info(f"GGUF base model '{base}' is a vision model")
except Exception as e:
@ -1444,27 +1513,30 @@ class ModelConfig:
gguf_is_vision = True
logger.info(f"Detected mmproj for vision: {mmproj_file}")
elif base_is_vision:
logger.warning(f"Base model is vision but no mmproj file found in {gguf_dir}")
logger.warning(
f"Base model is vision but no mmproj file found in {gguf_dir}"
)
return cls(
identifier=identifier,
display_name=display_name,
path=path,
is_local=True,
is_cached=True,
is_vision=gguf_is_vision,
is_lora=False,
is_gguf=True,
gguf_file=gguf_file,
gguf_mmproj_file=mmproj_file,
identifier = identifier,
display_name = display_name,
path = path,
is_local = True,
is_cached = True,
is_vision = gguf_is_vision,
is_lora = False,
is_gguf = True,
gguf_file = gguf_file,
gguf_mmproj_file = mmproj_file,
)
else:
# Check if the HF repo contains GGUF files
gguf_filename = detect_gguf_model_remote(identifier, hf_token=hf_token)
gguf_filename = detect_gguf_model_remote(identifier, hf_token = hf_token)
if gguf_filename:
# Preflight: verify llama-server binary exists BEFORE user waits
# for a multi-GB download that llama-server handles natively
from core.inference.llama_cpp import LlamaCppBackend
if not LlamaCppBackend._find_llama_server_binary():
raise RuntimeError(
"llama-server binary not found — cannot load GGUF models. "
@ -1472,7 +1544,7 @@ class ModelConfig:
)
# Use list_gguf_variants() to detect vision & resolve variant
variants, has_vision = list_gguf_variants(identifier, hf_token=hf_token)
variants, has_vision = list_gguf_variants(identifier, hf_token = hf_token)
variant = gguf_variant
if not variant:
# Auto-select best quantization
@ -1489,17 +1561,17 @@ class ModelConfig:
f"variant={variant}, vision={has_vision}"
)
return cls(
identifier=identifier,
display_name=display_name,
path=identifier,
is_local=False,
is_cached=False,
is_vision=has_vision,
is_lora=False,
is_gguf=True,
gguf_file=None,
gguf_hf_repo=identifier,
gguf_variant=variant,
identifier = identifier,
display_name = display_name,
path = identifier,
is_local = False,
is_cached = False,
is_vision = has_vision,
is_lora = False,
is_gguf = True,
gguf_file = None,
gguf_hf_repo = identifier,
gguf_variant = variant,
)
# Auto-detect LoRA for local paths (check adapter_config.json on disk)
@ -1507,20 +1579,25 @@ class ModelConfig:
detected_base = get_base_model_from_lora(path)
if detected_base:
is_lora = True
logger.info(f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})")
logger.info(
f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})"
)
# Auto-detect LoRA for remote HF models (check repo file listing)
if not is_lora and not is_local:
try:
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(identifier, token=hf_token)
info = hf_model_info(identifier, token = hf_token)
repo_files = [s.rfilename for s in info.siblings]
if "adapter_config.json" in repo_files:
is_lora = True
logger.info(f"Auto-detected remote LoRA adapter: '{identifier}'")
except Exception as e:
logger.debug(f"Could not check remote LoRA status for '{identifier}': {e}")
logger.debug(
f"Could not check remote LoRA status for '{identifier}': {e}"
)
# Handle LoRA adapters
base_model = None
if is_lora:
@ -1531,15 +1608,20 @@ class ModelConfig:
# Remote LoRA: download adapter_config.json from HF
try:
from huggingface_hub import hf_hub_download
config_path = hf_hub_download(identifier, "adapter_config.json", token=hf_token)
with open(config_path, 'r') as f:
config_path = hf_hub_download(
identifier, "adapter_config.json", token = hf_token
)
with open(config_path, "r") as f:
adapter_config = json.load(f)
base_model = adapter_config.get("base_model_name_or_path")
if base_model:
logger.info(f"Resolved remote LoRA base model: '{base_model}'")
except Exception as e:
logger.warning(f"Could not download adapter_config.json for '{identifier}': {e}")
logger.warning(
f"Could not download adapter_config.json for '{identifier}': {e}"
)
if not base_model:
logger.warning(f"Could not determine base model for LoRA '{path}'")
return None
@ -1547,34 +1629,35 @@ class ModelConfig:
else:
check_model = identifier
vision = is_vision_model(check_model, hf_token=hf_token)
audio_type_val = detect_audio_type(check_model, hf_token=hf_token)
vision = is_vision_model(check_model, hf_token = hf_token)
audio_type_val = detect_audio_type(check_model, hf_token = hf_token)
has_audio_in = is_audio_input_type(audio_type_val)
display_name = Path(path).name if is_local else identifier.split("/")[-1]
return cls(
identifier=identifier,
display_name=display_name,
path=path,
is_local=is_local,
is_cached=is_model_cached(identifier) if not is_local else True,
is_vision=vision,
is_lora=is_lora,
is_audio=audio_type_val is not None and audio_type_val != 'audio_vlm',
audio_type=audio_type_val,
has_audio_input=has_audio_in,
base_model=base_model,
identifier = identifier,
display_name = display_name,
path = path,
is_local = is_local,
is_cached = is_model_cached(identifier) if not is_local else True,
is_vision = vision,
is_lora = is_lora,
is_audio = audio_type_val is not None and audio_type_val != "audio_vlm",
audio_type = audio_type_val,
has_audio_input = has_audio_in,
base_model = base_model,
)
@classmethod
def from_ui_selection(cls,
dropdown_value: Optional[str],
search_value: Optional[str],
local_models: list = None,
hf_token: Optional[str] = None,
is_lora: bool = False) -> Optional['ModelConfig']:
def from_ui_selection(
cls,
dropdown_value: Optional[str],
search_value: Optional[str],
local_models: list = None,
hf_token: Optional[str] = None,
is_lora: bool = False,
) -> Optional["ModelConfig"]:
"""
Create a universal ModelConfig from UI dropdown/search selections.
Handles base models and LoRA adapters.
@ -1592,7 +1675,9 @@ class ModelConfig:
# Use the correct 'local_models' parameter to resolve display names
if " (Active)" in selected or " (Ready)" in selected:
clean_display_name = selected.replace(" (Active)", "").replace(" (Ready)", "")
clean_display_name = selected.replace(" (Active)", "").replace(
" (Ready)", ""
)
if local_models:
for local_display, local_path in local_models:
if local_display == clean_display_name:
@ -1621,25 +1706,28 @@ class ModelConfig:
# For a LoRA, we MUST find its base model.
base_model = get_base_model_from_lora(path)
if not base_model:
logger.warning(f"Could not determine base model for LoRA '{path}'. Cannot create config.")
return None # Cannot proceed without a base model
logger.warning(
f"Could not determine base model for LoRA '{path}'. Cannot create config."
)
return None # Cannot proceed without a base model
# A LoRA's vision capability is determined by its base model.
is_vision = is_vision_model(base_model, hf_token=hf_token)
is_vision = is_vision_model(base_model, hf_token = hf_token)
else:
# For a base model, just check its own vision status.
is_vision = is_vision_model(identifier, hf_token=hf_token)
is_vision = is_vision_model(identifier, hf_token = hf_token)
from utils.paths import is_model_cached
is_cached = is_model_cached(identifier) if not is_local else True
return cls(
identifier=identifier,
display_name=display_name,
path=path,
is_local=is_local,
is_cached=is_cached,
is_vision=is_vision,
is_lora=is_lora,
base_model=base_model, # This will be None for base models, and populated for LoRAs
identifier = identifier,
display_name = display_name,
path = path,
is_local = is_local,
is_cached = is_cached,
is_vision = is_vision,
is_lora = is_lora,
base_model = base_model, # This will be None for base models, and populated for LoRAs
)

Some files were not shown because too many files have changed in this diff Show more