From 47654cb91c0b80c17a058ebc3a5307e65f5549a1 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Thu, 12 Mar 2026 18:28:04 +0000 Subject: [PATCH] Final cleanup --- README.md | 240 -- cli/__init__.py | 6 +- cli/commands/__init__.py | 2 + cli/commands/export.py | 90 +- cli/commands/inference.py | 44 +- cli/commands/studio.py | 249 ++- cli/commands/train.py | 74 +- cli/commands/ui.py | 59 +- cli/config.py | 17 +- cli/options.py | 18 +- .../Unsloth_Studio_Colab.ipynb | 12 + studio/__init__.py | 2 + studio/backend/__init__.py | 2 + studio/backend/assets/__init__.py | 2 + studio/backend/assets/configs/__init__.py | 2 + studio/backend/auth/__init__.py | 4 +- studio/backend/auth/authentication.py | 20 +- studio/backend/auth/hashing.py | 6 +- studio/backend/auth/storage.py | 17 +- studio/backend/colab.py | 30 +- studio/backend/core/__init__.py | 142 +- .../backend/core/data_recipe/jobs/__init__.py | 1 - .../backend/core/data_recipe/jobs/manager.py | 80 +- studio/backend/core/data_recipe/jobs/parse.py | 90 +- studio/backend/core/data_recipe/jobs/types.py | 14 +- .../backend/core/data_recipe/jobs/worker.py | 56 +- studio/backend/core/data_recipe/jsonable.py | 2 +- .../data_recipe/local_callable_validators.py | 95 +- .../data_recipe/oxc-validator/validate.mjs | 3 + studio/backend/core/data_recipe/service.py | 61 +- studio/backend/core/export/__init__.py | 7 +- studio/backend/core/export/export.py | 287 ++- studio/backend/core/export/orchestrator.py | 111 +- studio/backend/core/export/worker.py | 341 +-- studio/backend/core/inference/__init__.py | 9 +- studio/backend/core/inference/audio_codecs.py | 73 +- studio/backend/core/inference/inference.py | 855 ++++--- studio/backend/core/inference/llama_cpp.py | 109 +- studio/backend/core/inference/orchestrator.py | 185 +- studio/backend/core/inference/worker.py | 493 +++-- studio/backend/core/training/__init__.py | 7 +- studio/backend/core/training/trainer.py | 1959 +++++++++++------ studio/backend/core/training/training.py | 151 +- studio/backend/core/training/worker.py | 658 +++--- studio/backend/loggers/__init__.py | 3 + studio/backend/loggers/config.py | 16 +- studio/backend/loggers/handlers.py | 25 +- studio/backend/main.py | 84 +- studio/backend/models/__init__.py | 1 + studio/backend/models/auth.py | 24 +- studio/backend/models/data_recipe.py | 34 +- studio/backend/models/datasets.py | 14 +- studio/backend/models/export.py | 55 +- studio/backend/models/inference.py | 220 +- studio/backend/models/models.py | 149 +- studio/backend/models/responses.py | 50 +- studio/backend/models/training.py | 207 +- studio/backend/models/users.py | 7 +- .../backend/plugins/__init__.py | 0 .../__init__.py | 2 + .../pyproject.toml | 3 + .../chunking.py | 32 +- .../data_designer_unstructured_seed/config.py | 8 +- .../data_designer_unstructured_seed/impl.py | 6 +- .../data_designer_unstructured_seed/plugin.py | 6 +- studio/backend/requirements/__init__.py | 2 + .../requirements/single-env/patch_metadata.py | 5 +- studio/backend/routes/auth.py | 68 +- studio/backend/routes/data_recipe/__init__.py | 2 +- studio/backend/routes/data_recipe/jobs.py | 49 +- studio/backend/routes/data_recipe/mcp.py | 30 +- studio/backend/routes/data_recipe/seed.py | 153 +- studio/backend/routes/data_recipe/validate.py | 28 +- studio/backend/routes/datasets.py | 216 +- studio/backend/routes/export.py | 158 +- studio/backend/routes/inference.py | 713 +++--- studio/backend/routes/models.py | 395 ++-- studio/backend/routes/training.py | 334 +-- studio/backend/run.py | 40 +- studio/backend/state/__init__.py | 2 + studio/backend/tests/__init__.py | 2 + studio/backend/tests/conftest.py | 1 + studio/backend/tests/test_utils.py | 112 +- studio/backend/utils/__init__.py | 2 + studio/backend/utils/cache_cleanup.py | 7 +- .../backend/utils/datasets/chat_templates.py | 28 +- .../backend/utils/datasets/data_collators.py | 52 +- .../backend/utils/datasets/dataset_utils.py | 393 ++-- .../utils/datasets/format_conversion.py | 308 ++- .../utils/datasets/format_detection.py | 300 ++- studio/backend/utils/datasets/llm_assist.py | 207 +- .../backend/utils/datasets/model_mappings.py | 5 - .../backend/utils/datasets/vlm_processing.py | 41 +- studio/backend/utils/hardware/__init__.py | 27 +- studio/backend/utils/hardware/hardware.py | 34 +- studio/backend/utils/inference/__init__.py | 2 +- .../utils/inference/inference_config.py | 26 +- studio/backend/utils/models/__init__.py | 33 +- studio/backend/utils/models/checkpoints.py | 17 +- studio/backend/utils/models/model_config.py | 512 +++-- studio/backend/utils/paths/__init__.py | 51 +- studio/backend/utils/paths/path_utils.py | 34 +- studio/backend/utils/paths/storage_roots.py | 39 +- studio/backend/utils/transformers_version.py | 59 +- studio/backend/utils/utils.py | 36 +- studio/frontend/.gitignore | 3 + studio/frontend/eslint.config.js | 3 + studio/frontend/index.html | 3 + studio/frontend/public/huggingface.svg | 3 + studio/frontend/public/vite.svg | 3 + studio/frontend/src/assets/react.svg | 3 + .../src/components/ui/collapsible.tsx | 80 +- studio/frontend/src/components/ui/tooltip.tsx | 124 +- studio/frontend/src/index.css | 4 +- studio/frontend/vite.config.ts | 3 + studio/install_python_stack.py | 122 +- studio/setup.bat | 3 + studio/setup.ps1 | 60 +- studio/setup.sh | 62 +- studio/tests/__init__.py | 0 studio/tests/test_openai_completions.py | 285 --- studio/tests/test_remote_lora_detection.py | 43 - studio/tests/test_sse_progress.py | 323 --- 123 files changed, 7475 insertions(+), 5806 deletions(-) delete mode 100644 README.md rename Unsloth_Studio_Colab.ipynb => studio/Unsloth_Studio_Colab.ipynb (90%) create mode 100644 studio/__init__.py create mode 100644 studio/backend/__init__.py create mode 100644 studio/backend/assets/__init__.py create mode 100644 studio/backend/assets/configs/__init__.py rename __init__.py => studio/backend/plugins/__init__.py (100%) create mode 100644 studio/backend/plugins/data-designer-unstructured-seed/__init__.py create mode 100644 studio/backend/requirements/__init__.py create mode 100644 studio/backend/utils/__init__.py delete mode 100644 studio/tests/__init__.py delete mode 100644 studio/tests/test_openai_completions.py delete mode 100644 studio/tests/test_remote_lora_detection.py delete mode 100644 studio/tests/test_sse_progress.py diff --git a/README.md b/README.md deleted file mode 100644 index 9656f1c634..0000000000 --- a/README.md +++ /dev/null @@ -1,240 +0,0 @@ -

- Unsloth Studio -

- -

πŸ¦₯ Unsloth Studio

- -

- A modern, full-stack web interface for fine-tuning, managing, and chatting with large language models β€” locally or in the cloud. -

- -

- Features β€’ - Quick Start β€’ - API Reference β€’ - Project Structure -

