ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/core/training/trainer.py
Roland Tannous 21ef22a9ff fix: skip streaming when dataset_slice_start > dataset_slice_end 
Prevents training on the wrong row range when start exceeds end by
falling back to full download where existing clamping handles it.
2026-03-10 20:21:34 +00:00

2820 lines
128 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0
# Copyright © 2025 Unsloth AI
"""
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
import json
import threading
import math
import logging
import time
from pathlib import Path
from typing import Optional, Callable
from dataclasses import dataclass
import pandas as pd
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 trl import SFTTrainer, SFTConfig
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
_BACKEND_ROOT = Path(__file__).resolve().parents[2]
_ASSETS_DATASETS_ROOT = _BACKEND_ROOT / "assets" / "datasets"
@dataclass
class TrainingProgress:
"""Training progress tracking"""
epoch: float = 0
step: int = 0
total_steps: int = 0
loss: float = 0.0
learning_rate: float = 0.0
is_training: bool = False
is_completed: bool = False
error: Optional[str] = None
status_message: str = "Ready to train" # Current stage message
elapsed_seconds: Optional[float] = None
eta_seconds: Optional[float] = None
grad_norm: Optional[float] = None
num_tokens: Optional[int] = None
eval_loss: Optional[float] = None
class UnslothTrainer:
"""
Unsloth Training Backend
"""
def __init__(self):
self.model = None
self.tokenizer = None
self.trainer = None
self.training_thread = None
self.training_progress = TrainingProgress()
self.progress_callbacks = []
self.is_training = False
self.should_stop = False
self.save_on_stop = True
self.load_in_4bit = True # Track quantization mode for metadata
# 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._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.model_name = None
# Training metrics tracking
self.training_start_time: Optional[float] = None
self.batch_size: Optional[int] = None
self.max_seq_length: Optional[int] = None
self.gradient_accumulation_steps: Optional[int] = None
# Thread safety
self._lock = threading.Lock()
# Store training context for later transfer
self.training_context = {
'base_model_name': None,
'output_dir': None,
'is_lora': True, # Default to LoRA
}
def add_progress_callback(self, callback: Callable[[TrainingProgress], None]):
"""Add callback for training progress updates"""
self.progress_callbacks.append(callback)
def _update_progress(self, **kwargs):
"""Update training progress and notify callbacks"""
with self._lock:
for key, value in kwargs.items():
if hasattr(self.training_progress, key):
setattr(self.training_progress, key, value)
# Notify all callbacks
for callback in self.progress_callbacks:
try:
callback(self.training_progress)
except Exception as e:
logger.error(f"Error in progress callback: {e}")
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):
if not logs:
return
loss_value = logs.get('loss', logs.get('train_loss', 0.0))
current_step = state.global_step
grad_norm = logs.get('grad_norm', None)
elapsed_seconds = None
if trainer_ref.training_start_time is not None:
elapsed_seconds = time.time() - trainer_ref.training_start_time
eta_seconds = None
if elapsed_seconds is not None and current_step > 0:
total_steps = trainer_ref.training_progress.total_steps
if total_steps > 0:
steps_remaining = total_steps - current_step
if steps_remaining > 0:
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="",
)
def on_epoch_end(self, args, state, control, **kwargs):
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:
print(f"Stop detected at step {state.global_step}\n")
control.should_training_stop = True
return control
return _ProgressCallback()
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)
return steps_per_epoch * num_epochs
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')
config = {
"per_device_train_batch_size": batch_size,
"gradient_accumulation_steps": gradient_accumulation_steps,
"warmup_steps": warmup_steps_val if warmup_steps_val is not None else 5,
"learning_rate": learning_rate,
"fp16": not is_bfloat16_supported(),
"bf16": is_bfloat16_supported(),
"logging_steps": 1,
"optim": optim_value,
"weight_decay": weight_decay,
"lr_scheduler_type": lr_scheduler_type,
"seed": random_seed,
"output_dir": output_dir,
"report_to": ["wandb"] if training_args.get('enable_wandb', False) else "none",
}
# max_steps vs epochs
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)
# save_steps
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"
# Apply per-branch overrides
if extra_args:
config.update(extra_args)
return config
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()
self.tokenizer.save_pretrained(output_dir)
self._patch_adapter_config(output_dir)
msg = f"{label} training stopped" if label else "Training stopped"
print(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}",
)
elif self.should_stop:
msg = f"{label} training cancelled" if label else "Training cancelled"
print(f"\n{msg}.\n")
self._update_progress(is_training=False, status_message="Training cancelled.")
else:
self.trainer.save_model()
self.tokenizer.save_pretrained(output_dir)
self._patch_adapter_config(output_dir)
msg = f"{label} training completed" if label else "Training completed"
print(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}",
)
def _cleanup_audio_artifacts(self):
"""Remove sys.path entries and sys.modules from previous audio preprocessing.
After audio training, cloned repo dirs (OuteTTS, Spark-TTS) remain on
sys.path and heavy audio modules (snac, whisper, sparktts, outetts) stay
in sys.modules. When the next training run calls dataset.map(num_proc=N),
forked child processes inherit this stale state and deadlock.
"""
import sys as _sys
# Remove cloned audio repo paths from sys.path
base_dir = os.path.dirname(os.path.abspath(__file__))
audio_paths = [
os.path.join(base_dir, "inference", "OuteTTS"), # DAC/OuteTTS
]
# 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")
audio_paths.append(spark_code_dir)
removed_paths = []
for path in audio_paths:
if path in _sys.path:
_sys.path.remove(path)
removed_paths.append(path)
# Remove stale audio modules from sys.modules
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:
print(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.
Returns:
dict with keys: audio_col, text_col, speaker_col (speaker_col may be None)
"""
cols = dataset.column_names
if custom_format_mapping:
audio_col = None
text_col = None
speaker_col = None
for col, role in custom_format_mapping.items():
if role == "audio":
audio_col = col
elif role == "text":
text_col = col
elif role == "speaker_id":
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}
# 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)
speaker_col = None
if "source" in cols:
speaker_col = "source"
elif "speaker_id" in cols:
speaker_col = "speaker_id"
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:
"""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
try:
if self.model is not None:
del self.model
if self.tokenizer is not None:
del self.tokenizer
if self.trainer is not None:
del self.trainer
print("\nClearing GPU memory before training...")
clear_gpu_cache()
# Clean up sys.path and sys.modules from previous audio preprocessing
# to prevent deadlocks when forking worker processes in dataset.map()
self._cleanup_audio_artifacts()
# Reload Unsloth-patched transformers modeling modules before clearing
# the compiled cache. unsloth_compile_transformers() sets __UNSLOTH_PATCHED__
# on each modeling module and replaces methods with exec'd code.
# clear_unsloth_compiled_cache() deletes the disk cache, but the flag
# prevents re-compilation — leaving missing cache files. Reloading
# 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__'):
try:
importlib.reload(_mod)
except Exception:
pass # Non-critical — Unsloth will handle stale modules
# 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':
self.is_audio = False
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
self.is_audio_vlm = False
if not self.is_audio and not self.is_audio_vlm:
self._cuda_audio_used = False
# VLM: vision model with image dataset (mutually exclusive with audio paths)
vision = is_vision_model(model_name) if not self.is_audio else False
self.is_vlm = not self.is_audio_vlm and vision and is_dataset_image
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"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
)
# 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')
self._update_progress(
status_message=f"Loading {model_type_label} model... {model_display}"
)
print(f"\nLoading {model_type_label} model: {model_name}")
# Set HF token if provided
if hf_token:
os.environ["HF_TOKEN"] = hf_token
# 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
try:
from huggingface_hub import model_info as hf_model_info
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:
friendly = (
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)
return False
except Exception as gate_err:
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)
return False
# Branch based on model type
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,
)
logger.info("Loaded CSM audio model")
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,
)
# Configure generation settings (notebook lines 100-105)
self.model.generation_config.language = "<|en|>"
self.model.generation_config.task = "transcribe"
self.model.config.suppress_tokens = []
self.model.generation_config.forced_decoder_ids = None
logger.info("Loaded Whisper audio model (FastModel)")
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,
)
logger.info(f"Loaded {self._audio_type} audio model (FastLanguageModel)")
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:
# "Spark-TTS-0.5B/LLM" (local-style, from YAML mapping)
# "unsloth/Spark-TTS-0.5B" (HF repo ID)
from unsloth import FastModel
from huggingface_hub import snapshot_download
if model_name.endswith("/LLM"):
# "Spark-TTS-0.5B/LLM" → parent="Spark-TTS-0.5B"
local_dir = model_name.rsplit("/", 1)[0]
hf_repo = f"unsloth/{local_dir}"
llm_path = model_name
else:
# "unsloth/Spark-TTS-0.5B" → local_dir="Spark-TTS-0.5B"
hf_repo = model_name
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
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,
)
logger.info("Loaded Spark-TTS (bicodec) model")
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,
)
logger.info("Loaded OuteTTS (dac) model (FastModel)")
elif self.is_audio_vlm:
# 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,
)
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,
)
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")
print(f"\n[VLM Diagnostic] FastVisionModel returned: {type(tok).__name__}")
print(f"[VLM Diagnostic] Is ProcessorMixin: {isinstance(tok, ProcessorMixin)}")
print(f"[VLM Diagnostic] Has image_processor: {hasattr(tok, 'image_processor')}")
print(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,
)
logger.info("Loaded text model")
if self.should_stop:
return False
self._update_progress(status_message="Model loaded successfully")
print("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):
# 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
# second attempt because the failed first call's partial
# imports clean up the stale state as a side effect.
self._source_code_retried = True
print(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)
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")):
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)
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")):
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)
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:
"""
Prepare model for training (with optional LoRA).
"""
try:
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")
print("Full finetuning mode - training all parameters\n")
return True
# LoRA/QLoRA mode - apply PEFT
# "all-linear" is a PEFT keyword that targets every linear layer
if isinstance(target_modules, list) and "all-linear" in target_modules:
if len(target_modules) == 1:
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"]
# 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":
use_gradient_checkpointing = "unsloth"
elif use_gradient_checkpointing in ("true", "1", "yes"):
use_gradient_checkpointing = True
elif use_gradient_checkpointing in ("false", "0", "no"):
use_gradient_checkpointing = False
else:
# Invalid value, default 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'")
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)
return False
# Check if model has the expected attributes
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)
return False
print(f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n")
print(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:
# Models using FastModel.get_peft_model (codec audio + audio VLM)
from unsloth import FastModel
label = self._audio_type or 'audio_vlm'
print(f"{label} LoRA configuration:")
print(f" - Target modules: {target_modules}")
if self.is_audio_vlm:
print(f" - Finetune vision layers: {finetune_vision_layers}")
print(f" - Finetune language layers: {finetune_language_layers}")
print(f" - Finetune attention modules: {finetune_attention_modules}")
print(f" - Finetune MLP modules: {finetune_mlp_modules}")
print()
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,
)
# 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,
)
self.model = FastModel.get_peft_model(self.model, **peft_kwargs)
elif self._audio_type == 'whisper':
# Phase 2: Whisper uses FastModel.get_peft_model with task_type=None
from unsloth import FastModel
print(f"Audio model (whisper) LoRA configuration:")
print(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,
)
elif self._audio_type == 'snac':
# Orpheus uses FastLanguageModel.get_peft_model
print(f"Audio model ({self._audio_type}) LoRA configuration:")
print(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,
)
elif self.is_vlm:
# Vision model LoRA
print(f"Vision model LoRA configuration:")
print(f" - Finetune vision layers: {finetune_vision_layers}")
print(f" - Finetune language layers: {finetune_language_layers}")
print(f" - Finetune attention modules: {finetune_attention_modules}")
print(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,
)
else:
# Text model LoRA
print(f"Text model LoRA configuration:")
print(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,
)
# Check if stopped during LoRA preparation
if self.should_stop:
print("Stopped during LoRA configuration\n")
return False
self._update_progress(status_message="LoRA adapters configured")
print("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)"
full_traceback = traceback.format_exc()
logger.error(f"Error preparing model: {error_details}")
logger.error(f"Full traceback:\n{full_traceback}")
print(f"\n[ERROR] Error preparing model: {error_details}", file=sys.stderr, flush=True)
print(f"[ERROR] Full traceback:\n{full_traceback}", file=sys.stderr, flush=True)
self._update_progress(error=error_details)
return False
def _apply_csm_forward_fix(self):
"""Monkey-patch CsmForConditionalGeneration.forward to fix depth decoder kwargs.
The original transformers forward passes raw **kwargs (num_items_in_batch,
causal_mask, etc.) from the Trainer/PEFT through to the depth decoder,
causing depth_decoder_loss=None and 'Tensor + NoneType' crash.
We patch at both instance AND class level for maximum reliability,
and strip non-TransformersKwargs params that Unsloth/PEFT inject.
"""
import types
import torch
import torch.nn as nn
from transformers.models.csm.modeling_csm import (
CsmForConditionalGeneration,
CsmOutputWithPast,
)
base_csm = self.model.base_model.model # CsmForConditionalGeneration
# Save original forward (the @can_return_tuple wrapped version)
_original_forward = CsmForConditionalGeneration.forward
# 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',
}
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,
):
# 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)
# Only keep recognized TransformersKwargs
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(
input_ids, input_values, input_values_cutoffs, labels
)
inputs_embeds = merged["inputs_embeds"]
labels = merged["labels"]
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,
**clean_kwargs,
)
backbone_hidden_states = backbone_outputs[0]
slice_indices = (
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, :])
loss = None
backbone_loss = None
depth_decoder_loss = None
depth_decoder_outputs = None
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,
)
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
)
train_idxs = train_mask.nonzero(as_tuple=True)
backbone_last_hidden_states = backbone_hidden_states[
train_idxs[0], train_idxs[1] - 1, :
]
depth_decoder_labels = labels[train_mask]
# 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)
)
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,
**dd_kwargs,
)
depth_decoder_loss = depth_decoder_outputs.loss
if depth_decoder_loss is None:
logger.warning(
"CSM depth_decoder_loss is None! "
f"labels shape={depth_decoder_labels.shape}, "
f"train_mask sum={train_mask.sum().item()}"
)
# Fallback: use only backbone loss instead of crashing
loss = backbone_loss
else:
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=(
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_hidden_states=(
depth_decoder_outputs.hidden_states if depth_decoder_outputs else None
),
depth_decoder_attentions=(
depth_decoder_outputs.attentions if depth_decoder_outputs else None
),
)
# Patch at BOTH instance and class level for maximum reliability.
# Instance-level: catches calls via BaseTuner.forward -> self.model.forward()
base_csm.forward = types.MethodType(_fixed_csm_forward, base_csm)
# Class-level: catches any path that resolves through the class dict
CsmForConditionalGeneration.forward = _fixed_csm_forward
print("Applied CSM forward fix (class + instance level)\n")
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
import torch
processor = AutoProcessor.from_pretrained(
self.model_name,
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)
# Resolve columns from user mapping or hardcoded fallback
resolved = self._resolve_audio_columns(dataset, custom_format_mapping)
audio_col = resolved["audio_col"]
text_col = resolved["text_col"]
speaker_key = resolved["speaker_col"]
if audio_col is None:
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}")
if speaker_key is None:
print("No speaker found, adding default 'source' of 0 for all examples\n")
dataset = dataset.add_column("source", ["0"] * len(dataset))
speaker_key = "source"
print(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))
required_keys = ["input_ids", "attention_mask", "labels", "input_values", "input_values_cutoffs"]
self._update_progress(status_message="Preprocessing CSM dataset...")
processed_examples = []
skipped = 0
for idx in range(len(dataset)):
if self.should_stop:
print("Stopped during CSM preprocessing\n")
break
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"]},
],
}]
# 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={
"padding": "max_length",
"max_length": 256,
"padding_side": "right",
},
audio_kwargs={
"sampling_rate": 24_000,
"max_length": 240001,
"padding": "max_length",
},
common_kwargs={"return_tensors": "pt"},
)
out = {}
for k in required_keys:
if k not in model_inputs:
raise KeyError(f"Missing required key '{k}' in model outputs")
out[k] = model_inputs[k][0]
if not all(isinstance(out[k], torch.Tensor) for k in out):
skipped += 1
continue
processed_examples.append(out)
except Exception as e:
logger.warning(f"Error processing CSM example {idx}: {e}")
skipped += 1
continue
if (idx + 1) % 100 == 0:
self._update_progress(
status_message=f"Preprocessing CSM... {idx + 1}/{len(dataset)}"
)
if not processed_examples:
raise ValueError(
f"No valid examples after CSM preprocessing (skipped {skipped})"
)
result_dataset = Dataset.from_list(processed_examples)
print(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):
"""Format dataset as audio chat messages for multimodal models (e.g. Gemma 3N).
Expects columns: audio (Audio), text (str).
Produces: messages column with system/user/assistant chat format.
"""
from datasets import Audio
resolved = self._resolve_audio_columns(dataset, custom_format_mapping)
audio_col = resolved["audio_col"]
text_col = resolved["text_col"]
if not audio_col or not text_col:
raise ValueError(
f"Audio VLM dataset needs 'audio' and 'text' columns, got: {dataset.column_names}"
)
# Store resolved audio column name for the collator closure
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))
def format_messages(samples):
formatted = {"messages": []}
for idx in range(len(samples[audio_col])):
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}
]},
]
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))
print(f"Audio VLM dataset formatted: {len(dataset)} examples\n")
return dataset
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
layers), interleave 7 codes per frame, wrap with Orpheus special tokens,
train on full sequence (no label masking).
"""
import torch
import torchaudio.transforms as T
SNAC_MODEL_NAME = "hubertsiuzdak/snac_24khz"
SNAC_SAMPLE_RATE = 24000
device = "cuda" if torch.cuda.is_available() else "cpu"
max_length = self.max_seq_length or 2048
tokenizer = self.tokenizer
# Orpheus special token IDs (hardcoded in tokenizer vocabulary)
START_OF_HUMAN = 128259
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
resolved = self._resolve_audio_columns(dataset, custom_format_mapping)
audio_col = resolved["audio_col"]
text_col = resolved["text_col"]
speaker_col = resolved["speaker_col"]
has_source = speaker_col is not None
if not audio_col or not text_col:
raise ValueError(
f"SNAC dataset needs 'audio' and 'text' columns, got: {dataset.column_names}"
)
# 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))
# 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
# Load SNAC codec model
self._update_progress(status_message="Loading SNAC codec model...")
print("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
self._update_progress(status_message="Encoding audio with SNAC...")
print(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
for idx in range(len(dataset)):
if self.should_stop:
print("Stopped during SNAC preprocessing\n")
break
example = dataset[idx]
try:
text = example.get(text_col)
if not text:
skipped += 1
continue
audio_data = example.get(audio_col)
if audio_data is None or audio_data.get("array") is None:
skipped += 1
continue
# --- Encode audio with SNAC (notebook lines 122-142) ---
waveform = torch.from_numpy(audio_data["array"]).unsqueeze(0).to(dtype=torch.float32)
if resample_transform is not None:
waveform = resample_transform(waveform)
waveform = waveform.unsqueeze(0).to(device)
with torch.inference_mode():
codes = snac_model.encode(waveform)
# Interleave 7 codes per frame with layer offsets (notebook lines 134-142)
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))
if len(all_codes) == 0:
skipped += 1
continue
# Deduplicate consecutive frames with same first code (notebook lines 185-207)
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])
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_ids.append(END_OF_TEXT)
# --- Build full input_ids (notebook lines 225-234) ---
input_ids = (
[START_OF_HUMAN]
+ text_ids
+ [END_OF_HUMAN]
+ [START_OF_AI]
+ [START_OF_SPEECH]
+ all_codes
+ [END_OF_SPEECH]
+ [END_OF_AI]
)
# Truncate to max_length
input_ids = input_ids[:max_length]
# Labels = input_ids (no masking — Orpheus trains on full sequence)
labels = list(input_ids)
attention_mask = [1] * len(input_ids)
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}")
skipped += 1
continue
# Progress update every 100 examples
if (idx + 1) % 100 == 0:
self._update_progress(
status_message=f"Encoding audio... {idx + 1}/{len(dataset)}"
)
# Free SNAC model from GPU
print("Freeing SNAC codec model from GPU...\n")
snac_model.to("cpu")
del snac_model
import gc
gc.collect()
torch.cuda.empty_cache()
self._cuda_audio_used = True
if not processed_examples:
raise ValueError(
f"No valid examples after SNAC preprocessing (skipped {skipped})"
)
result_dataset = Dataset.from_list(processed_examples)
print(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):
"""Preprocess dataset for Spark-TTS training with BiCodec tokenizer.
Mirrors Spark_TTS_(0_5B).ipynb: encode audio with BiCodec (semantic + global tokens),
format as special-token text strings for SFTTrainer with dataset_text_field="text".
"""
import sys
import torch
import numpy as np
import torchaudio.transforms as T
import subprocess
device = "cuda" if torch.cuda.is_available() else "cpu"
# 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")
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...")
print(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,
)
if spark_code_dir not in sys.path:
sys.path.insert(0, spark_code_dir)
from sparktts.models.audio_tokenizer import BiCodecTokenizer
from sparktts.utils.audio import audio_volume_normalize
# Resolve audio and text columns
resolved = self._resolve_audio_columns(dataset, custom_format_mapping)
audio_col = resolved["audio_col"]
text_col = resolved["text_col"]
speaker_col = resolved["speaker_col"]
has_source = speaker_col is not None
if not audio_col or not text_col:
raise ValueError(
f"BiCodec dataset needs 'audio' and 'text' columns, got: {dataset.column_names}"
)
# 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...")
print("Loading BiCodec tokenizer...\n")
audio_tokenizer = BiCodecTokenizer(self._spark_tts_repo_dir, device)
target_sr = audio_tokenizer.config['sample_rate']
self._update_progress(status_message="Encoding audio with BiCodec...")
print(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)."""
if wavs.shape[0] != 1:
raise ValueError(f"Expected batch size 1, but got shape {wavs.shape}")
wav_np = wavs.squeeze(0).cpu().numpy()
processed = audio_tokenizer.processor(
wav_np,
sampling_rate=16000,
return_tensors="pt",
padding=True,
)
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:
raise ValueError("Wav2Vec2Model did not return hidden states.")
feats_mix = (
model_output.hidden_states[11]
+ model_output.hidden_states[14]
+ model_output.hidden_states[16]
) / 3
return feats_mix
processed_examples = []
skipped = 0
for idx in range(len(dataset)):
if self.should_stop:
print("Stopped during BiCodec preprocessing\n")
break
example = dataset[idx]
try:
text = example.get(text_col)
if not text:
skipped += 1
continue
audio_data = example.get(audio_col)
if audio_data is None or audio_data.get("array") is None:
skipped += 1
continue
audio_array = audio_data["array"]
sampling_rate = audio_data.get("sampling_rate", target_sr)
# Resample if needed
if sampling_rate != 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()
# Volume normalize if configured
if audio_tokenizer.config.get("volume_normalize", False):
audio_array = audio_volume_normalize(audio_array)
# Get reference clip
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)
# Extract wav2vec2 features
feat = extract_wav2vec2_features(audio_tensor)
batch = {
"wav": audio_tensor,
"ref_wav": ref_wav_tensor,
"feat": feat.to(device),
}
# BiCodec tokenize
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()]
)
semantic_tokens = "".join(
[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
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})
except Exception as e:
logger.warning(f"Error processing BiCodec example {idx}: {e}")
skipped += 1
continue
# Progress update every 100 examples
if (idx + 1) % 100 == 0:
self._update_progress(
status_message=f"Encoding audio with BiCodec... {idx + 1}/{len(dataset)}"
)
# Free BiCodec model from GPU
print("Freeing BiCodec tokenizer from GPU...\n")
audio_tokenizer.model.cpu()
audio_tokenizer.feature_extractor.cpu()
del audio_tokenizer
import gc
gc.collect()
torch.cuda.empty_cache()
self._cuda_audio_used = True
if not processed_examples:
raise ValueError(
f"No valid examples after BiCodec preprocessing (skipped {skipped})"
)
result_dataset = Dataset.from_list(processed_examples)
print(f"BiCodec preprocessing complete: {len(result_dataset)} examples "
f"({skipped} skipped)\n")
# Debug: show first example text (truncated)
sample = result_dataset[0]["text"]
print(f"Sample text (first 200 chars): {sample[:200]}...\n")
print(f"Sample text length: {len(sample)} chars\n")
return result_dataset
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,
OuteTTS AudioProcessor for speaker representations, PromptProcessor for
training prompts. Outputs text strings for SFTTrainer with dataset_text_field="text".
"""
import sys
import io
import tempfile
import torch
import numpy as np
import soundfile as sf
from datasets import Dataset as HFDataset
device = "cuda" if torch.cuda.is_available() else "cpu"
# 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...")
print(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,
)
for fpath in [
os.path.join(outetts_pkg, "models", "gguf_model.py"),
os.path.join(outetts_pkg, "interface.py"),
os.path.join(outetts_pkg, "__init__.py"),
]:
if os.path.exists(fpath):
os.remove(fpath)
print(f"Removed {fpath}\n")
if outetts_code_dir not in sys.path:
sys.path.insert(0, outetts_code_dir)
from outetts.version.v3.audio_processor import AudioProcessor
from outetts.version.v3.prompt_processor import PromptProcessor
from outetts.models.config import ModelConfig as OuteTTSModelConfig
from outetts.utils.preprocessing import text_normalizations
# Resolve audio and text columns
resolved = self._resolve_audio_columns(dataset, custom_format_mapping)
audio_col = resolved["audio_col"]
text_col = resolved["text_col"]
if not audio_col or not text_col:
raise ValueError(
f"DAC dataset needs 'audio' and 'text' columns, got: {dataset.column_names}"
)
# 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))
print("Cast audio column to 24kHz\n")
# Load Whisper for word timings
self._update_progress(status_message="Loading Whisper model for word timings...")
print("Loading Whisper model for word timings...\n")
import whisper
whisper_model = whisper.load_model("turbo", device=device)
# Load OuteTTS AudioProcessor + PromptProcessor
self._update_progress(status_message="Loading OuteTTS AudioProcessor...")
print("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,
)
audio_processor = AudioProcessor(config=dummy_config)
prompt_processor = PromptProcessor(model_tokenizer_path)
self._update_progress(status_message="Preprocessing audio with OuteTTS...")
print(f"DAC preprocessing: audio_col='{audio_col}', text_col='{text_col}'\n")
processed_examples = []
skipped = 0
for idx in range(len(dataset)):
if self.should_stop:
print("Stopped during DAC preprocessing\n")
break
example = dataset[idx]
try:
text = example.get(text_col)
if not text or not isinstance(text, str):
skipped += 1
continue
audio_data = example.get(audio_col)
if audio_data is None or audio_data.get("array") is None:
skipped += 1
continue
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")
buf.seek(0)
audio_bytes = buf.getvalue()
# 1. Get word timings from Whisper
with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as tmp:
tmp.write(audio_bytes)
tmp.flush()
whisper_result = whisper_model.transcribe(tmp.name, word_timestamps=True)
normalized_transcript = text_normalizations(text)
words_with_timings = []
if whisper_result and "segments" in whisper_result:
for segment in whisper_result["segments"]:
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"]),
})
if not words_with_timings:
skipped += 1
continue
# 2. Create speaker representation with AudioProcessor
speaker_data_dict = {
"audio": {"bytes": audio_bytes},
"text": normalized_transcript,
"words": words_with_timings,
}
speaker = audio_processor.create_speaker_from_dict(speaker_data_dict)
if speaker is None:
skipped += 1
continue
# 3. Get training prompt from PromptProcessor
prompt = prompt_processor.get_training_prompt(speaker)
if prompt:
processed_examples.append({"text": prompt})
except Exception as e:
logger.warning(f"Error processing DAC example {idx}: {e}")
skipped += 1
continue
if (idx + 1) % 100 == 0:
self._update_progress(
status_message=f"Preprocessing audio with OuteTTS... {idx + 1}/{len(dataset)}"
)
# Free Whisper from GPU (notebook: data_processor.whisper_model.to('cpu'))
print("Moving Whisper model to CPU...\n")
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
if not processed_examples:
raise ValueError(
f"No valid examples after DAC preprocessing (skipped {skipped})"
)
result_dataset = HFDataset.from_list(processed_examples)
print(f"DAC preprocessing complete: {len(result_dataset)} examples "
f"({skipped} skipped)\n")
sample = result_dataset[0]["text"]
print(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):
"""Preprocess dataset for Whisper speech-to-text training.
Mirrors Whisper.ipynb: extract audio features with Whisper's feature
extractor, tokenize text labels. Returns (train_data, eval_data) where
each is a list of dicts with 'input_features' and 'labels'.
"""
from datasets import Audio
WHISPER_SAMPLE_RATE = 16000
resolved = self._resolve_audio_columns(dataset, custom_format_mapping)
audio_col = resolved["audio_col"]
text_col = resolved["text_col"]
if not audio_col or not text_col:
raise ValueError(
f"Whisper dataset needs 'audio' and 'text' columns, got: {dataset.column_names}"
)
# Cast audio to 16kHz (Whisper's expected 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)
dataset = splits["train"]
eval_dataset_raw = splits["test"]
self._update_progress(status_message="Processing audio for Whisper...")
print(f"Whisper preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
f"samples={len(dataset)}\n")
def process_split(ds, split_name="train"):
processed = []
skipped = 0
for idx in range(len(ds)):
if self.should_stop:
print(f"Stopped during Whisper {split_name} preprocessing\n")
break
example = ds[idx]
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:
skipped += 1
continue
# Extract audio features (notebook line 112-115)
features = self.tokenizer.feature_extractor(
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,
})
except Exception as 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)}"
)
print(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
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]:
"""
Load and prepare dataset for training.
Strategy: format first, then split — ensures both train and eval
portions are properly formatted and templated.
Returns:
Tuple of (dataset_info, eval_dataset) or None on error.
eval_dataset may be None if no eval split is available.
"""
try:
dataset = None
eval_dataset = None
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:
# Load local datasets using load_dataset() so the result is
# Arrow-backed (has cache files). Dataset.from_list() creates
# an in-memory dataset with no cache, which forces num_proc=1
# during tokenization/map because sharding requires Arrow files.
all_files: list[str] = []
for dataset_file in local_datasets:
# dataset_file may already be an absolute path from routes/training.py
if os.path.isabs(dataset_file):
file_path = dataset_file
else:
# Fallback: try relative to assets/datasets
file_path = str(_ASSETS_DATASETS_ROOT / dataset_file)
file_path_obj = Path(file_path)
if file_path_obj.is_dir():
parquet_dir = (
file_path_obj / "parquet-files"
if (file_path_obj / "parquet-files").exists()
else file_path_obj
)
parquet_files = sorted(parquet_dir.glob("*.parquet"))
if parquet_files:
all_files.extend(str(p) for p in parquet_files)
continue
# Fall through to single-file detection for dirs with json/csv
candidates: list[Path] = []
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}")
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'
else:
raise ValueError(f"Unsupported local dataset format: {all_files[0]}")
dataset = load_dataset(loader, data_files=all_files, split='train')
# Check if stopped during dataset loading
if self.should_stop:
print("Stopped during dataset loading\n")
return None
self._update_progress(status_message=f"Loaded {len(dataset)} samples from local files")
print(f"Loaded {len(dataset)} samples from local files\n")
print(f"[DEBUG] Dataset cache_files: {dataset.cache_files}\n")
elif dataset_source:
# Load from Hugging Face
split_name = train_split or "train"
load_kwargs = {"path": dataset_source, "split": split_name}
if subset:
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):
# Manual slice — stream only the rows we need instead of
# downloading the entire dataset.
rows_to_stream = dataset_slice_end + 1
print(
f"[dataset-slice] Manual slice specified "
f"(start={dataset_slice_start}, end={dataset_slice_end}), "
f"streaming {rows_to_stream} rows\n"
)
stream = load_dataset(**load_kwargs, streaming=True)
dataset = Dataset.from_list(list(stream.take(rows_to_stream)))
print(
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"
)
else:
dataset = load_dataset(**load_kwargs)
# Check if stopped during dataset loading
if self.should_stop:
print("Stopped during dataset loading\n")
return None
self._update_progress(status_message=f"Loaded dataset from HuggingFace: {dataset_source}")
print(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:
effective_train = train_split or "train"
if eval_split and eval_split != effective_train:
# Explicit eval split provided - load it directly
print(f"Loading explicit eval split: '{eval_split}'\n")
eval_load_kwargs = {"path": dataset_source, "split": eval_split}
if subset:
eval_load_kwargs["name"] = subset
eval_dataset = load_dataset(**eval_load_kwargs)
has_separate_eval_source = True
print(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
print(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,
)
if eval_dataset is not None:
has_separate_eval_source = True
else:
print("Eval disabled (eval_steps <= 0), skipping eval split detection\n")
if dataset is None:
raise ValueError("No dataset provided")
# Apply index range slicing if requested (inclusive on both ends)
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
# 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))
print(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:
print("Stopped before applying chat template\n")
return None
# ========== AUDIO MODELS: custom preprocessing ==========
if self._audio_type == 'csm':
processed = self._preprocess_csm_dataset(dataset, custom_format_mapping)
return (processed, None)
elif self._audio_type == 'whisper':
train_data, eval_data = self._preprocess_whisper_dataset(
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)
return (processed, None)
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':
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)
return (formatted, None)
# ========== FORMAT FIRST ==========
print(f"Formatting dataset with format_type='{format_type}'...\n")
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,
)
# Check if stopped during formatting
if self.should_stop:
print("Stopped during dataset formatting\n")
return None
# Abort if dataset formatting/conversion failed
if not dataset_info.get("success", True):
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)
return None
self._update_progress(status_message=f"Dataset formatted and ready for training")
print(f"Dataset formatted successfully\n")
# ========== THEN SPLIT ==========
if has_separate_eval_source and eval_dataset is not None:
# Eval came from a separate HF split — format it too
print(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,
)
eval_dataset = eval_info["dataset"]
print(f"Eval dataset formatted successfully\n")
elif eval_enabled and not has_separate_eval_source:
# No separate eval source — split the already-formatted dataset
formatted_dataset = dataset_info["dataset"]
split_result = self._resolve_eval_split_from_dataset(formatted_dataset)
if split_result is not None:
train_portion, eval_dataset = split_result
dataset_info["dataset"] = train_portion
return (dataset_info, eval_dataset)
except Exception as e:
logger.error(f"Error loading dataset: {e}")
self._update_progress(error=str(e))
return None
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
available_splits = get_dataset_split_names(**load_kwargs)
print(f"Available splits: {available_splits}\n")
# Check for common eval split names
for candidate in ["eval", "validation", "valid", "val", "test"]:
if candidate in available_splits:
eval_load_kwargs = {"path": dataset_source, "split": candidate}
if subset:
eval_load_kwargs["name"] = subset
candidate_ds = load_dataset(**eval_load_kwargs)
if len(candidate_ds) >= 16:
print(f"Auto-detected eval split '{candidate}' with {len(candidate_ds)} rows\n")
return candidate_ds
else:
print(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}")
# No separate HF eval split found — caller will handle programmatic splitting
return None
def _resolve_eval_split_from_dataset(self, dataset) -> Optional[tuple]:
"""Split a dataset into train and eval portions.
Returns:
Tuple of (train_dataset, eval_dataset), or None if dataset too small.
"""
MIN_EVAL_ROWS = 16
MIN_TOTAL_ROWS = 32 # Need at least 16 train + 16 eval
n = len(dataset)
if n < MIN_TOTAL_ROWS:
print(f"Dataset too small ({n} rows) for eval split, skipping eval\n")
return None
eval_size = max(MIN_EVAL_ROWS, min(128, int(0.05 * n)))
# Ensure we don't take more than half the dataset
eval_size = min(eval_size, n // 2)
print(f"Auto-splitting: {eval_size} rows for eval from {n} total\n")
split_result = dataset.train_test_split(test_size=eval_size, seed=3407)
print(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 = "./outputs",
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 = "runs",
**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")
return False
# Pre-import heavy transformers modules on the main thread.
# Unsloth's patched_import hook (deepseek_v3_moe.py) is not thread-safe
# with Python's importlib cache, causing KeyError: 'size' if these are
# first imported inside the worker thread.
import transformers # noqa: F401 – ensures submodules are cached
from transformers import ( # noqa: F401
Trainer as _HFTrainer,
TrainingArguments as _TrainingArguments,
TrainerCallback as _TrainerCallback,
)
if self._audio_type == 'whisper':
from transformers import ( # noqa: F401
Seq2SeqTrainer as _Seq2SeqTrainer,
Seq2SeqTrainingArguments as _Seq2SeqTrainingArguments,
)
# 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
}
)
self.should_stop = False
self.is_training = True
try:
self.training_thread.start()
return True
except Exception as e:
self.is_training = False
logger.error(f"Failed to start training thread: {e}")
return False
def _train_worker(self, dataset: Dataset, **training_args):
"""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)
# Set training start time
self.training_start_time = time.time()
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']
import wandb
wandb.init(project=training_args.get('wandb_project', 'unsloth-training'))
# Create output directory
output_dir = training_args.get('output_dir', './outputs')
os.makedirs(output_dir, exist_ok=True)
# ========== AUDIO TRAINER BRANCH ==========
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,
})
self.trainer = HFTrainer(
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)
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),
)
self._update_progress(total_steps=total, status_message="Starting CSM training...")
print(f"CSM training config: {config}\n")
self.trainer.train()
self._finalize_training(output_dir, "CSM")
return
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
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,
),
)
self.trainer.add_callback(self._create_progress_callback())
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),
)
self._update_progress(total_steps=total, status_message="Starting SNAC training...")
print(f"SNAC training config: {config}\n")
self.trainer.train()
self._finalize_training(output_dir, "SNAC")
return
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)
extra = {"remove_unused_columns": False, "label_names": ["labels"]}
if eval_dataset:
extra["eval_strategy"] = "steps"
extra["eval_steps"] = training_args.get('eval_steps', 5)
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),
"processing_class": self.tokenizer.feature_extractor,
"args": Seq2SeqTrainingArguments(**config),
}
if eval_dataset:
trainer_kwargs["eval_dataset"] = eval_dataset
self.trainer = Seq2SeqTrainer(**trainer_kwargs)
self.trainer.add_callback(self._create_progress_callback())
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),
)
self._update_progress(total_steps=total, status_message="Starting Whisper training...")
print(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'):
# bicodec/dac use the standard SFTTrainer text path below
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
print("Configuring data collator...\n")
data_collator = None # Default to built-in data collator
if is_deepseek_ocr:
# Special DeepSeek OCR collator - auto-install if needed
print("Detected DeepSeek OCR model\n")
# Ensure DeepSeek OCR module is installed
if not _ensure_deepseek_ocr_installed():
error_msg = (
"Failed to install DeepSeek OCR module. "
"Please install manually: "
"from huggingface_hub import snapshot_download; "
"snapshot_download('unsloth/DeepSeek-OCR', local_dir='deepseek_ocr')"
)
logger.error(error_msg)
self._update_progress(error=error_msg, is_training=False)
return
try:
from backend.data_utils import DeepSeekOCRDataCollator
print("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),
)
print("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)
return
elif self.is_audio_vlm:
# Audio VLM collator (e.g. Gemma 3N with audio data)
# Mirrors the collate_fn from Gemma3N_(4B)-Audio notebook
print("Configuring audio VLM data collator...\n")
processor = self.tokenizer # FastModel returns processor as tokenizer
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
).strip()
texts.append(text)
audios.append(example[audio_col_name]["array"])
batch = processor(
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'):
token_id = getattr(processor.tokenizer, attr, None)
if token_id is not None:
labels[labels == token_id] = -100
batch["labels"] = labels
return batch
data_collator = audio_vlm_collate_fn
print("Audio VLM data collator configured\n")
elif self.is_vlm:
# Standard VLM collator (images)
print("Using UnslothVisionDataCollator for vision model\n")
from unsloth.trainer import UnslothVisionDataCollator
FastVisionModel.for_training(self.model)
data_collator = UnslothVisionDataCollator(self.model, self.tokenizer)
print("Vision data collator configured\n")
# ========== 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)
print(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
"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),
"output_dir": output_dir,
"report_to": ["wandb"] if training_args.get('enable_wandb', False) else "none",
"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),
}
print(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
print(f"Using warmup_ratio: {warmup_ratio_val}\n")
elif warmup_steps_val is not None:
config_args["warmup_steps"] = warmup_steps_val
print(f"Using warmup_steps: {warmup_steps_val}\n")
else:
# Default to warmup_steps if neither provided
config_args["warmup_steps"] = 5
print(f"Using default warmup_steps: 5\n")
# Add save_steps if specified
save_steps_val = training_args.get('save_steps', 0)
if save_steps_val and save_steps_val > 0:
config_args["save_steps"] = save_steps_val
config_args["save_strategy"] = "steps"
# If max_steps is specified, use it instead of epochs
max_steps_val = training_args.get('max_steps', 0)
if max_steps_val and max_steps_val > 0:
del config_args["num_train_epochs"] # Remove epochs
config_args["max_steps"] = max_steps_val # Use steps instead
print(f"Training for {max_steps_val} steps\n")
else:
print(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)
if eval_dataset is not None:
if eval_steps_val > 0:
config_args["eval_strategy"] = "steps"
config_args["eval_steps"] = eval_steps_val
print(f"✅ Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n")
print(f"Eval dataset: {len(eval_dataset)} rows\n")
else:
print(f"⚠️ Eval dataset provided but eval_steps={eval_steps_val} (disabled)\n")
print("To enable evaluation, set eval_steps > 0.0\n")
else:
print("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")
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"
print(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),
})
else:
print("Configuring text model training parameters\n")
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)
config_args["packing"] = packing_enabled
print(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':
config_args["packing"] = False
print("Applied BiCodec overrides: packing=False\n")
elif self._audio_type == 'dac':
config_args["packing"] = False
print("Applied DAC overrides: packing=False\n")
print(f"The configuration is: {config_args}")
print("Training configuration prepared\n")
# ========== TRAINER INITIALIZATION ==========
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
trainer_kwargs = {
"model": self.model,
"train_dataset": train_dataset,
"processing_class": processing_class,
"data_collator": data_collator,
"args": SFTConfig(**config_args),
}
if eval_dataset is not None:
trainer_kwargs["eval_dataset"] = eval_dataset
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
trainer_kwargs = {
"model": self.model,
"train_dataset": train_dataset,
"processing_class": self.tokenizer,
"data_collator": data_collator,
"args": SFTConfig(**config_args),
}
if eval_dataset is not None:
trainer_kwargs["eval_dataset"] = eval_dataset
self.trainer = SFTTrainer(**trainer_kwargs)
else:
# For text-only training, if the tokenizer is actually a Processor
# (e.g., Gemma-3 returns a ProcessorMixin even for text), we must
# unwrap to the raw tokenizer. Otherwise Unsloth's SFTTrainer detects
# 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'):
print(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'],
"data_collator": data_collator,
"args": SFTConfig(**config_args),
}
if eval_dataset is not None:
trainer_kwargs["eval_dataset"] = eval_dataset
self.trainer = SFTTrainer(**trainer_kwargs)
print("Trainer initialized\n")
# ========== TRAIN ON RESPONSES ONLY ==========
# 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)
# 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'):
try:
print("Configuring train on responses only...\n")
# Get the template mapping for this model
model_name_lower = self.model_name.lower()
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
print(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"]
print(f"Instruction marker: {instruction_part[:50]}...\n")
print(f"Response marker: {response_part[:50]}...\n")
else:
print(f"No response mapping found for template: {template_name}\n")
train_on_responses_enabled = False
else:
print(f"No template mapping found for model: {self.model_name}\n")
train_on_responses_enabled = False
except Exception as e:
logger.warning(f"Could not configure train on responses: {e}")
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'):
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"],
)
print("Train on responses only configured successfully\n")
# ── Safety net: check if all samples were filtered out ──
# Unsloth's train_on_responses_only masks non-response
# tokens with -100. If max_seq_length is too short and the
# response portion gets truncated away, EVERY sample ends
# up with all labels == -100 and Unsloth removes them,
# leaving 0 usable training samples.
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
if filtered_len == 0 or drop_pct > 30:
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 "
f"only' — only {filtered_len} remain. This usually "
f"means max_seq_length ({max_seq}) is too short "
f"and the response portion is being truncated "
f"away. Try increasing max_seq_length (e.g. 8192) "
f"or disabling 'Train on completions'."
)
logger.error(error_msg)
self._update_progress(error=error_msg, is_training=False)
return
if dropped > 0:
print(
f"⚠️ {dropped}/{original_len} samples "
f"({drop_pct}%) were dropped (all labels "
f"masked). {filtered_len} samples remain.\n"
)
print(f"Post-filter dataset size: {filtered_len} samples\n")
except Exception as e:
logger.warning(f"Failed to apply train on responses only: {e}")
train_on_responses_enabled = False
else:
if train_on_responses_enabled and is_deepseek_ocr:
print("Train on responses handled by DeepSeek OCR collator\n")
else:
print("Training on full sequences (including prompts)\n")
# ========== 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),
)
self._update_progress(total_steps=total_steps)
# ========== START TRAINING ==========
self._update_progress(status_message="Starting training...")
print("Starting training...\n")
self.trainer.train()
# ========== SAVE MODEL ==========
self._finalize_training(output_dir)
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))
finally:
self.is_training = False
def _patch_adapter_config(self, output_dir: str) -> None:
"""Patch adapter_config.json with unsloth_training_method.
Values: 'qlora', 'lora', 'FT', 'CPT', 'DPO', 'GRPO', etc.
For LoRA/QLoRA, the distinction comes from load_in_4bit.
"""
config_path = os.path.join(output_dir, "adapter_config.json")
if not os.path.exists(config_path):
logger.info("No adapter_config.json found — skipping training method patch")
return
try:
with open(config_path, "r") as f:
config = json.load(f)
# Determine the training method
if self.load_in_4bit:
method = "qlora"
else:
method = "lora"
config["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)
except Exception as e:
logger.warning(f"Failed to patch adapter_config.json: {e}")
def stop_training(self, save: bool = True):
"""Stop ongoing training"""
print(f"\nStopping training (save={save})...")
self.should_stop = True
self.save_on_stop = save
stop_msg = (
"Stopping training and saving checkpoint..."
if save
else "Cancelling training..."
)
self._update_progress(status_message=stop_msg)
# If trainer exists, try to stop it gracefully
if self.trainer:
try:
# The callback will catch should_stop flag and stop the training loop
print("Training will stop at next step...\n")
except Exception as e:
logger.error(f"Error stopping trainer: {e}")
def get_training_progress(self) -> TrainingProgress:
"""Get current training progress"""
with self._lock:
return self.training_progress
def cleanup(self):
"""Cleanup resources"""
if self.trainer:
self.trainer = None
if self.model:
self.model = None
if self.tokenizer:
self.tokenizer = None
# Clear GPU memory
clear_gpu_cache()
def _ensure_deepseek_ocr_installed():
"""
Auto-install DeepSeek OCR module if not available.
Downloads from HuggingFace hub as a local module.
Returns:
bool: True if available (either already installed or just 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...")
print("\n Downloading DeepSeek OCR module from HuggingFace...\n")
from huggingface_hub import snapshot_download
import sys
import os
# Get the script directory to install locally
script_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(script_dir) # Go up to project root
# Download to project root as 'deepseek_ocr' folder
local_dir = os.path.join(parent_dir, "deepseek_ocr")
snapshot_download(
"unsloth/DeepSeek-OCR",
local_dir=local_dir,
local_dir_use_symlinks=False
)
# Add to sys.path if not already there
if parent_dir not in sys.path:
sys.path.insert(0, parent_dir)
# Try importing again
from deepseek_ocr.modeling_deepseekocr import format_messages
logger.info("DeepSeek OCR module installed successfully")
print("DeepSeek OCR module installed successfully!\n")
return True
except Exception as e:
logger.error(f"Failed to install DeepSeek OCR module: {e}")
print(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
if _trainer_instance is None:
_trainer_instance = UnslothTrainer()
return _trainer_instance
Reference in a new issue View git blame Copy permalink