- ---- - -## 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 -``` - -
-What does setup.sh do? - -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) - -
- ---- - -### 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 -``` - -
-What does setup.ps1 do? - -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) - -
- ---- - -### 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 ` 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. diff --git a/cli/__init__.py b/cli/__init__.py index 2779fca281..434ee4d2ed 100644 --- a/cli/__init__.py +++ b/cli/__init__.py @@ -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.") diff --git a/cli/commands/__init__.py b/cli/commands/__init__.py index e69de29bb2..32014236c6 100644 --- a/cli/commands/__init__.py +++ b/cli/commands/__init__.py @@ -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 diff --git a/cli/commands/export.py b/cli/commands/export.py index 14f52f724e..785f309e8e 100644 --- a/cli/commands/export.py +++ b/cli/commands/export.py @@ -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) diff --git a/cli/commands/inference.py b/cli/commands/inference.py index cf36a7ad4a..18de68c4d9 100644 --- a/cli/commands/inference.py +++ b/cli/commands/inference.py @@ -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() diff --git a/cli/commands/studio.py b/cli/commands/studio.py index 1d7198c766..4658ca5d3c 100644 --- a/cli/commands/studio.py +++ b/cli/commands/studio.py @@ -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: diff --git a/cli/commands/train.py b/cli/commands/train.py index 65f27002d4..665956a80c 100644 --- a/cli/commands/train.py +++ b/cli/commands/train.py @@ -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) diff --git a/cli/commands/ui.py b/cli/commands/ui.py index 2ee8340212..10db9e243e 100644 --- a/cli/commands/ui.py +++ b/cli/commands/ui.py @@ -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) diff --git a/cli/config.py b/cli/config.py index 65c7471ccc..d5362f5673 100644 --- a/cli/config.py +++ b/cli/config.py @@ -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) diff --git a/cli/options.py b/cli/options.py index b5009dc36b..f1ab586ee5 100644 --- a/cli/options.py +++ b/cli/options.py @@ -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): diff --git a/Unsloth_Studio_Colab.ipynb b/studio/Unsloth_Studio_Colab.ipynb similarity index 90% rename from Unsloth_Studio_Colab.ipynb rename to studio/Unsloth_Studio_Colab.ipynb index 7778b2254a..15bc7d4822 100644 --- a/Unsloth_Studio_Colab.ipynb +++ b/studio/Unsloth_Studio_Colab.ipynb @@ -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, diff --git a/studio/__init__.py b/studio/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/__init__.py @@ -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 diff --git a/studio/backend/__init__.py b/studio/backend/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/__init__.py @@ -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 diff --git a/studio/backend/assets/__init__.py b/studio/backend/assets/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/assets/__init__.py @@ -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 diff --git a/studio/backend/assets/configs/__init__.py b/studio/backend/assets/configs/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/assets/configs/__init__.py @@ -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 diff --git a/studio/backend/auth/__init__.py b/studio/backend/auth/__init__.py index 4e4dbc9d72..8e489f2a9f 100644 --- a/studio/backend/auth/__init__.py +++ b/studio/backend/auth/__init__.py @@ -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", ] - - - diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index 84048fb106..c41f60bea8 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -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", ) diff --git a/studio/backend/auth/hashing.py b/studio/backend/auth/hashing.py index 1b26ef4280..873381db04 100644 --- a/studio/backend/auth/hashing.py +++ b/studio/backend/auth/hashing.py @@ -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) - diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index 79d80de448..d246f18bf1 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -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) diff --git a/studio/backend/colab.py b/studio/backend/colab.py index d54634d726..8b4294f483 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -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"""
@@ -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) diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py index 84868cea0d..051f6615a7 100644 --- a/studio/backend/core/__init__.py +++ b/studio/backend/core/__init__.py @@ -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}") diff --git a/studio/backend/core/data_recipe/jobs/__init__.py b/studio/backend/core/data_recipe/jobs/__init__.py index 175c2c108f..cf03d62a3b 100644 --- a/studio/backend/core/data_recipe/jobs/__init__.py +++ b/studio/backend/core/data_recipe/jobs/__init__.py @@ -4,4 +4,3 @@ from .manager import JobManager, get_job_manager __all__ = ["JobManager", "get_job_manager"] - diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index 7188562bfb..61c3516b2a 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -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: diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py index b5567e7634..324b62a92e 100644 --- a/studio/backend/core/data_recipe/jobs/parse.py +++ b/studio/backend/core/data_recipe/jobs/parse.py @@ -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\d+) records across (?P\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, ) diff --git a/studio/backend/core/data_recipe/jobs/types.py b/studio/backend/core/data_recipe/jobs/types.py index 2ef22a387a..3079d76bdb 100644 --- a/studio/backend/core/data_recipe/jobs/types.py +++ b/studio/backend/core/data_recipe/jobs/types.py @@ -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) diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py index c67e93751c..63e38bd18d 100644 --- a/studio/backend/core/data_recipe/jobs/worker.py +++ b/studio/backend/core/data_recipe/jobs/worker.py @@ -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 diff --git a/studio/backend/core/data_recipe/jsonable.py b/studio/backend/core/data_recipe/jsonable.py index ba4d59d6dd..8bec60cf82 100644 --- a/studio/backend/core/data_recipe/jsonable.py +++ b/studio/backend/core/data_recipe/jsonable.py @@ -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", diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py index b5914c20f3..c32b2fccaf 100644 --- a/studio/backend/core/data_recipe/local_callable_validators.py +++ b/studio/backend/core/data_recipe/local_callable_validators.py @@ -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, diff --git a/studio/backend/core/data_recipe/oxc-validator/validate.mjs b/studio/backend/core/data_recipe/oxc-validator/validate.mjs index 7d2f206ce0..ad61fb5a9e 100644 --- a/studio/backend/core/data_recipe/oxc-validator/validate.mjs +++ b/studio/backend/core/data_recipe/oxc-validator/validate.mjs @@ -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"; diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index e90bb01c38..1c86ac42d9 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -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 diff --git a/studio/backend/core/export/__init__.py b/studio/backend/core/export/__init__.py index 3021d15a4d..16fba368ef 100644 --- a/studio/backend/core/export/__init__.py +++ b/studio/backend/core/export/__init__.py @@ -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", ] diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 2057f2a0ff..a6f67c5bae 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -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/ [](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 diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 49b89cc3b1..a9fbe659b3 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -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 ========== diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 197b9fa4a8..1128a5014a 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -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(), + }, + ) diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py index 97f5a03f79..35318f6357 100644 --- a/studio/backend/core/inference/__init__.py +++ b/studio/backend/core/inference/__init__.py @@ -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", ] diff --git a/studio/backend/core/inference/audio_codecs.py b/studio/backend/core/inference/audio_codecs.py index b22dd1deab..3a418d921d 100644 --- a/studio/backend/core/inference/audio_codecs.py +++ b/studio/backend/core/inference/audio_codecs.py @@ -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) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index a2ebbc0f5c..ff02816953 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -4,6 +4,7 @@ """ Core inference backend - streamlined """ + from unsloth import FastLanguageModel, FastVisionModel from unsloth.chat_templates import get_chat_template from transformers import TextStreamer @@ -24,9 +25,9 @@ import structlog from loggers import get_logger - logger = get_logger(__name__) + class InferenceBackend: """Unified inference backend supporting text, vision, and LoRA models""" @@ -52,6 +53,7 @@ class InferenceBackend: # concurrent compare-mode requests race on the GPU. The lock is # acquired by the *background generation thread*, not the event-loop. import threading + self._generation_lock = threading.Lock() self._model_state_lock = threading.Lock() @@ -62,13 +64,15 @@ class InferenceBackend: # API supports -1 as "disable top-k"; transformers expects 0 to disable. return 0 if top_k < 0 else top_k - def load_model(self, - config: ModelConfig, - max_seq_length: int = 2048, - dtype = None, - load_in_4bit: bool = True, - hf_token: Optional[str] = None, - trust_remote_code: bool = False) -> bool: + def load_model( + self, + config: ModelConfig, + max_seq_length: int = 2048, + dtype = None, + load_in_4bit: bool = True, + hf_token: Optional[str] = None, + trust_remote_code: bool = False, + ) -> bool: """ Load any model: base, LoRA adapter, text, or vision. """ @@ -104,18 +108,21 @@ class InferenceBackend: if config.is_audio: audio_type = config.audio_type adapter_info = " (LoRA adapter)" if config.is_lora else "" - logger.info(f"Loading audio ({audio_type}) model{adapter_info}: {model_name}") + logger.info( + f"Loading audio ({audio_type}) model{adapter_info}: {model_name}" + ) log_gpu_memory(f"Before loading {model_name}") if audio_type == "csm": from unsloth import FastModel from transformers import CsmForConditionalGeneration + model, processor = FastModel.from_pretrained( config.path, - auto_model=CsmForConditionalGeneration, - load_in_4bit=False, - token=hf_token if hf_token and hf_token.strip() else None, - trust_remote_code=trust_remote_code, + auto_model = CsmForConditionalGeneration, + load_in_4bit = False, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, ) FastModel.for_inference(model) self.models[model_name]["model"] = model @@ -135,34 +142,42 @@ class InferenceBackend: else: # base_model is an HF ID β€” download it from huggingface_hub import snapshot_download + local_dir = base_path.split("/")[-1] - repo_path = snapshot_download(base_path, local_dir=local_dir) + repo_path = snapshot_download( + base_path, local_dir = local_dir + ) abs_repo_path = os.path.abspath(repo_path) - logger.info(f"Spark-TTS LoRA: loading adapter from {config.path}, BiCodec from {abs_repo_path}") + logger.info( + f"Spark-TTS LoRA: loading adapter from {config.path}, BiCodec from {abs_repo_path}" + ) model, tokenizer = FastModel.from_pretrained( config.path, - dtype=torch.float32, - load_in_4bit=False, - token=hf_token if hf_token and hf_token.strip() else None, - trust_remote_code=trust_remote_code, + dtype = torch.float32, + load_in_4bit = False, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, ) else: # Base model: download full HF repo, then load from /LLM subfolder from huggingface_hub import snapshot_download + hf_repo = config.path local_dir = hf_repo.split("/")[-1] - repo_path = snapshot_download(hf_repo, local_dir=local_dir) + repo_path = snapshot_download(hf_repo, local_dir = local_dir) abs_repo_path = os.path.abspath(repo_path) llm_path = os.path.join(abs_repo_path, "LLM") - logger.info(f"Spark-TTS: downloaded repo to {repo_path}, loading LLM from {llm_path}") + logger.info( + f"Spark-TTS: downloaded repo to {repo_path}, loading LLM from {llm_path}" + ) model, tokenizer = FastModel.from_pretrained( llm_path, - dtype=torch.float32, - load_in_4bit=False, - token=hf_token if hf_token and hf_token.strip() else None, - trust_remote_code=trust_remote_code, + dtype = torch.float32, + load_in_4bit = False, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, ) FastModel.for_inference(model) @@ -172,12 +187,13 @@ class InferenceBackend: elif audio_type == "dac": # OuteTTS uses FastModel (not FastLanguageModel) from unsloth import FastModel + model, tokenizer = FastModel.from_pretrained( config.path, - max_seq_length=max_seq_length, - load_in_4bit=False, - token=hf_token if hf_token and hf_token.strip() else None, - trust_remote_code=trust_remote_code, + max_seq_length = max_seq_length, + load_in_4bit = False, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, ) FastModel.for_inference(model) self.models[model_name]["model"] = model @@ -186,28 +202,30 @@ class InferenceBackend: # Whisper ASR β€” uses FastModel with WhisperForConditionalGeneration from unsloth import FastModel from transformers import WhisperForConditionalGeneration + model, tokenizer = FastModel.from_pretrained( config.path, - auto_model=WhisperForConditionalGeneration, - whisper_language="English", - whisper_task="transcribe", - load_in_4bit=False, - token=hf_token if hf_token and hf_token.strip() else None, - trust_remote_code=trust_remote_code, + auto_model = WhisperForConditionalGeneration, + whisper_language = "English", + whisper_task = "transcribe", + load_in_4bit = False, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, ) FastModel.for_inference(model) model.eval() # Create ASR pipeline (per notebook) from transformers import pipeline as hf_pipeline + whisper_pipe = hf_pipeline( "automatic-speech-recognition", - model=model, - tokenizer=tokenizer.tokenizer, - feature_extractor=tokenizer.feature_extractor, - processor=tokenizer, - return_language=True, - torch_dtype=torch.float16, + model = model, + tokenizer = tokenizer.tokenizer, + feature_extractor = tokenizer.feature_extractor, + processor = tokenizer, + return_language = True, + torch_dtype = torch.float16, ) self.models[model_name]["model"] = model self.models[model_name]["tokenizer"] = tokenizer @@ -215,11 +233,11 @@ class InferenceBackend: else: # SNAC (Orpheus) uses FastLanguageModel model, tokenizer = FastLanguageModel.from_pretrained( - model_name=config.path, - max_seq_length=max_seq_length, - load_in_4bit=False, - token=hf_token if hf_token and hf_token.strip() else None, - trust_remote_code=trust_remote_code, + model_name = config.path, + max_seq_length = max_seq_length, + load_in_4bit = False, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, ) FastLanguageModel.for_inference(model) self.models[model_name]["model"] = model @@ -229,7 +247,9 @@ class InferenceBackend: # (Whisper is ASR, audio_vlm is audio input β€” neither needs a codec) if audio_type not in ("whisper", "audio_vlm"): model_repo_path = self.models[model_name].get("model_repo_path") - self._audio_codec_manager.load_codec(audio_type, self.device, model_repo_path=model_repo_path) + self._audio_codec_manager.load_codec( + audio_type, self.device, model_repo_path = model_repo_path + ) self.active_model_name = model_name self.loading_models.discard(model_name) @@ -238,7 +258,9 @@ class InferenceBackend: return True model_type = "vision" if config.is_vision else "text" - adapter_info = " (LoRA adapter)" if self.models[model_name]["is_lora"] else "" + adapter_info = ( + " (LoRA adapter)" if self.models[model_name]["is_lora"] else "" + ) logger.info(f"Loading {model_type} model{adapter_info}: {model_name}") log_gpu_memory(f"Before loading {model_name}") @@ -246,12 +268,12 @@ class InferenceBackend: if config.is_vision: # Vision model (or vision LoRA adapter) model, processor = FastVisionModel.from_pretrained( - model_name=config.path, # Can be base model OR LoRA adapter path - max_seq_length=max_seq_length, - dtype=dtype, - load_in_4bit=load_in_4bit, - token=hf_token if hf_token and hf_token.strip() else None, - trust_remote_code=trust_remote_code, + model_name = config.path, # Can be base model OR LoRA adapter path + max_seq_length = max_seq_length, + dtype = dtype, + load_in_4bit = load_in_4bit, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, ) # Apply inference optimization @@ -261,10 +283,16 @@ class InferenceBackend: # instead of a proper Processor for some models (e.g. Gemma-3). # In that case, load the real processor from the base model. from transformers import ProcessorMixin - if not (isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")): + + if not ( + isinstance(processor, ProcessorMixin) + or hasattr(processor, "image_processor") + ): # For LoRA adapters, use the base model. For local merged exports, # read export_metadata.json to find the original base model. - processor_source = config.base_model if config.is_lora else config.identifier + processor_source = ( + config.base_model if config.is_lora else config.identifier + ) if not config.is_lora and config.is_local: _meta_path = Path(config.path) / "export_metadata.json" try: @@ -279,12 +307,15 @@ class InferenceBackend: f"for '{model_name}' β€” loading proper processor from '{processor_source}'" ) from transformers import AutoProcessor + processor = AutoProcessor.from_pretrained( processor_source, - token=hf_token if hf_token and hf_token.strip() else None, - trust_remote_code=trust_remote_code, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, + ) + logger.info( + f"Loaded {type(processor).__name__} from {processor_source}" ) - logger.info(f"Loaded {type(processor).__name__} from {processor_source}") self.models[model_name]["model"] = model self.models[model_name]["tokenizer"] = processor @@ -293,12 +324,12 @@ class InferenceBackend: else: # Text model (or text LoRA adapter) model, tokenizer = FastLanguageModel.from_pretrained( - model_name=config.path, # Can be base model OR LoRA adapter path - max_seq_length=max_seq_length, - dtype=dtype, - load_in_4bit=load_in_4bit, - token=hf_token if hf_token and hf_token.strip() else None, - trust_remote_code=trust_remote_code, + model_name = config.path, # Can be base model OR LoRA adapter path + max_seq_length = max_seq_length, + dtype = dtype, + load_in_4bit = load_in_4bit, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, ) # Apply inference optimization @@ -351,6 +382,7 @@ class InferenceBackend: # Remove stale compiled cache so the next model gets a fresh one from utils.cache_cleanup import clear_unsloth_compiled_cache + clear_unsloth_compiled_cache() logger.info(f"Model '{model_name}' successfully unloaded.") @@ -359,7 +391,9 @@ class InferenceBackend: logger.error(f"Error while unloading model '{model_name}': {e}") return False else: - logger.warning(f"Attempted to unload model '{model_name}', but it was not found in the registry.") + logger.warning( + f"Attempted to unload model '{model_name}', but it was not found in the registry." + ) return True def revert_to_base_model(self, base_model_name: str) -> bool: @@ -384,7 +418,7 @@ class InferenceBackend: # After model.unload(), the base model may still carry a peft_config # attribute. Removing it ensures PeftModel.from_pretrained() gets # a clean base model without "multiple adapters" warnings. - if hasattr(model, 'peft_config'): + if hasattr(model, "peft_config"): del model.peft_config logger.info(f"Model '{base_model_name}' reverted to clean base state.") @@ -393,12 +427,18 @@ class InferenceBackend: except Exception as e: logger.error(f"Failed to revert model to base state: {e}") import traceback + logger.error(traceback.format_exc()) return False - def load_for_eval(self, lora_path: str, max_seq_length: int = 2048, - dtype = None, load_in_4bit: bool = True, - hf_token: Optional[str] = None) -> Tuple[bool, Optional[str], Optional[str]]: + def load_for_eval( + self, + lora_path: str, + max_seq_length: int = 2048, + dtype = None, + load_in_4bit: bool = True, + hf_token: Optional[str] = None, + ) -> Tuple[bool, Optional[str], Optional[str]]: """ Final Corrected Version: Ensures the base model and the specified adapter are loaded. @@ -406,6 +446,7 @@ class InferenceBackend: """ try: from utils.models import ModelConfig + lora_config = ModelConfig.from_lora_path(lora_path, hf_token) if not lora_config: return False, None, None @@ -413,10 +454,16 @@ class InferenceBackend: base_model_name = lora_config.base_model # 1. Load the base model if it's not already in memory - if base_model_name not in self.models or not self.models[base_model_name].get("model"): + if base_model_name not in self.models or not self.models[ + base_model_name + ].get("model"): logger.info(f"Base model '{base_model_name}' not loaded, loading now.") - base_config = ModelConfig.from_ui_selection(base_model_name, None, is_lora=False) - if not self.load_model(base_config, max_seq_length, dtype, load_in_4bit, hf_token): + base_config = ModelConfig.from_ui_selection( + base_model_name, None, is_lora = False + ) + if not self.load_model( + base_config, max_seq_length, dtype, load_in_4bit, hf_token + ): return False, None, None self.active_model_name = base_model_name @@ -427,9 +474,9 @@ class InferenceBackend: # 3. Call our robust load_adapter function to ensure this specific adapter is loaded. # It will only load from disk if the model doesn't already have it. adapter_success = self.load_adapter( - base_model_name=base_model_name, - adapter_path=lora_path, - adapter_name=adapter_name + base_model_name = base_model_name, + adapter_path = lora_path, + adapter_name = adapter_name, ) if not adapter_success: return False, base_model_name, None @@ -440,10 +487,13 @@ class InferenceBackend: except Exception as e: logger.error(f"Error during load_for_eval: {e}") import traceback + logger.error(traceback.format_exc()) return False, None, None - def load_adapter(self, base_model_name: str, adapter_path: str, adapter_name: str) -> bool: + def load_adapter( + self, base_model_name: str, adapter_path: str, adapter_name: str + ) -> bool: """ Loads an adapter onto the model ONLY if it's not already attached. """ @@ -451,20 +501,26 @@ class InferenceBackend: # Check if this adapter name is already part of the model's config. This is the most reliable check. if hasattr(model, "peft_config") and adapter_name in model.peft_config: - logger.info(f"Adapter '{adapter_name}' is already attached to the model. Skipping load.") + logger.info( + f"Adapter '{adapter_name}' is already attached to the model. Skipping load." + ) return True try: - logger.info(f"Loading new adapter '{adapter_name}' from '{adapter_path}' onto {base_model_name}") - model.load_adapter(adapter_path, adapter_name=adapter_name) + logger.info( + f"Loading new adapter '{adapter_name}' from '{adapter_path}' onto {base_model_name}" + ) + model.load_adapter(adapter_path, adapter_name = adapter_name) # Update our internal registry ONLY after a successful load. if "loaded_adapters" not in self.models[base_model_name]: self.models[base_model_name]["loaded_adapters"] = {} self.models[base_model_name]["loaded_adapters"][adapter_name] = adapter_path - total_adapters = len(getattr(model, 'peft_config', {})) - logger.info(f"Adapter '{adapter_name}' loaded successfully. (Total unique adapters on model: {total_adapters})") + total_adapters = len(getattr(model, "peft_config", {})) + logger.info( + f"Adapter '{adapter_name}' loaded successfully. (Total unique adapters on model: {total_adapters})" + ) return True except Exception as e: logger.error(f"Failed to load adapter '{adapter_name}': {e}") @@ -513,15 +569,21 @@ class InferenceBackend: if use_adapter is False: # Disable LoRA layers β†’ base model output if isinstance(model, (PeftModel, PeftModelForCausalLM)): - logger.info(f"Compare mode: disabling adapters on '{base}' for base model generation") + logger.info( + f"Compare mode: disabling adapters on '{base}' for base model generation" + ) model.base_model.disable_adapter_layers() else: - logger.info(f"Compare mode: model '{base}' is not a PeftModel, already base") + logger.info( + f"Compare mode: model '{base}' is not a PeftModel, already base" + ) elif use_adapter is True: # Re-enable LoRA layers β†’ adapter output if isinstance(model, (PeftModel, PeftModelForCausalLM)): - logger.info(f"Compare mode: enabling adapters on '{base}' for LoRA generation") + logger.info( + f"Compare mode: enabling adapters on '{base}' for LoRA generation" + ) model.base_model.enable_adapter_layers() else: logger.warning("use_adapter=true but model is not a PeftModel") @@ -529,16 +591,20 @@ class InferenceBackend: elif isinstance(use_adapter, str): # Enable adapters and set the specific one active if isinstance(model, (PeftModel, PeftModelForCausalLM)): - logger.info(f"Compare mode: enabling adapter '{use_adapter}' on '{base}'") + logger.info( + f"Compare mode: enabling adapter '{use_adapter}' on '{base}'" + ) model.base_model.enable_adapter_layers() self.set_active_adapter(base, use_adapter) else: - logger.warning(f"use_adapter='{use_adapter}' but model is not a PeftModel") + logger.warning( + f"use_adapter='{use_adapter}' but model is not a PeftModel" + ) 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]: """ @@ -554,49 +620,53 @@ class InferenceBackend: **gen_kwargs: Forwarded to generate_chat_response. """ yield from self._generate_chat_response_inner( - cancel_event=cancel_event, _adapter_state=use_adapter, **gen_kwargs + cancel_event = cancel_event, _adapter_state = use_adapter, **gen_kwargs ) - def generate_chat_response(self, - messages: list, - system_prompt: str, - 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) -> Generator[str, None, None]: + def generate_chat_response( + self, + messages: list, + system_prompt: str, + 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, + ) -> Generator[str, None, None]: """ Generate response for text or vision models. The generation lock is acquired by the background generation thread. """ yield from self._generate_chat_response_inner( - messages=messages, - system_prompt=system_prompt, - image=image, - temperature=temperature, - top_p=top_p, - top_k=top_k, - min_p=min_p, - max_new_tokens=max_new_tokens, - repetition_penalty=repetition_penalty, - cancel_event=cancel_event, + 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, ) - def _generate_chat_response_inner(self, - messages: list, - system_prompt: str = "", - image=None, - temperature: float = 0.7, - top_p: float = 0.9, - top_k: int = 40, - min_p: float = 0.0, - max_new_tokens: int = 256, - repetition_penalty: float = 1.1, - cancel_event=None, - _adapter_state=None) -> Generator[str, None, None]: + def _generate_chat_response_inner( + self, + messages: list, + system_prompt: str = "", + image = None, + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 40, + min_p: float = 0.0, + max_new_tokens: int = 256, + repetition_penalty: float = 1.1, + cancel_event = None, + _adapter_state = None, + ) -> Generator[str, None, None]: """ Inner generation logic. Called by both generate_chat_response and generate_with_adapter_control. @@ -621,16 +691,24 @@ class InferenceBackend: # FastVisionModel may return a raw tokenizer (e.g. GemmaTokenizerFast) # instead of a proper ProcessorMixin for some models (e.g. Gemma-3). from transformers import ProcessorMixin + processor = model_info.get("processor") - has_image_processing = ( - processor is not None - and (isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")) + has_image_processing = processor is not None and ( + isinstance(processor, ProcessorMixin) + or hasattr(processor, "image_processor") ) if has_image_processing: yield from self._generate_vision_response( - messages, system_prompt, image, - temperature, top_p, top_k, min_p, max_new_tokens, repetition_penalty, - cancel_event=cancel_event, + messages, + system_prompt, + image, + temperature, + top_p, + top_k, + min_p, + max_new_tokens, + repetition_penalty, + cancel_event = cancel_event, ) return else: @@ -645,31 +723,36 @@ class InferenceBackend: # Step 1: Apply get_chat_template if model is in mapper try: - from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, get_tokenizer_chat_template + from utils.datasets import ( + MODEL_TO_TEMPLATE_MAPPER, + get_tokenizer_chat_template, + ) model_name_lower = self.active_model_name.lower() # Check if model has a registered template if model_name_lower in MODEL_TO_TEMPLATE_MAPPER: template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower] - logger.info(f"Applying chat template '{template_name}' for {self.active_model_name}") + logger.info( + f"Applying chat template '{template_name}' for {self.active_model_name}" + ) # This modifies the tokenizer with the correct template tokenizer = get_chat_template( tokenizer, - chat_template=template_name, + chat_template = template_name, ) else: - logger.info(f"No registered template for {self.active_model_name}, using tokenizer default") + logger.info( + f"No registered template for {self.active_model_name}, using tokenizer default" + ) except Exception as e: logger.warning(f"Could not apply get_chat_template: {e}") # Step 2: Format with tokenizer.apply_chat_template() try: formatted_prompt = tokenizer.apply_chat_template( - messages, - tokenize=False, - add_generation_prompt=True + messages, tokenize = False, add_generation_prompt = True ) logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...") except Exception as e: @@ -679,14 +762,30 @@ class InferenceBackend: # Step 3: Generate yield from self.generate_stream( - formatted_prompt, temperature, top_p, top_k, min_p, max_new_tokens, repetition_penalty, - cancel_event=cancel_event, - _adapter_state=_adapter_state, + formatted_prompt, + temperature, + top_p, + top_k, + min_p, + max_new_tokens, + repetition_penalty, + cancel_event = cancel_event, + _adapter_state = _adapter_state, ) - def _generate_vision_response(self, messages, system_prompt, image, - temperature, top_p, top_k, min_p, max_new_tokens, - repetition_penalty, cancel_event=None) -> Generator[str, None, None]: + def _generate_vision_response( + self, + messages, + system_prompt, + image, + temperature, + top_p, + top_k, + min_p, + max_new_tokens, + repetition_penalty, + cancel_event = None, + ) -> Generator[str, None, None]: """Handle vision model generation with true token-by-token streaming.""" model_info = self.models[self.active_model_name] model = model_info["model"] @@ -699,8 +798,9 @@ class InferenceBackend: user_message = "" if messages and messages[-1]["role"] == "user": import re + user_message = messages[-1]["content"] - user_message = re.sub(r']*>', '', user_message).strip() + user_message = re.sub(r"]*>", "", user_message).strip() if not user_message: user_message = "Describe this image." if image else "Hello" @@ -712,22 +812,26 @@ class InferenceBackend: "role": "user", "content": [ {"type": "image"}, - {"type": "text", "text": user_message} + {"type": "text", "text": user_message}, ], } ] - input_text = processor.apply_chat_template(vision_messages, add_generation_prompt=True, tokenize=False) + input_text = processor.apply_chat_template( + vision_messages, add_generation_prompt = True, tokenize = False + ) inputs = processor( image, input_text, - add_special_tokens=False, - return_tensors="pt", + add_special_tokens = False, + return_tensors = "pt", ).to(self.device) else: # Text-only for vision model formatted_prompt = self.format_chat_prompt(messages, system_prompt) - inputs = raw_tokenizer(formatted_prompt, return_tensors="pt").to(self.device) + inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to( + self.device + ) # Stream with TextIteratorStreamer + background thread try: @@ -736,21 +840,21 @@ class InferenceBackend: streamer = TextIteratorStreamer( raw_tokenizer, - skip_prompt=True, - skip_special_tokens=True, - timeout=0.2, + skip_prompt = True, + skip_special_tokens = True, + timeout = 0.2, ) generation_kwargs = dict( **inputs, - streamer=streamer, - max_new_tokens=max_new_tokens, - use_cache=True, - do_sample=temperature > 0, - temperature=temperature, - top_p=top_p, - top_k=top_k, - min_p=min_p, + streamer = streamer, + max_new_tokens = max_new_tokens, + use_cache = True, + do_sample = temperature > 0, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, ) err: dict[str, str] = {} @@ -768,11 +872,12 @@ class InferenceBackend: except Exception: pass - thread = threading.Thread(target=generate_fn) + thread = threading.Thread(target = generate_fn) thread.start() output = "" from queue import Empty + try: while True: if cancel_event is not None and cancel_event.is_set(): @@ -792,9 +897,11 @@ class InferenceBackend: finally: if cancel_event is not None: cancel_event.set() - thread.join(timeout=10) + thread.join(timeout = 10) if thread.is_alive(): - logger.warning("Vision generation thread did not exit after cancel/join timeout") + logger.warning( + "Vision generation thread did not exit after cancel/join timeout" + ) if err.get("msg"): yield f"Error: {err['msg']}" @@ -803,10 +910,19 @@ class InferenceBackend: logger.error(f"Vision generation error: {e}") yield f"Error: {str(e)}" - def generate_audio_input_response(self, messages, system_prompt, audio_array, - temperature, top_p, top_k, min_p, - max_new_tokens, repetition_penalty, - cancel_event=None) -> Generator[str, None, None]: + def generate_audio_input_response( + self, + messages, + system_prompt, + audio_array, + temperature, + top_p, + top_k, + min_p, + max_new_tokens, + repetition_penalty, + cancel_event = None, + ) -> Generator[str, None, None]: """Handle audio input (ASR) generation β€” accepts audio numpy array, streams text output. Uses processor.apply_chat_template with audio embedded in messages (Gemma 3n pattern). @@ -846,11 +962,11 @@ class InferenceBackend: # apply_chat_template handles audio embedding + tokenization in one step inputs = processor.apply_chat_template( audio_messages, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - return_tensors="pt", - truncation=False, + add_generation_prompt = True, + tokenize = True, + return_dict = True, + return_tensors = "pt", + truncation = False, ).to(self.device) try: @@ -859,18 +975,18 @@ class InferenceBackend: streamer = TextIteratorStreamer( raw_tokenizer, - skip_prompt=True, - skip_special_tokens=True, - timeout=0.2, + skip_prompt = True, + skip_special_tokens = True, + timeout = 0.2, ) # Notebook uses do_sample=False for ASR (greedy decoding for accuracy) generation_kwargs = dict( **inputs, - streamer=streamer, - max_new_tokens=max_new_tokens, - use_cache=True, - do_sample=False, + streamer = streamer, + max_new_tokens = max_new_tokens, + use_cache = True, + do_sample = False, ) err: dict[str, str] = {} @@ -888,7 +1004,7 @@ class InferenceBackend: except Exception: pass - thread = threading.Thread(target=generate_fn) + thread = threading.Thread(target = generate_fn) thread.start() output = "" @@ -910,9 +1026,11 @@ class InferenceBackend: finally: if cancel_event is not None: cancel_event.set() - thread.join(timeout=10) + thread.join(timeout = 10) if thread.is_alive(): - logger.warning("Audio input generation thread did not exit after cancel/join timeout") + logger.warning( + "Audio input generation thread did not exit after cancel/join timeout" + ) if err.get("msg"): yield f"Error: {err['msg']}" @@ -921,7 +1039,9 @@ class InferenceBackend: logger.error(f"Audio input generation error: {e}") yield f"Error: {str(e)}" - def generate_whisper_response(self, audio_array, cancel_event=None) -> Generator[str, None, None]: + def generate_whisper_response( + self, audio_array, cancel_event = None + ) -> Generator[str, None, None]: """Whisper ASR β€” takes audio numpy array, yields transcribed text. Uses the pre-built transformers pipeline (created during model loading). @@ -943,16 +1063,18 @@ class InferenceBackend: logger.error(f"Whisper ASR error: {e}") yield f"Error: {str(e)}" - def generate_stream(self, - prompt: str, - 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, - _adapter_state=None) -> Generator[str, None, None]: + def generate_stream( + self, + prompt: str, + 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, + _adapter_state = None, + ) -> Generator[str, None, None]: """Generate streaming text response (text models only). _adapter_state: if not None, the background thread toggles adapters @@ -971,30 +1093,32 @@ class InferenceBackend: tokenizer = getattr(tokenizer, "tokenizer", tokenizer) try: - inputs = tokenizer(prompt, return_tensors="pt").to(model.device) + inputs = tokenizer(prompt, return_tensors = "pt").to(model.device) from transformers import TextIteratorStreamer import threading streamer = TextIteratorStreamer( tokenizer, - skip_prompt=True, - skip_special_tokens=True, - timeout=0.2, + skip_prompt = True, + skip_special_tokens = True, + timeout = 0.2, ) generation_kwargs = dict( **inputs, - streamer=streamer, - max_new_tokens=max_new_tokens, - temperature=temperature, - top_p=top_p, - top_k=top_k, - min_p=min_p, - repetition_penalty=repetition_penalty, - do_sample=temperature > 0, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.eos_token_id if tokenizer.pad_token_id is None else tokenizer.pad_token_id, + streamer = streamer, + max_new_tokens = max_new_tokens, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + repetition_penalty = repetition_penalty, + do_sample = temperature > 0, + eos_token_id = tokenizer.eos_token_id, + pad_token_id = tokenizer.eos_token_id + if tokenizer.pad_token_id is None + else tokenizer.pad_token_id, ) if cancel_event is not None: from transformers.generation.stopping_criteria import ( @@ -1029,11 +1153,12 @@ class InferenceBackend: pass err: dict[str, str] = {} - thread = threading.Thread(target=generate_fn) + thread = threading.Thread(target = generate_fn) thread.start() output = "" from queue import Empty + try: while True: if cancel_event is not None and cancel_event.is_set(): @@ -1053,9 +1178,11 @@ class InferenceBackend: finally: if cancel_event is not None: cancel_event.set() - thread.join(timeout=10) + thread.join(timeout = 10) if thread.is_alive(): - logger.warning("Generation thread did not exit after cancel/join timeout") + logger.warning( + "Generation thread did not exit after cancel/join timeout" + ) if err.get("msg"): yield f"Error: {err['msg']}" @@ -1100,83 +1227,139 @@ class InferenceBackend: self._apply_adapter_state(use_adapter) if audio_type == "snac": - return self._generate_snac(model, tokenizer, text, temperature, top_p, max_new_tokens, repetition_penalty) + return self._generate_snac( + model, + tokenizer, + text, + temperature, + top_p, + max_new_tokens, + repetition_penalty, + ) elif audio_type == "csm": processor = model_info.get("processor", tokenizer) return self._generate_csm(model, processor, text, max_new_tokens) elif audio_type == "bicodec": - return self._generate_bicodec(model, tokenizer, text, temperature, top_k, max_new_tokens) + return self._generate_bicodec( + model, tokenizer, text, temperature, top_k, max_new_tokens + ) elif audio_type == "dac": - return self._generate_dac(model, tokenizer, text, temperature, top_k, top_p, min_p, max_new_tokens, repetition_penalty) + return self._generate_dac( + model, + tokenizer, + text, + temperature, + top_k, + top_p, + min_p, + max_new_tokens, + repetition_penalty, + ) else: raise RuntimeError(f"Unknown audio_type: {audio_type}") - def _generate_snac(self, model, tokenizer, text, temperature, top_p, max_new_tokens, repetition_penalty): + def _generate_snac( + self, + model, + tokenizer, + text, + temperature, + top_p, + max_new_tokens, + repetition_penalty, + ): """Generate audio using SNAC codec (Orpheus).""" device = model.device - start_token = torch.tensor([[128259]], device=device) # START_OF_HUMAN - end_tokens = torch.tensor([[128009, 128260]], device=device) # EOT, END_OF_HUMAN - text_ids = tokenizer(text, return_tensors="pt").input_ids.to(device) - input_ids = torch.cat([start_token, text_ids, end_tokens], dim=1) + start_token = torch.tensor([[128259]], device = device) # START_OF_HUMAN + end_tokens = torch.tensor( + [[128009, 128260]], device = device + ) # EOT, END_OF_HUMAN + text_ids = tokenizer(text, return_tensors = "pt").input_ids.to(device) + input_ids = torch.cat([start_token, text_ids, end_tokens], dim = 1) attention_mask = torch.ones_like(input_ids) generated = model.generate( - input_ids=input_ids, - attention_mask=attention_mask, - max_new_tokens=max_new_tokens, - do_sample=True, - temperature=temperature, - top_p=top_p, - repetition_penalty=repetition_penalty, - eos_token_id=128258, # END_OF_SPEECH - use_cache=True, + input_ids = input_ids, + attention_mask = attention_mask, + max_new_tokens = max_new_tokens, + do_sample = True, + temperature = temperature, + top_p = top_p, + repetition_penalty = repetition_penalty, + eos_token_id = 128258, # END_OF_SPEECH + use_cache = True, ) return self._audio_codec_manager.decode_snac(generated, str(device)) def _generate_csm(self, model, processor, text, max_new_tokens): """Generate audio using CSM (Sesame).""" speaker_id = 0 - inputs = processor(f"[{speaker_id}]{text}", add_special_tokens=True, return_tensors="pt").to(model.device) - audio_values = model.generate(**inputs, max_new_tokens=max_new_tokens, output_audio=True) + inputs = processor( + f"[{speaker_id}]{text}", add_special_tokens = True, return_tensors = "pt" + ).to(model.device) + audio_values = model.generate( + **inputs, max_new_tokens = max_new_tokens, output_audio = True + ) return self._audio_codec_manager.decode_csm(audio_values) - def _generate_bicodec(self, model, tokenizer, text, temperature, top_k, max_new_tokens): + def _generate_bicodec( + self, model, tokenizer, text, temperature, top_k, max_new_tokens + ): """Generate audio using BiCodec (Spark-TTS).""" - prompt = "<|task_tts|><|start_content|>" + text + "<|end_content|><|start_global_token|>" - inputs = tokenizer([prompt], return_tensors="pt").to(model.device) + prompt = ( + "<|task_tts|><|start_content|>" + + text + + "<|end_content|><|start_global_token|>" + ) + inputs = tokenizer([prompt], return_tensors = "pt").to(model.device) generated = model.generate( **inputs, - max_new_tokens=max_new_tokens, - do_sample=True, - temperature=temperature, - top_k=top_k, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id, + max_new_tokens = max_new_tokens, + do_sample = True, + temperature = temperature, + top_k = top_k, + eos_token_id = tokenizer.eos_token_id, + pad_token_id = tokenizer.pad_token_id, ) - new_tokens = generated[:, inputs.input_ids.shape[1]:] - decoded_text = tokenizer.batch_decode(new_tokens, skip_special_tokens=False)[0] + new_tokens = generated[:, inputs.input_ids.shape[1] :] + decoded_text = tokenizer.batch_decode(new_tokens, skip_special_tokens = False)[0] return self._audio_codec_manager.decode_bicodec(decoded_text, str(model.device)) - def _generate_dac(self, model, tokenizer, text, temperature, top_k, top_p, min_p, max_new_tokens, repetition_penalty): + def _generate_dac( + self, + model, + tokenizer, + text, + temperature, + top_k, + top_p, + min_p, + max_new_tokens, + repetition_penalty, + ): """Generate audio using DAC (OuteTTS). Follows Oute_TTS_(1B).ipynb exactly.""" # Monkey-patch RepetitionPenaltyLogitsProcessor with a 64-token penalty # window (same as the OuteTTS notebook) to avoid degenerate repetition. self._patch_repetition_penalty_processor() - prompt = "<|im_start|>\n<|text_start|>" + text + "<|text_end|>\n<|audio_start|><|global_features_start|>\n" + prompt = ( + "<|im_start|>\n<|text_start|>" + + text + + "<|text_end|>\n<|audio_start|><|global_features_start|>\n" + ) with torch.inference_mode(): - with torch.amp.autocast('cuda', dtype=model.dtype): - inputs = tokenizer([prompt], return_tensors="pt").to(model.device) + with torch.amp.autocast("cuda", dtype = model.dtype): + inputs = tokenizer([prompt], return_tensors = "pt").to(model.device) generated = model.generate( **inputs, - temperature=temperature, - top_k=top_k, - top_p=top_p, - min_p=min_p, - repetition_penalty=repetition_penalty, - max_new_tokens=max_new_tokens, + temperature = temperature, + top_k = top_k, + top_p = top_p, + min_p = min_p, + repetition_penalty = repetition_penalty, + max_new_tokens = max_new_tokens, ) - decoded_text = tokenizer.batch_decode(generated, skip_special_tokens=False)[0] + decoded_text = tokenizer.batch_decode(generated, skip_special_tokens = False)[0] return self._audio_codec_manager.decode_dac(decoded_text, str(model.device)) _repetition_penalty_patched = False @@ -1199,11 +1382,15 @@ class InferenceBackend: def __init__(self, penalty: float): self.penalty_last_n = 64 if not isinstance(penalty, float) or penalty <= 0: - raise ValueError(f"`penalty` has to be a positive float, but is {penalty}") + raise ValueError( + f"`penalty` has to be a positive float, but is {penalty}" + ) self.penalty = penalty @torch.no_grad() - def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + def __call__( + self, input_ids: torch.LongTensor, scores: torch.FloatTensor + ) -> torch.FloatTensor: if self.penalty_last_n == 0 or self.penalty == 1.0: return scores batch_size, seq_len = input_ids.shape @@ -1217,11 +1404,17 @@ class InferenceBackend: if token_id >= vocab_size: continue logit = scores[b, token_id] - scores[b, token_id] = logit * self.penalty if logit <= 0 else logit / self.penalty + scores[b, token_id] = ( + logit * self.penalty if logit <= 0 else logit / self.penalty + ) return scores - generation_utils.RepetitionPenaltyLogitsProcessor = RepetitionPenaltyLogitsProcessorPatch - logger.info("Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS") + generation_utils.RepetitionPenaltyLogitsProcessor = ( + RepetitionPenaltyLogitsProcessorPatch + ) + logger.info( + "Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS" + ) def format_chat_prompt(self, messages: list, system_prompt: str = None) -> str: if not self.active_model_name or self.active_model_name not in self.models: @@ -1232,7 +1425,9 @@ class InferenceBackend: logger.error("Tokenizer not loaded for active model") return "" - chat_template_info = self.models[self.active_model_name].get("chat_template_info", {}) + chat_template_info = self.models[self.active_model_name].get( + "chat_template_info", {} + ) tokenizer = self.models[self.active_model_name]["tokenizer"] tokenizer = getattr(tokenizer, "tokenizer", tokenizer) @@ -1249,12 +1444,15 @@ class InferenceBackend: if role in ["system", "user", "assistant"] and content.strip(): if role == last_role: - logger.debug(f"Skipping consecutive {role} message to maintain alternation") + logger.debug( + f"Skipping consecutive {role} message to maintain alternation" + ) continue if role == "user": import re - clean_content = re.sub(r'<[^>]+>', '', content).strip() + + clean_content = re.sub(r"<[^>]+>", "", content).strip() if clean_content: chat_messages.append({"role": role, "content": clean_content}) last_role = role @@ -1265,7 +1463,9 @@ class InferenceBackend: continue if chat_messages and chat_messages[-1]["role"] == "assistant": - logger.debug("Removing final assistant message to ensure proper alternation") + logger.debug( + "Removing final assistant message to ensure proper alternation" + ) chat_messages.pop() logger.info(f"Sending {len(chat_messages)} messages to tokenizer:") @@ -1274,31 +1474,44 @@ class InferenceBackend: try: formatted_prompt = tokenizer.apply_chat_template( - chat_messages, - tokenize=False, - add_generation_prompt=True + chat_messages, tokenize = False, add_generation_prompt = True ) logger.info(f"Successfully applied tokenizer's native chat template") return formatted_prompt except Exception as e: error_msg = str(e).lower() - if "chat_template is not set" in error_msg or "no template argument" in error_msg: - logger.info(f"Base model detected - no built-in chat template available, using fallback formatting") + if ( + "chat_template is not set" in error_msg + or "no template argument" in error_msg + ): + logger.info( + f"Base model detected - no built-in chat template available, using fallback formatting" + ) else: logger.warning(f"Failed to apply tokenizer chat template: {e}") - logger.debug(f"""Failed with messages: {[f"{m['role']}: {m['content'][:30]}..." for m in chat_messages]}""") + logger.debug( + f"""Failed with messages: {[f"{m['role']}: {m['content'][:30]}..." for m in chat_messages]}""" + ) if chat_template_info.get("has_template", False): - logger.info("Falling back to manual template formatting based on detected patterns") + logger.info( + "Falling back to manual template formatting based on detected patterns" + ) template_type = chat_template_info.get("format_type", "generic") - manual_prompt = self._format_chat_manual(chat_messages, template_type, chat_template_info.get("special_tokens", {})) + manual_prompt = self._format_chat_manual( + chat_messages, + template_type, + chat_template_info.get("special_tokens", {}), + ) logger.info(f"Manual template result: {manual_prompt[:200]}...") return manual_prompt else: logger.info("Using generic chat formatting for base model") return self._format_generic_template(chat_messages, {}) - def _format_chat_manual(self, messages: list, template_type: str, special_tokens: dict) -> str: + def _format_chat_manual( + self, messages: list, template_type: str, special_tokens: dict + ) -> str: """ Manual chat formatting fallback for when tokenizer template fails @@ -1329,7 +1542,9 @@ class InferenceBackend: for msg in messages: role = msg["role"] content = msg["content"] - formatted += f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>" + formatted += ( + f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>" + ) formatted += "<|start_header_id|>assistant<|end_header_id|>\n\n" return formatted @@ -1358,7 +1573,10 @@ class InferenceBackend: formatted += f"[INST] {user_content} [/INST]" - if i + 1 < len(conversation) and conversation[i + 1]["role"] == "assistant": + if ( + i + 1 < len(conversation) + and conversation[i + 1]["role"] == "assistant" + ): formatted += f" {conversation[i + 1]['content']}" i += 2 else: @@ -1435,11 +1653,11 @@ class InferenceBackend: try: # This is a common pattern for Unsloth/Hugging Face models - if hasattr(model, 'past_key_values'): + if hasattr(model, "past_key_values"): model.past_key_values = None - if hasattr(model, 'generation_config'): - if hasattr(model.generation_config, 'past_key_values'): - model.generation_config.past_key_values = None + if hasattr(model, "generation_config"): + if hasattr(model.generation_config, "past_key_values"): + model.generation_config.past_key_values = None logger.debug(f"Reset generation state for model: {model_name}") except Exception as e: @@ -1456,6 +1674,7 @@ class InferenceBackend: logger.debug("Cleared GPU cache") import gc + gc.collect() logger.info("Performed comprehensive generation state reset") @@ -1468,8 +1687,9 @@ class InferenceBackend: 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)) + + 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) return img @@ -1483,7 +1703,9 @@ class InferenceBackend: return text.strip() def _load_chat_template_info(self, model_name: str): - if model_name not in self.models or not self.models[model_name].get("tokenizer"): + if model_name not in self.models or not self.models[model_name].get( + "tokenizer" + ): return tokenizer = self.models[model_name]["tokenizer"] @@ -1497,29 +1719,43 @@ class InferenceBackend: try: from utils.datasets import MODEL_TO_TEMPLATE_MAPPER - #Try exact match first + + # Try exact match first model_name_lower = model_name.lower() if model_name_lower in MODEL_TO_TEMPLATE_MAPPER: - chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[model_name_lower] - logger.info(f"Detected template '{chat_template_info['template_name']}' for {model_name} from mapper") + chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[ + model_name_lower + ] + logger.info( + f"Detected template '{chat_template_info['template_name']}' for {model_name} from mapper" + ) else: # Try partial match (for variants like model_name-bnb-4bit) for key in MODEL_TO_TEMPLATE_MAPPER: if key in model_name_lower or model_name_lower in key: - chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[key] - logger.info(f"Detected template '{chat_template_info['template_name']}' for {model_name} (partial match)") + chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[ + key + ] + logger.info( + f"Detected template '{chat_template_info['template_name']}' for {model_name} (partial match)" + ) break except Exception as e: - logger.warning(f"Could not detect template from mapper for {model_name}: {e}") + logger.warning( + f"Could not detect template from mapper for {model_name}: {e}" + ) try: - if hasattr(tokenizer, 'chat_template') and tokenizer.chat_template: + if hasattr(tokenizer, "chat_template") and tokenizer.chat_template: chat_template_info["has_template"] = True chat_template_info["template"] = tokenizer.chat_template template_str = tokenizer.chat_template.lower() - if "start_header_id" in template_str and "end_header_id" in template_str: + if ( + "start_header_id" in template_str + and "end_header_id" in template_str + ): chat_template_info["format_type"] = "llama3" elif "[inst]" in template_str and "[/inst]" in template_str: chat_template_info["format_type"] = "mistral" @@ -1530,21 +1766,25 @@ class InferenceBackend: else: chat_template_info["format_type"] = "custom" - logger.info(f"Loaded chat template for {model_name} (detected as {chat_template_info['format_type']} format)") + logger.info( + f"Loaded chat template for {model_name} (detected as {chat_template_info['format_type']} format)" + ) logger.debug(f"Template preview: {tokenizer.chat_template[:200]}...") special_tokens = {} - if hasattr(tokenizer, 'bos_token') and tokenizer.bos_token: + if hasattr(tokenizer, "bos_token") and tokenizer.bos_token: special_tokens["bos_token"] = tokenizer.bos_token - if hasattr(tokenizer, 'eos_token') and tokenizer.eos_token: + if hasattr(tokenizer, "eos_token") and tokenizer.eos_token: special_tokens["eos_token"] = tokenizer.eos_token - if hasattr(tokenizer, 'pad_token') and tokenizer.pad_token: + if hasattr(tokenizer, "pad_token") and tokenizer.pad_token: special_tokens["pad_token"] = tokenizer.pad_token chat_template_info["special_tokens"] = special_tokens else: - logger.info(f"No chat template found for {model_name}, will use generic formatting") + logger.info( + f"No chat template found for {model_name}, will use generic formatting" + ) except Exception as e: logger.error(f"Error loading chat template info for {model_name}: {e}") @@ -1552,10 +1792,13 @@ class InferenceBackend: self.models[model_name]["chat_template_info"] = chat_template_info if chat_template_info["has_template"]: - logger.info(f"Chat template loaded for {model_name}: {chat_template_info['format_type']} format") + logger.info( + f"Chat template loaded for {model_name}: {chat_template_info['format_type']} format" + ) else: - logger.info(f"No built-in chat template for {model_name}, will use generic formatting") - + logger.info( + f"No built-in chat template for {model_name}, will use generic formatting" + ) def get_current_model(self) -> Optional[str]: """Get currently active model name""" @@ -1569,11 +1812,13 @@ class InferenceBackend: """Get name of currently loading model""" return next(iter(self.loading_models)) if self.loading_models else None - def load_model_simple(self, - model_path: str, - hf_token: Optional[str] = None, - max_seq_length: int = 2048, - load_in_4bit: bool = True) -> bool: + def load_model_simple( + self, + model_path: str, + hf_token: Optional[str] = None, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + ) -> bool: """ Simple model loading wrapper for chat interface. Accepts model path as string and handles ModelConfig creation internally. @@ -1591,17 +1836,17 @@ class InferenceBackend: # Create config from string path config = ModelConfig.from_ui_selection( model_path, - lora_path=None, # No LoRA for chat - is_lora=False + lora_path = None, # No LoRA for chat + is_lora = False, ) # Call existing load_model with config return self.load_model( - config=config, - max_seq_length=max_seq_length, - dtype=None, # Auto-detect - load_in_4bit=load_in_4bit, - hf_token=hf_token + config = config, + max_seq_length = max_seq_length, + dtype = None, # Auto-detect + load_in_4bit = load_in_4bit, + hf_token = hf_token, ) except Exception as e: @@ -1609,9 +1854,9 @@ class InferenceBackend: return False - # Global inference backend instance inference_backend = InferenceBackend() + def get_inference_backend() -> InferenceBackend: return inference_backend diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 13973b17f8..21a25d51b2 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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'(? 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]: diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 3b878b348f..b266f5e0d4 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -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(), + }, + ) diff --git a/studio/backend/core/training/__init__.py b/studio/backend/core/training/__init__.py index 3b92aa704d..6e19805eb0 100644 --- a/studio/backend/core/training/__init__.py +++ b/studio/backend/core/training/__init__.py @@ -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", ] diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index f933b5582e..cee84a8809 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -5,13 +5,16 @@ Unsloth Training Backend Integrates Unsloth training capabilities with the FastAPI backend """ + import os import sys + # Prevent tokenizer parallelism deadlocks when datasets uses multiprocessing fork os.environ["TOKENIZERS_PARALLELISM"] = "false" import torch from utils.hardware import clear_gpu_cache, safe_num_proc + torch._dynamo.config.recompile_limit = 64 from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported from unsloth.chat_templates import get_chat_template @@ -31,13 +34,17 @@ from datasets import Dataset, load_dataset from utils.models import is_vision_model, detect_audio_type from utils.datasets import format_and_template_dataset from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER -from utils.paths import ensure_dir, resolve_dataset_path, resolve_output_dir, resolve_tensorboard_dir +from utils.paths import ( + ensure_dir, + resolve_dataset_path, + resolve_output_dir, + resolve_tensorboard_dir, +) from trl import SFTTrainer, SFTConfig logger = get_logger(__name__) - def _build_report_targets(training_args) -> list[str] | str: report_to: list[str] = [] if training_args.get("enable_wandb", False): @@ -50,6 +57,7 @@ def _build_report_targets(training_args) -> list[str] | str: @dataclass class TrainingProgress: """Training progress tracking""" + epoch: float = 0 step: int = 0 total_steps: int = 0 @@ -65,6 +73,7 @@ class TrainingProgress: num_tokens: Optional[int] = None eval_loss: Optional[float] = None + class UnslothTrainer: """ Unsloth Training Backend @@ -85,10 +94,16 @@ class UnslothTrainer: # Model state tracking self.is_vlm = False self.is_audio = False - self.is_audio_vlm = False # Multimodal model (e.g. Gemma 3N) trained on audio data + self.is_audio_vlm = ( + False # Multimodal model (e.g. Gemma 3N) trained on audio data + ) self._audio_type = None # 'csm', 'whisper', 'snac', 'bicodec', 'dac' - self._cuda_audio_used = False # Set once after audio CUDA preprocessing; never cleared - self._spark_tts_repo_dir = None # Path to downloaded Spark-TTS repo (for BiCodecTokenizer) + self._cuda_audio_used = ( + False # Set once after audio CUDA preprocessing; never cleared + ) + self._spark_tts_repo_dir = ( + None # Path to downloaded Spark-TTS repo (for BiCodecTokenizer) + ) self.model_name = None # Training metrics tracking @@ -102,9 +117,9 @@ class UnslothTrainer: # Store training context for later transfer self.training_context = { - 'base_model_name': None, - 'output_dir': None, - 'is_lora': True, # Default to LoRA + "base_model_name": None, + "output_dir": None, + "is_lora": True, # Default to LoRA } def pre_detect_and_load_tokenizer( @@ -136,7 +151,7 @@ class UnslothTrainer: # --- Detect audio type (reads config.json only, no VRAM) --- self._audio_type = detect_audio_type(model_name, hf_token) - if self._audio_type == 'audio_vlm': + if self._audio_type == "audio_vlm": self.is_audio = False self.is_audio_vlm = is_dataset_audio self._audio_type = None @@ -153,21 +168,30 @@ class UnslothTrainer: logger.info( "pre_detect: audio_type=%s, is_audio=%s, is_audio_vlm=%s, is_vlm=%s", - self._audio_type, self.is_audio, self.is_audio_vlm, self.is_vlm, + self._audio_type, + self.is_audio, + self.is_audio_vlm, + self.is_vlm, ) # --- Load lightweight tokenizer/processor (CPU only, no VRAM) --- # Whisper needs AutoProcessor (has feature_extractor + tokenizer). # All others work with AutoTokenizer (CSM loads its own processor inline). - if self._audio_type == 'whisper': + if self._audio_type == "whisper": from transformers import AutoProcessor + self.tokenizer = AutoProcessor.from_pretrained( - model_name, trust_remote_code=trust_remote_code, token=hf_token, + model_name, + trust_remote_code = trust_remote_code, + token = hf_token, ) else: from transformers import AutoTokenizer + self.tokenizer = AutoTokenizer.from_pretrained( - model_name, trust_remote_code=trust_remote_code, token=hf_token, + model_name, + trust_remote_code = trust_remote_code, + token = hf_token, ) logger.info("Pre-loaded tokenizer for %s", model_name) @@ -193,15 +217,16 @@ class UnslothTrainer: def _create_progress_callback(self): """Create a TrainerCallback for progress tracking. Reused by all training branches.""" from transformers import TrainerCallback + trainer_ref = self class _ProgressCallback(TrainerCallback): - 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)) + loss_value = logs.get("loss", logs.get("train_loss", 0.0)) current_step = state.global_step - grad_norm = logs.get('grad_norm', None) + grad_norm = logs.get("grad_norm", None) elapsed_seconds = None if trainer_ref.training_start_time is not None: @@ -213,25 +238,27 @@ class UnslothTrainer: if total_steps > 0: steps_remaining = total_steps - current_step if steps_remaining > 0: - eta_seconds = (elapsed_seconds / current_step) * steps_remaining + eta_seconds = ( + elapsed_seconds / current_step + ) * steps_remaining num_tokens = getattr(state, "num_input_tokens_seen", None) trainer_ref._update_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), - elapsed_seconds=elapsed_seconds, - eta_seconds=eta_seconds, - grad_norm=grad_norm, - num_tokens=num_tokens, - eval_loss=logs.get('eval_loss', None), - status_message="", + step = current_step, + epoch = round(state.epoch, 2) if state.epoch else 0, + loss = loss_value, + learning_rate = logs.get("learning_rate", 0.0), + elapsed_seconds = elapsed_seconds, + eta_seconds = eta_seconds, + grad_norm = grad_norm, + num_tokens = num_tokens, + eval_loss = logs.get("eval_loss", None), + status_message = "", ) def on_epoch_end(self, args, state, control, **kwargs): - trainer_ref._update_progress(epoch=state.epoch, step=state.global_step) + trainer_ref._update_progress(epoch = state.epoch, step = state.global_step) def on_step_end(self, args, state, control, **kwargs): if trainer_ref.should_stop: @@ -241,29 +268,35 @@ class UnslothTrainer: return _ProgressCallback() - def _calculate_total_steps(self, num_samples, batch_size, grad_accum, num_epochs, max_steps): + def _calculate_total_steps( + self, num_samples, batch_size, grad_accum, num_epochs, max_steps + ): """Calculate total training steps from dataset size and training params.""" if max_steps and max_steps > 0: return max_steps len_dataloader = math.ceil(num_samples / batch_size) - steps_per_epoch = max(len_dataloader // grad_accum + int(len_dataloader % grad_accum > 0), 1) + steps_per_epoch = max( + len_dataloader // grad_accum + int(len_dataloader % grad_accum > 0), 1 + ) return steps_per_epoch * num_epochs - def _build_audio_training_args(self, training_args, output_dir, *, extra_args=None): + def _build_audio_training_args(self, training_args, output_dir, *, extra_args = None): """Build training args dict for audio branches. Constructs the common config (batch size, lr, warmup, fp16/bf16, etc.) and applies per-branch overrides via extra_args. """ - batch_size = training_args.get('batch_size', 2) - gradient_accumulation_steps = training_args.get('gradient_accumulation_steps', 4) - warmup_steps_val = training_args.get('warmup_steps', 5) - max_steps_val = training_args.get('max_steps', 0) - learning_rate = training_args.get('learning_rate', 2e-4) - weight_decay = training_args.get('weight_decay', 0.001) - lr_scheduler_type = training_args.get('lr_scheduler_type', 'linear') - random_seed = training_args.get('random_seed', 3407) - optim_value = training_args.get('optim', 'adamw_8bit') + batch_size = training_args.get("batch_size", 2) + gradient_accumulation_steps = training_args.get( + "gradient_accumulation_steps", 4 + ) + warmup_steps_val = training_args.get("warmup_steps", 5) + max_steps_val = training_args.get("max_steps", 0) + learning_rate = training_args.get("learning_rate", 2e-4) + weight_decay = training_args.get("weight_decay", 0.001) + lr_scheduler_type = training_args.get("lr_scheduler_type", "linear") + random_seed = training_args.get("random_seed", 3407) + optim_value = training_args.get("optim", "adamw_8bit") config = { "per_device_train_batch_size": batch_size, @@ -290,10 +323,10 @@ class UnslothTrainer: if max_steps_val and max_steps_val > 0: config["max_steps"] = max_steps_val else: - config["num_train_epochs"] = training_args.get('num_epochs', 3) + config["num_train_epochs"] = training_args.get("num_epochs", 3) # save_steps - save_steps_val = training_args.get('save_steps', 0) + save_steps_val = training_args.get("save_steps", 0) if save_steps_val and save_steps_val > 0: config["save_steps"] = save_steps_val config["save_strategy"] = "steps" @@ -304,7 +337,7 @@ class UnslothTrainer: return config - def _finalize_training(self, output_dir, label=""): + def _finalize_training(self, output_dir, label = ""): """Save model after training and update progress. Used by all training branches.""" if self.should_stop and self.save_on_stop: self.trainer.save_model() @@ -313,13 +346,15 @@ class UnslothTrainer: msg = f"{label} training stopped" if label else "Training stopped" logger.info(f"\n{msg}. Model saved to {output_dir}\n") self._update_progress( - is_training=False, - status_message=f"Training stopped. Model saved to {output_dir}", + is_training = False, + status_message = f"Training stopped. Model saved to {output_dir}", ) elif self.should_stop: msg = f"{label} training cancelled" if label else "Training cancelled" logger.info(f"\n{msg}.\n") - self._update_progress(is_training=False, status_message="Training cancelled.") + self._update_progress( + is_training = False, status_message = "Training cancelled." + ) else: self.trainer.save_model() self.tokenizer.save_pretrained(output_dir) @@ -327,9 +362,9 @@ class UnslothTrainer: msg = f"{label} training completed" if label else "Training completed" logger.info(f"\n{msg}! Model saved to {output_dir}\n") self._update_progress( - is_training=False, - is_completed=True, - status_message=f"Training completed! Model saved to {output_dir}", + is_training = False, + is_completed = True, + status_message = f"Training completed! Model saved to {output_dir}", ) def _cleanup_audio_artifacts(self): @@ -349,7 +384,9 @@ class UnslothTrainer: ] # Spark-TTS path is relative to the downloaded repo if self._spark_tts_repo_dir: - spark_code_dir = os.path.join(os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS") + spark_code_dir = os.path.join( + os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS" + ) audio_paths.append(spark_code_dir) removed_paths = [] @@ -359,14 +396,16 @@ class UnslothTrainer: removed_paths.append(path) # Remove stale audio modules from sys.modules - prefixes = ('snac', 'whisper', 'sparktts', 'outetts') + prefixes = ("snac", "whisper", "sparktts", "outetts") removed_modules = [key for key in _sys.modules if key.startswith(prefixes)] for key in removed_modules: del _sys.modules[key] if removed_paths or removed_modules: - logger.info(f"Cleaned up audio artifacts: {len(removed_paths)} paths, " - f"{len(removed_modules)} modules\n") + logger.info( + f"Cleaned up audio artifacts: {len(removed_paths)} paths, " + f"{len(removed_modules)} modules\n" + ) def _resolve_audio_columns(self, dataset, custom_format_mapping: dict = None): """Resolve audio, text, and speaker columns from user mapping or hardcoded fallback. @@ -389,11 +428,22 @@ class UnslothTrainer: speaker_col = col # Use mapping if both required columns exist in the dataset if audio_col and audio_col in cols and text_col and text_col in cols: - return {"audio_col": audio_col, "text_col": text_col, "speaker_col": speaker_col} + return { + "audio_col": audio_col, + "text_col": text_col, + "speaker_col": speaker_col, + } # Hardcoded fallback (existing behavior) audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None) - text_col = next((c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), None) + text_col = next( + ( + c + for c in cols + if c.lower() in ("text", "sentence", "transcript", "transcription") + ), + None, + ) speaker_col = None if "source" in cols: @@ -401,20 +451,27 @@ class UnslothTrainer: elif "speaker_id" in cols: speaker_col = "speaker_id" - return {"audio_col": audio_col, "text_col": text_col, "speaker_col": speaker_col} + return { + "audio_col": audio_col, + "text_col": text_col, + "speaker_col": speaker_col, + } - - def load_model(self, - model_name: str, - max_seq_length: int = 2048, - load_in_4bit: bool = True, - hf_token: Optional[str] = None, - is_dataset_image: bool = False, - is_dataset_audio: bool = False, - trust_remote_code: bool = False) -> bool: + def load_model( + self, + model_name: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + hf_token: Optional[str] = None, + is_dataset_image: bool = False, + is_dataset_audio: bool = False, + trust_remote_code: bool = False, + ) -> bool: """Load model for training (supports both text and vision models)""" self.load_in_4bit = load_in_4bit # Store for training_meta.json - self.trust_remote_code = trust_remote_code # For AutoProcessor etc. used during training + self.trust_remote_code = ( + trust_remote_code # For AutoProcessor etc. used during training + ) try: if self.model is not None: del self.model @@ -439,9 +496,10 @@ class UnslothTrainer: # restores original class definitions so Unsloth can re-compile cleanly. import sys as _sys import importlib + for _key, _mod in list(_sys.modules.items()): - if 'transformers.models.' in _key and '.modeling_' in _key: - if hasattr(_mod, '__UNSLOTH_PATCHED__'): + if "transformers.models." in _key and ".modeling_" in _key: + if hasattr(_mod, "__UNSLOTH_PATCHED__"): try: importlib.reload(_mod) except Exception: @@ -449,13 +507,16 @@ class UnslothTrainer: # Remove stale compiled cache so the new model gets a fresh one from utils.cache_cleanup import clear_unsloth_compiled_cache + clear_unsloth_compiled_cache() # Detect audio model type dynamically (config.json + tokenizer) self._audio_type = detect_audio_type(model_name, hf_token) # audio_vlm is detected as an audio_type now, handle it separately - if self._audio_type == 'audio_vlm': + if self._audio_type == "audio_vlm": self.is_audio = False - self.is_audio_vlm = is_dataset_audio # Only use audio VLM path if dataset has audio + self.is_audio_vlm = ( + is_dataset_audio # Only use audio VLM path if dataset has audio + ) self._audio_type = None else: self.is_audio = self._audio_type is not None @@ -470,25 +531,33 @@ class UnslothTrainer: self.model_name = model_name self.max_seq_length = max_seq_length - logger.info(f"Audio type: {self._audio_type}, is_audio: {self.is_audio}, is_audio_vlm: {self.is_audio_vlm}") - logger.info(f"Dataset has images: {is_dataset_image}, audio: {is_dataset_audio}") + logger.info( + f"Audio type: {self._audio_type}, is_audio: {self.is_audio}, is_audio_vlm: {self.is_audio_vlm}" + ) + logger.info( + f"Dataset has images: {is_dataset_image}, audio: {is_dataset_audio}" + ) logger.info(f"Using VLM path: {self.is_vlm}") # Reset training state for new run self._update_progress( - is_training=True, - is_completed=False, - error=None, - step=0, - loss=0.0, - epoch=0 + is_training = True, + is_completed = False, + error = None, + step = 0, + loss = 0.0, + epoch = 0, ) # Update UI immediately with loading message - model_display = model_name.split('/')[-1] if '/' in model_name else model_name - model_type_label = 'audio' if self.is_audio else ('vision' if self.is_vlm else 'text') + model_display = ( + model_name.split("/")[-1] if "/" in model_name else model_name + ) + model_type_label = ( + "audio" if self.is_audio else ("vision" if self.is_vlm else "text") + ) self._update_progress( - status_message=f"Loading {model_type_label} model... {model_display}" + status_message = f"Loading {model_type_label} model... {model_display}" ) logger.info(f"\nLoading {model_type_label} model: {model_name}") @@ -499,10 +568,11 @@ class UnslothTrainer: # Proactive gated-model check: verify access BEFORE from_pretrained. # Catches ALL gated/private models (text, vision, audio) globally. - if '/' in model_name: # Only check HF repo IDs, not local paths + if "/" in model_name: # Only check HF repo IDs, not local paths try: from huggingface_hub import model_info as hf_model_info - info = hf_model_info(model_name, token=hf_token or None) + + info = hf_model_info(model_name, token = hf_token or None) # model_info succeeds even for gated repos (metadata is public), # but info.gated tells us if files require acceptance/token. if info.gated and not hf_token: @@ -510,49 +580,57 @@ class UnslothTrainer: f"Access denied for '{model_name}'. This model is gated. " f"Please add a Hugging Face token with access and try again." ) - logger.error(f"Model '{model_name}' is gated (gated={info.gated}) and no HF token provided") - self._update_progress(error=friendly, is_training=False) + logger.error( + f"Model '{model_name}' is gated (gated={info.gated}) and no HF token provided" + ) + self._update_progress(error = friendly, is_training = False) return False except Exception as gate_err: - from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError + from huggingface_hub.utils import ( + GatedRepoError, + RepositoryNotFoundError, + ) + if isinstance(gate_err, (GatedRepoError, RepositoryNotFoundError)): friendly = ( f"Access denied for '{model_name}'. This model is gated or private. " f"Please add a Hugging Face token with access and try again." ) logger.error(f"Gated model check failed: {gate_err}") - self._update_progress(error=friendly, is_training=False) + self._update_progress(error = friendly, is_training = False) return False # Branch based on model type - if self._audio_type == 'csm': + if self._audio_type == "csm": # CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False from unsloth import FastModel from transformers import CsmForConditionalGeneration + self.model, self.tokenizer = FastModel.from_pretrained( - model_name=model_name, - max_seq_length=max_seq_length, - dtype=None, - auto_model=CsmForConditionalGeneration, - load_in_4bit=False, - token=hf_token, - trust_remote_code=trust_remote_code, + model_name = model_name, + max_seq_length = max_seq_length, + dtype = None, + auto_model = CsmForConditionalGeneration, + load_in_4bit = False, + token = hf_token, + trust_remote_code = trust_remote_code, ) logger.info("Loaded CSM audio model") - elif self._audio_type == 'whisper': + elif self._audio_type == "whisper": # Whisper: FastModel + auto_model=WhisperForConditionalGeneration + load_in_4bit=False from unsloth import FastModel from transformers import WhisperForConditionalGeneration + self.model, self.tokenizer = FastModel.from_pretrained( - model_name=model_name, - dtype=None, - load_in_4bit=False, - auto_model=WhisperForConditionalGeneration, - whisper_language="English", - whisper_task="transcribe", - token=hf_token, - trust_remote_code=trust_remote_code, + model_name = model_name, + dtype = None, + load_in_4bit = False, + auto_model = WhisperForConditionalGeneration, + whisper_language = "English", + whisper_task = "transcribe", + token = hf_token, + trust_remote_code = trust_remote_code, ) # Configure generation settings (notebook lines 100-105) self.model.generation_config.language = "<|en|>" @@ -561,19 +639,21 @@ class UnslothTrainer: self.model.generation_config.forced_decoder_ids = None logger.info("Loaded Whisper audio model (FastModel)") - elif self._audio_type == 'snac': + elif self._audio_type == "snac": # Orpheus: language model with audio codec tokens self.model, self.tokenizer = FastLanguageModel.from_pretrained( - model_name=model_name, - max_seq_length=max_seq_length, - dtype=None, - load_in_4bit=load_in_4bit, - token=hf_token, - trust_remote_code=trust_remote_code, + model_name = model_name, + max_seq_length = max_seq_length, + dtype = None, + load_in_4bit = load_in_4bit, + token = hf_token, + trust_remote_code = trust_remote_code, + ) + logger.info( + f"Loaded {self._audio_type} audio model (FastLanguageModel)" ) - logger.info(f"Loaded {self._audio_type} audio model (FastLanguageModel)") - elif self._audio_type == 'bicodec': + elif self._audio_type == "bicodec": # Spark-TTS: download full repo (contains sparktts package + BiCodec weights), # then load only the LLM subfolder with FastModel. # model_name may be: @@ -593,29 +673,32 @@ class UnslothTrainer: local_dir = model_name.split("/")[-1] llm_path = f"{local_dir}/LLM" - repo_path = snapshot_download(hf_repo, local_dir=local_dir) - self._spark_tts_repo_dir = os.path.abspath(repo_path) # Absolute path for sys.path + repo_path = snapshot_download(hf_repo, local_dir = local_dir) + self._spark_tts_repo_dir = os.path.abspath( + repo_path + ) # Absolute path for sys.path llm_path = os.path.join(self._spark_tts_repo_dir, "LLM") self.model, self.tokenizer = FastModel.from_pretrained( - model_name=llm_path, - max_seq_length=max_seq_length, - dtype=torch.float32, # Spark-TTS requires float32 - load_in_4bit=False, - token=hf_token, - trust_remote_code=trust_remote_code, + model_name = llm_path, + max_seq_length = max_seq_length, + dtype = torch.float32, # Spark-TTS requires float32 + load_in_4bit = False, + token = hf_token, + trust_remote_code = trust_remote_code, ) logger.info("Loaded Spark-TTS (bicodec) model") - elif self._audio_type == 'dac': + elif self._audio_type == "dac": # OuteTTS: uses FastModel (not FastLanguageModel) with load_in_4bit=False from unsloth import FastModel + self.model, self.tokenizer = FastModel.from_pretrained( model_name, - max_seq_length=max_seq_length, - load_in_4bit=False, - token=hf_token, - trust_remote_code=trust_remote_code, + max_seq_length = max_seq_length, + load_in_4bit = False, + token = hf_token, + trust_remote_code = trust_remote_code, ) logger.info("Loaded OuteTTS (dac) model (FastModel)") @@ -623,57 +706,71 @@ class UnslothTrainer: # Audio VLM: multimodal model trained on audio (e.g. Gemma 3N) # Uses FastModel (general loader) β€” returns (model, processor) from unsloth import FastModel + self.model, self.tokenizer = FastModel.from_pretrained( - model_name=model_name, - max_seq_length=max_seq_length, - dtype=None, - load_in_4bit=load_in_4bit, - token=hf_token, - trust_remote_code=trust_remote_code, + model_name = model_name, + max_seq_length = max_seq_length, + dtype = None, + load_in_4bit = load_in_4bit, + token = hf_token, + trust_remote_code = trust_remote_code, ) logger.info("Loaded audio VLM model (FastModel)") elif self.is_vlm: # Load vision model - returns (model, tokenizer) self.model, self.tokenizer = FastVisionModel.from_pretrained( - model_name=model_name, - max_seq_length=max_seq_length, - dtype=None, # Auto-detect - load_in_4bit=load_in_4bit, - token=hf_token, - trust_remote_code=trust_remote_code, + model_name = model_name, + max_seq_length = max_seq_length, + dtype = None, # Auto-detect + load_in_4bit = load_in_4bit, + token = hf_token, + trust_remote_code = trust_remote_code, ) logger.info("Loaded vision model") # Diagnostic: check if FastVisionModel returned a real Processor or a raw tokenizer from transformers import ProcessorMixin + tok = self.tokenizer - has_image_proc = isinstance(tok, ProcessorMixin) or hasattr(tok, "image_processor") - logger.info(f"\n[VLM Diagnostic] FastVisionModel returned: {type(tok).__name__}") - logger.info(f"[VLM Diagnostic] Is ProcessorMixin: {isinstance(tok, ProcessorMixin)}") - logger.info(f"[VLM Diagnostic] Has image_processor: {hasattr(tok, 'image_processor')}") - logger.info(f"[VLM Diagnostic] Usable as vision processor: {has_image_proc}\n") + has_image_proc = isinstance(tok, ProcessorMixin) or hasattr( + tok, "image_processor" + ) + logger.info( + f"\n[VLM Diagnostic] FastVisionModel returned: {type(tok).__name__}" + ) + logger.info( + f"[VLM Diagnostic] Is ProcessorMixin: {isinstance(tok, ProcessorMixin)}" + ) + logger.info( + f"[VLM Diagnostic] Has image_processor: {hasattr(tok, 'image_processor')}" + ) + logger.info( + f"[VLM Diagnostic] Usable as vision processor: {has_image_proc}\n" + ) else: # Load text model - returns (model, tokenizer) self.model, self.tokenizer = FastLanguageModel.from_pretrained( - model_name=model_name, - max_seq_length=max_seq_length, - dtype=None, # Auto-detect - load_in_4bit=load_in_4bit, - token=hf_token, - trust_remote_code=trust_remote_code, + model_name = model_name, + max_seq_length = max_seq_length, + dtype = None, # Auto-detect + load_in_4bit = load_in_4bit, + token = hf_token, + trust_remote_code = trust_remote_code, ) logger.info("Loaded text model") if self.should_stop: return False - self._update_progress(status_message="Model loaded successfully") + self._update_progress(status_message = "Model loaded successfully") logger.info("Model loaded successfully") return True except OSError as e: - if "could not get source code" in str(e) and not getattr(self, '_source_code_retried', False): + if "could not get source code" in str(e) and not getattr( + self, "_source_code_retried", False + ): # Unsloth's patching can leave stale state that makes # inspect.getsource() fail when switching model families # (e.g. gemma3 β†’ gemma3n). The load always succeeds on a @@ -681,48 +778,77 @@ class UnslothTrainer: # imports clean up the stale state as a side effect. self._source_code_retried = True logger.info(f"\n'could not get source code' β€” retrying once...\n") - return self.load_model(model_name, max_seq_length, load_in_4bit, hf_token, - is_dataset_image, is_dataset_audio, trust_remote_code) + return self.load_model( + model_name, + max_seq_length, + load_in_4bit, + hf_token, + is_dataset_image, + is_dataset_audio, + trust_remote_code, + ) error_msg = str(e) error_lower = error_msg.lower() - if any(k in error_lower for k in ("gated repo", "access to it at", "401", "403", "unauthorized", "forbidden")): + if any( + k in error_lower + for k in ( + "gated repo", + "access to it at", + "401", + "403", + "unauthorized", + "forbidden", + ) + ): error_msg = ( f"Access denied for '{model_name}'. This model is gated or private. " f"Please add a Hugging Face token with access and try again." ) logger.error(f"Error loading model: {e}") - self._update_progress(error=error_msg, is_training=False) + self._update_progress(error = error_msg, is_training = False) return False except Exception as e: error_msg = str(e) # Catch gated/auth errors and surface a friendly message error_lower = error_msg.lower() - if any(k in error_lower for k in ("gated repo", "access to it at", "401", "403", "unauthorized", "forbidden")): + if any( + k in error_lower + for k in ( + "gated repo", + "access to it at", + "401", + "403", + "unauthorized", + "forbidden", + ) + ): error_msg = ( f"Access denied for '{model_name}'. This model is gated or private. " f"Please add a Hugging Face token with access and try again." ) logger.error(f"Error loading model: {e}") - self._update_progress(error=error_msg, is_training=False) + self._update_progress(error = error_msg, is_training = False) return False finally: self._source_code_retried = False - def prepare_model_for_training(self, - use_lora: bool = True, - # Vision-specific LoRA parameters (only used if is_vlm=True) - finetune_vision_layers: bool = True, - finetune_language_layers: bool = True, - finetune_attention_modules: bool = True, - finetune_mlp_modules: bool = True, - # Standard LoRA parameters - target_modules: list = None, - lora_r: int = 16, - lora_alpha: int = 16, - lora_dropout: float = 0.0, - use_gradient_checkpointing: str = "unsloth", - use_rslora: bool = False, - use_loftq: bool = False) -> bool: + def prepare_model_for_training( + self, + use_lora: bool = True, + # Vision-specific LoRA parameters (only used if is_vlm=True) + finetune_vision_layers: bool = True, + finetune_language_layers: bool = True, + finetune_attention_modules: bool = True, + finetune_mlp_modules: bool = True, + # Standard LoRA parameters + target_modules: list = None, + lora_r: int = 16, + lora_alpha: int = 16, + lora_dropout: float = 0.0, + use_gradient_checkpointing: str = "unsloth", + use_rslora: bool = False, + use_loftq: bool = False, + ) -> bool: """ Prepare model for training (with optional LoRA). """ @@ -730,10 +856,11 @@ class UnslothTrainer: if self.model is None: raise ValueError("Model not loaded. Call load_model() first.") - # Full finetuning mode - skip PEFT entirely if not use_lora: - self._update_progress(status_message="Full finetuning mode - no LoRA adapters") + self._update_progress( + status_message = "Full finetuning mode - no LoRA adapters" + ) logger.info("Full finetuning mode - training all parameters\n") return True @@ -744,15 +871,27 @@ class UnslothTrainer: target_modules = "all-linear" else: target_modules = [m for m in target_modules if m != "all-linear"] - elif target_modules is None or (isinstance(target_modules, list) and len(target_modules) == 0): - target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", - "gate_proj", "up_proj", "down_proj"] + elif target_modules is None or ( + isinstance(target_modules, list) and len(target_modules) == 0 + ): + target_modules = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ] # Validate and normalize gradient_checkpointing # Must be one of: True, False, or "unsloth" if isinstance(use_gradient_checkpointing, str): use_gradient_checkpointing = use_gradient_checkpointing.strip().lower() - if use_gradient_checkpointing == "" or use_gradient_checkpointing == "unsloth": + if ( + use_gradient_checkpointing == "" + or use_gradient_checkpointing == "unsloth" + ): use_gradient_checkpointing = "unsloth" elif use_gradient_checkpointing in ("true", "1", "yes"): use_gradient_checkpointing = True @@ -760,102 +899,122 @@ class UnslothTrainer: use_gradient_checkpointing = False else: # Invalid value, default to "unsloth" - logger.warning(f"Invalid gradient_checkpointing value: {use_gradient_checkpointing}, defaulting to 'unsloth'") + logger.warning( + f"Invalid gradient_checkpointing value: {use_gradient_checkpointing}, defaulting to 'unsloth'" + ) use_gradient_checkpointing = "unsloth" elif use_gradient_checkpointing not in (True, False, "unsloth"): # Invalid type or value, default to "unsloth" - logger.warning(f"Invalid gradient_checkpointing type/value: {use_gradient_checkpointing}, defaulting to 'unsloth'") + logger.warning( + f"Invalid gradient_checkpointing type/value: {use_gradient_checkpointing}, defaulting to 'unsloth'" + ) use_gradient_checkpointing = "unsloth" # Verify model is loaded if self.model is None: error_msg = "Model is None - model was not loaded properly" logger.error(error_msg) - self._update_progress(error=error_msg) + self._update_progress(error = error_msg) return False # Check if model has the expected attributes - if not hasattr(self.model, 'config'): + if not hasattr(self.model, "config"): error_msg = "Model does not have config attribute - model may not be loaded correctly" logger.error(error_msg) - self._update_progress(error=error_msg) + self._update_progress(error = error_msg) return False - logger.info(f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n") - logger.info(f"Gradient checkpointing: {use_gradient_checkpointing} (type: {type(use_gradient_checkpointing).__name__})\n") + logger.info( + f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n" + ) + logger.info( + f"Gradient checkpointing: {use_gradient_checkpointing} (type: {type(use_gradient_checkpointing).__name__})\n" + ) # Branch based on model type: audio, audio_vlm, vision, or text - if self._audio_type in ('csm', 'bicodec', 'dac') or self.is_audio_vlm: + if self._audio_type in ("csm", "bicodec", "dac") or self.is_audio_vlm: # Models using FastModel.get_peft_model (codec audio + audio VLM) from unsloth import FastModel - label = self._audio_type or 'audio_vlm' + + label = self._audio_type or "audio_vlm" logger.info(f"{label} LoRA configuration:") logger.info(f" - Target modules: {target_modules}") if self.is_audio_vlm: logger.info(f" - Finetune vision layers: {finetune_vision_layers}") - logger.info(f" - Finetune language layers: {finetune_language_layers}") - logger.info(f" - Finetune attention modules: {finetune_attention_modules}") + logger.info( + f" - Finetune language layers: {finetune_language_layers}" + ) + logger.info( + f" - Finetune attention modules: {finetune_attention_modules}" + ) logger.info(f" - Finetune MLP modules: {finetune_mlp_modules}") logger.info() peft_kwargs = dict( - r=lora_r, - target_modules=target_modules, - lora_alpha=lora_alpha, - lora_dropout=lora_dropout, - bias="none", - use_gradient_checkpointing=use_gradient_checkpointing, - random_state=3407, - use_rslora=use_rslora, - loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + r = lora_r, + target_modules = target_modules, + lora_alpha = lora_alpha, + lora_dropout = lora_dropout, + bias = "none", + use_gradient_checkpointing = use_gradient_checkpointing, + random_state = 3407, + use_rslora = use_rslora, + loftq_config = {"loftq_bits": 4, "loftq_iter": 1} + if use_loftq + else None, ) # Audio VLM models support VLM-style layer selection if self.is_audio_vlm: peft_kwargs.update( - finetune_vision_layers=finetune_vision_layers, - finetune_language_layers=finetune_language_layers, - finetune_attention_modules=finetune_attention_modules, - finetune_mlp_modules=finetune_mlp_modules, + finetune_vision_layers = finetune_vision_layers, + finetune_language_layers = finetune_language_layers, + finetune_attention_modules = finetune_attention_modules, + finetune_mlp_modules = finetune_mlp_modules, ) self.model = FastModel.get_peft_model(self.model, **peft_kwargs) - elif self._audio_type == 'whisper': + elif self._audio_type == "whisper": # Phase 2: Whisper uses FastModel.get_peft_model with task_type=None from unsloth import FastModel + logger.info(f"Audio model (whisper) LoRA configuration:") logger.info(f" - Target modules: {target_modules}\n") self.model = FastModel.get_peft_model( self.model, - r=lora_r, - target_modules=target_modules, - lora_alpha=lora_alpha, - lora_dropout=lora_dropout, - bias="none", - use_gradient_checkpointing=use_gradient_checkpointing, - random_state=3407, - use_rslora=use_rslora, - loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, - task_type=None, + r = lora_r, + target_modules = target_modules, + lora_alpha = lora_alpha, + lora_dropout = lora_dropout, + bias = "none", + use_gradient_checkpointing = use_gradient_checkpointing, + random_state = 3407, + use_rslora = use_rslora, + loftq_config = {"loftq_bits": 4, "loftq_iter": 1} + if use_loftq + else None, + task_type = None, ) - elif self._audio_type == 'snac': + elif self._audio_type == "snac": # Orpheus uses FastLanguageModel.get_peft_model logger.info(f"Audio model ({self._audio_type}) LoRA configuration:") logger.info(f" - Target modules: {target_modules}\n") self.model = FastLanguageModel.get_peft_model( self.model, - r=lora_r, - target_modules=target_modules, - lora_alpha=lora_alpha, - lora_dropout=lora_dropout, - bias="none", - use_gradient_checkpointing=use_gradient_checkpointing, - random_state=3407, - use_rslora=use_rslora, - loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + r = lora_r, + target_modules = target_modules, + lora_alpha = lora_alpha, + lora_dropout = lora_dropout, + bias = "none", + use_gradient_checkpointing = use_gradient_checkpointing, + random_state = 3407, + use_rslora = use_rslora, + loftq_config = {"loftq_bits": 4, "loftq_iter": 1} + if use_loftq + else None, ) elif self.is_vlm: @@ -863,24 +1022,28 @@ class UnslothTrainer: logger.info(f"Vision model LoRA configuration:") logger.info(f" - Finetune vision layers: {finetune_vision_layers}") logger.info(f" - Finetune language layers: {finetune_language_layers}") - logger.info(f" - Finetune attention modules: {finetune_attention_modules}") + logger.info( + f" - Finetune attention modules: {finetune_attention_modules}" + ) logger.info(f" - Finetune MLP modules: {finetune_mlp_modules}\n") self.model = FastVisionModel.get_peft_model( self.model, - finetune_vision_layers=finetune_vision_layers, - finetune_language_layers=finetune_language_layers, - finetune_attention_modules=finetune_attention_modules, - finetune_mlp_modules=finetune_mlp_modules, - r=lora_r, - target_modules=target_modules, - lora_alpha=lora_alpha, - lora_dropout=lora_dropout, - bias="none", - use_gradient_checkpointing=use_gradient_checkpointing, - random_state=3407, - use_rslora=use_rslora, - loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + finetune_vision_layers = finetune_vision_layers, + finetune_language_layers = finetune_language_layers, + finetune_attention_modules = finetune_attention_modules, + finetune_mlp_modules = finetune_mlp_modules, + r = lora_r, + target_modules = target_modules, + lora_alpha = lora_alpha, + lora_dropout = lora_dropout, + bias = "none", + use_gradient_checkpointing = use_gradient_checkpointing, + random_state = 3407, + use_rslora = use_rslora, + loftq_config = {"loftq_bits": 4, "loftq_iter": 1} + if use_loftq + else None, ) else: # Text model LoRA @@ -889,15 +1052,17 @@ class UnslothTrainer: self.model = FastLanguageModel.get_peft_model( self.model, - r=lora_r, - target_modules=target_modules, - lora_alpha=lora_alpha, - lora_dropout=lora_dropout, - bias="none", - use_gradient_checkpointing=use_gradient_checkpointing, - random_state=3407, - use_rslora=use_rslora, - loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + r = lora_r, + target_modules = target_modules, + lora_alpha = lora_alpha, + lora_dropout = lora_dropout, + bias = "none", + use_gradient_checkpointing = use_gradient_checkpointing, + random_state = 3407, + use_rslora = use_rslora, + loftq_config = {"loftq_bits": 4, "loftq_iter": 1} + if use_loftq + else None, ) # Check if stopped during LoRA preparation @@ -905,20 +1070,25 @@ class UnslothTrainer: logger.info("Stopped during LoRA configuration\n") return False - self._update_progress(status_message="LoRA adapters configured") + self._update_progress(status_message = "LoRA adapters configured") logger.info("LoRA adapters configured successfully\n") return True except Exception as e: import traceback import sys - error_details = f"{type(e).__name__}: {str(e)}" if str(e) else f"{type(e).__name__} (no message)" + + error_details = ( + f"{type(e).__name__}: {str(e)}" + if str(e) + else f"{type(e).__name__} (no message)" + ) full_traceback = traceback.format_exc() logger.error(f"Error preparing model: {error_details}") logger.error(f"Full traceback:\n{full_traceback}") logger.info(f"\n[ERROR] Error preparing model: {error_details}") logger.info(f"[ERROR] Full traceback:\n{full_traceback}") - self._update_progress(error=error_details) + self._update_progress(error = error_details) return False def _apply_csm_forward_fix(self): @@ -946,29 +1116,44 @@ class UnslothTrainer: # Keys that the depth decoder and its sub-layers actually understand _TRANSFORMERS_KWARGS = { - 'num_items_in_batch', 'output_hidden_states', 'output_attentions', - 'output_router_logits', 'cu_seq_lens_q', 'cu_seq_lens_k', - 'max_length_q', 'max_length_k', + "num_items_in_batch", + "output_hidden_states", + "output_attentions", + "output_router_logits", + "cu_seq_lens_q", + "cu_seq_lens_k", + "max_length_q", + "max_length_k", } def _fixed_csm_forward( self, - input_ids=None, input_values=None, attention_mask=None, - input_values_cutoffs=None, position_ids=None, past_key_values=None, - inputs_embeds=None, labels=None, use_cache=None, - cache_position=None, logits_to_keep=0, **kwargs, + input_ids = None, + input_values = None, + attention_mask = None, + input_values_cutoffs = None, + position_ids = None, + past_key_values = None, + inputs_embeds = None, + labels = None, + use_cache = None, + cache_position = None, + logits_to_keep = 0, + **kwargs, ): # Strip non-standard kwargs injected by Unsloth/PEFT (causal_mask, # num_logits_to_keep, task_ids, return_dict, etc.) - output_attentions = kwargs.pop('output_attentions', None) - output_hidden_states = kwargs.pop('output_hidden_states', None) - kwargs.pop('return_dict', None) - kwargs.pop('causal_mask', None) - kwargs.pop('num_logits_to_keep', None) - kwargs.pop('task_ids', None) + output_attentions = kwargs.pop("output_attentions", None) + output_hidden_states = kwargs.pop("output_hidden_states", None) + kwargs.pop("return_dict", None) + kwargs.pop("causal_mask", None) + kwargs.pop("num_logits_to_keep", None) + kwargs.pop("task_ids", None) # Only keep recognized TransformersKwargs - clean_kwargs = {k: v for k, v in kwargs.items() if k in _TRANSFORMERS_KWARGS} + clean_kwargs = { + k: v for k, v in kwargs.items() if k in _TRANSFORMERS_KWARGS + } if input_ids is not None and input_ids.ndim == 2: merged = self._merge_input_ids_with_input_values( @@ -979,18 +1164,22 @@ class UnslothTrainer: input_ids = None backbone_outputs = self.backbone_model( - input_ids=input_ids, attention_mask=attention_mask, - position_ids=position_ids, past_key_values=past_key_values, - inputs_embeds=inputs_embeds, use_cache=use_cache, - cache_position=cache_position, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, + input_ids = input_ids, + attention_mask = attention_mask, + position_ids = position_ids, + past_key_values = past_key_values, + inputs_embeds = inputs_embeds, + use_cache = use_cache, + cache_position = cache_position, + output_attentions = output_attentions, + output_hidden_states = output_hidden_states, **clean_kwargs, ) backbone_hidden_states = backbone_outputs[0] slice_indices = ( - slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) + slice(-logits_to_keep, None) + if isinstance(logits_to_keep, int) else logits_to_keep ) backbone_logits = self.lm_head(backbone_hidden_states[:, slice_indices, :]) @@ -1002,17 +1191,21 @@ class UnslothTrainer: if labels is not None: backbone_labels = labels[:, :, 0] backbone_loss = self.loss_function( - logits=backbone_logits, labels=backbone_labels, - vocab_size=self.config.vocab_size, **clean_kwargs, + logits = backbone_logits, + labels = backbone_labels, + vocab_size = self.config.vocab_size, + **clean_kwargs, ) - train_mask = ~(labels[:, :, 1:] == -100).all(dim=-1) - depth_decoder_input_ids = labels[train_mask][..., :self.config.num_codebooks - 1] + train_mask = ~(labels[:, :, 1:] == -100).all(dim = -1) + depth_decoder_input_ids = labels[train_mask][ + ..., : self.config.num_codebooks - 1 + ] depth_decoder_input_ids = nn.functional.pad( - depth_decoder_input_ids, (1, 0), value=0 + depth_decoder_input_ids, (1, 0), value = 0 ) - train_idxs = train_mask.nonzero(as_tuple=True) + train_idxs = train_mask.nonzero(as_tuple = True) backbone_last_hidden_states = backbone_hidden_states[ train_idxs[0], train_idxs[1] - 1, : ] @@ -1021,18 +1214,19 @@ class UnslothTrainer: # Build clean kwargs for depth decoder dd_kwargs = clean_kwargs.copy() # Scale num_items_in_batch for depth decoder (31 codebooks) - if 'num_items_in_batch' in dd_kwargs: - dd_kwargs['num_items_in_batch'] = ( - dd_kwargs['num_items_in_batch'] * (self.config.num_codebooks - 1) - ) + if "num_items_in_batch" in dd_kwargs: + dd_kwargs["num_items_in_batch"] = dd_kwargs[ + "num_items_in_batch" + ] * (self.config.num_codebooks - 1) depth_decoder_outputs = self.depth_decoder( - input_ids=depth_decoder_input_ids, - backbone_last_hidden_state=backbone_last_hidden_states, - use_cache=False, return_dict=True, - labels=depth_decoder_labels, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, + input_ids = depth_decoder_input_ids, + backbone_last_hidden_state = backbone_last_hidden_states, + use_cache = False, + return_dict = True, + labels = depth_decoder_labels, + output_attentions = output_attentions, + output_hidden_states = output_hidden_states, **dd_kwargs, ) @@ -1049,21 +1243,27 @@ class UnslothTrainer: loss = backbone_loss + depth_decoder_loss return CsmOutputWithPast( - loss=loss, backbone_loss=backbone_loss, - depth_decoder_loss=depth_decoder_loss, logits=backbone_logits, - past_key_values=backbone_outputs.past_key_values, - hidden_states=backbone_outputs.hidden_states, - attentions=backbone_outputs.attentions, - depth_decoder_logits=( + loss = loss, + backbone_loss = backbone_loss, + depth_decoder_loss = depth_decoder_loss, + logits = backbone_logits, + past_key_values = backbone_outputs.past_key_values, + hidden_states = backbone_outputs.hidden_states, + attentions = backbone_outputs.attentions, + depth_decoder_logits = ( depth_decoder_outputs.logits if depth_decoder_outputs else None ), - depth_decoder_past_key_values=( - depth_decoder_outputs.past_key_values if depth_decoder_outputs else None + depth_decoder_past_key_values = ( + depth_decoder_outputs.past_key_values + if depth_decoder_outputs + else None ), - depth_decoder_hidden_states=( - depth_decoder_outputs.hidden_states if depth_decoder_outputs else None + depth_decoder_hidden_states = ( + depth_decoder_outputs.hidden_states + if depth_decoder_outputs + else None ), - depth_decoder_attentions=( + depth_decoder_attentions = ( depth_decoder_outputs.attentions if depth_decoder_outputs else None ), ) @@ -1075,7 +1275,7 @@ class UnslothTrainer: CsmForConditionalGeneration.forward = _fixed_csm_forward logger.info("Applied CSM forward fix (class + instance level)\n") - def _preprocess_csm_dataset(self, dataset, custom_format_mapping=None): + def _preprocess_csm_dataset(self, dataset, custom_format_mapping = None): """Preprocess dataset for CSM TTS training (exact notebook copy).""" from transformers import AutoProcessor from datasets import Audio @@ -1083,13 +1283,13 @@ class UnslothTrainer: processor = AutoProcessor.from_pretrained( self.model_name, - trust_remote_code=getattr(self, "trust_remote_code", False), + trust_remote_code = getattr(self, "trust_remote_code", False), ) # Strip pad_to_multiple_of from tokenizer init_kwargs β€” fine-tuned models # (e.g. keanteng/sesame-csm-elise) save it in tokenizer_config.json, and # _merge_kwargs leaks it into audio_kwargs where EncodecFeatureExtractor rejects it. - processor.tokenizer.init_kwargs.pop('pad_to_multiple_of', None) + processor.tokenizer.init_kwargs.pop("pad_to_multiple_of", None) # Resolve columns from user mapping or hardcoded fallback resolved = self._resolve_audio_columns(dataset, custom_format_mapping) @@ -1098,21 +1298,35 @@ class UnslothTrainer: speaker_key = resolved["speaker_col"] if audio_col is None: - raise ValueError(f"No audio column found in dataset. Columns: {dataset.column_names}") + raise ValueError( + f"No audio column found in dataset. Columns: {dataset.column_names}" + ) if text_col is None: - raise ValueError(f"No text column found in dataset. Columns: {dataset.column_names}") + raise ValueError( + f"No text column found in dataset. Columns: {dataset.column_names}" + ) if speaker_key is None: - logger.info("No speaker found, adding default 'source' of 0 for all examples\n") + logger.info( + "No speaker found, adding default 'source' of 0 for all examples\n" + ) dataset = dataset.add_column("source", ["0"] * len(dataset)) speaker_key = "source" - logger.info(f"CSM preprocessing: audio_col='{audio_col}', text_col='{text_col}', speaker_key='{speaker_key}'\n") + logger.info( + f"CSM preprocessing: audio_col='{audio_col}', text_col='{text_col}', speaker_key='{speaker_key}'\n" + ) - dataset = dataset.cast_column(audio_col, Audio(sampling_rate=24000)) + dataset = dataset.cast_column(audio_col, Audio(sampling_rate = 24000)) - required_keys = ["input_ids", "attention_mask", "labels", "input_values", "input_values_cutoffs"] + required_keys = [ + "input_ids", + "attention_mask", + "labels", + "input_values", + "input_values_cutoffs", + ] - self._update_progress(status_message="Preprocessing CSM dataset...") + self._update_progress(status_message = "Preprocessing CSM dataset...") processed_examples = [] skipped = 0 for idx in range(len(dataset)): @@ -1122,31 +1336,33 @@ class UnslothTrainer: example = dataset[idx] try: - conversation = [{ - "role": str(example[speaker_key]), - "content": [ - {"type": "text", "text": example.get(text_col, "")}, - {"type": "audio", "path": example[audio_col]["array"]}, - ], - }] + conversation = [ + { + "role": str(example[speaker_key]), + "content": [ + {"type": "text", "text": example.get(text_col, "")}, + {"type": "audio", "path": example[audio_col]["array"]}, + ], + } + ] # NOTE: pad_to_multiple_of intentionally omitted from text_kwargs β€” # CsmProcessor._merge_kwargs leaks it to EncodecFeatureExtractor which rejects it. model_inputs = processor.apply_chat_template( conversation, - tokenize=True, - return_dict=True, - output_labels=True, - text_kwargs={ + tokenize = True, + return_dict = True, + output_labels = True, + text_kwargs = { "padding": "max_length", "max_length": 256, "padding_side": "right", }, - audio_kwargs={ + audio_kwargs = { "sampling_rate": 24_000, "max_length": 240001, "padding": "max_length", }, - common_kwargs={"return_tensors": "pt"}, + common_kwargs = {"return_tensors": "pt"}, ) out = {} @@ -1168,7 +1384,7 @@ class UnslothTrainer: if (idx + 1) % 100 == 0: self._update_progress( - status_message=f"Preprocessing CSM... {idx + 1}/{len(dataset)}" + status_message = f"Preprocessing CSM... {idx + 1}/{len(dataset)}" ) if not processed_examples: @@ -1177,11 +1393,13 @@ class UnslothTrainer: ) result_dataset = Dataset.from_list(processed_examples) - logger.info(f"CSM preprocessing complete: {len(result_dataset)} examples " - f"({skipped} skipped)\n") + logger.info( + f"CSM preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n" + ) return result_dataset - def _format_audio_vlm_dataset(self, dataset, custom_format_mapping=None): + def _format_audio_vlm_dataset(self, dataset, custom_format_mapping = None): """Format dataset as audio chat messages for multimodal models (e.g. Gemma 3N). Expects columns: audio (Audio), text (str). @@ -1201,7 +1419,7 @@ class UnslothTrainer: self._audio_vlm_audio_col = audio_col # Cast audio to 16kHz (standard for speech models) - dataset = dataset.cast_column(audio_col, Audio(sampling_rate=16000)) + dataset = dataset.cast_column(audio_col, Audio(sampling_rate = 16000)) def format_messages(samples): formatted = {"messages": []} @@ -1209,26 +1427,35 @@ class UnslothTrainer: audio = samples[audio_col][idx]["array"] label = str(samples[text_col][idx]) message = [ - {"role": "system", "content": [ - {"type": "text", "text": "You are an assistant that transcribes speech accurately."} - ]}, - {"role": "user", "content": [ - {"type": "audio", "audio": audio}, - {"type": "text", "text": "Please transcribe this audio."} - ]}, - {"role": "assistant", "content": [ - {"type": "text", "text": label} - ]}, + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are an assistant that transcribes speech accurately.", + } + ], + }, + { + "role": "user", + "content": [ + {"type": "audio", "audio": audio}, + {"type": "text", "text": "Please transcribe this audio."}, + ], + }, + {"role": "assistant", "content": [{"type": "text", "text": label}]}, ] formatted["messages"].append(message) return formatted - self._update_progress(status_message="Formatting audio VLM dataset...") - dataset = dataset.map(format_messages, batched=True, batch_size=4, num_proc=safe_num_proc(4)) + self._update_progress(status_message = "Formatting audio VLM dataset...") + dataset = dataset.map( + format_messages, batched = True, batch_size = 4, num_proc = safe_num_proc(4) + ) logger.info(f"Audio VLM dataset formatted: {len(dataset)} examples\n") return dataset - def _preprocess_snac_dataset(self, dataset, custom_format_mapping=None): + def _preprocess_snac_dataset(self, dataset, custom_format_mapping = None): """Preprocess dataset for Orpheus TTS training with SNAC codec. Mirrors Orpheus_(3B)-TTS.ipynb: encode audio with SNAC (24kHz, 3 hierarchical @@ -1246,13 +1473,13 @@ class UnslothTrainer: # Orpheus special token IDs (hardcoded in tokenizer vocabulary) START_OF_HUMAN = 128259 - END_OF_HUMAN = 128260 - START_OF_AI = 128261 - END_OF_AI = 128262 + END_OF_HUMAN = 128260 + START_OF_AI = 128261 + END_OF_AI = 128262 START_OF_SPEECH = 128257 - END_OF_SPEECH = 128258 - END_OF_TEXT = 128009 - AUDIO_OFFSET = 128266 + END_OF_SPEECH = 128258 + END_OF_TEXT = 128009 + AUDIO_OFFSET = 128266 resolved = self._resolve_audio_columns(dataset, custom_format_mapping) audio_col = resolved["audio_col"] @@ -1266,25 +1493,37 @@ class UnslothTrainer: # Cast audio column so datasets 4.x AudioDecoder objects are decoded to dicts from datasets import Audio - dataset = dataset.cast_column(audio_col, Audio(sampling_rate=SNAC_SAMPLE_RATE)) + + dataset = dataset.cast_column(audio_col, Audio(sampling_rate = SNAC_SAMPLE_RATE)) # Get dataset sample rate from first example (after cast, always SNAC_SAMPLE_RATE) first_audio = dataset[0][audio_col] - ds_sample_rate = first_audio.get("sampling_rate", SNAC_SAMPLE_RATE) if isinstance(first_audio, dict) else SNAC_SAMPLE_RATE + ds_sample_rate = ( + first_audio.get("sampling_rate", SNAC_SAMPLE_RATE) + if isinstance(first_audio, dict) + else SNAC_SAMPLE_RATE + ) # Load SNAC codec model - self._update_progress(status_message="Loading SNAC codec model...") + self._update_progress(status_message = "Loading SNAC codec model...") logger.info("Loading SNAC codec model...\n") from snac import SNAC + snac_model = SNAC.from_pretrained(SNAC_MODEL_NAME) snac_model = snac_model.to(device).eval() # Resample transform (created once) - resample_transform = T.Resample(orig_freq=ds_sample_rate, new_freq=SNAC_SAMPLE_RATE) if ds_sample_rate != SNAC_SAMPLE_RATE else None + resample_transform = ( + T.Resample(orig_freq = ds_sample_rate, new_freq = SNAC_SAMPLE_RATE) + if ds_sample_rate != SNAC_SAMPLE_RATE + else None + ) - self._update_progress(status_message="Encoding audio with SNAC...") - logger.info(f"SNAC preprocessing: audio_col='{audio_col}', text_col='{text_col}', " - f"has_source={has_source}, ds_sample_rate={ds_sample_rate}\n") + self._update_progress(status_message = "Encoding audio with SNAC...") + logger.info( + f"SNAC preprocessing: audio_col='{audio_col}', text_col='{text_col}', " + f"has_source={has_source}, ds_sample_rate={ds_sample_rate}\n" + ) processed_examples = [] skipped = 0 @@ -1306,7 +1545,11 @@ class UnslothTrainer: continue # --- Encode audio with SNAC (notebook lines 122-142) --- - waveform = torch.from_numpy(audio_data["array"]).unsqueeze(0).to(dtype=torch.float32) + waveform = ( + torch.from_numpy(audio_data["array"]) + .unsqueeze(0) + .to(dtype = torch.float32) + ) if resample_transform is not None: waveform = resample_transform(waveform) @@ -1318,12 +1561,22 @@ class UnslothTrainer: all_codes = [] for i in range(codes[0].shape[1]): all_codes.append(codes[0][0][i].item() + AUDIO_OFFSET) - all_codes.append(codes[1][0][2*i].item() + AUDIO_OFFSET + 4096) - all_codes.append(codes[2][0][4*i].item() + AUDIO_OFFSET + (2*4096)) - all_codes.append(codes[2][0][(4*i)+1].item() + AUDIO_OFFSET + (3*4096)) - all_codes.append(codes[1][0][(2*i)+1].item() + AUDIO_OFFSET + (4*4096)) - all_codes.append(codes[2][0][(4*i)+2].item() + AUDIO_OFFSET + (5*4096)) - all_codes.append(codes[2][0][(4*i)+3].item() + AUDIO_OFFSET + (6*4096)) + all_codes.append(codes[1][0][2 * i].item() + AUDIO_OFFSET + 4096) + all_codes.append( + codes[2][0][4 * i].item() + AUDIO_OFFSET + (2 * 4096) + ) + all_codes.append( + codes[2][0][(4 * i) + 1].item() + AUDIO_OFFSET + (3 * 4096) + ) + all_codes.append( + codes[1][0][(2 * i) + 1].item() + AUDIO_OFFSET + (4 * 4096) + ) + all_codes.append( + codes[2][0][(4 * i) + 2].item() + AUDIO_OFFSET + (5 * 4096) + ) + all_codes.append( + codes[2][0][(4 * i) + 3].item() + AUDIO_OFFSET + (6 * 4096) + ) if len(all_codes) == 0: skipped += 1 @@ -1333,12 +1586,16 @@ class UnslothTrainer: deduped = all_codes[:7] for i in range(7, len(all_codes), 7): if all_codes[i] != deduped[-7]: - deduped.extend(all_codes[i:i+7]) + deduped.extend(all_codes[i : i + 7]) all_codes = deduped # --- Build text tokens (notebook lines 217-224) --- - text_prompt = f"{example[speaker_col]}: {text}" if has_source and example.get(speaker_col) else text - text_ids = tokenizer.encode(text_prompt, add_special_tokens=True) + text_prompt = ( + f"{example[speaker_col]}: {text}" + if has_source and example.get(speaker_col) + else text + ) + text_ids = tokenizer.encode(text_prompt, add_special_tokens = True) text_ids.append(END_OF_TEXT) # --- Build full input_ids (notebook lines 225-234) --- @@ -1360,11 +1617,13 @@ class UnslothTrainer: labels = list(input_ids) attention_mask = [1] * len(input_ids) - processed_examples.append({ - "input_ids": input_ids, - "labels": labels, - "attention_mask": attention_mask, - }) + processed_examples.append( + { + "input_ids": input_ids, + "labels": labels, + "attention_mask": attention_mask, + } + ) except Exception as e: logger.warning(f"Error processing SNAC example {idx}: {e}") @@ -1374,7 +1633,7 @@ class UnslothTrainer: # Progress update every 100 examples if (idx + 1) % 100 == 0: self._update_progress( - status_message=f"Encoding audio... {idx + 1}/{len(dataset)}" + status_message = f"Encoding audio... {idx + 1}/{len(dataset)}" ) # Free SNAC model from GPU @@ -1382,6 +1641,7 @@ class UnslothTrainer: snac_model.to("cpu") del snac_model import gc + gc.collect() torch.cuda.empty_cache() self._cuda_audio_used = True @@ -1392,11 +1652,13 @@ class UnslothTrainer: ) result_dataset = Dataset.from_list(processed_examples) - logger.info(f"SNAC preprocessing complete: {len(result_dataset)} examples " - f"({skipped} skipped)\n") + logger.info( + f"SNAC preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n" + ) return result_dataset - def _preprocess_bicodec_dataset(self, dataset, custom_format_mapping=None): + def _preprocess_bicodec_dataset(self, dataset, custom_format_mapping = None): """Preprocess dataset for Spark-TTS training with BiCodec tokenizer. Mirrors Spark_TTS_(0_5B).ipynb: encode audio with BiCodec (semantic + global tokens), @@ -1413,14 +1675,23 @@ class UnslothTrainer: # The sparktts Python package lives in the SparkAudio/Spark-TTS GitHub repo, # NOT in the unsloth/Spark-TTS-0.5B HF model repo. Clone it if needed. - spark_code_dir = os.path.join(os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS") + spark_code_dir = os.path.join( + os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS" + ) sparktts_pkg = os.path.join(spark_code_dir, "sparktts") if not os.path.isdir(sparktts_pkg): - self._update_progress(status_message="Cloning Spark-TTS code repo...") + self._update_progress(status_message = "Cloning Spark-TTS code repo...") logger.info(f"Cloning SparkAudio/Spark-TTS to {spark_code_dir}...\n") 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: @@ -1443,18 +1714,21 @@ class UnslothTrainer: # Cast audio column so datasets 4.x AudioDecoder objects are decoded to dicts. # Don't resample here β€” BiCodec's target_sr may differ; the loop handles resampling. from datasets import Audio + dataset = dataset.cast_column(audio_col, Audio()) # Load BiCodec tokenizer - self._update_progress(status_message="Loading BiCodec tokenizer...") + self._update_progress(status_message = "Loading BiCodec tokenizer...") logger.info("Loading BiCodec tokenizer...\n") audio_tokenizer = BiCodecTokenizer(self._spark_tts_repo_dir, device) - target_sr = audio_tokenizer.config['sample_rate'] + target_sr = audio_tokenizer.config["sample_rate"] - self._update_progress(status_message="Encoding audio with BiCodec...") - logger.info(f"BiCodec preprocessing: audio_col='{audio_col}', text_col='{text_col}', " - f"has_source={has_source}, target_sr={target_sr}\n") + self._update_progress(status_message = "Encoding audio with BiCodec...") + logger.info( + f"BiCodec preprocessing: audio_col='{audio_col}', text_col='{text_col}', " + f"has_source={has_source}, target_sr={target_sr}\n" + ) def extract_wav2vec2_features(wavs: torch.Tensor) -> torch.Tensor: """Extract wav2vec2 features (average of layers 11, 14, 16).""" @@ -1464,11 +1738,13 @@ class UnslothTrainer: processed = audio_tokenizer.processor( wav_np, - sampling_rate=16000, - return_tensors="pt", - padding=True, + sampling_rate = 16000, + return_tensors = "pt", + padding = True, + ) + input_values = processed.input_values.to( + audio_tokenizer.feature_extractor.device ) - input_values = processed.input_values.to(audio_tokenizer.feature_extractor.device) model_output = audio_tokenizer.feature_extractor(input_values) if model_output.hidden_states is None: @@ -1505,7 +1781,7 @@ class UnslothTrainer: # Resample if needed if sampling_rate != target_sr: - resampler = T.Resample(orig_freq=sampling_rate, new_freq=target_sr) + resampler = T.Resample(orig_freq = sampling_rate, new_freq = target_sr) audio_tensor_temp = torch.from_numpy(audio_array).float() audio_array = resampler(audio_tensor_temp).numpy() @@ -1517,8 +1793,12 @@ class UnslothTrainer: ref_wav_np = audio_tokenizer.get_ref_clip(audio_array) # Prepare tensors - audio_tensor = torch.from_numpy(audio_array).unsqueeze(0).float().to(device) - ref_wav_tensor = torch.from_numpy(ref_wav_np).unsqueeze(0).float().to(device) + audio_tensor = ( + torch.from_numpy(audio_array).unsqueeze(0).float().to(device) + ) + ref_wav_tensor = ( + torch.from_numpy(ref_wav_np).unsqueeze(0).float().to(device) + ) # Extract wav2vec2 features feat = extract_wav2vec2_features(audio_tensor) @@ -1530,31 +1810,45 @@ class UnslothTrainer: } # BiCodec tokenize - semantic_token_ids, global_token_ids = audio_tokenizer.model.tokenize(batch) + semantic_token_ids, global_token_ids = audio_tokenizer.model.tokenize( + batch + ) global_tokens = "".join( - [f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze().cpu().numpy()] + [ + f"<|bicodec_global_{i}|>" + for i in global_token_ids.squeeze().cpu().numpy() + ] ) semantic_tokens = "".join( - [f"<|bicodec_semantic_{i}|>" for i in semantic_token_ids.squeeze().cpu().numpy()] + [ + f"<|bicodec_semantic_{i}|>" + for i in semantic_token_ids.squeeze().cpu().numpy() + ] ) # Format text with source prefix if available - text_content = f"{example[speaker_col]}: {text}" if has_source and example.get(speaker_col) else text + text_content = ( + f"{example[speaker_col]}: {text}" + if has_source and example.get(speaker_col) + else text + ) - formatted = "".join([ - "<|task_tts|>", - "<|start_content|>", - text_content, - "<|end_content|>", - "<|start_global_token|>", - global_tokens, - "<|end_global_token|>", - "<|start_semantic_token|>", - semantic_tokens, - "<|end_semantic_token|>", - "<|im_end|>", - ]) + formatted = "".join( + [ + "<|task_tts|>", + "<|start_content|>", + text_content, + "<|end_content|>", + "<|start_global_token|>", + global_tokens, + "<|end_global_token|>", + "<|start_semantic_token|>", + semantic_tokens, + "<|end_semantic_token|>", + "<|im_end|>", + ] + ) processed_examples.append({"text": formatted}) @@ -1566,7 +1860,7 @@ class UnslothTrainer: # Progress update every 100 examples if (idx + 1) % 100 == 0: self._update_progress( - status_message=f"Encoding audio with BiCodec... {idx + 1}/{len(dataset)}" + status_message = f"Encoding audio with BiCodec... {idx + 1}/{len(dataset)}" ) # Free BiCodec model from GPU @@ -1575,6 +1869,7 @@ class UnslothTrainer: audio_tokenizer.feature_extractor.cpu() del audio_tokenizer import gc + gc.collect() torch.cuda.empty_cache() self._cuda_audio_used = True @@ -1585,15 +1880,17 @@ class UnslothTrainer: ) result_dataset = Dataset.from_list(processed_examples) - logger.info(f"BiCodec preprocessing complete: {len(result_dataset)} examples " - f"({skipped} skipped)\n") + logger.info( + f"BiCodec preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n" + ) # Debug: show first example text (truncated) sample = result_dataset[0]["text"] logger.info(f"Sample text (first 200 chars): {sample[:200]}...\n") logger.info(f"Sample text length: {len(sample)} chars\n") return result_dataset - def _preprocess_dac_dataset(self, dataset, custom_format_mapping=None): + def _preprocess_dac_dataset(self, dataset, custom_format_mapping = None): """Preprocess dataset for OuteTTS training with DAC codec. Mirrors Oute_TTS_(1B).ipynb DataCreationV3: uses Whisper for word timings, @@ -1613,15 +1910,23 @@ class UnslothTrainer: # Clone OuteTTS repo (same as audio_codecs._load_dac) import subprocess + base_dir = os.path.dirname(os.path.abspath(__file__)) outetts_code_dir = os.path.join(base_dir, "inference", "OuteTTS") outetts_pkg = os.path.join(outetts_code_dir, "outetts") if not os.path.isdir(outetts_pkg): - self._update_progress(status_message="Cloning OuteTTS code repo...") + self._update_progress(status_message = "Cloning OuteTTS code repo...") logger.info(f"Cloning edwko/OuteTTS to {outetts_code_dir}...\n") 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, ) for fpath in [ os.path.join(outetts_pkg, "models", "gguf_model.py"), @@ -1651,29 +1956,35 @@ class UnslothTrainer: # Cast audio to 24kHz (notebook: dataset.cast_column("audio", Audio(sampling_rate=24000))) from datasets import Audio - dataset = dataset.cast_column(audio_col, Audio(sampling_rate=24000)) + + dataset = dataset.cast_column(audio_col, Audio(sampling_rate = 24000)) logger.info("Cast audio column to 24kHz\n") # Load Whisper for word timings - self._update_progress(status_message="Loading Whisper model for word timings...") + self._update_progress( + status_message = "Loading Whisper model for word timings..." + ) logger.info("Loading Whisper model for word timings...\n") import whisper - whisper_model = whisper.load_model("turbo", device=device) + + whisper_model = whisper.load_model("turbo", device = device) # Load OuteTTS AudioProcessor + PromptProcessor - self._update_progress(status_message="Loading OuteTTS AudioProcessor...") + self._update_progress(status_message = "Loading OuteTTS AudioProcessor...") logger.info("Loading OuteTTS AudioProcessor...\n") model_tokenizer_path = "OuteAI/Llama-OuteTTS-1.0-1B" dummy_config = OuteTTSModelConfig( - tokenizer_path=model_tokenizer_path, - device=device, - audio_codec_path=None, + tokenizer_path = model_tokenizer_path, + device = device, + audio_codec_path = None, ) - audio_processor = AudioProcessor(config=dummy_config) + audio_processor = AudioProcessor(config = dummy_config) prompt_processor = PromptProcessor(model_tokenizer_path) - self._update_progress(status_message="Preprocessing audio with OuteTTS...") - logger.info(f"DAC preprocessing: audio_col='{audio_col}', text_col='{text_col}'\n") + self._update_progress(status_message = "Preprocessing audio with OuteTTS...") + logger.info( + f"DAC preprocessing: audio_col='{audio_col}', text_col='{text_col}'\n" + ) processed_examples = [] skipped = 0 @@ -1694,28 +2005,30 @@ class UnslothTrainer: skipped += 1 continue - audio_array = np.array(audio_data["array"], dtype=np.float32) + audio_array = np.array(audio_data["array"], dtype = np.float32) sampling_rate = audio_data.get("sampling_rate", 24000) # Convert to WAV bytes (Whisper needs a file path) buf = io.BytesIO() - sf.write(buf, audio_array, sampling_rate, format="WAV", subtype="FLOAT") + sf.write(buf, audio_array, sampling_rate, format = "WAV", subtype = "FLOAT") buf.seek(0) audio_bytes = buf.getvalue() # 1. Get word timings from Whisper with tempfile.NamedTemporaryFile( - suffix=".wav", - delete=False, - dir=str(ensure_dir(tmp_root())), + suffix = ".wav", + delete = False, + dir = str(ensure_dir(tmp_root())), ) as tmp: tmp.write(audio_bytes) tmp.flush() tmp_path = tmp.name try: - whisper_result = whisper_model.transcribe(tmp_path, word_timestamps=True) + whisper_result = whisper_model.transcribe( + tmp_path, word_timestamps = True + ) finally: - Path(tmp_path).unlink(missing_ok=True) + Path(tmp_path).unlink(missing_ok = True) normalized_transcript = text_normalizations(text) words_with_timings = [] @@ -1724,11 +2037,13 @@ class UnslothTrainer: for word_info in segment.get("words", []): cleaned = word_info["word"].strip() if cleaned: - words_with_timings.append({ - "word": cleaned, - "start": float(word_info["start"]), - "end": float(word_info["end"]), - }) + words_with_timings.append( + { + "word": cleaned, + "start": float(word_info["start"]), + "end": float(word_info["end"]), + } + ) if not words_with_timings: skipped += 1 @@ -1757,16 +2072,17 @@ class UnslothTrainer: if (idx + 1) % 100 == 0: self._update_progress( - status_message=f"Preprocessing audio with OuteTTS... {idx + 1}/{len(dataset)}" + status_message = f"Preprocessing audio with OuteTTS... {idx + 1}/{len(dataset)}" ) # Free Whisper from GPU (notebook: data_processor.whisper_model.to('cpu')) logger.info("Moving Whisper model to CPU...\n") - whisper_model.to('cpu') + whisper_model.to("cpu") del whisper_model del audio_processor del prompt_processor import gc + gc.collect() torch.cuda.empty_cache() self._cuda_audio_used = True @@ -1777,13 +2093,17 @@ class UnslothTrainer: ) result_dataset = HFDataset.from_list(processed_examples) - logger.info(f"DAC preprocessing complete: {len(result_dataset)} examples " - f"({skipped} skipped)\n") + logger.info( + f"DAC preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n" + ) sample = result_dataset[0]["text"] logger.info(f"Sample text (first 200 chars): {sample[:200]}...\n") return result_dataset - def _preprocess_whisper_dataset(self, dataset, eval_split=None, custom_format_mapping=None): + def _preprocess_whisper_dataset( + self, dataset, eval_split = None, custom_format_mapping = None + ): """Preprocess dataset for Whisper speech-to-text training. Mirrors Whisper.ipynb: extract audio features with Whisper's feature @@ -1803,20 +2123,24 @@ class UnslothTrainer: ) # Cast audio to 16kHz (Whisper's expected sample rate) - dataset = dataset.cast_column(audio_col, Audio(sampling_rate=WHISPER_SAMPLE_RATE)) + dataset = dataset.cast_column( + audio_col, Audio(sampling_rate = WHISPER_SAMPLE_RATE) + ) # Train/eval split (notebook does dataset.train_test_split) eval_dataset_raw = None if eval_split: - splits = dataset.train_test_split(test_size=0.06, seed=42) + splits = dataset.train_test_split(test_size = 0.06, seed = 42) dataset = splits["train"] eval_dataset_raw = splits["test"] - self._update_progress(status_message="Processing audio for Whisper...") - logger.info(f"Whisper preprocessing: audio_col='{audio_col}', text_col='{text_col}', " - f"samples={len(dataset)}\n") + self._update_progress(status_message = "Processing audio for Whisper...") + logger.info( + f"Whisper preprocessing: audio_col='{audio_col}', text_col='{text_col}', " + f"samples={len(dataset)}\n" + ) - def process_split(ds, split_name="train"): + def process_split(ds, split_name = "train"): processed = [] skipped = 0 for idx in range(len(ds)): @@ -1828,53 +2152,67 @@ class UnslothTrainer: try: audio_data = example.get(audio_col) text = example.get(text_col) - if audio_data is None or audio_data.get("array") is None or not text: + if ( + audio_data is None + or audio_data.get("array") is None + or not text + ): skipped += 1 continue # Extract audio features (notebook line 112-115) features = self.tokenizer.feature_extractor( - audio_data["array"], sampling_rate=audio_data["sampling_rate"] + audio_data["array"], sampling_rate = audio_data["sampling_rate"] ) # Tokenize text (notebook line 116) tokenized_text = self.tokenizer.tokenizer(text) - processed.append({ - "input_features": features.input_features[0], - "labels": tokenized_text.input_ids, - }) + processed.append( + { + "input_features": features.input_features[0], + "labels": tokenized_text.input_ids, + } + ) except Exception as e: - logger.warning(f"Error processing Whisper {split_name} example {idx}: {e}") + logger.warning( + f"Error processing Whisper {split_name} example {idx}: {e}" + ) skipped += 1 continue if (idx + 1) % 100 == 0: self._update_progress( - status_message=f"Processing {split_name} audio... {idx + 1}/{len(ds)}" + status_message = f"Processing {split_name} audio... {idx + 1}/{len(ds)}" ) - logger.info(f"Whisper {split_name} preprocessing: {len(processed)} examples ({skipped} skipped)\n") + logger.info( + f"Whisper {split_name} preprocessing: {len(processed)} examples ({skipped} skipped)\n" + ) return processed train_data = process_split(dataset, "train") - eval_data = process_split(eval_dataset_raw, "eval") if eval_dataset_raw else None + eval_data = ( + process_split(eval_dataset_raw, "eval") if eval_dataset_raw else None + ) if not train_data: raise ValueError("No valid examples after Whisper preprocessing") return (train_data, eval_data) - def load_and_format_dataset(self, - dataset_source: str, - format_type: str = "auto", - local_datasets: list = None, - custom_format_mapping: dict = None, - subset: str = None, - train_split: str = "train", - eval_split: str = None, - eval_steps: float = 0.00, - dataset_slice_start: int = None, - dataset_slice_end: int = None) -> Optional[tuple]: + def load_and_format_dataset( + self, + dataset_source: str, + format_type: str = "auto", + local_datasets: list = None, + custom_format_mapping: dict = None, + subset: str = None, + train_split: str = "train", + eval_split: str = None, + eval_steps: float = 0.00, + dataset_slice_start: int = None, + dataset_slice_end: int = None, + ) -> Optional[tuple]: """ Load and prepare dataset for training. @@ -1888,7 +2226,9 @@ class UnslothTrainer: try: dataset = None eval_dataset = None - has_separate_eval_source = False # True if eval comes from a separate HF split + has_separate_eval_source = ( + False # True if eval comes from a separate HF split + ) eval_enabled = eval_steps is not None and eval_steps > 0 if local_datasets: @@ -1919,35 +2259,41 @@ class UnslothTrainer: continue # Fall through to single-file detection for dirs with json/csv candidates: list[Path] = [] - for ext in ('.json', '.jsonl', '.csv', '.parquet'): + for ext in (".json", ".jsonl", ".csv", ".parquet"): candidates.extend(sorted(file_path_obj.glob(f"*{ext}"))) 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(str(file_path_obj)) if all_files: # Determine loader type from the first file extension first_ext = Path(all_files[0]).suffix.lower() - if first_ext in ('.json', '.jsonl'): - loader = 'json' - elif first_ext == '.csv': - loader = 'csv' - elif first_ext == '.parquet': - loader = 'parquet' + if first_ext in (".json", ".jsonl"): + loader = "json" + elif first_ext == ".csv": + loader = "csv" + elif first_ext == ".parquet": + loader = "parquet" else: - raise ValueError(f"Unsupported local dataset format: {all_files[0]}") + raise ValueError( + f"Unsupported local dataset format: {all_files[0]}" + ) - dataset = load_dataset(loader, data_files=all_files, split='train') + dataset = load_dataset(loader, data_files = all_files, split = "train") # Check if stopped during dataset loading if self.should_stop: logger.info("Stopped during dataset loading\n") return None - self._update_progress(status_message=f"Loaded {len(dataset)} samples from local files") + self._update_progress( + status_message = f"Loaded {len(dataset)} samples from local files" + ) logger.info(f"Loaded {len(dataset)} samples from local files\n") logger.info(f"[DEBUG] Dataset cache_files: {dataset.cache_files}\n") @@ -1959,9 +2305,11 @@ class UnslothTrainer: load_kwargs["name"] = subset _slice_start = dataset_slice_start or 0 - if (dataset_slice_end is not None - and dataset_slice_end >= 0 - and dataset_slice_end >= _slice_start): + if ( + dataset_slice_end is not None + and dataset_slice_end >= 0 + and dataset_slice_end >= _slice_start + ): # Manual slice β€” stream only the rows we need instead of # downloading the entire dataset. rows_to_stream = dataset_slice_end + 1 @@ -1970,14 +2318,14 @@ class UnslothTrainer: f"(start={dataset_slice_start}, end={dataset_slice_end}), " f"streaming {rows_to_stream} rows\n" ) - stream = load_dataset(**load_kwargs, streaming=True) + stream = load_dataset(**load_kwargs, streaming = True) dataset = Dataset.from_list(list(stream.take(rows_to_stream))) logger.info( f"[dataset-slice] Downloaded {len(dataset)} rows " f"(requested {rows_to_stream})\n" ) self._update_progress( - status_message=f"Streamed {len(dataset)} rows from HuggingFace" + status_message = f"Streamed {len(dataset)} rows from HuggingFace" ) else: dataset = load_dataset(**load_kwargs) @@ -1987,8 +2335,12 @@ class UnslothTrainer: logger.info("Stopped during dataset loading\n") return None - self._update_progress(status_message=f"Loaded dataset from HuggingFace: {dataset_source}") - logger.info(f"Loaded dataset from Hugging Face: {dataset_source} ({len(dataset)} rows)\n") + self._update_progress( + status_message = f"Loaded dataset from HuggingFace: {dataset_source}" + ) + logger.info( + f"Loaded dataset from Hugging Face: {dataset_source} ({len(dataset)} rows)\n" + ) # Resolve eval split from a separate HF split (explicit or auto-detected) if eval_enabled: @@ -2001,20 +2353,26 @@ class UnslothTrainer: eval_load_kwargs["name"] = subset eval_dataset = load_dataset(**eval_load_kwargs) has_separate_eval_source = True - logger.info(f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n") + logger.info( + f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n" + ) elif eval_split and eval_split == effective_train: # Same split as training β€” will do 80/20 split after formatting - logger.info(f"Eval split '{eval_split}' is the same as train split β€” will split 80/20\n") + logger.info( + f"Eval split '{eval_split}' is the same as train split β€” will split 80/20\n" + ) else: # Auto-detect eval split from HF (returns a separate dataset, or None) eval_dataset = self._auto_detect_eval_split_from_hf( - dataset_source=dataset_source, - subset=subset, + dataset_source = dataset_source, + subset = subset, ) if eval_dataset is not None: has_separate_eval_source = True else: - logger.info("Eval disabled (eval_steps <= 0), skipping eval split detection\n") + logger.info( + "Eval disabled (eval_steps <= 0), skipping eval split detection\n" + ) if dataset is None: raise ValueError("No dataset provided") @@ -2023,13 +2381,21 @@ class UnslothTrainer: if dataset_slice_start is not None or dataset_slice_end is not None: total_rows = len(dataset) start = dataset_slice_start if dataset_slice_start is not None else 0 - end = dataset_slice_end if dataset_slice_end is not None else total_rows - 1 + end = ( + dataset_slice_end + if dataset_slice_end is not None + else total_rows - 1 + ) # Clamp to valid range start = max(0, min(start, total_rows - 1)) end = max(start, min(end, total_rows - 1)) dataset = dataset.select(range(start, end + 1)) - logger.info(f"Sliced dataset to rows [{start}, {end}]: {len(dataset)} of {total_rows} rows\n") - self._update_progress(status_message=f"Sliced dataset to {len(dataset)} rows (indices {start}-{end})") + logger.info( + f"Sliced dataset to rows [{start}, {end}]: {len(dataset)} of {total_rows} rows\n" + ) + self._update_progress( + status_message = f"Sliced dataset to {len(dataset)} rows (indices {start}-{end})" + ) # Check if stopped before applying template if self.should_stop: @@ -2037,30 +2403,38 @@ class UnslothTrainer: return None # ========== AUDIO MODELS: custom preprocessing ========== - if self._audio_type == 'csm': + if self._audio_type == "csm": processed = self._preprocess_csm_dataset(dataset, custom_format_mapping) return (processed, None) - elif self._audio_type == 'whisper': + elif self._audio_type == "whisper": train_data, eval_data = self._preprocess_whisper_dataset( - dataset, eval_split=eval_split, custom_format_mapping=custom_format_mapping + dataset, + eval_split = eval_split, + custom_format_mapping = custom_format_mapping, ) return (train_data, eval_data) - elif self._audio_type == 'snac': - processed = self._preprocess_snac_dataset(dataset, custom_format_mapping) + elif self._audio_type == "snac": + processed = self._preprocess_snac_dataset( + dataset, custom_format_mapping + ) return (processed, None) - elif self._audio_type == 'bicodec': - processed = self._preprocess_bicodec_dataset(dataset, custom_format_mapping) + elif self._audio_type == "bicodec": + processed = self._preprocess_bicodec_dataset( + dataset, custom_format_mapping + ) return ({"dataset": processed, "final_format": "audio_bicodec"}, None) - elif self._audio_type == 'dac': + elif self._audio_type == "dac": processed = self._preprocess_dac_dataset(dataset, custom_format_mapping) return ({"dataset": processed, "final_format": "audio_dac"}, None) elif self.is_audio_vlm: - formatted = self._format_audio_vlm_dataset(dataset, custom_format_mapping) + formatted = self._format_audio_vlm_dataset( + dataset, custom_format_mapping + ) return (formatted, None) # ========== FORMAT FIRST ========== @@ -2068,13 +2442,13 @@ class UnslothTrainer: dataset_info = format_and_template_dataset( dataset, - model_name=self.model_name, - tokenizer=self.tokenizer, - is_vlm=self.is_vlm, - format_type=format_type, - dataset_name=dataset_source, - custom_format_mapping=custom_format_mapping, - progress_callback=self._update_progress, + model_name = self.model_name, + tokenizer = self.tokenizer, + is_vlm = self.is_vlm, + format_type = format_type, + dataset_name = dataset_source, + custom_format_mapping = custom_format_mapping, + progress_callback = self._update_progress, ) # Check if stopped during formatting @@ -2087,10 +2461,12 @@ class UnslothTrainer: errors = dataset_info.get("errors", []) error_msg = "; ".join(errors) if errors else "Dataset formatting failed" logger.error(f"Dataset conversion failed: {error_msg}") - self._update_progress(error=error_msg) + self._update_progress(error = error_msg) return None - self._update_progress(status_message=f"Dataset formatted and ready for training") + self._update_progress( + status_message = f"Dataset formatted and ready for training" + ) logger.info(f"Dataset formatted successfully\n") # ========== THEN SPLIT ========== @@ -2099,12 +2475,12 @@ class UnslothTrainer: logger.info(f"Formatting eval dataset ({len(eval_dataset)} rows)...\n") eval_info = format_and_template_dataset( eval_dataset, - model_name=self.model_name, - tokenizer=self.tokenizer, - is_vlm=self.is_vlm, - format_type=format_type, - dataset_name=dataset_source, - custom_format_mapping=custom_format_mapping, + model_name = self.model_name, + tokenizer = self.tokenizer, + is_vlm = self.is_vlm, + format_type = format_type, + dataset_name = dataset_source, + custom_format_mapping = custom_format_mapping, ) eval_dataset = eval_info["dataset"] logger.info(f"Eval dataset formatted successfully\n") @@ -2120,14 +2496,16 @@ class UnslothTrainer: except Exception as e: logger.error(f"Error loading dataset: {e}") - self._update_progress(error=str(e)) + self._update_progress(error = str(e)) return None - def _auto_detect_eval_split_from_hf(self, dataset_source: str, - subset: str) -> Optional[Dataset]: + def _auto_detect_eval_split_from_hf( + self, dataset_source: str, subset: str + ) -> Optional[Dataset]: """Auto-detect an eval split from HF dataset (separate named split only).""" try: from datasets import get_dataset_split_names + load_kwargs = {"path": dataset_source} if subset: load_kwargs["config_name"] = subset @@ -2142,10 +2520,14 @@ class UnslothTrainer: eval_load_kwargs["name"] = subset candidate_ds = load_dataset(**eval_load_kwargs) if len(candidate_ds) >= 16: - logger.info(f"Auto-detected eval split '{candidate}' with {len(candidate_ds)} rows\n") + logger.info( + f"Auto-detected eval split '{candidate}' with {len(candidate_ds)} rows\n" + ) return candidate_ds else: - logger.info(f"Found eval split '{candidate}' but only {len(candidate_ds)} rows (< 16), skipping\n") + logger.info( + f"Found eval split '{candidate}' but only {len(candidate_ds)} rows (< 16), skipping\n" + ) except Exception as e: logger.warning(f"Could not check dataset splits: {e}") @@ -2172,42 +2554,45 @@ class UnslothTrainer: eval_size = min(eval_size, n // 2) logger.info(f"Auto-splitting: {eval_size} rows for eval from {n} total\n") - split_result = dataset.train_test_split(test_size=eval_size, seed=3407) - logger.info(f"Split complete: {len(split_result['train'])} train, {len(split_result['test'])} eval\n") - return (split_result['train'], split_result['test']) + split_result = dataset.train_test_split(test_size = eval_size, seed = 3407) + logger.info( + f"Split complete: {len(split_result['train'])} train, {len(split_result['test'])} eval\n" + ) + return (split_result["train"], split_result["test"]) - def start_training(self, - dataset: Dataset, - eval_dataset: Dataset = None, - eval_steps: float = 0.00, - output_dir: str | None = None, - num_epochs: int = 3, - learning_rate: float = 5e-5, - batch_size: int = 2, - gradient_accumulation_steps: int = 4, - warmup_steps: int = None, - warmup_ratio: float = None, - max_steps: int = 0, - save_steps: int = 0, - weight_decay: float = 0.01, - random_seed: int = 3407, - packing: bool = False, - train_on_completions: bool = False, - enable_wandb: bool = False, - wandb_project: str = "unsloth-training", - wandb_token: str = None, - enable_tensorboard: bool = False, - tensorboard_dir: str | None = None, - **kwargs) -> bool: + def start_training( + self, + dataset: Dataset, + eval_dataset: Dataset = None, + eval_steps: float = 0.00, + output_dir: str | None = None, + num_epochs: int = 3, + learning_rate: float = 5e-5, + batch_size: int = 2, + gradient_accumulation_steps: int = 4, + warmup_steps: int = None, + warmup_ratio: float = None, + max_steps: int = 0, + save_steps: int = 0, + weight_decay: float = 0.01, + random_seed: int = 3407, + packing: bool = False, + train_on_completions: bool = False, + enable_wandb: bool = False, + wandb_project: str = "unsloth-training", + wandb_token: str = None, + enable_tensorboard: bool = False, + tensorboard_dir: str | None = None, + **kwargs, + ) -> bool: """Start training in a separate thread""" if self.is_training: logger.warning("Training already in progress") return False - if self.model is None or self.tokenizer is None: - self._update_progress(error="Model not loaded") + self._update_progress(error = "Model not loaded") return False # Pre-import heavy transformers modules on the main thread. @@ -2220,7 +2605,8 @@ class UnslothTrainer: TrainingArguments as _TrainingArguments, TrainerCallback as _TrainerCallback, ) - if self._audio_type == 'whisper': + + if self._audio_type == "whisper": from transformers import ( # noqa: F401 Seq2SeqTrainer as _Seq2SeqTrainer, Seq2SeqTrainingArguments as _Seq2SeqTrainingArguments, @@ -2228,31 +2614,31 @@ class UnslothTrainer: # Start training in separate thread self.training_thread = threading.Thread( - target=self._train_worker, - args=(dataset,), - kwargs={ - 'output_dir': output_dir, - 'num_epochs': num_epochs, - 'learning_rate': learning_rate, - 'batch_size': batch_size, - 'gradient_accumulation_steps': gradient_accumulation_steps, - 'warmup_steps': warmup_steps, - 'warmup_ratio': warmup_ratio, - 'max_steps': max_steps, - 'save_steps': save_steps, - 'weight_decay': weight_decay, - 'random_seed': random_seed, - 'packing': packing, - 'train_on_completions': train_on_completions, - 'enable_wandb': enable_wandb, - 'wandb_project': wandb_project, - 'wandb_token': wandb_token, - 'enable_tensorboard': enable_tensorboard, - 'tensorboard_dir': tensorboard_dir, - 'eval_dataset': eval_dataset, - 'eval_steps': eval_steps, - **kwargs - } + target = self._train_worker, + args = (dataset,), + kwargs = { + "output_dir": output_dir, + "num_epochs": num_epochs, + "learning_rate": learning_rate, + "batch_size": batch_size, + "gradient_accumulation_steps": gradient_accumulation_steps, + "warmup_steps": warmup_steps, + "warmup_ratio": warmup_ratio, + "max_steps": max_steps, + "save_steps": save_steps, + "weight_decay": weight_decay, + "random_seed": random_seed, + "packing": packing, + "train_on_completions": train_on_completions, + "enable_wandb": enable_wandb, + "wandb_project": wandb_project, + "wandb_token": wandb_token, + "enable_tensorboard": enable_tensorboard, + "tensorboard_dir": tensorboard_dir, + "eval_dataset": eval_dataset, + "eval_steps": eval_steps, + **kwargs, + }, ) self.should_stop = False @@ -2269,100 +2655,130 @@ class UnslothTrainer: """Worker function for training (runs in separate thread)""" try: # Store training parameters for metrics calculation - self.batch_size = training_args.get('batch_size', 2) - self.max_seq_length = training_args.get('max_seq_length', 2048) - self.gradient_accumulation_steps = training_args.get('gradient_accumulation_steps', 4) - + self.batch_size = training_args.get("batch_size", 2) + self.max_seq_length = training_args.get("max_seq_length", 2048) + self.gradient_accumulation_steps = training_args.get( + "gradient_accumulation_steps", 4 + ) + # Set training start time self.training_start_time = time.time() - - self._update_progress(is_training=True, error=None) + + self._update_progress(is_training = True, error = None) # Setup logging - if training_args.get('enable_wandb', False) and training_args.get('wandb_token'): - os.environ["WANDB_API_KEY"] = training_args['wandb_token'] + if training_args.get("enable_wandb", False) and training_args.get( + "wandb_token" + ): + os.environ["WANDB_API_KEY"] = training_args["wandb_token"] import wandb - wandb.init(project=training_args.get('wandb_project', 'unsloth-training')) + + wandb.init( + project = training_args.get("wandb_project", "unsloth-training") + ) # Create output directory output_dir = str(resolve_output_dir(training_args.get("output_dir"))) ensure_dir(Path(output_dir)) # ========== AUDIO TRAINER BRANCH ========== - if self._audio_type == 'csm': + if self._audio_type == "csm": # CSM uses plain HF Trainer (NOT SFTTrainer) # Needs remove_unused_columns=False for depth decoder (input_values + cutoffs) from transformers import Trainer as HFTrainer, TrainingArguments + self._apply_csm_forward_fix() - config = self._build_audio_training_args(training_args, output_dir, extra_args={ - "remove_unused_columns": False, - }) + config = self._build_audio_training_args( + training_args, + output_dir, + extra_args = { + "remove_unused_columns": False, + }, + ) self.trainer = HFTrainer( - model=self.model, train_dataset=dataset, - args=TrainingArguments(**config), + model = self.model, + train_dataset = dataset, + args = TrainingArguments(**config), ) self.trainer.add_callback(self._create_progress_callback()) - batch_size = training_args.get('batch_size', 2) + batch_size = training_args.get("batch_size", 2) total = self._calculate_total_steps( - len(dataset), batch_size, - training_args.get('gradient_accumulation_steps', 4), - training_args.get('num_epochs', 3), - training_args.get('max_steps', 0), + len(dataset), + batch_size, + training_args.get("gradient_accumulation_steps", 4), + training_args.get("num_epochs", 3), + training_args.get("max_steps", 0), + ) + self._update_progress( + total_steps = total, status_message = "Starting CSM training..." ) - self._update_progress(total_steps=total, status_message="Starting CSM training...") logger.info(f"CSM training config: {config}\n") self.trainer.train() self._finalize_training(output_dir, "CSM") return - elif self._audio_type == 'snac': + elif self._audio_type == "snac": # Orpheus: language model with SNAC codec tokens β€” plain HF Trainer # DataCollatorForSeq2Seq dynamically pads variable-length sequences per batch # (text + audio codes vary in length) and pads labels with -100. - from transformers import Trainer as HFTrainer, TrainingArguments, DataCollatorForSeq2Seq + from transformers import ( + Trainer as HFTrainer, + TrainingArguments, + DataCollatorForSeq2Seq, + ) config = self._build_audio_training_args(training_args, output_dir) self.trainer = HFTrainer( - model=self.model, train_dataset=dataset, - args=TrainingArguments(**config), - data_collator=DataCollatorForSeq2Seq( - tokenizer=self.tokenizer, padding=True, pad_to_multiple_of=8, + model = self.model, + train_dataset = dataset, + args = TrainingArguments(**config), + data_collator = DataCollatorForSeq2Seq( + tokenizer = self.tokenizer, + padding = True, + pad_to_multiple_of = 8, ), ) self.trainer.add_callback(self._create_progress_callback()) - batch_size = training_args.get('batch_size', 2) + batch_size = training_args.get("batch_size", 2) total = self._calculate_total_steps( - len(dataset), batch_size, - training_args.get('gradient_accumulation_steps', 4), - training_args.get('num_epochs', 3), - training_args.get('max_steps', 0), + len(dataset), + batch_size, + training_args.get("gradient_accumulation_steps", 4), + training_args.get("num_epochs", 3), + training_args.get("max_steps", 0), + ) + self._update_progress( + total_steps = total, status_message = "Starting SNAC training..." ) - self._update_progress(total_steps=total, status_message="Starting SNAC training...") logger.info(f"SNAC training config: {config}\n") self.trainer.train() self._finalize_training(output_dir, "SNAC") return - elif self._audio_type == 'whisper': + elif self._audio_type == "whisper": # Whisper: Seq2SeqTrainer with custom speech collator from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments from utils.datasets import DataCollatorSpeechSeq2SeqWithPadding - eval_dataset = training_args.get('eval_dataset', None) + eval_dataset = training_args.get("eval_dataset", None) extra = {"remove_unused_columns": False, "label_names": ["labels"]} if eval_dataset: extra["eval_strategy"] = "steps" - extra["eval_steps"] = training_args.get('eval_steps', 5) + extra["eval_steps"] = training_args.get("eval_steps", 5) - config = self._build_audio_training_args(training_args, output_dir, extra_args=extra) + config = self._build_audio_training_args( + training_args, output_dir, extra_args = extra + ) trainer_kwargs = { "model": self.model, "train_dataset": dataset, - "data_collator": DataCollatorSpeechSeq2SeqWithPadding(processor=self.tokenizer), + "data_collator": DataCollatorSpeechSeq2SeqWithPadding( + processor = self.tokenizer + ), "processing_class": self.tokenizer.feature_extractor, "args": Seq2SeqTrainingArguments(**config), } @@ -2372,27 +2788,37 @@ class UnslothTrainer: self.trainer = Seq2SeqTrainer(**trainer_kwargs) self.trainer.add_callback(self._create_progress_callback()) - batch_size = training_args.get('batch_size', 2) + batch_size = training_args.get("batch_size", 2) total = self._calculate_total_steps( - len(dataset), batch_size, - training_args.get('gradient_accumulation_steps', 4), - training_args.get('num_epochs', 3), - training_args.get('max_steps', 0), + len(dataset), + batch_size, + training_args.get("gradient_accumulation_steps", 4), + training_args.get("num_epochs", 3), + training_args.get("max_steps", 0), + ) + self._update_progress( + total_steps = total, status_message = "Starting Whisper training..." ) - self._update_progress(total_steps=total, status_message="Starting Whisper training...") logger.info(f"Whisper training config: {config}\n") self.trainer.train() self._finalize_training(output_dir, "Whisper") return - elif self._audio_type is not None and self._audio_type not in ('bicodec', 'dac'): + elif self._audio_type is not None and self._audio_type not in ( + "bicodec", + "dac", + ): # bicodec/dac use the standard SFTTrainer text path below - raise NotImplementedError(f"Audio training for '{self._audio_type}' not yet implemented") + raise NotImplementedError( + f"Audio training for '{self._audio_type}' not yet implemented" + ) # ========== DATA COLLATOR SELECTION ========== # Detect special model types model_name_lower = self.model_name.lower() - is_deepseek_ocr = "deepseek" in model_name_lower and "ocr" in model_name_lower + is_deepseek_ocr = ( + "deepseek" in model_name_lower and "ocr" in model_name_lower + ) logger.info("Configuring data collator...\n") @@ -2409,7 +2835,7 @@ class UnslothTrainer: "snapshot_download('unsloth/DeepSeek-OCR', local_dir='deepseek_ocr')" ) logger.error(error_msg) - self._update_progress(error=error_msg, is_training=False) + self._update_progress(error = error_msg, is_training = False) return try: @@ -2418,19 +2844,21 @@ class UnslothTrainer: logger.info("Configuring DeepSeek OCR data collator...\n") FastVisionModel.for_training(self.model) data_collator = DeepSeekOCRDataCollator( - tokenizer=self.tokenizer, - model=self.model, - image_size=640, - base_size=1024, - crop_mode=True, - train_on_responses_only=training_args.get('train_on_completions', False), + tokenizer = self.tokenizer, + model = self.model, + image_size = 640, + base_size = 1024, + crop_mode = True, + train_on_responses_only = training_args.get( + "train_on_completions", False + ), ) logger.info("DeepSeek OCR data collator configured successfully\n") except Exception as e: logger.error(f"Failed to configure DeepSeek OCR collator: {e}") error_msg = f"Error configuring DeepSeek OCR: {str(e)}" - self._update_progress(error=error_msg, is_training=False) + self._update_progress(error = error_msg, is_training = False) return elif self.is_audio_vlm: @@ -2439,26 +2867,33 @@ class UnslothTrainer: logger.info("Configuring audio VLM data collator...\n") processor = self.tokenizer # FastModel returns processor as tokenizer - audio_col_name = getattr(self, '_audio_vlm_audio_col', 'audio') + audio_col_name = getattr(self, "_audio_vlm_audio_col", "audio") def audio_vlm_collate_fn(examples): texts = [] audios = [] for example in examples: text = processor.apply_chat_template( - example["messages"], tokenize=False, add_generation_prompt=False + example["messages"], + tokenize = False, + add_generation_prompt = False, ).strip() texts.append(text) audios.append(example[audio_col_name]["array"]) batch = processor( - text=texts, audio=audios, return_tensors="pt", padding=True + text = texts, audio = audios, return_tensors = "pt", padding = True ) # Labels = input_ids with special tokens masked labels = batch["input_ids"].clone() labels[labels == processor.tokenizer.pad_token_id] = -100 - for attr in ('audio_token_id', 'image_token_id', 'boi_token_id', 'eoi_token_id'): + for attr in ( + "audio_token_id", + "image_token_id", + "boi_token_id", + "eoi_token_id", + ): token_id = getattr(processor.tokenizer, attr, None) if token_id is not None: labels[labels == token_id] = -100 @@ -2479,41 +2914,52 @@ class UnslothTrainer: # ========== TRAINING CONFIGURATION ========== # Handle warmup_steps vs warmup_ratio - warmup_steps_val = training_args.get('warmup_steps', None) - warmup_ratio_val = training_args.get('warmup_ratio', None) - - lr_value = training_args.get('learning_rate', 2e-4) - logger.info(f"[DEBUG] learning_rate from training_args: {lr_value} (type: {type(lr_value).__name__})\n") + warmup_steps_val = training_args.get("warmup_steps", None) + warmup_ratio_val = training_args.get("warmup_ratio", None) + + lr_value = training_args.get("learning_rate", 2e-4) + logger.info( + f"[DEBUG] learning_rate from training_args: {lr_value} (type: {type(lr_value).__name__})\n" + ) config_args = { - "per_device_train_batch_size": training_args.get('batch_size', 2), - "gradient_accumulation_steps": training_args.get('gradient_accumulation_steps', 4), - "num_train_epochs": training_args.get('num_epochs', 3), # Default to epochs + "per_device_train_batch_size": training_args.get("batch_size", 2), + "gradient_accumulation_steps": training_args.get( + "gradient_accumulation_steps", 4 + ), + "num_train_epochs": training_args.get( + "num_epochs", 3 + ), # Default to epochs "learning_rate": lr_value, "fp16": not is_bfloat16_supported(), "bf16": is_bfloat16_supported(), "logging_steps": 1, - "weight_decay": training_args.get('weight_decay', 0.01), - "seed": training_args.get('random_seed', 3407), + "weight_decay": training_args.get("weight_decay", 0.01), + "seed": training_args.get("random_seed", 3407), "output_dir": output_dir, "report_to": _build_report_targets(training_args), "include_num_input_tokens_seen": True, # Enable token counting - "dataset_num_proc": 1 if (self.is_audio or self.is_audio_vlm or self._cuda_audio_used) else safe_num_proc(max(1, os.cpu_count() // 4)), - "max_seq_length": training_args.get('max_seq_length', 2048), + "dataset_num_proc": 1 + if (self.is_audio or self.is_audio_vlm or self._cuda_audio_used) + else safe_num_proc(max(1, os.cpu_count() // 4)), + "max_seq_length": training_args.get("max_seq_length", 2048), } if training_args.get("enable_tensorboard", False): config_args["logging_dir"] = str( resolve_tensorboard_dir(training_args.get("tensorboard_dir")) ) - logger.info(f"[DEBUG] dataset_num_proc={config_args['dataset_num_proc']} (is_audio={self.is_audio}, is_audio_vlm={self.is_audio_vlm}, _cuda_audio_used={self._cuda_audio_used})") + logger.info( + f"[DEBUG] dataset_num_proc={config_args['dataset_num_proc']} (is_audio={self.is_audio}, is_audio_vlm={self.is_audio_vlm}, _cuda_audio_used={self._cuda_audio_used})" + ) # On Windows with transformers 5.x, disable DataLoader multiprocessing # to avoid issues with modified sys.path (.venv_t5) in spawned workers. if sys.platform == "win32": import transformers as _tf + if _tf.__version__.startswith("5."): config_args["dataloader_num_workers"] = 0 - + # Add warmup parameter - use warmup_ratio if provided, otherwise warmup_steps if warmup_ratio_val is not None: config_args["warmup_ratio"] = warmup_ratio_val @@ -2527,13 +2973,13 @@ class UnslothTrainer: logger.info(f"Using default warmup_steps: 5\n") # Add save_steps if specified - save_steps_val = training_args.get('save_steps', 0) + save_steps_val = training_args.get("save_steps", 0) if save_steps_val and save_steps_val > 0: config_args["save_steps"] = save_steps_val config_args["save_strategy"] = "steps" # If max_steps is specified, use it instead of epochs - max_steps_val = training_args.get('max_steps', 0) + max_steps_val = training_args.get("max_steps", 0) if max_steps_val and max_steps_val > 0: del config_args["num_train_epochs"] # Remove epochs config_args["max_steps"] = max_steps_val # Use steps instead @@ -2542,62 +2988,74 @@ class UnslothTrainer: logger.info(f"Training for {config_args['num_train_epochs']} epochs\n") # ========== EVAL CONFIGURATION ========== - eval_dataset = training_args.get('eval_dataset', None) - eval_steps_val = training_args.get('eval_steps', 0.00) + eval_dataset = training_args.get("eval_dataset", None) + eval_steps_val = training_args.get("eval_steps", 0.00) if eval_dataset is not None: if eval_steps_val > 0: config_args["eval_strategy"] = "steps" config_args["eval_steps"] = eval_steps_val - logger.info(f"βœ… Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n") + logger.info( + f"βœ… Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n" + ) logger.info(f"Eval dataset: {len(eval_dataset)} rows\n") else: - logger.info(f"⚠️ Eval dataset provided but eval_steps={eval_steps_val} (disabled)\n") + logger.info( + f"⚠️ Eval dataset provided but eval_steps={eval_steps_val} (disabled)\n" + ) logger.info("To enable evaluation, set eval_steps > 0.0\n") else: logger.info("No eval dataset β€” evaluation disabled\n") # Add model-specific parameters # Use optim and lr_scheduler_type from training_args if provided, otherwise use defaults - optim_value = training_args.get('optim', "adamw_8bit") - lr_scheduler_type_value = training_args.get('lr_scheduler_type', "linear") - + optim_value = training_args.get("optim", "adamw_8bit") + lr_scheduler_type_value = training_args.get("lr_scheduler_type", "linear") + if self.is_vlm or self.is_audio_vlm: # Vision / audio VLM config (both need skip_prepare_dataset + remove_unused_columns) label = "audio VLM" if self.is_audio_vlm else "vision" logger.info(f"Configuring {label} model training parameters\n") # Use provided values or defaults for vision models - optim_value = training_args.get('optim', "adamw_torch_fused") - lr_scheduler_type_value = training_args.get('lr_scheduler_type', "cosine") - config_args.update({ - "optim": optim_value, - "lr_scheduler_type": lr_scheduler_type_value, - "gradient_checkpointing": True, - "gradient_checkpointing_kwargs": {"use_reentrant": False}, - "max_grad_norm": 0.3, - "remove_unused_columns": False, - "dataset_text_field": "", - "dataset_kwargs": {"skip_prepare_dataset": True}, - "max_length": training_args.get('max_seq_length', 2048), - }) + optim_value = training_args.get("optim", "adamw_torch_fused") + lr_scheduler_type_value = training_args.get( + "lr_scheduler_type", "cosine" + ) + config_args.update( + { + "optim": optim_value, + "lr_scheduler_type": lr_scheduler_type_value, + "gradient_checkpointing": True, + "gradient_checkpointing_kwargs": {"use_reentrant": False}, + "max_grad_norm": 0.3, + "remove_unused_columns": False, + "dataset_text_field": "", + "dataset_kwargs": {"skip_prepare_dataset": True}, + "max_length": training_args.get("max_seq_length", 2048), + } + ) else: logger.info("Configuring text model training parameters\n") - config_args.update({ - "optim": optim_value, - "lr_scheduler_type": lr_scheduler_type_value, - "dataset_text_field": "text", - }) + config_args.update( + { + "optim": optim_value, + "lr_scheduler_type": lr_scheduler_type_value, + "dataset_text_field": "text", + } + ) # Only add packing for text models (not DeepSeek OCR which is VLM) if not is_deepseek_ocr: - packing_enabled = training_args.get('packing', False) + packing_enabled = training_args.get("packing", False) config_args["packing"] = packing_enabled - logger.info(f"Sequence packing: {'enabled' if packing_enabled else 'disabled'}\n") + logger.info( + f"Sequence packing: {'enabled' if packing_enabled else 'disabled'}\n" + ) # Audio codec overrides β€” BiCodec/DAC use the text SFTTrainer path - if self._audio_type == 'bicodec': + if self._audio_type == "bicodec": config_args["packing"] = False logger.info("Applied BiCodec overrides: packing=False\n") - elif self._audio_type == 'dac': + elif self._audio_type == "dac": config_args["packing"] = False logger.info("Applied DAC overrides: packing=False\n") @@ -2608,8 +3066,14 @@ class UnslothTrainer: if self.is_audio_vlm: # Audio VLM (e.g. Gemma 3N + audio): raw Dataset from _format_audio_vlm_dataset # Notebook uses processing_class=processor.tokenizer (text tokenizer only) - train_dataset = dataset if isinstance(dataset, Dataset) else dataset['dataset'] - processing_class = self.tokenizer.tokenizer if hasattr(self.tokenizer, 'tokenizer') else self.tokenizer + train_dataset = ( + dataset if isinstance(dataset, Dataset) else dataset["dataset"] + ) + processing_class = ( + self.tokenizer.tokenizer + if hasattr(self.tokenizer, "tokenizer") + else self.tokenizer + ) trainer_kwargs = { "model": self.model, "train_dataset": train_dataset, @@ -2622,7 +3086,9 @@ class UnslothTrainer: self.trainer = SFTTrainer(**trainer_kwargs) elif self.is_vlm: # Image VLM: dataset is dict wrapper from format_and_template_dataset - train_dataset = dataset['dataset'] if isinstance(dataset, dict) else dataset + train_dataset = ( + dataset["dataset"] if isinstance(dataset, dict) else dataset + ) trainer_kwargs = { "model": self.model, "train_dataset": train_dataset, @@ -2640,15 +3106,20 @@ class UnslothTrainer: # ProcessorMixin β†’ sets _is_vlm=True β†’ skips _prepare_dataset entirely, # and the 'text' column never gets tokenized to 'input_ids'. from transformers import ProcessorMixin + sft_tokenizer = self.tokenizer - if isinstance(self.tokenizer, ProcessorMixin) and hasattr(self.tokenizer, 'tokenizer'): - logger.info(f" ⚠️ Unwrapping Processor β†’ raw tokenizer for text-only SFTTrainer") + if isinstance(self.tokenizer, ProcessorMixin) and hasattr( + self.tokenizer, "tokenizer" + ): + logger.info( + f" ⚠️ Unwrapping Processor β†’ raw tokenizer for text-only SFTTrainer" + ) sft_tokenizer = self.tokenizer.tokenizer trainer_kwargs = { "model": self.model, "tokenizer": sft_tokenizer, - "train_dataset": dataset['dataset'], + "train_dataset": dataset["dataset"], "data_collator": data_collator, "args": SFTConfig(**config_args), } @@ -2665,11 +3136,18 @@ class UnslothTrainer: # Determine if we should train on responses only instruction_part = None response_part = None - train_on_responses_enabled = training_args.get('train_on_completions', False) + train_on_responses_enabled = training_args.get( + "train_on_completions", False + ) # DeepSeek OCR handles this internally in its collator, so skip # Audio VLM handles label masking in its collator, so skip - if train_on_responses_enabled and not self.is_audio_vlm and not self.is_audio and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'): + if ( + train_on_responses_enabled + and not self.is_audio_vlm + and not self.is_audio + and not (is_deepseek_ocr or dataset["final_format"].lower() == "alpaca") + ): try: logger.info("Configuring train on responses only...\n") @@ -2681,16 +3159,26 @@ class UnslothTrainer: logger.info(f"Detected template: {template_name}\n") if template_name in TEMPLATE_TO_RESPONSES_MAPPER: - instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["instruction"] - response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"] + instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[ + template_name + ]["instruction"] + response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name][ + "response" + ] - logger.info(f"Instruction marker: {instruction_part[:50]}...\n") + logger.info( + f"Instruction marker: {instruction_part[:50]}...\n" + ) logger.info(f"Response marker: {response_part[:50]}...\n") else: - logger.info(f"No response mapping found for template: {template_name}\n") + logger.info( + f"No response mapping found for template: {template_name}\n" + ) train_on_responses_enabled = False else: - logger.info(f"No template mapping found for model: {self.model_name}\n") + logger.info( + f"No template mapping found for model: {self.model_name}\n" + ) train_on_responses_enabled = False except Exception as e: @@ -2698,15 +3186,22 @@ class UnslothTrainer: train_on_responses_enabled = False # Apply train on responses only if we have valid parts - if train_on_responses_enabled and instruction_part and response_part and not self.is_audio_vlm and not self.is_audio and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'): + if ( + train_on_responses_enabled + and instruction_part + and response_part + and not self.is_audio_vlm + and not self.is_audio + and not (is_deepseek_ocr or dataset["final_format"].lower() == "alpaca") + ): try: from unsloth.chat_templates import train_on_responses_only self.trainer = train_on_responses_only( self.trainer, - instruction_part=instruction_part, - response_part=response_part, - num_proc=config_args["dataset_num_proc"], + instruction_part = instruction_part, + response_part = response_part, + num_proc = config_args["dataset_num_proc"], ) logger.info("Train on responses only configured successfully\n") @@ -2719,10 +3214,14 @@ class UnslothTrainer: filtered_len = len(self.trainer.train_dataset) original_len = len(dataset["dataset"]) dropped = original_len - filtered_len - drop_pct = round(100 * dropped / original_len, 1) if original_len > 0 else 0 + drop_pct = ( + round(100 * dropped / original_len, 1) + if original_len > 0 + else 0 + ) if filtered_len == 0 or drop_pct > 30: - max_seq = training_args.get('max_seq_length', 2048) + max_seq = training_args.get("max_seq_length", 2048) error_msg = ( f"{dropped}/{original_len} samples ({drop_pct}%) " f"were dropped after applying 'train on responses " @@ -2733,7 +3232,7 @@ class UnslothTrainer: f"or disabling 'Train on completions'." ) logger.error(error_msg) - self._update_progress(error=error_msg, is_training=False) + self._update_progress(error = error_msg, is_training = False) return if dropped > 0: @@ -2747,12 +3246,23 @@ class UnslothTrainer: # [DEBUG] Decode first sample AFTER train_on_completions applied try: _row = self.trainer.train_dataset[0] - _space = self.tokenizer(" ", add_special_tokens=False).input_ids[0] - print("[DEBUG] === After train_on_completions ===", flush=True) - print(f"[DEBUG] input_ids decoded:\n{self.tokenizer.decode(_row['input_ids'])}\n", flush=True) - print(f"[DEBUG] labels decoded (-100 β†’ space):\n{self.tokenizer.decode([_space if x == -100 else x for x in _row['labels']])}\n", flush=True) + _space = self.tokenizer( + " ", add_special_tokens = False + ).input_ids[0] + print("[DEBUG] === After train_on_completions ===", flush = True) + print( + f"[DEBUG] input_ids decoded:\n{self.tokenizer.decode(_row['input_ids'])}\n", + flush = True, + ) + print( + f"[DEBUG] labels decoded (-100 β†’ space):\n{self.tokenizer.decode([_space if x == -100 else x for x in _row['labels']])}\n", + flush = True, + ) except Exception as _dbg_e: - print(f"[DEBUG] Could not decode post-completions sample: {_dbg_e}", flush=True) + print( + f"[DEBUG] Could not decode post-completions sample: {_dbg_e}", + flush = True, + ) except Exception as e: logger.warning(f"Failed to apply train on responses only: {e}") @@ -2766,18 +3276,21 @@ class UnslothTrainer: # ========== PROGRESS TRACKING ========== self.trainer.add_callback(self._create_progress_callback()) - num_samples = len(dataset['dataset'] if isinstance(dataset, dict) else dataset) - batch_size = training_args.get('batch_size', 2) - total_steps = self._calculate_total_steps( - num_samples, batch_size, - training_args.get('gradient_accumulation_steps', 4), - training_args.get('num_epochs', 3), - training_args.get('max_steps', 0), + num_samples = len( + dataset["dataset"] if isinstance(dataset, dict) else dataset ) - self._update_progress(total_steps=total_steps) + batch_size = training_args.get("batch_size", 2) + total_steps = self._calculate_total_steps( + num_samples, + batch_size, + training_args.get("gradient_accumulation_steps", 4), + training_args.get("num_epochs", 3), + training_args.get("max_steps", 0), + ) + self._update_progress(total_steps = total_steps) # ========== START TRAINING ========== - self._update_progress(status_message="Starting training...") + self._update_progress(status_message = "Starting training...") logger.info("Starting training...\n") self.trainer.train() @@ -2786,9 +3299,10 @@ class UnslothTrainer: except Exception as e: import traceback + logger.error(f"Training error: {e}") logger.error(f"Full traceback:\n{traceback.format_exc()}") - self._update_progress(is_training=False, error=str(e)) + self._update_progress(is_training = False, error = str(e)) finally: self.is_training = False @@ -2815,10 +3329,12 @@ class UnslothTrainer: method = "lora" config["unsloth_training_method"] = method - logger.info(f"Patching adapter_config.json with unsloth_training_method='{method}'") + logger.info( + f"Patching adapter_config.json with unsloth_training_method='{method}'" + ) with open(config_path, "w") as f: - json.dump(config, f, indent=2) + json.dump(config, f, indent = 2) except Exception as e: logger.warning(f"Failed to patch adapter_config.json: {e}") @@ -2833,7 +3349,7 @@ class UnslothTrainer: if save else "Cancelling training..." ) - self._update_progress(status_message=stop_msg) + self._update_progress(status_message = stop_msg) # If trainer exists, try to stop it gracefully if self.trainer: @@ -2872,13 +3388,16 @@ def _ensure_deepseek_ocr_installed(): try: # Try importing to see if already available from deepseek_ocr.modeling_deepseekocr import format_messages + logger.info("DeepSeek OCR module already available") return True except ImportError: pass try: - logger.info("DeepSeek OCR module not found. Auto-installing from HuggingFace...") + logger.info( + "DeepSeek OCR module not found. Auto-installing from HuggingFace..." + ) logger.info("\n Downloading DeepSeek OCR module from HuggingFace...\n") from huggingface_hub import snapshot_download @@ -2893,9 +3412,7 @@ def _ensure_deepseek_ocr_installed(): local_dir = os.path.join(parent_dir, "deepseek_ocr") snapshot_download( - "unsloth/DeepSeek-OCR", - local_dir=local_dir, - local_dir_use_symlinks=False + "unsloth/DeepSeek-OCR", local_dir = local_dir, local_dir_use_symlinks = False ) # Add to sys.path if not already there @@ -2914,9 +3431,11 @@ def _ensure_deepseek_ocr_installed(): logger.info(f"\n❌ Failed to install DeepSeek OCR module: {e}\n") return False + # Global trainer instance _trainer_instance = None + def get_trainer() -> UnslothTrainer: """Get global trainer instance""" global _trainer_instance diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 649e48b96f..96a92af942 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -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 diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 0eedeb5135..b5cf984bf9 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -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(), + } + ) diff --git a/studio/backend/loggers/__init__.py b/studio/backend/loggers/__init__.py index 81ae4948b3..721dde4688 100644 --- a/studio/backend/loggers/__init__.py +++ b/studio/backend/loggers/__init__.py @@ -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"] diff --git a/studio/backend/loggers/config.py b/studio/backend/loggers/config.py index d6247a139a..0d32a64657 100644 --- a/studio/backend/loggers/config.py +++ b/studio/backend/loggers/config.py @@ -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) diff --git a/studio/backend/loggers/handlers.py b/studio/backend/loggers/handlers.py index 08f692c256..3add92ea1e 100644 --- a/studio/backend/loggers/handlers.py +++ b/studio/backend/loggers/handlers.py @@ -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 diff --git a/studio/backend/main.py b/studio/backend/main.py index f0ba4be4f4..46615bb5f0 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -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 - diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py index 13ab08c50b..4a53418439 100644 --- a/studio/backend/models/__init__.py +++ b/studio/backend/models/__init__.py @@ -4,6 +4,7 @@ """ Pydantic models for API request/response schemas """ + from .training import ( TrainingStartRequest, TrainingJobResponse, diff --git a/studio/backend/models/auth.py b/studio/backend/models/auth.py index 47428a791d..c12d15617e 100644 --- a/studio/backend/models/auth.py +++ b/studio/backend/models/auth.py @@ -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") diff --git a/studio/backend/models/data_recipe.py b/studio/backend/models/data_recipe.py index d7758cec36..6992572b00 100644 --- a/studio/backend/models/data_recipe.py +++ b/studio/backend/models/data_recipe.py @@ -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) diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index 06d555eabd..f20d6f2d15 100644 --- a/studio/backend/models/datasets.py +++ b/studio/backend/models/datasets.py @@ -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) diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 9625123721..a86596f199 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -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 - - diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index e9e581d1bd..46925fb3db 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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) diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 0f87d4543c..6634f69385 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -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", ) diff --git a/studio/backend/models/responses.py b/studio/backend/models/responses.py index 9e426d2158..3081f67422 100644 --- a/studio/backend/models/responses.py +++ b/studio/backend/models/responses.py @@ -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" + ) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 8ea9eeec42..35bd9744a9 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -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" + ) diff --git a/studio/backend/models/users.py b/studio/backend/models/users.py index 7d3fbb4eba..8e982ed9f9 100644 --- a/studio/backend/models/users.py +++ b/studio/backend/models/users.py @@ -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'") diff --git a/__init__.py b/studio/backend/plugins/__init__.py similarity index 100% rename from __init__.py rename to studio/backend/plugins/__init__.py diff --git a/studio/backend/plugins/data-designer-unstructured-seed/__init__.py b/studio/backend/plugins/data-designer-unstructured-seed/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/__init__.py @@ -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 diff --git a/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml b/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml index 5b83e5468a..f826ffd992 100644 --- a/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml +++ b/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml @@ -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" diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py index 9e81a967b7..80f51b2a24 100644 --- a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py @@ -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) diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py index 1d6defa4df..e0a0392a69 100644 --- a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py @@ -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) diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py index acde1909c2..8a3deb9b92 100644 --- a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py @@ -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) diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py index de8c00d8b7..6f0d7ffd49 100644 --- a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py @@ -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, ) diff --git a/studio/backend/requirements/__init__.py b/studio/backend/requirements/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/requirements/__init__.py @@ -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 diff --git a/studio/backend/requirements/single-env/patch_metadata.py b/studio/backend/requirements/single-env/patch_metadata.py index b5d4bfcaad..7bcaa56ac8 100644 --- a/studio/backend/requirements/single-env/patch_metadata.py +++ b/studio/backend/requirements/single-env/patch_metadata.py @@ -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 diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index f7508cb6ff..83cc41e83b 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -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", ) - diff --git a/studio/backend/routes/data_recipe/__init__.py b/studio/backend/routes/data_recipe/__init__.py index 6ec530751e..c596f189d4 100644 --- a/studio/backend/routes/data_recipe/__init__.py +++ b/studio/backend/routes/data_recipe/__init__.py @@ -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) diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py index 02b8befab6..4661615338 100644 --- a/studio/backend/routes/data_recipe/jobs.py +++ b/studio/backend/routes/data_recipe/jobs.py @@ -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") diff --git a/studio/backend/routes/data_recipe/mcp.py b/studio/backend/routes/data_recipe/mcp.py index a106810a15..1f5c0f34e0 100644 --- a/studio/backend/routes/data_recipe/mcp.py +++ b/studio/backend/routes/data_recipe/mcp.py @@ -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, ) diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index cb8130e698..76e05be954 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -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, ) diff --git a/studio/backend/routes/data_recipe/validate.py b/studio/backend/routes/data_recipe/validate.py index f2858f97bd..a793a3b172 100644 --- a/studio/backend/routes/data_recipe/validate.py +++ b/studio/backend/routes/data_recipe/validate.py @@ -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) diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 937de73512..415c30cd72 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -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)}") diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 4daca71143..3e60eaaf20 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -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)}", ) - - diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2f6626d7fb..7de7336946 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4,6 +4,7 @@ """ Inference API routes for model loading and text generation. """ + import sys import time import uuid @@ -18,7 +19,6 @@ import asyncio import threading - # Add backend directory to path backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: @@ -69,15 +69,15 @@ router = APIRouter() logger = get_logger(__name__) - # GGUF inference backend (llama-server) _llama_cpp_backend = LlamaCppBackend() + def get_llama_cpp_backend() -> LlamaCppBackend: return _llama_cpp_backend -@router.post("/load", response_model=LoadResponse) +@router.post("/load", response_model = LoadResponse) async def load_model( request: LoadRequest, current_subject: str = Depends(get_current_subject), @@ -98,15 +98,15 @@ async def load_model( # Create config using clean factory method # is_lora is auto-detected from adapter_config.json on disk/HF config = ModelConfig.from_identifier( - model_id=request.model_path, - hf_token=request.hf_token, - gguf_variant=request.gguf_variant, + model_id = request.model_path, + hf_token = request.hf_token, + gguf_variant = request.gguf_variant, ) if not config: raise HTTPException( - status_code=400, - detail=f"Invalid model identifier: {request.model_path}" + status_code = 400, + detail = f"Invalid model identifier: {request.model_path}", ) # ── GGUF path: load via llama-server ────────────────────── @@ -116,34 +116,36 @@ async def load_model( # Unload any active Unsloth model first to free VRAM if unsloth_backend.active_model_name: - logger.info(f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF") + logger.info( + f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF" + ) unsloth_backend.unload_model(unsloth_backend.active_model_name) # Route to HF mode or local mode based on config if config.gguf_hf_repo: # HF mode: llama-server downloads via -hf "repo:quant" success = llama_backend.load_model( - hf_repo=config.gguf_hf_repo, - hf_variant=config.gguf_variant, - hf_token=request.hf_token, - model_identifier=config.identifier, - is_vision=config.is_vision, - n_ctx=request.max_seq_length, + hf_repo = config.gguf_hf_repo, + hf_variant = config.gguf_variant, + hf_token = request.hf_token, + model_identifier = config.identifier, + is_vision = config.is_vision, + n_ctx = request.max_seq_length, ) else: # Local mode: llama-server loads via -m success = llama_backend.load_model( - gguf_path=config.gguf_file, - mmproj_path=config.gguf_mmproj_file, - model_identifier=config.identifier, - is_vision=config.is_vision, - n_ctx=request.max_seq_length, + gguf_path = config.gguf_file, + mmproj_path = config.gguf_mmproj_file, + model_identifier = config.identifier, + is_vision = config.is_vision, + n_ctx = request.max_seq_length, ) if not success: raise HTTPException( - status_code=500, - detail=f"Failed to load GGUF model: {config.display_name}" + status_code = 500, + detail = f"Failed to load GGUF model: {config.display_name}", ) logger.info(f"Loaded GGUF model via llama-server: {config.identifier}") @@ -151,13 +153,13 @@ async def load_model( inference_config = load_inference_config(config.identifier) return LoadResponse( - status="loaded", - model=config.identifier, - display_name=config.display_name, - is_vision=config.is_vision, - is_lora=False, - is_gguf=True, - inference=inference_config, + status = "loaded", + model = config.identifier, + display_name = config.display_name, + is_vision = config.is_vision, + is_lora = False, + is_gguf = True, + inference = inference_config, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -172,9 +174,12 @@ async def load_model( # Shut down any export subprocess to free VRAM 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 inference") + logger.info( + "Shutting down export subprocess to free GPU memory for inference" + ) exp_backend._shutdown_subprocess() exp_backend.current_checkpoint = None exp_backend.is_vision = False @@ -189,6 +194,7 @@ async def load_model( if config.is_lora and config.path: import json from pathlib import Path + adapter_cfg_path = Path(config.path) / "adapter_config.json" if adapter_cfg_path.exists(): try: @@ -208,10 +214,16 @@ async def load_model( ) load_in_4bit = True elif training_method: - logger.info(f"Training method: {training_method}, load_in_4bit={load_in_4bit}") + logger.info( + f"Training method: {training_method}, load_in_4bit={load_in_4bit}" + ) else: # No unsloth_training_method β€” fallback to base model name - if config.base_model and "-bnb-4bit" not in config.base_model.lower() and load_in_4bit: + if ( + config.base_model + and "-bnb-4bit" not in config.base_model.lower() + and load_in_4bit + ): logger.info( f"No unsloth_training_method in adapter_config.json. " f"Base model '{config.base_model}' has no -bnb-4bit suffix β€” " @@ -223,29 +235,30 @@ async def load_model( # Load the model success = backend.load_model( - config=config, - max_seq_length=request.max_seq_length, - load_in_4bit=load_in_4bit, - hf_token=request.hf_token, - trust_remote_code=request.trust_remote_code, + config = config, + max_seq_length = request.max_seq_length, + load_in_4bit = load_in_4bit, + hf_token = request.hf_token, + trust_remote_code = request.trust_remote_code, ) if not success: # Check if YAML says this model needs trust_remote_code if not request.trust_remote_code: model_defaults = load_model_defaults(config.identifier) - yaml_trust = model_defaults.get("inference", {}).get("trust_remote_code", False) + yaml_trust = model_defaults.get("inference", {}).get( + "trust_remote_code", False + ) if yaml_trust: raise HTTPException( - status_code=400, - detail=( + status_code = 400, + detail = ( f"Model '{config.display_name}' requires trust_remote_code to be enabled. " f"Please enable 'Trust remote code' in Chat Settings and try again." ), ) raise HTTPException( - status_code=500, - detail=f"Failed to load model: {config.display_name}" + status_code = 500, detail = f"Failed to load model: {config.display_name}" ) logger.info(f"Loaded model: {config.identifier}") @@ -254,29 +267,26 @@ async def load_model( inference_config = load_inference_config(config.identifier) return LoadResponse( - status="loaded", - model=config.identifier, - display_name=config.display_name, - is_vision=config.is_vision, - is_lora=config.is_lora, - is_gguf=False, - is_audio=config.is_audio, - audio_type=config.audio_type, - has_audio_input=config.has_audio_input, - inference=inference_config, + status = "loaded", + model = config.identifier, + display_name = config.display_name, + is_vision = config.is_vision, + is_lora = config.is_lora, + is_gguf = False, + is_audio = config.is_audio, + audio_type = config.audio_type, + has_audio_input = config.has_audio_input, + inference = inference_config, ) except HTTPException: raise except Exception as e: - logger.error(f"Error loading model: {e}", exc_info=True) - raise HTTPException( - status_code=500, - detail=f"Failed to load model: {str(e)}" - ) + logger.error(f"Error loading model: {e}", exc_info = True) + raise HTTPException(status_code = 500, detail = f"Failed to load model: {str(e)}") -@router.post("/validate", response_model=ValidateModelResponse) +@router.post("/validate", response_model = ValidateModelResponse) async def validate_model( request: ValidateModelRequest, current_subject: str = Depends(get_current_subject), @@ -289,38 +299,41 @@ async def validate_model( """ try: config = ModelConfig.from_identifier( - model_id=request.model_path, - hf_token=request.hf_token, - gguf_variant=request.gguf_variant, + model_id = request.model_path, + hf_token = request.hf_token, + gguf_variant = request.gguf_variant, ) if not config: raise HTTPException( - status_code=400, - detail=f"Invalid model identifier: {request.model_path}", + status_code = 400, + detail = f"Invalid model identifier: {request.model_path}", ) return ValidateModelResponse( - valid=True, - message="Model identifier is valid.", - identifier=config.identifier, - display_name=getattr(config, "display_name", config.identifier), - is_gguf=getattr(config, "is_gguf", False), - is_lora=getattr(config, "is_lora", False), - is_vision=getattr(config, "is_vision", False), + valid = True, + message = "Model identifier is valid.", + identifier = config.identifier, + display_name = getattr(config, "display_name", config.identifier), + is_gguf = getattr(config, "is_gguf", False), + is_lora = getattr(config, "is_lora", False), + is_vision = getattr(config, "is_vision", False), ) except HTTPException: raise except Exception as e: - logger.error(f"Error validating model identifier '{request.model_path}': {e}", exc_info=True) + logger.error( + f"Error validating model identifier '{request.model_path}': {e}", + exc_info = True, + ) raise HTTPException( - status_code=400, - detail=f"Invalid model: {str(e)}", + status_code = 400, + detail = f"Invalid model: {str(e)}", ) -@router.post("/unload", response_model=UnloadResponse) +@router.post("/unload", response_model = UnloadResponse) async def unload_model( request: UnloadRequest, current_subject: str = Depends(get_current_subject), @@ -332,23 +345,23 @@ async def unload_model( try: # Check if the GGUF backend has this model loaded llama_backend = get_llama_cpp_backend() - if llama_backend.is_loaded and llama_backend.model_identifier == request.model_path: + if ( + llama_backend.is_loaded + and llama_backend.model_identifier == request.model_path + ): llama_backend.unload_model() logger.info(f"Unloaded GGUF model: {request.model_path}") - return UnloadResponse(status="unloaded", model=request.model_path) + return UnloadResponse(status = "unloaded", model = request.model_path) # Otherwise, unload from Unsloth backend backend = get_inference_backend() backend.unload_model(request.model_path) logger.info(f"Unloaded model: {request.model_path}") - return UnloadResponse(status="unloaded", model=request.model_path) + return UnloadResponse(status = "unloaded", model = request.model_path) except Exception as e: - logger.error(f"Error unloading model: {e}", exc_info=True) - raise HTTPException( - status_code=500, - detail=f"Failed to unload model: {str(e)}" - ) + logger.error(f"Error unloading model: {e}", exc_info = True) + raise HTTPException(status_code = 500, detail = f"Failed to unload model: {str(e)}") @router.post("/generate/stream") @@ -358,17 +371,16 @@ async def generate_stream( ): """ Generate a chat response with Server-Sent Events (SSE) streaming. - + For vision models, provide image_base64 with the base64-encoded image. """ backend = get_inference_backend() - + if not backend.active_model_name: raise HTTPException( - status_code=400, - detail="No model loaded. Call POST /inference/load first." + status_code = 400, detail = "No model loaded. Call POST /inference/load first." ) - + # Decode image if provided (for vision models) image = None if request.image_base64: @@ -376,58 +388,57 @@ async def generate_stream( import base64 from PIL import Image from io import BytesIO - + # Check if current model supports vision model_info = backend.models.get(backend.active_model_name, {}) if not model_info.get("is_vision"): raise HTTPException( - status_code=400, - detail="Image provided but current model is text-only. Load a vision model." + status_code = 400, + detail = "Image provided but current model is text-only. Load a vision model.", ) - + image_data = base64.b64decode(request.image_base64) image = Image.open(BytesIO(image_data)) image = backend.resize_image(image) - + except HTTPException: raise except Exception as e: raise HTTPException( - status_code=400, - detail=f"Failed to decode image: {str(e)}" + status_code = 400, detail = f"Failed to decode image: {str(e)}" ) - + async def stream(): try: for chunk in backend.generate_chat_response( - messages=request.messages, - system_prompt=request.system_prompt, - image=image, - temperature=request.temperature, - top_p=request.top_p, - top_k=request.top_k, - max_new_tokens=request.max_new_tokens, - repetition_penalty=request.repetition_penalty, + messages = request.messages, + system_prompt = request.system_prompt, + image = image, + temperature = request.temperature, + top_p = request.top_p, + top_k = request.top_k, + max_new_tokens = request.max_new_tokens, + repetition_penalty = request.repetition_penalty, ): yield f"data: {json.dumps({'content': chunk})}\n\n" yield "data: [DONE]\n\n" - + except Exception as e: backend.reset_generation_state() - logger.error(f"Error during generation: {e}", exc_info=True) - yield f"data: {json.dumps({'error': str(e)})}\n\n" - + logger.error(f"Error during generation: {e}", exc_info = True) + yield f"data: {json.dumps({'error': 'An internal error occurred'})}\n\n" + return StreamingResponse( stream(), - media_type="text/event-stream", - headers={ + media_type = "text/event-stream", + headers = { "Cache-Control": "no-cache", "Connection": "keep-alive", - } + }, ) -@router.get("/status", response_model=InferenceStatusResponse) +@router.get("/status", response_model = InferenceStatusResponse) async def get_status( current_subject: str = Depends(get_current_subject), ): @@ -441,12 +452,12 @@ async def get_status( # If a GGUF model is loaded via llama-server, report that if llama_backend.is_loaded: return InferenceStatusResponse( - active_model=llama_backend.model_identifier, - is_vision=llama_backend.is_vision, - is_gguf=True, - gguf_variant=llama_backend.hf_variant, - loading=[], - loaded=[llama_backend.model_identifier], + active_model = llama_backend.model_identifier, + is_vision = llama_backend.is_vision, + is_gguf = True, + gguf_variant = llama_backend.hf_variant, + loading = [], + loaded = [llama_backend.model_identifier], ) # Otherwise, report Unsloth backend status @@ -464,22 +475,19 @@ async def get_status( has_audio_input = model_info.get("has_audio_input", False) return InferenceStatusResponse( - active_model=backend.active_model_name, - is_vision=is_vision, - is_gguf=False, - is_audio=is_audio, - audio_type=audio_type, - has_audio_input=has_audio_input, - loading=list(getattr(backend, 'loading_models', set())), - loaded=list(backend.models.keys()), + active_model = backend.active_model_name, + is_vision = is_vision, + is_gguf = False, + is_audio = is_audio, + audio_type = audio_type, + has_audio_input = has_audio_input, + loading = list(getattr(backend, "loading_models", set())), + loaded = list(backend.models.keys()), ) except Exception as e: - logger.error(f"Error getting status: {e}", exc_info=True) - raise HTTPException( - status_code=500, - detail=f"Failed to get status: {str(e)}" - ) + logger.error(f"Error getting status: {e}", exc_info = True) + raise HTTPException(status_code = 500, detail = f"Failed to get status: {str(e)}") # ===================================================================== @@ -488,7 +496,11 @@ async def get_status( @router.post("/audio/generate") -async def generate_audio(payload: ChatCompletionRequest, request: Request, current_subject: str = Depends(get_current_subject)): +async def generate_audio( + payload: ChatCompletionRequest, + request: Request, + current_subject: str = Depends(get_current_subject), +): """ Generate audio (TTS) from the latest user message. Returns a JSON response with base64-encoded WAV audio. @@ -498,22 +510,24 @@ async def generate_audio(payload: ChatCompletionRequest, request: Request, curre backend = get_inference_backend() if not backend.active_model_name: - raise HTTPException(status_code=400, detail="No model loaded.") + raise HTTPException(status_code = 400, detail = "No model loaded.") model_info = backend.models.get(backend.active_model_name, {}) if not model_info.get("is_audio"): - raise HTTPException(status_code=400, detail="Active model is not an audio model.") + raise HTTPException( + status_code = 400, detail = "Active model is not an audio model." + ) # Extract text from the last user message _, chat_messages, _ = _extract_content_parts(payload.messages) if not chat_messages: - raise HTTPException(status_code=400, detail="No messages provided.") + raise HTTPException(status_code = 400, detail = "No messages provided.") last_user_msg = next( (m for m in reversed(chat_messages) if m["role"] == "user"), None ) if not last_user_msg: - raise HTTPException(status_code=400, detail="No user message found.") + raise HTTPException(status_code = 400, detail = "No user message found.") text = last_user_msg["content"] @@ -521,42 +535,46 @@ async def generate_audio(payload: ChatCompletionRequest, request: Request, curre wav_bytes, sample_rate = await asyncio.get_event_loop().run_in_executor( None, lambda: backend.generate_audio_response( - text=text, - temperature=payload.temperature, - top_p=payload.top_p, - top_k=payload.top_k, - min_p=payload.min_p, - max_new_tokens=payload.max_tokens or 2048, - repetition_penalty=payload.repetition_penalty, - use_adapter=payload.use_adapter, + text = text, + temperature = payload.temperature, + top_p = payload.top_p, + top_k = payload.top_k, + min_p = payload.min_p, + max_new_tokens = payload.max_tokens or 2048, + repetition_penalty = payload.repetition_penalty, + use_adapter = payload.use_adapter, ), ) audio_b64 = base64.b64encode(wav_bytes).decode("ascii") completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" - return JSONResponse(content={ - "id": completion_id, - "object": "chat.completion.audio", - "model": backend.active_model_name, - "audio": { - "data": audio_b64, - "format": "wav", - "sample_rate": sample_rate, - }, - "choices": [{ - "index": 0, - "message": { - "role": "assistant", - "content": f"[Generated audio from: \"{text[:100]}\"]", + return JSONResponse( + content = { + "id": completion_id, + "object": "chat.completion.audio", + "model": backend.active_model_name, + "audio": { + "data": audio_b64, + "format": "wav", + "sample_rate": sample_rate, }, - "finish_reason": "stop", - }], - }) + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": f'[Generated audio from: "{text[:100]}"]', + }, + "finish_reason": "stop", + } + ], + } + ) except Exception as e: - logger.error(f"Audio generation error: {e}", exc_info=True) - raise HTTPException(status_code=500, detail=str(e)) + logger.error(f"Audio generation error: {e}", exc_info = True) + raise HTTPException(status_code = 500, detail = str(e)) # ===================================================================== @@ -576,9 +594,9 @@ def _decode_audio_base64(b64: str) -> np.ndarray: # torchaudio.load needs a file path or file-like object with format hint # Write to a temp file so torchaudio can auto-detect the format with tempfile.NamedTemporaryFile( - suffix=".audio", - delete=False, - dir=str(ensure_dir(tmp_root())), + suffix = ".audio", + delete = False, + dir = str(ensure_dir(tmp_root())), ) as tmp: tmp.write(raw) tmp_path = tmp.name @@ -589,11 +607,11 @@ def _decode_audio_base64(b64: str) -> np.ndarray: # Convert to mono if stereo if waveform.shape[0] > 1: - waveform = waveform.mean(dim=0, keepdim=True) + waveform = waveform.mean(dim = 0, keepdim = True) # Resample to 16kHz if needed if sr != 16000: - resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000) + resampler = torchaudio.transforms.Resample(orig_freq = sr, new_freq = 16000) waveform = resampler(waveform) return waveform.squeeze(0).numpy() @@ -683,8 +701,8 @@ async def openai_chat_completions( backend = get_inference_backend() if not backend.active_model_name: raise HTTPException( - status_code=400, - detail="No model loaded. Call POST /inference/load first.", + status_code = 400, + detail = "No model loaded. Call POST /inference/load first.", ) model_name = backend.active_model_name or payload.model @@ -697,8 +715,8 @@ async def openai_chat_completions( # ── Whisper without audio: return clear error ── if model_info.get("audio_type") == "whisper" and not payload.audio_base64: raise HTTPException( - status_code=400, - detail="Whisper models require audio input. Please upload an audio file.", + status_code = 400, + detail = "Whisper models require audio input. Please upload an audio file.", ) # ── Audio INPUT path: decode WAV and route to audio input generation ── @@ -712,30 +730,38 @@ async def openai_chat_completions( def audio_input_generate(): if model_info.get("audio_type") == "whisper": return backend.generate_whisper_response( - audio_array=audio_array, - cancel_event=cancel_event, + audio_array = audio_array, + cancel_event = cancel_event, ) return backend.generate_audio_input_response( - messages=chat_messages, - system_prompt=system_prompt, - audio_array=audio_array, - temperature=payload.temperature, - top_p=payload.top_p, - top_k=payload.top_k, - min_p=payload.min_p, - max_new_tokens=payload.max_tokens or 2048, - repetition_penalty=payload.repetition_penalty, - cancel_event=cancel_event, + messages = chat_messages, + system_prompt = system_prompt, + audio_array = audio_array, + temperature = payload.temperature, + top_p = payload.top_p, + top_k = payload.top_k, + min_p = payload.min_p, + max_new_tokens = payload.max_tokens or 2048, + repetition_penalty = payload.repetition_penalty, + cancel_event = cancel_event, ) if payload.stream: + async def audio_input_stream(): try: first_chunk = ChatCompletionChunk( - id=completion_id, created=created, model=model_name, - choices=[ChunkChoice(delta=ChoiceDelta(role="assistant"), finish_reason=None)], + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(role = "assistant"), + finish_reason = None, + ) + ], ) - yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n" + yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" for chunk_text in audio_input_generate(): if await request.is_disconnected(): @@ -743,36 +769,60 @@ async def openai_chat_completions( return if chunk_text: chunk = ChatCompletionChunk( - id=completion_id, created=created, model=model_name, - choices=[ChunkChoice(delta=ChoiceDelta(content=chunk_text), finish_reason=None)], + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(content = chunk_text), + finish_reason = None, + ) + ], ) - yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" + yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" final_chunk = ChatCompletionChunk( - id=completion_id, created=created, model=model_name, - choices=[ChunkChoice(delta=ChoiceDelta(), finish_reason="stop")], + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice(delta = ChoiceDelta(), finish_reason = "stop") + ], ) - yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n" + yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() raise except Exception as e: - logger.error(f"Error during audio input streaming: {e}", exc_info=True) - yield f"data: {json.dumps({'error': {'message': str(e), 'type': 'server_error'}})}\n\n" + logger.error( + f"Error during audio input streaming: {e}", exc_info = True + ) + yield f"data: {json.dumps({'error': {'message': 'An internal error occurred', 'type': 'server_error'}})}\n\n" return StreamingResponse( audio_input_stream(), - media_type="text/event-stream", - headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"}, + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, ) else: full_text = "".join(audio_input_generate()) response = ChatCompletion( - id=completion_id, created=created, model=model_name, - choices=[CompletionChoice(message=CompletionMessage(content=full_text), finish_reason="stop")], + id = completion_id, + created = created, + model = model_name, + choices = [ + CompletionChoice( + message = CompletionMessage(content = full_text), + finish_reason = "stop", + ) + ], ) - return JSONResponse(content=response.model_dump()) + return JSONResponse(content = response.model_dump()) # ── Parse messages (handles multimodal content parts) ───── system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts( @@ -781,8 +831,8 @@ async def openai_chat_completions( if not chat_messages: raise HTTPException( - status_code=400, - detail="At least one non-system message is required.", + status_code = 400, + detail = "At least one non-system message is required.", ) # ── GGUF path: proxy to llama-server /v1/chat/completions ── @@ -791,8 +841,8 @@ async def openai_chat_completions( image_b64 = extracted_image_b64 or payload.image_base64 if image_b64 and not llama_backend.is_vision: raise HTTPException( - status_code=400, - detail="Image provided but current GGUF model does not support vision.", + status_code = 400, + detail = "Image provided but current GGUF model does not support vision.", ) # Build message list with system prompt prepended @@ -808,31 +858,34 @@ async def openai_chat_completions( def gguf_generate(): return llama_backend.generate_chat_completion( - messages=gguf_messages, - image_b64=image_b64, - temperature=payload.temperature, - top_p=payload.top_p, - top_k=payload.top_k, - min_p=payload.min_p, - max_tokens=payload.max_tokens or 2048, - repetition_penalty=payload.repetition_penalty, - cancel_event=cancel_event, + messages = gguf_messages, + image_b64 = image_b64, + temperature = payload.temperature, + top_p = payload.top_p, + top_k = payload.top_k, + min_p = payload.min_p, + max_tokens = payload.max_tokens or 2048, + repetition_penalty = payload.repetition_penalty, + cancel_event = cancel_event, ) if payload.stream: + async def gguf_stream_chunks(): try: # First chunk: role first_chunk = ChatCompletionChunk( - id=completion_id, - created=created, - model=model_name, - choices=[ChunkChoice( - delta=ChoiceDelta(role="assistant"), - finish_reason=None, - )], + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(role = "assistant"), + finish_reason = None, + ) + ], ) - yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n" + yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" # Content chunks β€” llama backend yields cumulative text prev_text = "" @@ -840,48 +893,55 @@ async def openai_chat_completions( if await request.is_disconnected(): cancel_event.set() return - new_text = cumulative[len(prev_text):] + new_text = cumulative[len(prev_text) :] prev_text = cumulative if not new_text: continue chunk = ChatCompletionChunk( - id=completion_id, - created=created, - model=model_name, - choices=[ChunkChoice( - delta=ChoiceDelta(content=new_text), - finish_reason=None, - )], + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(content = new_text), + finish_reason = None, + ) + ], ) - yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" + yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" # Final chunk final_chunk = ChatCompletionChunk( - id=completion_id, - created=created, - model=model_name, - choices=[ChunkChoice( - delta=ChoiceDelta(), - finish_reason="stop", - )], + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(), + finish_reason = "stop", + ) + ], ) - yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n" + yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() raise except Exception as e: - logger.error(f"Error during GGUF streaming: {e}", exc_info=True) + logger.error(f"Error during GGUF streaming: {e}", exc_info = True) error_chunk = { - "error": {"message": str(e), "type": "server_error"}, + "error": { + "message": "An internal error occurred", + "type": "server_error", + }, } yield f"data: {json.dumps(error_chunk)}\n\n" return StreamingResponse( gguf_stream_chunks(), - media_type="text/event-stream", - headers={ + media_type = "text/event-stream", + headers = { "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", @@ -894,19 +954,21 @@ async def openai_chat_completions( full_text = token response = ChatCompletion( - id=completion_id, - created=created, - model=model_name, - choices=[CompletionChoice( - message=CompletionMessage(content=full_text), - finish_reason="stop", - )], + id = completion_id, + created = created, + model = model_name, + choices = [ + CompletionChoice( + message = CompletionMessage(content = full_text), + finish_reason = "stop", + ) + ], ) - return JSONResponse(content=response.model_dump()) + return JSONResponse(content = response.model_dump()) except Exception as e: - logger.error(f"Error during GGUF completion: {e}", exc_info=True) - raise HTTPException(status_code=500, detail=str(e)) + logger.error(f"Error during GGUF completion: {e}", exc_info = True) + raise HTTPException(status_code = 500, detail = str(e)) # ── Standard Unsloth path ───────────────────────────────── @@ -923,8 +985,8 @@ async def openai_chat_completions( model_info = backend.models.get(backend.active_model_name, {}) if not model_info.get("is_vision"): raise HTTPException( - status_code=400, - detail="Image provided but current model is text-only. Load a vision model.", + status_code = 400, + detail = "Image provided but current model is text-only. Load a vision model.", ) image_data = base64.b64decode(image_b64) @@ -934,52 +996,59 @@ async def openai_chat_completions( except HTTPException: raise except Exception as e: - raise HTTPException(status_code=400, detail=f"Failed to decode image: {e}") + raise HTTPException(status_code = 400, detail = f"Failed to decode image: {e}") # Shared generation kwargs gen_kwargs = dict( - messages=chat_messages, - system_prompt=system_prompt, - image=image, - temperature=payload.temperature, - top_p=payload.top_p, - top_k=payload.top_k, - min_p=payload.min_p, - max_new_tokens=payload.max_tokens or 2048, - repetition_penalty=payload.repetition_penalty, + messages = chat_messages, + system_prompt = system_prompt, + image = image, + temperature = payload.temperature, + top_p = payload.top_p, + top_k = payload.top_k, + min_p = payload.min_p, + max_new_tokens = payload.max_tokens or 2048, + repetition_penalty = payload.repetition_penalty, ) # Choose generation path (adapter-controlled or standard) cancel_event = threading.Event() if payload.use_adapter is not None: + def generate(): return backend.generate_with_adapter_control( - use_adapter=payload.use_adapter, - cancel_event=cancel_event, + use_adapter = payload.use_adapter, + cancel_event = cancel_event, **gen_kwargs, ) else: + def generate(): - return backend.generate_chat_response(cancel_event=cancel_event, **gen_kwargs) + return backend.generate_chat_response( + cancel_event = cancel_event, **gen_kwargs + ) completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) # ── Streaming response ──────────────────────────────────────── if payload.stream: + async def stream_chunks(): try: first_chunk = ChatCompletionChunk( - id=completion_id, - created=created, - model=model_name, - choices=[ChunkChoice( - delta=ChoiceDelta(role="assistant"), - finish_reason=None, - )], + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(role = "assistant"), + finish_reason = None, + ) + ], ) - yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n" + yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" prev_text = "" for cumulative in generate(): @@ -987,31 +1056,35 @@ async def openai_chat_completions( cancel_event.set() backend.reset_generation_state() return - new_text = cumulative[len(prev_text):] + new_text = cumulative[len(prev_text) :] prev_text = cumulative if not new_text: continue chunk = ChatCompletionChunk( - id=completion_id, - created=created, - model=model_name, - choices=[ChunkChoice( - delta=ChoiceDelta(content=new_text), - finish_reason=None, - )], + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(content = new_text), + finish_reason = None, + ) + ], ) - yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" + yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" final_chunk = ChatCompletionChunk( - id=completion_id, - created=created, - model=model_name, - choices=[ChunkChoice( - delta=ChoiceDelta(), - finish_reason="stop", - )], + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(), + finish_reason = "stop", + ) + ], ) - yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n" + yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -1020,16 +1093,19 @@ async def openai_chat_completions( raise except Exception as e: backend.reset_generation_state() - logger.error(f"Error during OpenAI streaming: {e}", exc_info=True) + logger.error(f"Error during OpenAI streaming: {e}", exc_info = True) error_chunk = { - "error": {"message": str(e), "type": "server_error"}, + "error": { + "message": "An internal error occurred", + "type": "server_error", + }, } yield f"data: {json.dumps(error_chunk)}\n\n" return StreamingResponse( stream_chunks(), - media_type="text/event-stream", - headers={ + media_type = "text/event-stream", + headers = { "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", @@ -1044,26 +1120,29 @@ async def openai_chat_completions( full_text = token response = ChatCompletion( - id=completion_id, - created=created, - model=model_name, - choices=[CompletionChoice( - message=CompletionMessage(content=full_text), - finish_reason="stop", - )], + id = completion_id, + created = created, + model = model_name, + choices = [ + CompletionChoice( + message = CompletionMessage(content = full_text), + finish_reason = "stop", + ) + ], ) - return JSONResponse(content=response.model_dump()) + return JSONResponse(content = response.model_dump()) except Exception as e: backend.reset_generation_state() - logger.error(f"Error during OpenAI completion: {e}", exc_info=True) - raise HTTPException(status_code=500, detail=str(e)) + logger.error(f"Error during OpenAI completion: {e}", exc_info = True) + raise HTTPException(status_code = 500, detail = str(e)) # ===================================================================== # OpenAI-Compatible Models Listing (/models β†’ /v1/models) # ===================================================================== + @router.get("/models") async def openai_list_models( current_subject: str = Depends(get_current_subject), @@ -1079,19 +1158,23 @@ async def openai_list_models( # Check GGUF backend llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded: - models.append({ - "id": llama_backend.model_identifier, - "object": "model", - "owned_by": "local", - }) + models.append( + { + "id": llama_backend.model_identifier, + "object": "model", + "owned_by": "local", + } + ) # Check Unsloth backend backend = get_inference_backend() if backend.active_model_name: - models.append({ - "id": backend.active_model_name, - "object": "model", - "owned_by": "local", - }) + models.append( + { + "id": backend.active_model_name, + "object": "model", + "owned_by": "local", + } + ) return {"object": "list", "data": models} diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index b6b36de346..99b42b491e 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -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)}", ) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index c4010835ac..43b5212010 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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", - } + }, ) diff --git a/studio/backend/run.py b/studio/backend/run.py index 32aae9831a..dde1335eae 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -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) - diff --git a/studio/backend/state/__init__.py b/studio/backend/state/__init__.py index e69de29bb2..32014236c6 100644 --- a/studio/backend/state/__init__.py +++ b/studio/backend/state/__init__.py @@ -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 diff --git a/studio/backend/tests/__init__.py b/studio/backend/tests/__init__.py index e69de29bb2..32014236c6 100644 --- a/studio/backend/tests/__init__.py +++ b/studio/backend/tests/__init__.py @@ -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 diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py index 4b2bc38cb0..053e9b85d9 100644 --- a/studio/backend/tests/conftest.py +++ b/studio/backend/tests/conftest.py @@ -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 diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index e896c9965c..3c33b33cb3 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -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() diff --git a/studio/backend/utils/__init__.py b/studio/backend/utils/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/utils/__init__.py @@ -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 diff --git a/studio/backend/utils/cache_cleanup.py b/studio/backend/utils/cache_cleanup.py index 187f1c02cd..b3ffbcfc05 100644 --- a/studio/backend/utils/cache_cleanup.py +++ b/studio/backend/utils/cache_cleanup.py @@ -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) diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index ca2e7a8c98..23fc856772 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -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: diff --git a/studio/backend/utils/datasets/data_collators.py b/studio/backend/utils/datasets/data_collators.py index 3c0be0706c..04b24bbae3 100644 --- a/studio/backend/utils/datasets/data_collators.py +++ b/studio/backend/utils/datasets/data_collators.py @@ -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 diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 9629c5aec4..9d15b86ca1 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -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 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 diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py index be707e6127..264789c41c 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -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} diff --git a/studio/backend/utils/datasets/format_detection.py b/studio/backend/utils/datasets/format_detection.py index e03e85e594..7b70ff3a76 100644 --- a/studio/backend/utils/datasets/format_detection.py +++ b/studio/backend/utils/datasets/format_detection.py @@ -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 diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 556caba2f5..6e8c25cba5 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -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, diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index 31c2da28dc..96b0c50876 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -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|>", }, } - diff --git a/studio/backend/utils/datasets/vlm_processing.py b/studio/backend/utils/datasets/vlm_processing.py index 63aeec514e..7b63152ede 100644 --- a/studio/backend/utils/datasets/vlm_processing.py +++ b/studio/backend/utils/datasets/vlm_processing.py @@ -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] = "" 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 ===== diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index 600084f5f2..c1f8b62010 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -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", ] diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index be8e2eb9c0..c743dcd897 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -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()) diff --git a/studio/backend/utils/inference/__init__.py b/studio/backend/utils/inference/__init__.py index 15645ef89a..9dc9c08767 100644 --- a/studio/backend/utils/inference/__init__.py +++ b/studio/backend/utils/inference/__init__.py @@ -4,7 +4,7 @@ """ Inference utility functions """ + from utils.inference.inference_config import load_inference_config __all__ = ["load_inference_config"] - diff --git a/studio/backend/utils/inference/inference_config.py b/studio/backend/utils/inference/inference_config.py index 3e4727e4e3..5b92efce6e 100644 --- a/studio/backend/utils/inference/inference_config.py +++ b/studio/backend/utils/inference/inference_config.py @@ -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 diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 46ec1362ff..82236d8013 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -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", ] diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py index f33d170816..a7cb80f338 100644 --- a/studio/backend/utils/models/checkpoints.py +++ b/studio/backend/utils/models/checkpoints.py @@ -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 diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 21aa1daae1..960e263cb2 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -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: '' 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(' 10000, + "csm": lambda tokens: "<|AUDIO|>" in tokens and "<|audio_eos|>" in tokens, + "whisper": lambda tokens: "<|startoftranscript|>" in tokens, + "audio_vlm": lambda tokens: "" 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(" 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 ) diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py index 3850b57638..507fb1106b 100644 --- a/studio/backend/utils/paths/__init__.py +++ b/studio/backend/utils/paths/__init__.py @@ -4,6 +4,7 @@ """ Path utilities for model and dataset handling """ + from .path_utils import normalize_path, is_local_path, is_model_cached, get_cache_path from .storage_roots import ( studio_root, @@ -21,6 +22,7 @@ from .storage_roots import ( oxc_validator_tmp_root, tensorboard_root, ensure_dir, + ensure_studio_directories, resolve_under_root, resolve_output_dir, resolve_export_dir, @@ -29,28 +31,29 @@ from .storage_roots import ( ) __all__ = [ - 'normalize_path', - 'is_local_path', - 'is_model_cached', - 'get_cache_path', - 'studio_root', - 'assets_root', - 'datasets_root', - 'dataset_uploads_root', - 'recipe_datasets_root', - 'outputs_root', - 'exports_root', - 'auth_root', - 'auth_db_path', - 'tmp_root', - 'seed_uploads_root', - 'unstructured_seed_cache_root', - 'oxc_validator_tmp_root', - 'tensorboard_root', - 'ensure_dir', - 'resolve_under_root', - 'resolve_output_dir', - 'resolve_export_dir', - 'resolve_tensorboard_dir', - 'resolve_dataset_path', + "normalize_path", + "is_local_path", + "is_model_cached", + "get_cache_path", + "studio_root", + "assets_root", + "datasets_root", + "dataset_uploads_root", + "recipe_datasets_root", + "outputs_root", + "exports_root", + "auth_root", + "auth_db_path", + "tmp_root", + "seed_uploads_root", + "unstructured_seed_cache_root", + "oxc_validator_tmp_root", + "tensorboard_root", + "ensure_dir", + "ensure_studio_directories", + "resolve_under_root", + "resolve_output_dir", + "resolve_export_dir", + "resolve_tensorboard_dir", + "resolve_dataset_path", ] diff --git a/studio/backend/utils/paths/path_utils.py b/studio/backend/utils/paths/path_utils.py index f93262c2cb..1d6a952399 100644 --- a/studio/backend/utils/paths/path_utils.py +++ b/studio/backend/utils/paths/path_utils.py @@ -4,6 +4,7 @@ """ Path utilities for model and dataset handling """ + import os from pathlib import Path from typing import Optional @@ -25,14 +26,14 @@ def normalize_path(path: str) -> str: return path # Handle Windows drive letters (C:\\ or c:\\) - if len(path) >= 3 and path[1] == ':' and path[2] in ('\\', '/'): + if len(path) >= 3 and path[1] == ":" and path[2] in ("\\", "/"): drive = path[0].lower() - rest = path[3:].replace('\\', '/') - return f'/mnt/{drive}/{rest}' + rest = path[3:].replace("\\", "/") + return f"/mnt/{drive}/{rest}" # Already Unix-style or relative - return path.replace('\\', '/') -pass + return path.replace("\\", "/") + def is_local_path(path: str) -> bool: """ @@ -53,26 +54,26 @@ def is_local_path(path: str) -> bool: pass # Obvious HF patterns - if path.count('/') == 1 and not path.startswith(('/', '.', '~')): + if path.count("/") == 1 and not path.startswith(("/", ".", "~")): return False # Looks like org/model format # Filesystem indicators return ( - path.startswith(('/', '.', '~')) or # Unix absolute/relative - ':' in path or # Windows drive or URL - '\\' in path or # Windows separator - os.path.isabs(path) # System-absolute + path.startswith(("/", ".", "~")) # Unix absolute/relative + or ":" in path # Windows drive or URL + or "\\" in path # Windows separator + or os.path.isabs(path) # System-absolute ) -pass + def get_cache_path(model_name: str) -> Optional[Path]: """Get HuggingFace cache path for a model if it exists.""" - cache_dir = Path.home() / '.cache' / 'huggingface' / 'hub' + cache_dir = Path.home() / ".cache" / "huggingface" / "hub" model_cache_name = model_name.replace("/", "--") - model_cache_path = cache_dir / f'models--{model_cache_name}' + model_cache_path = cache_dir / f"models--{model_cache_name}" return model_cache_path if model_cache_path.exists() else None -pass + def is_model_cached(model_name: str) -> bool: """Check if model is downloaded in HuggingFace cache.""" @@ -81,9 +82,8 @@ def is_model_cached(model_name: str) -> bool: return False # Check for actual model files - for suffix in ['.safetensors', '.bin', '.json']: - if list(cache_path.rglob(f'*{suffix}')): + for suffix in [".safetensors", ".bin", ".json"]: + if list(cache_path.rglob(f"*{suffix}")): return True return False -pass diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index a005d604a8..d9adce2105 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -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 __future__ import annotations from pathlib import Path @@ -61,11 +64,29 @@ def tensorboard_root() -> Path: def ensure_dir(path: Path) -> Path: - path.mkdir(parents=True, exist_ok=True) + path.mkdir(parents = True, exist_ok = True) return path -def _clean_relative_path(path_value: str, *, strip_prefixes: tuple[str, ...] = ()) -> Path: +def ensure_studio_directories() -> None: + """Create all standard studio directories on startup.""" + for dir_fn in ( + studio_root, + assets_root, + datasets_root, + dataset_uploads_root, + recipe_datasets_root, + outputs_root, + exports_root, + auth_root, + tensorboard_root, + ): + ensure_dir(dir_fn()) + + +def _clean_relative_path( + path_value: str, *, strip_prefixes: tuple[str, ...] = () +) -> Path: path = Path(path_value).expanduser() parts = [part for part in path.parts if part not in ("", ".")] while parts and parts[0] in strip_prefixes: @@ -86,31 +107,31 @@ def resolve_under_root( if path.is_absolute(): return path - cleaned = _clean_relative_path(str(path), strip_prefixes=strip_prefixes) + cleaned = _clean_relative_path(str(path), strip_prefixes = strip_prefixes) return root / cleaned def resolve_output_dir(path_value: str | None = None) -> Path: return resolve_under_root( path_value, - root=outputs_root(), - strip_prefixes=("outputs",), + root = outputs_root(), + strip_prefixes = ("outputs",), ) def resolve_export_dir(path_value: str | None = None) -> Path: return resolve_under_root( path_value, - root=exports_root(), - strip_prefixes=("exports",), + root = exports_root(), + strip_prefixes = ("exports",), ) def resolve_tensorboard_dir(path_value: str | None = None) -> Path: return resolve_under_root( path_value, - root=tensorboard_root(), - strip_prefixes=("runs", "tensorboard"), + root = tensorboard_root(), + strip_prefixes = ("runs", "tensorboard"), ) diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 4cc7217ccb..58fbc35e33 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -33,7 +33,6 @@ from pathlib import Path logger = get_logger(__name__) - # --------------------------------------------------------------------------- # Detection # --------------------------------------------------------------------------- @@ -41,12 +40,12 @@ logger = get_logger(__name__) # Lowercase substrings β€” if ANY appears anywhere in the lowered model name, # we need transformers 5.x. TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = ( - "ministral-3-", # Ministral-3-{3,8,14}B-{Instruct,Reasoning,Base}-2512 - "glm-4.7-flash", # GLM-4.7-Flash - "qwen3-30b-a3b", # Qwen3-30B-A3B-Instruct-2507 and variants - "qwen3.5", # Qwen3.5 family (35B-A3B, etc.) - "qwen3-next", # Qwen3-Next and variants - "tiny_qwen3_moe", # imdatta0/tiny_qwen3_moe_2.8B_0.7B + "ministral-3-", # Ministral-3-{3,8,14}B-{Instruct,Reasoning,Base}-2512 + "glm-4.7-flash", # GLM-4.7-Flash + "qwen3-30b-a3b", # Qwen3-30B-A3B-Instruct-2507 and variants + "qwen3.5", # Qwen3.5 family (35B-A3B, etc.) + "qwen3-next", # Qwen3-Next and variants + "tiny_qwen3_moe", # imdatta0/tiny_qwen3_moe_2.8B_0.7B ) # Versions @@ -77,7 +76,8 @@ def _resolve_base_model(model_name: str) -> str: if base: logger.info( "Resolved LoRA adapter '%s' β†’ base model '%s'", - model_name, base, + model_name, + base, ) return base except Exception as exc: @@ -87,18 +87,21 @@ def _resolve_base_model(model_name: str) -> str: if local_path.is_dir(): try: from utils.models import get_base_model_from_lora + base = get_base_model_from_lora(model_name) if base: logger.info( "Resolved LoRA adapter '%s' β†’ base model '%s' " "(via get_base_model_from_lora)", - model_name, base, + model_name, + base, ) return base except Exception as exc: logger.debug( "get_base_model_from_lora failed for '%s': %s", - model_name, exc, + model_name, + exc, ) return model_name @@ -115,6 +118,7 @@ def needs_transformers_5(model_name: str) -> bool: # Version switching (in-process β€” used only by export) # --------------------------------------------------------------------------- + def _get_in_memory_version() -> str | None: """Return the transformers version currently loaded in this process.""" tf = sys.modules.get("transformers") @@ -153,7 +157,8 @@ def _purge_modules() -> int: """ importlib.invalidate_caches() to_remove = [ - k for k in list(sys.modules.keys()) + k + for k in list(sys.modules.keys()) if any(k == p or k.startswith(p + ".") for p in _PURGE_PREFIXES) ] for key in to_remove: @@ -167,15 +172,21 @@ def _ensure_venv_t5_exists() -> bool: return True logger.warning(".venv_t5 not found at %s β€” installing at runtime", _VENV_T5_DIR) - os.makedirs(_VENV_T5_DIR, exist_ok=True) + os.makedirs(_VENV_T5_DIR, exist_ok = True) for pkg in (f"transformers=={TRANSFORMERS_5_VERSION}", "huggingface_hub==1.3.0"): cmd = [ - sys.executable, "-m", "pip", "install", - "--target", _VENV_T5_DIR, + sys.executable, + "-m", + "pip", + "install", + "--target", + _VENV_T5_DIR, "--no-deps", pkg, ] - result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + result = subprocess.run( + cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True + ) if result.returncode != 0: logger.error("pip install failed:\n%s", result.stdout) return False @@ -186,7 +197,9 @@ def _ensure_venv_t5_exists() -> bool: def _activate_5x() -> None: """Prepend .venv_t5/ to sys.path, purge stale modules, reimport.""" if not _ensure_venv_t5_exists(): - raise RuntimeError(f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}") + raise RuntimeError( + f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}" + ) if _VENV_T5_DIR not in sys.path: sys.path.insert(0, _VENV_T5_DIR) @@ -196,6 +209,7 @@ def _activate_5x() -> None: logger.info("Purged %d cached modules", count) import transformers + logger.info("Loaded transformers %s", transformers.__version__) @@ -209,6 +223,7 @@ def _deactivate_5x() -> None: logger.info("Purged %d cached modules", count) import transformers + logger.info("Reverted to transformers %s", transformers.__version__) @@ -236,7 +251,10 @@ def ensure_transformers_version(model_name: str) -> None: logger.info( "Version check for '%s' (resolved: '%s'): need=%s, in_memory=%s", - model_name, resolved, target_version, in_memory, + model_name, + resolved, + target_version, + in_memory, ) # --- Already correct? --------------------------------------------------- @@ -245,7 +263,8 @@ def ensure_transformers_version(model_name: str) -> None: if in_memory_major == target_major: logger.info( "transformers %s already loaded β€” correct for '%s'", - in_memory, model_name, + in_memory, + model_name, ) return @@ -254,7 +273,9 @@ def ensure_transformers_version(model_name: str) -> None: logger.info("Activating transformers %s via .venv_t5…", TRANSFORMERS_5_VERSION) _activate_5x() else: - logger.info("Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION) + logger.info( + "Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION + ) _deactivate_5x() final = _get_in_memory_version() diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 0cf3abfebf..4e61a5b969 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -4,6 +4,7 @@ """ Shared backend utilities """ + import os import structlog from loggers import get_logger @@ -28,26 +29,26 @@ def without_hf_auth(): """ # Save environment variables saved_env = {} - env_vars = ['HF_TOKEN', 'HUGGINGFACE_HUB_TOKEN', 'HF_HOME'] + env_vars = ["HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_HOME"] for var in env_vars: if var in os.environ: saved_env[var] = os.environ[var] del os.environ[var] # Save disable flag - saved_disable = os.environ.get('HF_HUB_DISABLE_IMPLICIT_TOKEN') - os.environ['HF_HUB_DISABLE_IMPLICIT_TOKEN'] = '1' + saved_disable = os.environ.get("HF_HUB_DISABLE_IMPLICIT_TOKEN") + os.environ["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "1" # Move token files temporarily token_files = [] token_locations = [ - Path.home() / '.cache' / 'huggingface' / 'token', - Path.home() / '.huggingface' / 'token' + Path.home() / ".cache" / "huggingface" / "token", + Path.home() / ".huggingface" / "token", ] for token_loc in token_locations: if token_loc.exists(): - temp = tempfile.NamedTemporaryFile(delete=False) + temp = tempfile.NamedTemporaryFile(delete = False) temp.close() shutil.move(str(token_loc), temp.name) token_files.append((token_loc, temp.name)) @@ -58,7 +59,7 @@ def without_hf_auth(): # Restore tokens for original, temp in token_files: try: - original.parent.mkdir(parents=True, exist_ok=True) + original.parent.mkdir(parents = True, exist_ok = True) shutil.move(temp, str(original)) except Exception as e: logger.error(f"Failed to restore token {original}: {e}") @@ -68,10 +69,10 @@ def without_hf_auth(): os.environ[var] = value if saved_disable is not None: - os.environ['HF_HUB_DISABLE_IMPLICIT_TOKEN'] = saved_disable + os.environ["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = saved_disable else: - os.environ.pop('HF_HUB_DISABLE_IMPLICIT_TOKEN', None) -pass + os.environ.pop("HF_HUB_DISABLE_IMPLICIT_TOKEN", None) + def format_error_message(error: Exception, model_name: str) -> str: """ @@ -85,7 +86,7 @@ def format_error_message(error: Exception, model_name: str) -> str: User-friendly error string """ error_str = str(error).lower() - model_short = model_name.split('/')[-1] if '/' in model_name else model_name + model_short = model_name.split("/")[-1] if "/" in model_name else model_name if "repository not found" in error_str or "404" in error_str: return f"Model '{model_short}' not found. Check the model name." @@ -99,12 +100,19 @@ def format_error_message(error: Exception, model_name: str) -> str: if "invalid user token" in error_str: return "Invalid HF token. Please check your token and try again." - if "memory" in error_str or "cuda" in error_str or "mlx" in error_str or "out of memory" in error_str: + if ( + "memory" in error_str + or "cuda" in error_str + or "mlx" in error_str + or "out of memory" in error_str + ): from utils.hardware import get_device + device = get_device() - device_label = {"cuda": "GPU", "mlx": "Apple Silicon GPU", "cpu": "system"}.get(device.value, "GPU") + device_label = {"cuda": "GPU", "mlx": "Apple Silicon GPU", "cpu": "system"}.get( + device.value, "GPU" + ) return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory." # Generic fallback return str(error) -pass diff --git a/studio/frontend/.gitignore b/studio/frontend/.gitignore index 3483430dcf..bf7ac45ef1 100644 --- a/studio/frontend/.gitignore +++ b/studio/frontend/.gitignore @@ -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 + # Logs logs *.log diff --git a/studio/frontend/eslint.config.js b/studio/frontend/eslint.config.js index 88dec236a8..e76e018697 100644 --- a/studio/frontend/eslint.config.js +++ b/studio/frontend/eslint.config.js @@ -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 js from "@eslint/js"; import reactHooks from "eslint-plugin-react-hooks"; import reactRefresh from "eslint-plugin-react-refresh"; diff --git a/studio/frontend/index.html b/studio/frontend/index.html index c55ae77f62..4f81ffd4ff 100644 --- a/studio/frontend/index.html +++ b/studio/frontend/index.html @@ -1,4 +1,7 @@ + + + diff --git a/studio/frontend/public/huggingface.svg b/studio/frontend/public/huggingface.svg index 45459331bb..c48757ff87 100644 --- a/studio/frontend/public/huggingface.svg +++ b/studio/frontend/public/huggingface.svg @@ -1 +1,4 @@ + + + HuggingFace \ No newline at end of file diff --git a/studio/frontend/public/vite.svg b/studio/frontend/public/vite.svg index e7b8dfb1b2..8f4c4a2a10 100644 --- a/studio/frontend/public/vite.svg +++ b/studio/frontend/public/vite.svg @@ -1 +1,4 @@ + + + \ No newline at end of file diff --git a/studio/frontend/src/assets/react.svg b/studio/frontend/src/assets/react.svg index 6c87de9bb3..2b67796fea 100644 --- a/studio/frontend/src/assets/react.svg +++ b/studio/frontend/src/assets/react.svg @@ -1 +1,4 @@ + + + \ No newline at end of file diff --git a/studio/frontend/src/components/ui/collapsible.tsx b/studio/frontend/src/components/ui/collapsible.tsx index 3cc4953cf3..3566eb9859 100644 --- a/studio/frontend/src/components/ui/collapsible.tsx +++ b/studio/frontend/src/components/ui/collapsible.tsx @@ -1,40 +1,40 @@ -// 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 { cn } from "@/lib/utils"; -import { Collapsible as CollapsiblePrimitive } from "radix-ui"; - -function Collapsible({ - ...props -}: React.ComponentProps) { - return ; -} - -function CollapsibleTrigger({ - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function CollapsibleContent({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -export { Collapsible, CollapsibleTrigger, CollapsibleContent }; +// 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 { cn } from "@/lib/utils"; +import { Collapsible as CollapsiblePrimitive } from "radix-ui"; + +function Collapsible({ + ...props +}: React.ComponentProps) { + return ; +} + +function CollapsibleTrigger({ + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function CollapsibleContent({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +export { Collapsible, CollapsibleTrigger, CollapsibleContent }; diff --git a/studio/frontend/src/components/ui/tooltip.tsx b/studio/frontend/src/components/ui/tooltip.tsx index d5f1497041..f58d9f5309 100644 --- a/studio/frontend/src/components/ui/tooltip.tsx +++ b/studio/frontend/src/components/ui/tooltip.tsx @@ -1,62 +1,62 @@ -// 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 { Tooltip as TooltipPrimitive } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; - -function TooltipProvider({ - delayDuration = 400, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function Tooltip({ - ...props -}: React.ComponentProps) { - return ( - - - - ); -} - -function TooltipTrigger({ - ...props -}: React.ComponentProps) { - return ; -} - -function TooltipContent({ - className, - sideOffset = 0, - children, - ...props -}: React.ComponentProps) { - return ( - - - {children} - - - - ); -} - -export { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger }; +// 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 { Tooltip as TooltipPrimitive } from "radix-ui"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; + +function TooltipProvider({ + delayDuration = 400, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function Tooltip({ + ...props +}: React.ComponentProps) { + return ( + + + + ); +} + +function TooltipTrigger({ + ...props +}: React.ComponentProps) { + return ; +} + +function TooltipContent({ + className, + sideOffset = 0, + children, + ...props +}: React.ComponentProps) { + return ( + + + {children} + + + + ); +} + +export { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger }; diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css index 419883ebf6..397853a7b7 100644 --- a/studio/frontend/src/index.css +++ b/studio/frontend/src/index.css @@ -1,5 +1,5 @@ -/* SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0 */ -/* Copyright Β© 2025 Unsloth AI */ +/* 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 "tailwindcss"; @import "tw-animate-css"; diff --git a/studio/frontend/vite.config.ts b/studio/frontend/vite.config.ts index c8c9028b2e..6e9a03df95 100644 --- a/studio/frontend/vite.config.ts +++ b/studio/frontend/vite.config.ts @@ -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 path from "node:path"; import tailwindcss from "@tailwindcss/vite"; import react from "@vitejs/plugin-react"; diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index a4fed429a2..9ce16cd7bb 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1,4 +1,5 @@ #!/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 @@ -31,6 +32,7 @@ LOCAL_DD_UNSTRUCTURED_PLUGIN = ( # ── Color support ────────────────────────────────────────────────────── + def _enable_colors() -> bool: """Try to enable ANSI color support. Returns True if available.""" if not hasattr(sys.stdout, "fileno"): @@ -43,6 +45,7 @@ def _enable_colors() -> bool: if IS_WINDOWS: try: import ctypes + kernel32 = ctypes.windll.kernel32 # Enable ENABLE_VIRTUAL_TERMINAL_PROCESSING (0x0004) on stdout handle = kernel32.GetStdHandle(-11) # STD_OUTPUT_HANDLE @@ -54,14 +57,18 @@ def _enable_colors() -> bool: return False return True # Unix terminals support ANSI by default + _HAS_COLOR = _enable_colors() + def _green(msg: str) -> str: return f"\033[92m{msg}\033[0m" if _HAS_COLOR else msg + def _cyan(msg: str) -> str: return f"\033[96m{msg}\033[0m" if _HAS_COLOR else msg + def _red(msg: str) -> str: return f"\033[91m{msg}\033[0m" if _HAS_COLOR else msg @@ -71,29 +78,33 @@ def run(label: str, cmd: list[str], *, quiet: bool = True) -> None: print(_cyan(f" {label}...")) result = subprocess.run( cmd, - stdout=subprocess.PIPE if quiet else None, - stderr=subprocess.STDOUT if quiet else None, + stdout = subprocess.PIPE if quiet else None, + stderr = subprocess.STDOUT if quiet else None, ) if result.returncode != 0: print(_red(f"❌ {label} failed (exit code {result.returncode}):")) if result.stdout: - print(result.stdout.decode(errors="replace")) + print(result.stdout.decode(errors = "replace")) sys.exit(result.returncode) # Packages to skip on Windows (require special build steps) -WINDOWS_SKIP_PACKAGES = {"open_spiel"} +WINDOWS_SKIP_PACKAGES = {"open_spiel", "triton_kernels"} def _filter_requirements(req: Path, skip: set[str]) -> Path: """Return a temp copy of a requirements file with certain packages removed.""" - lines = req.read_text(encoding="utf-8").splitlines(keepends=True) + lines = req.read_text(encoding = "utf-8").splitlines(keepends = True) filtered = [ - line for line in lines + line + for line in lines if not any(line.strip().lower().startswith(pkg) for pkg in skip) ] tmp = tempfile.NamedTemporaryFile( - mode="w", suffix=".txt", delete=False, encoding="utf-8", + mode = "w", + suffix = ".txt", + delete = False, + encoding = "utf-8", ) tmp.writelines(filtered) tmp.close() @@ -121,8 +132,7 @@ def pip_install( finally: # Clean up temp file if we created one if actual_req is not None and actual_req != req: - actual_req.unlink(missing_ok=True) - + actual_req.unlink(missing_ok = True) def download_file(url: str, dest: Path) -> None: @@ -134,7 +144,8 @@ def patch_package_file(package_name: str, relative_path: str, url: str) -> None: """Download a file from url and overwrite a file inside an installed package.""" result = subprocess.run( [sys.executable, "-m", "pip", "show", package_name], - capture_output=True, text=True, + capture_output = True, + text = True, ) if result.returncode != 0: print(_red(f" ⚠️ Could not find package {package_name}, skipping patch")) @@ -157,6 +168,7 @@ def patch_package_file(package_name: str, relative_path: str, url: str) -> None: # ── Main install sequence ───────────────────────────────────────────── + def install_python_stack() -> int: print(_cyan("── Installing Python stack ──")) @@ -165,80 +177,85 @@ def install_python_stack() -> int: # 2. Core packages: unsloth-zoo + unsloth pip_install( - "Installing unsloth-zoo + unsloth", + "Installing base packages", "--no-cache-dir", - req=REQ_ROOT / "base.txt", + req = REQ_ROOT / "base.txt", ) # 3. Extra dependencies pip_install( "Installing additional unsloth dependencies", "--no-cache-dir", - req=REQ_ROOT / "extras.txt", + req = REQ_ROOT / "extras.txt", ) # 3b. Extra dependencies (no-deps) β€” audio model support etc. pip_install( "Installing extras (no-deps)", - "--no-deps", "--no-cache-dir", - req=REQ_ROOT / "extras-no-deps.txt", + "--no-deps", + "--no-cache-dir", + req = REQ_ROOT / "extras-no-deps.txt", ) # 4. Overrides (torchao, transformers) β€” force-reinstall pip_install( - "Installing torchao + transformers overrides", - "--force-reinstall", "--no-cache-dir", - req=REQ_ROOT / "overrides.txt", + "Installing dependency overrides", + "--force-reinstall", + "--no-cache-dir", + req = REQ_ROOT / "overrides.txt", ) # 5. Triton kernels (no-deps, from source) - pip_install( - "Installing triton kernels", - "--no-deps", "--no-cache-dir", - req=REQ_ROOT / "triton-kernels.txt", - constrain=False, - ) + if not IS_WINDOWS: + pip_install( + "Installing triton kernels", + "--no-deps", + "--no-cache-dir", + req = REQ_ROOT / "triton-kernels.txt", + constrain = False, + ) - # 6. Patch: override llama_cpp.py with fix from unsloth-zoo feature/llama-cpp-windows-support branch - patch_package_file( - "unsloth-zoo", - os.path.join("unsloth_zoo", "llama_cpp.py"), - "https://raw.githubusercontent.com/unslothai/unsloth-zoo/refs/heads/main/unsloth_zoo/llama_cpp.py", - ) + # # 6. Patch: override llama_cpp.py with fix from unsloth-zoo feature/llama-cpp-windows-support branch + # patch_package_file( + # "unsloth-zoo", + # os.path.join("unsloth_zoo", "llama_cpp.py"), + # "https://raw.githubusercontent.com/unslothai/unsloth-zoo/refs/heads/main/unsloth_zoo/llama_cpp.py", + # ) - # 7a. Patch: override vision.py with fix from unsloth PR #4091 - patch_package_file( - "unsloth", - os.path.join("unsloth", "models", "vision.py"), - "https://raw.githubusercontent.com/unslothai/unsloth/80e0108a684c882965a02a8ed851e3473c1145ab/unsloth/models/vision.py", - ) + # # 7a. Patch: override vision.py with fix from unsloth PR #4091 + # patch_package_file( + # "unsloth", + # os.path.join("unsloth", "models", "vision.py"), + # "https://raw.githubusercontent.com/unslothai/unsloth/80e0108a684c882965a02a8ed851e3473c1145ab/unsloth/models/vision.py", + # ) - # 7b. Patch : override save.py with fix from feature/llama-cpp-windows-support - patch_package_file( - "unsloth", - os.path.join("unsloth", "save.py"), - "https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/main/unsloth/save.py", - ) + # # 7b. Patch : override save.py with fix from feature/llama-cpp-windows-support + # patch_package_file( + # "unsloth", + # os.path.join("unsloth", "save.py"), + # "https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/main/unsloth/save.py", + # ) # 8. Studio dependencies pip_install( "Installing studio dependencies", "--no-cache-dir", - req=REQ_ROOT / "studio.txt", + req = REQ_ROOT / "studio.txt", ) # 9. Data-designer dependencies pip_install( - "Installing data-designer dependencies", + "Installing data-designer base dependencies", "--no-cache-dir", - req=SINGLE_ENV / "data-designer-deps.txt", + req = SINGLE_ENV / "data-designer-deps.txt", ) # 10. Data-designer packages (no-deps to avoid conflicts) pip_install( "Installing data-designer", - "--no-cache-dir", "--no-deps", - req=SINGLE_ENV / "data-designer.txt", + "--no-cache-dir", + "--no-deps", + req = SINGLE_ENV / "data-designer.txt", ) # 11. Local Data Designer seed plugin @@ -251,9 +268,11 @@ def install_python_stack() -> int: return 1 pip_install( "Installing local data-designer unstructured plugin", - "--no-cache-dir", "--no-deps", - "-e", str(LOCAL_DD_UNSTRUCTURED_PLUGIN), - constrain=False, + "--no-cache-dir", + "--no-deps", + "-e", + str(LOCAL_DD_UNSTRUCTURED_PLUGIN), + constrain = False, ) # 12. Patch metadata for single-env compatibility @@ -265,7 +284,8 @@ def install_python_stack() -> int: # 13. Final check (silent; third-party conflicts are expected) subprocess.run( [sys.executable, "-m", "pip", "check"], - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + stdout = subprocess.DEVNULL, + stderr = subprocess.DEVNULL, ) print(_green("βœ… Python dependencies installed")) diff --git a/studio/setup.bat b/studio/setup.bat index ef16abd263..72d1acb141 100644 --- a/studio/setup.bat +++ b/studio/setup.bat @@ -1,2 +1,5 @@ @echo off +REM SPDX-License-Identifier: AGPL-3.0-only +REM Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + powershell -ExecutionPolicy Bypass -File "%~dp0setup.ps1" %* diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 8b662e1998..bf0ba754f0 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1,6 +1,6 @@ #Requires -Version 5.1 -# SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0 -# Copyright Β© 2025 Unsloth AI +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 <# .SYNOPSIS Full environment setup for Unsloth Studio on Windows (bundled version). @@ -653,7 +653,7 @@ Write-Host "[OK] Using $PythonCmd ($(& $PythonCmd --version 2>&1))" -ForegroundC # Always create a .venv for isolation -- even for pip installs. # Created in the repo root (parent of studio/). -$VenvDir = Join-Path (Split-Path -Parent $PSScriptRoot) ".venv" +$VenvDir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv" if (-not (Test-Path $VenvDir)) { Write-Host " Creating virtual environment at $VenvDir..." -ForegroundColor Cyan & $PythonCmd -m venv $VenvDir @@ -726,7 +726,7 @@ $ErrorActionPreference = $prevEAP # The training subprocess just prepends .venv_t5/ to sys.path β€” instant switch. Write-Host "" Write-Host " Pre-installing transformers 5.x for newer model support..." -ForegroundColor Cyan -$VenvT5Dir = Join-Path (Split-Path -Parent $PSScriptRoot) ".venv_t5" +$VenvT5Dir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv_t5" if (Test-Path $VenvT5Dir) { Remove-Item -Recurse -Force $VenvT5Dir } New-Item -ItemType Directory -Path $VenvT5Dir -Force | Out-Null $prevEAP_t5 = $ErrorActionPreference @@ -956,55 +956,6 @@ if (Test-Path $LlamaServerBin) { } } -# ============================================ -# Add shell aliases (PowerShell profile + cmd batch files) -# ============================================ -Write-Host "" -$RepoDir = Split-Path -Parent $PSScriptRoot -$VenvPython = Join-Path $RepoDir ".venv\Scripts\python.exe" -$CliScript = Join-Path $RepoDir "cli.py" -$FrontendDist = Join-Path $PSScriptRoot "frontend\dist" -$AliasAdded = $false - -# --- PowerShell profile: add functions --- -$ProfileDir = Split-Path $PROFILE -Parent -if (-not (Test-Path $ProfileDir)) { New-Item -ItemType Directory -Path $ProfileDir -Force | Out-Null } -if (-not (Test-Path $PROFILE)) { New-Item -ItemType File -Path $PROFILE -Force | Out-Null } - -if (-not (Select-String -Path $PROFILE -Pattern "unsloth-studio" -Quiet -ErrorAction SilentlyContinue)) { - $block = @" - -# Unsloth Studio launcher -function unsloth-studio { & "$VenvPython" "$CliScript" studio -f "$FrontendDist" @args } -function unsloth-ui { & "$VenvPython" "$CliScript" studio -f "$FrontendDist" @args } -"@ - Add-Content -Path $PROFILE -Value $block - Write-Host "[OK] Aliases 'unsloth-studio' and 'unsloth-ui' added to $PROFILE" -ForegroundColor Green - $AliasAdded = $true -} else { - Write-Host "[OK] Aliases 'unsloth-studio' and 'unsloth-ui' already exist in $PROFILE" -ForegroundColor Green -} - -# --- cmd.exe: create batch files and ensure they're on PATH --- -$BatDir = Join-Path $RepoDir ".venv\Scripts" -foreach ($name in @("unsloth-studio", "unsloth-ui")) { - $batPath = Join-Path $BatDir "$name.bat" - if (-not (Test-Path $batPath)) { - Set-Content -Path $batPath -Value "@echo off`r`n`"$VenvPython`" `"$CliScript`" studio -f `"$FrontendDist`" %*" - } -} -# Persist .venv\Scripts to User PATH so commands work in new cmd.exe terminals without activation -$userPath = [Environment]::GetEnvironmentVariable('Path', 'User') -if (-not $userPath -or $userPath -notlike "*$BatDir*") { - if ($userPath) { - [Environment]::SetEnvironmentVariable('Path', "$BatDir;$userPath", 'User') - } else { - [Environment]::SetEnvironmentVariable('Path', "$BatDir", 'User') - } - Write-Host " Persisted $BatDir to User PATH" -ForegroundColor Gray -} -Write-Host "[OK] Batch launchers created (works from any new cmd.exe or PowerShell)" -ForegroundColor Green - # ============================================ # Done # ============================================ @@ -1012,8 +963,7 @@ Write-Host "" Write-Host "+===============================================+" -ForegroundColor Green Write-Host "| Setup Complete! |" -ForegroundColor Green Write-Host "| |" -ForegroundColor Green -Write-Host "| IMPORTANT: Open a NEW terminal, then run: |" -ForegroundColor Yellow -Write-Host "| |" -ForegroundColor Green +Write-Host "| Launch with: |" -ForegroundColor Green Write-Host "| unsloth studio -H 0.0.0.0 -p 8000 |" -ForegroundColor Green Write-Host "| |" -ForegroundColor Green Write-Host "+===============================================+" -ForegroundColor Green diff --git a/studio/setup.sh b/studio/setup.sh index 45fedc8882..d42f772397 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -1,6 +1,6 @@ #!/usr/bin/env bash -# SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0 -# Copyright Β© 2025 Unsloth AI +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 set -euo pipefail @@ -322,57 +322,6 @@ rm -rf "$LLAMA_CPP_DIR" fi } -# ── 9. Add shell alias (skip in Colab) ── -# Note: venv activation does NOT persist across terminal sessions. -# This alias hardcodes the venv python path so users don't need to activate. -if [ "$IS_COLAB" = false ]; then -echo "" -REPO_DIR="$REPO_ROOT" -STUDIO_VENV_PY="$HOME/.unsloth/studio/.venv/bin/python" - -# Detect the user's default shell and pick the right rc file -USER_SHELL="$(basename "${SHELL:-/bin/bash}")" -case "$USER_SHELL" in - zsh) - SHELL_RC="$HOME/.zshrc" - ALIAS_BLOCK="alias unsloth-studio='${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio' -alias unsloth-ui='${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio'" - ;; - fish) - SHELL_RC="$HOME/.config/fish/config.fish" - ALIAS_BLOCK="alias unsloth-studio '${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio' -alias unsloth-ui '${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio'" - ;; - ksh) - SHELL_RC="$HOME/.kshrc" - ALIAS_BLOCK="alias unsloth-studio='${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio' -alias unsloth-ui='${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio'" - ;; - *) - SHELL_RC="$HOME/.bashrc" - ALIAS_BLOCK="alias unsloth-studio='${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio' -alias unsloth-ui='${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio'" - ;; -esac - -echo " Detected shell: $USER_SHELL β†’ $SHELL_RC" - -ALIAS_ADDED=false -if ! grep -qF "unsloth-studio" "$SHELL_RC" 2>/dev/null; then - mkdir -p "$(dirname "$SHELL_RC")" # needed for fish's nested config path - cat >> "$SHELL_RC" < list[dict | str]: - """Extract `data:` payloads from raw SSE text. Returns dicts or raw strings.""" - results = [] - for line in raw.split("\n"): - if line.startswith("data: "): - payload = line[len("data: "):] - if payload == "[DONE]": - results.append("[DONE]") - else: - try: - results.append(json.loads(payload)) - except json.JSONDecodeError: - results.append(payload) - return results - - -@pytest.fixture() -def client(): - yield TestClient(app) - - -# ===================================================================== -# Streaming tests -# ===================================================================== - - -class TestStreamingChunkFormat: - """Each SSE chunk must match the OpenAI chat.completion.chunk schema.""" - - def test_chunks_have_required_fields(self, client: TestClient): - mock_backend = _make_mock_backend(tokens=["Hi"]) - with patch("routes.inference.get_inference_backend", return_value=mock_backend): - resp = client.post( - "/api/inference/chat/completions", - json={ - "messages": [{"role": "user", "content": "Hello"}], - "stream": True, - }, - ) - - assert resp.status_code == 200 - chunks = _parse_sse_data(resp.text) - - # Filter to actual chunk dicts (not [DONE]) - json_chunks = [c for c in chunks if isinstance(c, dict) and "choices" in c] - assert len(json_chunks) >= 2 # role chunk + content chunk(s) + final - - for chunk in json_chunks: - assert "id" in chunk - assert chunk["object"] == "chat.completion.chunk" - assert "created" in chunk - assert "model" in chunk - assert len(chunk["choices"]) == 1 - assert "delta" in chunk["choices"][0] - - def test_first_chunk_has_role(self, client: TestClient): - mock_backend = _make_mock_backend(tokens=["Hi"]) - with patch("routes.inference.get_inference_backend", return_value=mock_backend): - resp = client.post( - "/api/inference/chat/completions", - json={"messages": [{"role": "user", "content": "Hello"}]}, - ) - - chunks = [c for c in _parse_sse_data(resp.text) if isinstance(c, dict) and "choices" in c] - first = chunks[0] - assert first["choices"][0]["delta"].get("role") == "assistant" - - def test_last_chunk_has_stop_finish_reason(self, client: TestClient): - mock_backend = _make_mock_backend(tokens=["Done"]) - with patch("routes.inference.get_inference_backend", return_value=mock_backend): - resp = client.post( - "/api/inference/chat/completions", - json={"messages": [{"role": "user", "content": "Hello"}]}, - ) - - chunks = [c for c in _parse_sse_data(resp.text) if isinstance(c, dict) and "choices" in c] - last = chunks[-1] - assert last["choices"][0]["finish_reason"] == "stop" - # Delta should be empty on the final chunk - assert last["choices"][0]["delta"].get("content") is None - - def test_stream_ends_with_done(self, client: TestClient): - mock_backend = _make_mock_backend(tokens=["x"]) - with patch("routes.inference.get_inference_backend", return_value=mock_backend): - resp = client.post( - "/api/inference/chat/completions", - json={"messages": [{"role": "user", "content": "Hello"}]}, - ) - - all_data = _parse_sse_data(resp.text) - assert all_data[-1] == "[DONE]" - - def test_consistent_id_across_chunks(self, client: TestClient): - mock_backend = _make_mock_backend(tokens=["a", "b", "c"]) - with patch("routes.inference.get_inference_backend", return_value=mock_backend): - resp = client.post( - "/api/inference/chat/completions", - json={"messages": [{"role": "user", "content": "Hello"}]}, - ) - - chunks = [c for c in _parse_sse_data(resp.text) if isinstance(c, dict) and "choices" in c] - ids = set(c["id"] for c in chunks) - assert len(ids) == 1, "All chunks should share the same completion ID" - - -class TestStreamingHeaders: - """Verify response headers for SSE proxy compatibility.""" - - def test_headers(self, client: TestClient): - mock_backend = _make_mock_backend(tokens=["x"]) - with patch("routes.inference.get_inference_backend", return_value=mock_backend): - resp = client.post( - "/api/inference/chat/completions", - json={"messages": [{"role": "user", "content": "Hello"}]}, - ) - - assert resp.headers["content-type"].startswith("text/event-stream") - assert resp.headers.get("cache-control") == "no-cache" - assert resp.headers.get("x-accel-buffering") == "no" - - -# ===================================================================== -# Non-streaming tests -# ===================================================================== - - -class TestNonStreaming: - """When stream=false, return a single ChatCompletion JSON object.""" - - def test_returns_json_object(self, client: TestClient): - mock_backend = _make_mock_backend(tokens=["Full response text"]) - with patch("routes.inference.get_inference_backend", return_value=mock_backend): - resp = client.post( - "/api/inference/chat/completions", - json={ - "messages": [{"role": "user", "content": "Hello"}], - "stream": False, - }, - ) - - assert resp.status_code == 200 - body = resp.json() - assert body["object"] == "chat.completion" - assert body["choices"][0]["message"]["role"] == "assistant" - assert body["choices"][0]["message"]["content"] == "Full response text" - assert body["choices"][0]["finish_reason"] == "stop" - - def test_non_streaming_has_model(self, client: TestClient): - mock_backend = _make_mock_backend(tokens=["x"], active_model="my-model") - with patch("routes.inference.get_inference_backend", return_value=mock_backend): - resp = client.post( - "/api/inference/chat/completions", - json={ - "messages": [{"role": "user", "content": "Hi"}], - "stream": False, - }, - ) - - body = resp.json() - assert body["model"] == "my-model" - - -# ===================================================================== -# System prompt extraction -# ===================================================================== - - -class TestSystemPromptExtraction: - """System messages should be extracted and passed as system_prompt.""" - - def test_system_message_extracted(self, client: TestClient): - mock_backend = _make_mock_backend(tokens=["ok"]) - with patch("routes.inference.get_inference_backend", return_value=mock_backend): - client.post( - "/api/inference/chat/completions", - json={ - "messages": [ - {"role": "system", "content": "You are a pirate."}, - {"role": "user", "content": "Hello"}, - ], - "stream": False, - }, - ) - - # Check that generate_chat_response was called with the correct system_prompt - call_kwargs = mock_backend.generate_chat_response.call_args[1] - assert call_kwargs["system_prompt"] == "You are a pirate." - # System message should NOT be in the chat_messages list - assert all(m["role"] != "system" for m in call_kwargs["messages"]) - - def test_default_system_prompt_when_none(self, client: TestClient): - mock_backend = _make_mock_backend(tokens=["ok"]) - with patch("routes.inference.get_inference_backend", return_value=mock_backend): - client.post( - "/api/inference/chat/completions", - json={ - "messages": [{"role": "user", "content": "Hello"}], - "stream": False, - }, - ) - - call_kwargs = mock_backend.generate_chat_response.call_args[1] - assert call_kwargs["system_prompt"] == "You are a helpful AI assistant." - - -# ===================================================================== -# Error handling -# ===================================================================== - - -class TestErrorHandling: - """Validate error responses for bad requests.""" - - def test_no_model_loaded(self, client: TestClient): - mock_backend = _make_mock_backend() - mock_backend.active_model_name = None - with patch("routes.inference.get_inference_backend", return_value=mock_backend): - resp = client.post( - "/api/inference/chat/completions", - json={"messages": [{"role": "user", "content": "Hi"}]}, - ) - - assert resp.status_code == 400 - assert "No model loaded" in resp.json()["detail"] - - def test_only_system_messages_rejected(self, client: TestClient): - mock_backend = _make_mock_backend() - with patch("routes.inference.get_inference_backend", return_value=mock_backend): - resp = client.post( - "/api/inference/chat/completions", - json={ - "messages": [{"role": "system", "content": "You are a bot."}], - }, - ) - - assert resp.status_code == 400 - assert "non-system message" in resp.json()["detail"] diff --git a/studio/tests/test_remote_lora_detection.py b/studio/tests/test_remote_lora_detection.py deleted file mode 100644 index d14ad25d6d..0000000000 --- a/studio/tests/test_remote_lora_detection.py +++ /dev/null @@ -1,43 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -""" -Test remote LoRA adapter detection via HuggingFace Hub API. - -Verifies that we can detect whether a remote HF model is a LoRA adapter -by checking for adapter_config.json in the repo file listing. -""" -import pytest -from huggingface_hub import model_info - - -def is_remote_lora_adapter(model_id: str, hf_token: str = None) -> bool: - """ - Check if a remote HuggingFace model is a LoRA adapter - by looking for adapter_config.json in its repo files. - """ - try: - info = model_info(model_id, token=hf_token) - filenames = [s.rfilename for s in info.siblings] - return "adapter_config.json" in filenames - except Exception: - return False - - -class TestRemoteLoRADetection: - """Test remote LoRA adapter detection via HF Hub API.""" - - def test_lora_adapter_detected(self): - """edbeeching/llama-se-rl-adapter is a known LoRA adapter on HF.""" - result = is_remote_lora_adapter("edbeeching/llama-se-rl-adapter") - assert result is True, "Expected edbeeching/llama-se-rl-adapter to be detected as a LoRA adapter" - - def test_base_model_not_detected_as_lora(self): - """google/gemma-3-4b-it is a full base model, not a LoRA adapter.""" - result = is_remote_lora_adapter("google/gemma-3-4b-it") - assert result is False, "Expected google/gemma-3-4b-it to NOT be detected as a LoRA adapter" - - def test_nonexistent_model_returns_false(self): - """A nonexistent model should return False, not raise.""" - result = is_remote_lora_adapter("this-org-does-not-exist/fake-model-12345") - assert result is False, "Expected nonexistent model to return False" diff --git a/studio/tests/test_sse_progress.py b/studio/tests/test_sse_progress.py deleted file mode 100644 index b7ee9c9b8f..0000000000 --- a/studio/tests/test_sse_progress.py +++ /dev/null @@ -1,323 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -""" -Tests for the SSE training progress endpoint and status fallback. - -Validates: - - SSE spec compliance: `retry:`, `id:`, `event:` fields - - Named event types: progress, heartbeat, complete, error - - Last-Event-ID reconnection and history replay - - /status metric_history fallback (Option B) - -All tests mock the training backend and bypass auth. -""" -import sys -from pathlib import Path -from typing import Optional -from unittest.mock import MagicMock, patch, PropertyMock -import re - -import pytest - -# ── Path setup ──────────────────────────────────────────────────── -# Add backend root so bare `from routes…`, `from models…` etc. resolve. -_backend_root = Path(__file__).resolve().parent.parent / "backend" -if str(_backend_root) not in sys.path: - sys.path.insert(0, str(_backend_root)) - -from fastapi.testclient import TestClient -from main import app -from auth.authentication import get_current_subject - - -# ── Fixtures ────────────────────────────────────────────────────── - -def _bypass_auth(): - """Dependency override that skips real JWT validation.""" - return "test-user" - - -def _make_mock_backend( - *, - is_active: bool = False, - step_history: list | None = None, - loss_history: list | None = None, - lr_history: list | None = None, - total_steps: int = 100, - epoch: int | None = 1, - job_id: str = "job_test_001", -): - """Build a lightweight mock that quacks like TrainingBackend.""" - backend = MagicMock() - backend.current_job_id = job_id - backend.step_history = step_history or [] - backend.loss_history = loss_history or [] - backend.lr_history = lr_history or [] - backend.is_training_active.return_value = is_active - backend._training_thread = None - - # trainer.training_progress / get_training_progress() - tp = MagicMock() - tp.total_steps = total_steps - tp.epoch = epoch - tp.step = step_history[-1] if step_history else 0 - tp.loss = loss_history[-1] if loss_history else 0.0 - tp.learning_rate = lr_history[-1] if lr_history else 0.0 - tp.status_message = "Training..." - tp.error = None - tp.is_completed = not is_active and bool(step_history) - - backend.trainer = MagicMock() - backend.trainer.training_progress = tp - backend.trainer.get_training_progress.return_value = tp - - return backend - - -@pytest.fixture() -def client(): - """TestClient with auth bypassed.""" - app.dependency_overrides[get_current_subject] = _bypass_auth - yield TestClient(app) - app.dependency_overrides.clear() - - -# ── SSE Parsing Helpers ─────────────────────────────────────────── - -def parse_sse_events(raw: str) -> list[dict]: - """ - Parse raw SSE text into a list of event dicts. - - Each dict has optional keys: 'id', 'event', 'data', 'retry'. - """ - events: list[dict] = [] - current: dict = {} - - for line in raw.split("\n"): - if line.startswith("retry:"): - # retry is a standalone directive, not part of a normal event - events.append({"retry": line.split(":", 1)[1].strip()}) - continue - if line.startswith("id:"): - current["id"] = line.split(":", 1)[1].strip() - elif line.startswith("event:"): - current["event"] = line.split(":", 1)[1].strip() - elif line.startswith("data:"): - current["data"] = line.split(":", 1)[1].strip() - elif line == "" and current: - events.append(current) - current = {} - - if current: - events.append(current) - return events - - -# ===================================================================== -# Option A β€” /api/train/progress (SSE) -# ===================================================================== - - -class TestSSERetryDirective: - """The first thing the stream emits must be `retry: 3000`.""" - - def test_retry_is_first_event(self, client: TestClient): - mock_backend = _make_mock_backend(is_active=False) - with patch("routes.training.get_training_backend", return_value=mock_backend): - resp = client.get("/api/train/progress") - - assert resp.status_code == 200 - assert resp.headers["content-type"].startswith("text/event-stream") - - events = parse_sse_events(resp.text) - assert len(events) >= 1 - assert events[0] == {"retry": "3000"} - - -class TestSSEEventFields: - """Every non-retry event must include `id:`, `event:`, and `data:` fields.""" - - def test_events_have_id_and_event_type(self, client: TestClient): - mock_backend = _make_mock_backend( - is_active=False, - step_history=[1, 2, 3], - loss_history=[2.0, 1.5, 1.0], - lr_history=[1e-4, 1e-4, 1e-4], - total_steps=3, - ) - with patch("routes.training.get_training_backend", return_value=mock_backend): - resp = client.get("/api/train/progress") - - events = parse_sse_events(resp.text) - data_events = [e for e in events if "data" in e] - - assert len(data_events) >= 1 - for evt in data_events: - assert "id" in evt, f"Missing `id:` field in event: {evt}" - assert "event" in evt, f"Missing `event:` field in event: {evt}" - assert "data" in evt - - -class TestSSENamedEventTypes: - """Events use the correct named types: progress, complete, heartbeat, error.""" - - def test_idle_sends_progress_then_complete(self, client: TestClient): - mock_backend = _make_mock_backend( - is_active=False, - step_history=[10], - loss_history=[1.5], - lr_history=[1e-4], - total_steps=10, - ) - with patch("routes.training.get_training_backend", return_value=mock_backend): - resp = client.get("/api/train/progress") - - events = parse_sse_events(resp.text) - data_events = [e for e in events if "event" in e and e.get("event") != "retry"] - - event_types = [e["event"] for e in data_events] - assert "progress" in event_types - assert "complete" in event_types - - def test_no_history_sends_complete(self, client: TestClient): - mock_backend = _make_mock_backend(is_active=False) - with patch("routes.training.get_training_backend", return_value=mock_backend): - resp = client.get("/api/train/progress") - - events = parse_sse_events(resp.text) - data_events = [e for e in events if "event" in e] - assert any(e["event"] == "complete" for e in data_events) - - -class TestSSELastEventIDResume: - """When `Last-Event-ID` header is sent, the server replays missed steps.""" - - def test_replays_steps_after_last_event_id(self, client: TestClient): - mock_backend = _make_mock_backend( - is_active=False, - step_history=[1, 2, 3, 4, 5], - loss_history=[2.5, 2.0, 1.5, 1.2, 1.0], - lr_history=[1e-4, 1e-4, 1e-4, 1e-4, 1e-4], - total_steps=5, - ) - with patch("routes.training.get_training_backend", return_value=mock_backend): - resp = client.get( - "/api/train/progress", - headers={"Last-Event-ID": "2"}, - ) - - events = parse_sse_events(resp.text) - # Filter to progress events (replayed ones) - progress_events = [e for e in events if e.get("event") == "progress"] - - # Steps 3, 4, 5 should have been replayed - replayed_ids = [int(e["id"]) for e in progress_events] - assert 3 in replayed_ids - assert 4 in replayed_ids - assert 5 in replayed_ids - # Steps 1, 2 should NOT be replayed - assert 1 not in replayed_ids - assert 2 not in replayed_ids - - def test_no_replay_without_header(self, client: TestClient): - """Without Last-Event-ID, should start fresh (initial progress event).""" - mock_backend = _make_mock_backend( - is_active=False, - step_history=[1, 2, 3], - loss_history=[2.0, 1.5, 1.0], - lr_history=[1e-4, 1e-4, 1e-4], - total_steps=3, - ) - with patch("routes.training.get_training_backend", return_value=mock_backend): - resp = client.get("/api/train/progress") - - events = parse_sse_events(resp.text) - progress_events = [e for e in events if e.get("event") == "progress"] - - # Should have initial step=0 progress event - assert any(e.get("id") == "0" for e in progress_events) - - def test_invalid_last_event_id_treated_as_fresh(self, client: TestClient): - """Non-integer Last-Event-ID should be ignored gracefully.""" - mock_backend = _make_mock_backend(is_active=False) - with patch("routes.training.get_training_backend", return_value=mock_backend): - resp = client.get( - "/api/train/progress", - headers={"Last-Event-ID": "not-a-number"}, - ) - - assert resp.status_code == 200 - events = parse_sse_events(resp.text) - # Should still work β€” treated as a fresh connection - assert any(e.get("event") == "progress" or e.get("event") == "complete" for e in events) - - -class TestSSEResponseHeaders: - """Verify SSE response headers for proxy compatibility.""" - - def test_headers(self, client: TestClient): - mock_backend = _make_mock_backend(is_active=False) - with patch("routes.training.get_training_backend", return_value=mock_backend): - resp = client.get("/api/train/progress") - - assert resp.headers["content-type"].startswith("text/event-stream") - assert resp.headers.get("cache-control") == "no-cache" - assert resp.headers.get("x-accel-buffering") == "no" - - -# ===================================================================== -# Option B β€” /api/train/status (metric_history fallback) -# ===================================================================== - - -class TestStatusMetricHistory: - """The /status endpoint returns metric_history for chart recovery.""" - - def test_metric_history_populated_when_history_exists(self, client: TestClient): - mock_backend = _make_mock_backend( - is_active=True, - step_history=[1, 2, 3, 4, 5], - loss_history=[2.5, 2.0, 1.5, 1.2, 1.0], - lr_history=[1e-4, 1e-4, 1e-4, 1e-4, 1e-4], - total_steps=10, - ) - with patch("routes.training.get_training_backend", return_value=mock_backend): - resp = client.get("/api/train/status") - - assert resp.status_code == 200 - body = resp.json() - - assert "metric_history" in body - mh = body["metric_history"] - assert mh is not None - assert mh["steps"] == [1, 2, 3, 4, 5] - assert mh["loss"] == [2.5, 2.0, 1.5, 1.2, 1.0] - assert mh["lr"] == [1e-4, 1e-4, 1e-4, 1e-4, 1e-4] - - def test_metric_history_null_when_no_history(self, client: TestClient): - mock_backend = _make_mock_backend(is_active=False) - with patch("routes.training.get_training_backend", return_value=mock_backend): - resp = client.get("/api/train/status") - - assert resp.status_code == 200 - body = resp.json() - assert body["metric_history"] is None - - def test_status_still_returns_phase_and_details(self, client: TestClient): - """Ensure adding metric_history didn't break existing fields.""" - mock_backend = _make_mock_backend( - is_active=True, - step_history=[5], - loss_history=[1.5], - lr_history=[1e-4], - total_steps=100, - ) - with patch("routes.training.get_training_backend", return_value=mock_backend): - resp = client.get("/api/train/status") - - body = resp.json() - assert body["phase"] == "training" - assert body["is_training_running"] is True - assert body["job_id"] == "job_test_001" - assert "details" in body