@@ -65,32 +65,32 @@ def show_link(port: int = 8000):
def start(port: int = 8000):
"""
Start Unsloth Studio server in Colab and display the URL.
-
+
Usage:
from colab import start
start()
"""
import sys
-
+
logger.info("π¦₯ Starting Unsloth Studio...")
-
+
logger.info(" Loading backend...")
from run import run_server
-
+
# Auto-detect frontend path
repo_root = Path(__file__).parent.parent
frontend_path = repo_root / "frontend" / "dist"
-
+
if not frontend_path.exists():
logger.info("β Frontend not built! Please run the setup cell first.")
return
-
+
logger.info(" Starting server...")
# Start server silently
- run_server(host="0.0.0.0", port=port, frontend_path=frontend_path, silent=True)
-
+ run_server(host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True)
+
logger.info(" Server started!")
-
+
# Show the clickable link with real URL
show_link(port)
diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py
index 84868cea0d..051f6615a7 100644
--- a/studio/backend/core/__init__.py
+++ b/studio/backend/core/__init__.py
@@ -12,101 +12,123 @@ code has a chance to run.
__all__ = [
# Inference
- 'InferenceBackend',
- 'get_inference_backend',
-
+ "InferenceBackend",
+ "get_inference_backend",
# Training
- 'get_training_backend',
- 'TrainingBackend',
- 'TrainingProgress',
-
+ "get_training_backend",
+ "TrainingBackend",
+ "TrainingProgress",
# Config
- 'ModelConfig',
- 'is_vision_model',
- 'scan_trained_loras',
- 'load_model_defaults',
- 'get_base_model_from_lora',
-
+ "ModelConfig",
+ "is_vision_model",
+ "scan_trained_loras",
+ "load_model_defaults",
+ "get_base_model_from_lora",
# Utils
- 'format_and_template_dataset',
- 'normalize_path',
- 'is_local_path',
- 'is_model_cached',
- 'without_hf_auth',
- 'format_error_message',
- 'get_gpu_memory_info',
- 'log_gpu_memory',
- 'get_device',
- 'is_apple_silicon',
- 'clear_gpu_cache',
- 'DeviceType',
+ "format_and_template_dataset",
+ "normalize_path",
+ "is_local_path",
+ "is_model_cached",
+ "without_hf_auth",
+ "format_error_message",
+ "get_gpu_memory_info",
+ "log_gpu_memory",
+ "get_device",
+ "is_apple_silicon",
+ "clear_gpu_cache",
+ "DeviceType",
]
def __getattr__(name):
# Inference
- if name in ('InferenceBackend', 'get_inference_backend'):
+ if name in ("InferenceBackend", "get_inference_backend"):
from .inference import InferenceBackend, get_inference_backend
- globals()['InferenceBackend'] = InferenceBackend
- globals()['get_inference_backend'] = get_inference_backend
+
+ globals()["InferenceBackend"] = InferenceBackend
+ globals()["get_inference_backend"] = get_inference_backend
return globals()[name]
# Training
- if name in ('TrainingBackend', 'get_training_backend', 'TrainingProgress'):
+ if name in ("TrainingBackend", "get_training_backend", "TrainingProgress"):
from .training import TrainingBackend, get_training_backend, TrainingProgress
- globals()['TrainingBackend'] = TrainingBackend
- globals()['get_training_backend'] = get_training_backend
- globals()['TrainingProgress'] = TrainingProgress
+
+ globals()["TrainingBackend"] = TrainingBackend
+ globals()["get_training_backend"] = get_training_backend
+ globals()["TrainingProgress"] = TrainingProgress
return globals()[name]
# Config (from utils.models)
- if name in ('is_vision_model', 'ModelConfig', 'scan_trained_loras',
- 'load_model_defaults', 'get_base_model_from_lora'):
+ if name in (
+ "is_vision_model",
+ "ModelConfig",
+ "scan_trained_loras",
+ "load_model_defaults",
+ "get_base_model_from_lora",
+ ):
from utils.models import (
- is_vision_model, ModelConfig, scan_trained_loras,
- load_model_defaults, get_base_model_from_lora,
+ is_vision_model,
+ ModelConfig,
+ scan_trained_loras,
+ load_model_defaults,
+ get_base_model_from_lora,
)
- globals()['is_vision_model'] = is_vision_model
- globals()['ModelConfig'] = ModelConfig
- globals()['scan_trained_loras'] = scan_trained_loras
- globals()['load_model_defaults'] = load_model_defaults
- globals()['get_base_model_from_lora'] = get_base_model_from_lora
+
+ globals()["is_vision_model"] = is_vision_model
+ globals()["ModelConfig"] = ModelConfig
+ globals()["scan_trained_loras"] = scan_trained_loras
+ globals()["load_model_defaults"] = load_model_defaults
+ globals()["get_base_model_from_lora"] = get_base_model_from_lora
return globals()[name]
# Paths
- if name in ('normalize_path', 'is_local_path', 'is_model_cached'):
+ if name in ("normalize_path", "is_local_path", "is_model_cached"):
from utils.paths import normalize_path, is_local_path, is_model_cached
- globals()['normalize_path'] = normalize_path
- globals()['is_local_path'] = is_local_path
- globals()['is_model_cached'] = is_model_cached
+
+ globals()["normalize_path"] = normalize_path
+ globals()["is_local_path"] = is_local_path
+ globals()["is_model_cached"] = is_model_cached
return globals()[name]
# Utils
- if name in ('without_hf_auth', 'format_error_message'):
+ if name in ("without_hf_auth", "format_error_message"):
from utils.utils import without_hf_auth, format_error_message
- globals()['without_hf_auth'] = without_hf_auth
- globals()['format_error_message'] = format_error_message
+
+ globals()["without_hf_auth"] = without_hf_auth
+ globals()["format_error_message"] = format_error_message
return globals()[name]
# Hardware
- if name in ('get_device', 'is_apple_silicon', 'clear_gpu_cache',
- 'get_gpu_memory_info', 'log_gpu_memory', 'DeviceType'):
+ if name in (
+ "get_device",
+ "is_apple_silicon",
+ "clear_gpu_cache",
+ "get_gpu_memory_info",
+ "log_gpu_memory",
+ "DeviceType",
+ ):
from utils.hardware import (
- get_device, is_apple_silicon, clear_gpu_cache,
- get_gpu_memory_info, log_gpu_memory, DeviceType,
+ get_device,
+ is_apple_silicon,
+ clear_gpu_cache,
+ get_gpu_memory_info,
+ log_gpu_memory,
+ DeviceType,
)
- globals()['get_device'] = get_device
- globals()['is_apple_silicon'] = is_apple_silicon
- globals()['clear_gpu_cache'] = clear_gpu_cache
- globals()['get_gpu_memory_info'] = get_gpu_memory_info
- globals()['log_gpu_memory'] = log_gpu_memory
- globals()['DeviceType'] = DeviceType
+
+ globals()["get_device"] = get_device
+ globals()["is_apple_silicon"] = is_apple_silicon
+ globals()["clear_gpu_cache"] = clear_gpu_cache
+ globals()["get_gpu_memory_info"] = get_gpu_memory_info
+ globals()["log_gpu_memory"] = log_gpu_memory
+ globals()["DeviceType"] = DeviceType
return globals()[name]
# Datasets
- if name == 'format_and_template_dataset':
+ if name == "format_and_template_dataset":
from utils.datasets import format_and_template_dataset
- globals()['format_and_template_dataset'] = format_and_template_dataset
+
+ globals()["format_and_template_dataset"] = format_and_template_dataset
return format_and_template_dataset
raise AttributeError(f"module 'core' has no attribute {name!r}")
diff --git a/studio/backend/core/data_recipe/jobs/__init__.py b/studio/backend/core/data_recipe/jobs/__init__.py
index 175c2c108f..cf03d62a3b 100644
--- a/studio/backend/core/data_recipe/jobs/__init__.py
+++ b/studio/backend/core/data_recipe/jobs/__init__.py
@@ -4,4 +4,3 @@
from .manager import JobManager, get_job_manager
__all__ = ["JobManager", "get_job_manager"]
-
diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py
index 7188562bfb..61c3516b2a 100644
--- a/studio/backend/core/data_recipe/jobs/manager.py
+++ b/studio/backend/core/data_recipe/jobs/manager.py
@@ -52,12 +52,10 @@ class Subscription:
if event_id is None:
self._next_id += 1
event_id = self._next_id
- body = json.dumps(event, separators=(",", ":"), ensure_ascii=False)
+ body = json.dumps(event, separators = (",", ":"), ensure_ascii = False)
event_type = event.get("type") or "message"
return (
- f"id: {event_id}\n"
- f"event: {event_type}\n"
- f"data: {body}\n\n"
+ f"id: {event_id}\n" f"event: {event_type}\n" f"data: {body}\n\n"
).encode("utf-8")
@@ -68,7 +66,7 @@ class JobManager:
self._job: Job | None = None
self._proc: mp.Process | None = None
self._mp_q: Any | None = None
- self._events: deque[dict] = deque(maxlen=5000)
+ self._events: deque[dict] = deque(maxlen = 5000)
self._subs: list[queue.Queue] = []
self._pump_thread: threading.Thread | None = None
self._seq: int = 0
@@ -92,7 +90,7 @@ class JobManager:
raise RuntimeError("job already running")
job_id = uuid.uuid4().hex
- self._job = Job(job_id=job_id, status="pending", started_at=time.time())
+ self._job = Job(job_id = job_id, status = "pending", started_at = time.time())
self._job.progress_columns_total = llm_column_count
self._events.clear()
self._seq = 0
@@ -101,18 +99,20 @@ class JobManager:
run_payload["_job_id"] = job_id
mp_q = _CTX.Queue()
proc = _CTX.Process(
- target=run_job_process,
- kwargs={"event_queue": mp_q, "recipe": recipe, "run": run_payload},
- daemon=True,
+ target = run_job_process,
+ kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload},
+ daemon = True,
)
proc.start()
self._mp_q = mp_q
self._proc = proc
- self._pump_thread = threading.Thread(target=self._pump_loop, daemon=True)
+ self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread.start()
- self._emit({"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id})
+ self._emit(
+ {"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id}
+ )
return job_id
def cancel(self, job_id: str) -> bool:
@@ -123,7 +123,9 @@ class JobManager:
if self._proc is None or not self._proc.is_alive():
return True
self._job.status = "cancelling"
- self._emit({"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id})
+ self._emit(
+ {"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id}
+ )
try:
self._proc.terminate()
except (AttributeError, OSError):
@@ -225,7 +227,7 @@ class JobManager:
if in_memory_dataset is not None:
total = len(in_memory_dataset)
- rows = in_memory_dataset[offset:offset + limit]
+ rows = in_memory_dataset[offset : offset + limit]
return {"dataset": rows, "total": total}
if not artifact_path:
if job_status in {"completed", "error", "cancelled"}:
@@ -238,7 +240,9 @@ class JobManager:
if not parquet_dir.exists():
return {"error": f"dataset path missing: {parquet_dir}"}
- return self._load_dataset_page(parquet_dir=parquet_dir, limit=limit, offset=offset)
+ return self._load_dataset_page(
+ parquet_dir = parquet_dir, limit = limit, offset = offset
+ )
except Exception as exc:
return {"error": f"dataset load failed: {exc}"}
@@ -250,16 +254,16 @@ class JobManager:
offset: int,
) -> dict[str, Any]:
dataset_page = JobManager._load_dataset_page_with_duckdb(
- parquet_dir=parquet_dir,
- limit=limit,
- offset=offset,
+ parquet_dir = parquet_dir,
+ limit = limit,
+ offset = offset,
)
if dataset_page is not None:
return dataset_page
return JobManager._load_dataset_page_with_data_designer(
- parquet_dir=parquet_dir,
- limit=limit,
- offset=offset,
+ parquet_dir = parquet_dir,
+ limit = limit,
+ offset = offset,
)
@staticmethod
@@ -299,9 +303,9 @@ class JobManager:
for helper_col in ("filename", "__row_num__"):
if helper_col in dataframe.columns:
- dataframe = dataframe.drop(columns=[helper_col])
+ dataframe = dataframe.drop(columns = [helper_col])
- rows = dataframe.to_dict(orient="records")
+ rows = dataframe.to_dict(orient = "records")
return {"dataset": to_preview_jsonable(rows), "total": total}
@staticmethod
@@ -315,21 +319,23 @@ class JobManager:
dataframe = read_parquet_dataset(parquet_dir)
total = int(len(dataframe.index))
- rows = dataframe.iloc[offset:offset + limit].to_dict(orient="records")
+ rows = dataframe.iloc[offset : offset + limit].to_dict(orient = "records")
return {"dataset": to_preview_jsonable(rows), "total": total}
- def subscribe(self, job_id: str, *, after_seq: int | None = None) -> Subscription | None:
+ def subscribe(
+ self, job_id: str, *, after_seq: int | None = None
+ ) -> Subscription | None:
"""SSE subscribe: get replay buffer + live events stream."""
with self._lock:
if self._job is None or self._job.job_id != job_id:
return None
- q: queue.Queue = queue.Queue(maxsize=2000)
+ q: queue.Queue = queue.Queue(maxsize = 2000)
self._subs.append(q)
if after_seq is None:
replay = list(self._events)
else:
replay = [e for e in self._events if int(e.get("seq") or 0) > after_seq]
- return Subscription(replay=replay, _q=q)
+ return Subscription(replay = replay, _q = q)
def unsubscribe(self, sub: Subscription) -> None:
"""Drop SSE subscriber (client disconnected)."""
@@ -361,7 +367,7 @@ class JobManager:
def _read_queue_with_timeout(q: Any, *, timeout_sec: float) -> dict | None:
"""Try read 1 event from mp queue. Timeout = pump stays responsive."""
try:
- return coerce_event(q.get(timeout=timeout_sec))
+ return coerce_event(q.get(timeout = timeout_sec))
except queue.Empty:
return None
except (EOFError, OSError, ValueError):
@@ -387,7 +393,7 @@ class JobManager:
return
job, proc, mp_q = snap
- event = self._read_queue_with_timeout(mp_q, timeout_sec=0.25)
+ event = self._read_queue_with_timeout(mp_q, timeout_sec = 0.25)
if event is not None:
self._handle_event(job, event)
continue
@@ -399,7 +405,11 @@ class JobManager:
self._handle_event(job, e)
with self._lock:
- if self._job and self._job.status in {"pending", "active", "cancelling"}:
+ if self._job and self._job.status in {
+ "pending",
+ "active",
+ "cancelling",
+ }:
if self._job.status == "cancelling":
self._job.status = "cancelled"
else:
@@ -407,9 +417,17 @@ class JobManager:
self._job.error = self._job.error or "process exited"
self._job.finished_at = time.time()
event_type = (
- EVENT_JOB_CANCELLED if self._job.status == "cancelled" else EVENT_JOB_ERROR
+ EVENT_JOB_CANCELLED
+ if self._job.status == "cancelled"
+ else EVENT_JOB_ERROR
+ )
+ self._emit(
+ {
+ "type": event_type,
+ "ts": time.time(),
+ "job_id": self._job.job_id,
+ }
)
- self._emit({"type": event_type, "ts": time.time(), "job_id": self._job.job_id})
return
def _handle_event(self, job: Job, event: dict) -> None:
diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py
index b5567e7634..324b62a92e 100644
--- a/studio/backend/core/data_recipe/jobs/parse.py
+++ b/studio/backend/core/data_recipe/jobs/parse.py
@@ -22,7 +22,7 @@ from .constants import (
from .types import Job, ModelUsage, Progress
-@dataclass(frozen=True)
+@dataclass(frozen = True)
class ParsedUpdate:
stage: str | None = None
current_column: str | None = None
@@ -42,6 +42,7 @@ class ParsedUpdate:
usage_rpm: float | None = None
usage_section_start: bool | None = None
+
# kinda of a bummber but currently only option, Best effort parser from data-designer logs -> structured status for UI.
_RE_SAMPLERS = re.compile(
r"Preparing samplers to generate (?P
\d+) records across (?P\d+) columns"
@@ -66,76 +67,76 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
m = _RE_SAMPLERS.search(msg)
if m:
return ParsedUpdate(
- stage=STAGE_SAMPLING,
- rows=int(m.group("rows")),
- cols=int(m.group("cols")),
+ stage = STAGE_SAMPLING,
+ rows = int(m.group("rows")),
+ cols = int(m.group("cols")),
)
if "Sorting column configs into a Directed Acyclic Graph" in msg:
- return ParsedUpdate(stage=STAGE_DAG)
+ return ParsedUpdate(stage = STAGE_DAG)
if "Running health checks for models" in msg:
- return ParsedUpdate(stage=STAGE_HEALTHCHECK)
+ return ParsedUpdate(stage = STAGE_HEALTHCHECK)
if "Preview generation in progress" in msg:
- return ParsedUpdate(stage=STAGE_PREVIEW)
+ return ParsedUpdate(stage = STAGE_PREVIEW)
if "Creating Data Designer dataset" in msg:
- return ParsedUpdate(stage=STAGE_CREATE)
+ return ParsedUpdate(stage = STAGE_CREATE)
if "Measuring dataset column statistics" in msg:
- return ParsedUpdate(stage=STAGE_PROFILING)
+ return ParsedUpdate(stage = STAGE_PROFILING)
m = _RE_COLCFG.search(msg)
if m:
col = m.group("col")
- return ParsedUpdate(stage=STAGE_COLUMN_CONFIG, current_column=col)
+ return ParsedUpdate(stage = STAGE_COLUMN_CONFIG, current_column = col)
m = _RE_PROCESSING_COL.search(msg)
if m:
col = m.group("col")
- return ParsedUpdate(stage=STAGE_GENERATING, current_column=col)
+ return ParsedUpdate(stage = STAGE_GENERATING, current_column = col)
m = _RE_PROGRESS.search(msg)
if m:
p = Progress(
- done=int(m.group("done")),
- total=int(m.group("total")),
- percent=float(m.group("pct")),
- ok=int(m.group("ok")),
- failed=int(m.group("failed")),
- rate=float(m.group("rate")),
- eta_sec=float(m.group("eta")),
+ done = int(m.group("done")),
+ total = int(m.group("total")),
+ percent = float(m.group("pct")),
+ ok = int(m.group("ok")),
+ failed = int(m.group("failed")),
+ rate = float(m.group("rate")),
+ eta_sec = float(m.group("eta")),
)
- return ParsedUpdate(stage=STAGE_GENERATING, progress=p)
+ return ParsedUpdate(stage = STAGE_GENERATING, progress = p)
m = _RE_BATCH.search(msg)
if m:
return ParsedUpdate(
- stage=STAGE_BATCH,
- batch_idx=int(m.group("idx")),
- batch_total=int(m.group("total")),
+ stage = STAGE_BATCH,
+ batch_idx = int(m.group("idx")),
+ batch_total = int(m.group("total")),
)
if "Model usage summary" in msg:
- return ParsedUpdate(usage_section_start=True)
+ return ParsedUpdate(usage_section_start = True)
m = _RE_USAGE_MODEL.search(msg)
if m and "|-- model:" in msg:
- return ParsedUpdate(usage_model=str(m.group("model")).strip())
+ return ParsedUpdate(usage_model = str(m.group("model")).strip())
m = _RE_USAGE_TOKENS.search(msg)
if m:
return ParsedUpdate(
- usage_input_tokens=int(m.group("input")),
- usage_output_tokens=int(m.group("output")),
- usage_total_tokens=int(m.group("total")),
- usage_tps=float(m.group("tps")),
+ usage_input_tokens = int(m.group("input")),
+ usage_output_tokens = int(m.group("output")),
+ usage_total_tokens = int(m.group("total")),
+ usage_tps = float(m.group("tps")),
)
m = _RE_USAGE_REQUESTS.search(msg)
if m:
return ParsedUpdate(
- usage_requests_success=int(m.group("success")),
- usage_requests_failed=int(m.group("failed")),
- usage_requests_total=int(m.group("total")),
- usage_rpm=float(m.group("rpm")),
+ usage_requests_success = int(m.group("success")),
+ usage_requests_failed = int(m.group("failed")),
+ usage_requests_total = int(m.group("total")),
+ usage_rpm = float(m.group("rpm")),
)
return None
@@ -146,7 +147,10 @@ def apply_update(job: Job, update: ParsedUpdate) -> None:
job.stage = update.stage
if update.current_column is not None:
job.current_column = update.current_column
- if update.stage == STAGE_GENERATING and update.current_column not in job._seen_generation_columns:
+ if (
+ update.stage == STAGE_GENERATING
+ and update.current_column not in job._seen_generation_columns
+ ):
job._seen_generation_columns.append(update.current_column)
if update.rows is not None:
job.rows = update.rows
@@ -185,7 +189,7 @@ def apply_update(job: Job, update: ParsedUpdate) -> None:
name = update.usage_model.strip().strip("'").strip('"')
job._current_usage_model = name
if name not in job.model_usage:
- job.model_usage[name] = ModelUsage(model=name)
+ job.model_usage[name] = ModelUsage(model = name)
if job._current_usage_model is None:
return
@@ -227,7 +231,9 @@ def _compute_overall_progress(job: Job, column_progress: Progress) -> Progress:
if len(job._column_done) == 0:
done = current_done
else:
- sum_done = sum(max(0, min(value, total_rows)) for value in job._column_done.values())
+ sum_done = sum(
+ max(0, min(value, total_rows)) for value in job._column_done.values()
+ )
done = int(sum_done / total_columns)
prev_done = int(job.progress.done or 0)
@@ -241,13 +247,13 @@ def _compute_overall_progress(job: Job, column_progress: Progress) -> Progress:
percent = prev_percent
return Progress(
- done=done,
- total=total_rows,
- percent=percent,
- eta_sec=column_progress.eta_sec,
- rate=column_progress.rate,
- ok=column_progress.ok,
- failed=column_progress.failed,
+ done = done,
+ total = total_rows,
+ percent = percent,
+ eta_sec = column_progress.eta_sec,
+ rate = column_progress.rate,
+ ok = column_progress.ok,
+ failed = column_progress.failed,
)
diff --git a/studio/backend/core/data_recipe/jobs/types.py b/studio/backend/core/data_recipe/jobs/types.py
index 2ef22a387a..3079d76bdb 100644
--- a/studio/backend/core/data_recipe/jobs/types.py
+++ b/studio/backend/core/data_recipe/jobs/types.py
@@ -54,9 +54,9 @@ class Job:
status: JobStatus = "created"
stage: str | None = None
current_column: str | None = None
- progress: Progress = field(default_factory=Progress)
- column_progress: Progress = field(default_factory=Progress)
- batch: BatchProgress = field(default_factory=BatchProgress)
+ progress: Progress = field(default_factory = Progress)
+ column_progress: Progress = field(default_factory = Progress)
+ batch: BatchProgress = field(default_factory = BatchProgress)
rows: int | None = None
cols: int | None = None
error: str | None = None
@@ -67,10 +67,10 @@ class Job:
artifact_path: str | None = None
dataset: list[dict[str, Any]] | None = None
processor_artifacts: dict[str, Any] | None = None
- model_usage: dict[str, ModelUsage] = field(default_factory=dict)
+ model_usage: dict[str, ModelUsage] = field(default_factory = dict)
progress_columns_total: int | None = None
- completed_columns: list[str] = field(default_factory=list)
+ completed_columns: list[str] = field(default_factory = list)
_current_usage_model: str | None = None
_in_usage_summary: bool = False
- _seen_generation_columns: list[str] = field(default_factory=list)
- _column_done: dict[str, int] = field(default_factory=dict)
+ _seen_generation_columns: list[str] = field(default_factory = list)
+ _column_done: dict[str, int] = field(default_factory = dict)
diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py
index c67e93751c..63e38bd18d 100644
--- a/studio/backend/core/data_recipe/jobs/worker.py
+++ b/studio/backend/core/data_recipe/jobs/worker.py
@@ -51,7 +51,9 @@ def _slugify_run_name(value: str) -> str:
return slug[:80].strip("-")
-def _build_dataset_name(*, run_name: str | None, job_id: str, artifact_root: Path) -> str:
+def _build_dataset_name(
+ *, run_name: str | None, job_id: str, artifact_root: Path
+) -> str:
fallback = f"recipe_{job_id}"
slug = _slugify_run_name(run_name or "")
base_name = f"recipe_{slug}" if slug else fallback
@@ -74,16 +76,20 @@ def run_job_process(
Sends events to `event_queue`.
"""
import os
- os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
+
+ os.environ["PYTHONWARNINGS"] = (
+ "ignore" # Suppress warnings at C-level before imports
+ )
import warnings
from loggers.config import LogConfig
+
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
-
+
LogConfig.setup_logging(
- service_name="unsloth-studio-data-worker",
- env=os.getenv("ENVIRONMENT_TYPE", "production"),
+ service_name = "unsloth-studio-data-worker",
+ env = os.getenv("ENVIRONMENT_TYPE", "production"),
)
event_queue.put({"type": EVENT_JOB_STARTED, "ts": time.time()})
@@ -98,16 +104,16 @@ def run_job_process(
run_name_raw = run.get("run_name")
run_name = run_name_raw if isinstance(run_name_raw, str) else None
dataset_name = _build_dataset_name(
- run_name=run_name,
- job_id=job_id,
- artifact_root=_ARTIFACT_ROOT,
+ run_name = run_name,
+ job_id = job_id,
+ artifact_root = _ARTIFACT_ROOT,
)
merge_batches = bool(run.get("merge_batches"))
ensure_dir(_ARTIFACT_ROOT)
run_config_raw = run.get("run_config") or {}
builder = build_config_builder(recipe)
- designer = create_data_designer(recipe, artifact_path=str(_ARTIFACT_ROOT))
+ designer = create_data_designer(recipe, artifact_path = str(_ARTIFACT_ROOT))
# DataDesigner configures root logging in DataDesigner.__init__.
# Attach queue logger directly to `data_designer` so parser events survive root resets.
@@ -123,16 +129,16 @@ def run_job_process(
execution_type = str(run.get("execution_type") or "full").strip().lower()
if execution_type == "preview":
- results = designer.preview(builder, num_records=rows)
+ results = designer.preview(builder, num_records = rows)
analysis = (
None
if results.analysis is None
- else to_jsonable(results.analysis.model_dump(mode="json"))
+ else to_jsonable(results.analysis.model_dump(mode = "json"))
)
dataset = (
[]
if results.dataset is None
- else to_preview_jsonable(results.dataset.to_dict(orient="records"))
+ else to_preview_jsonable(results.dataset.to_dict(orient = "records"))
)
processor_artifacts = (
None
@@ -151,10 +157,14 @@ def run_job_process(
}
)
else:
- results = designer.create(builder, num_records=rows, dataset_name=dataset_name)
- analysis = to_jsonable(results.load_analysis().model_dump(mode="json"))
+ results = designer.create(
+ builder, num_records = rows, dataset_name = dataset_name
+ )
+ analysis = to_jsonable(results.load_analysis().model_dump(mode = "json"))
if merge_batches:
- _merge_batches_to_single_parquet(results.artifact_storage.base_dataset_path)
+ _merge_batches_to_single_parquet(
+ results.artifact_storage.base_dataset_path
+ )
artifact_path = str(results.artifact_storage.base_dataset_path)
event_queue.put(
{
@@ -171,7 +181,7 @@ def run_job_process(
"type": EVENT_JOB_ERROR,
"ts": time.time(),
"error": str(exc),
- "stack": traceback.format_exc(limit=20),
+ "stack": traceback.format_exc(limit = 20),
}
)
@@ -189,12 +199,12 @@ def _merge_batches_to_single_parquet(base_dataset_path: Path) -> None:
dataframe = read_parquet_dataset(parquet_dir)
shutil.rmtree(parquet_dir)
- parquet_dir.mkdir(parents=True, exist_ok=True)
+ parquet_dir.mkdir(parents = True, exist_ok = True)
merged_file = parquet_dir / "batch_00000.parquet"
- dataframe.to_parquet(merged_file, index=False)
+ dataframe.to_parquet(merged_file, index = False)
_rewrite_merged_metadata(
- base_dataset_path=base_dataset_path,
- parquet_file=merged_file,
+ base_dataset_path = base_dataset_path,
+ parquet_file = merged_file,
)
@@ -204,7 +214,7 @@ def _rewrite_merged_metadata(*, base_dataset_path: Path, parquet_file: Path) ->
return
try:
- metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+ metadata = json.loads(metadata_path.read_text(encoding = "utf-8"))
except (OSError, TypeError, ValueError):
return
@@ -222,8 +232,8 @@ def _rewrite_merged_metadata(*, base_dataset_path: Path, parquet_file: Path) ->
try:
metadata_path.write_text(
- json.dumps(metadata, indent=2, sort_keys=True),
- encoding="utf-8",
+ json.dumps(metadata, indent = 2, sort_keys = True),
+ encoding = "utf-8",
)
except OSError:
return
diff --git a/studio/backend/core/data_recipe/jsonable.py b/studio/backend/core/data_recipe/jsonable.py
index ba4d59d6dd..8bec60cf82 100644
--- a/studio/backend/core/data_recipe/jsonable.py
+++ b/studio/backend/core/data_recipe/jsonable.py
@@ -11,7 +11,7 @@ from typing import Any
def _pil_to_preview_payload(image: Any) -> dict[str, Any]:
buffer = io.BytesIO()
- image.convert("RGB").save(buffer, format="JPEG", quality=85)
+ image.convert("RGB").save(buffer, format = "JPEG", quality = 85)
return {
"type": "image",
"mime": "image/jpeg",
diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py
index b5914c20f3..c32b2fccaf 100644
--- a/studio/backend/core/data_recipe/local_callable_validators.py
+++ b/studio/backend/core/data_recipe/local_callable_validators.py
@@ -33,7 +33,7 @@ _OXC_TOOL_DIR = Path(__file__).resolve().parent / "oxc-validator"
_OXC_RUNNER_PATH = _OXC_TOOL_DIR / "validate.mjs"
-@dataclass(frozen=True)
+@dataclass(frozen = True)
class OxcLocalCallableValidatorSpec:
name: str
drop: bool
@@ -64,7 +64,7 @@ def split_oxc_local_callable_validators(
kept_columns.append(column)
continue
- maybe_spec = _parse_oxc_spec(column=column)
+ maybe_spec = _parse_oxc_spec(column = column)
if maybe_spec is None:
kept_columns.append(column)
continue
@@ -96,14 +96,14 @@ def register_oxc_local_callable_validators(
)
builder.add_column(
ValidationColumnConfig(
- name=spec.name,
- drop=spec.drop,
- target_columns=spec.target_columns,
- validator_type=ValidatorType.LOCAL_CALLABLE,
- validator_params=LocalCallableValidatorParams(
- validation_function=validation_function,
+ name = spec.name,
+ drop = spec.drop,
+ target_columns = spec.target_columns,
+ validator_type = ValidatorType.LOCAL_CALLABLE,
+ validator_params = LocalCallableValidatorParams(
+ validation_function = validation_function,
),
- batch_size=spec.batch_size,
+ batch_size = spec.batch_size,
)
)
@@ -132,7 +132,11 @@ def _parse_oxc_spec(
target_columns_raw = column.get("target_columns")
target_columns = (
- [value.strip() for value in target_columns_raw if isinstance(value, str) and value.strip()]
+ [
+ value.strip()
+ for value in target_columns_raw
+ if isinstance(value, str) and value.strip()
+ ]
if isinstance(target_columns_raw, list)
else []
)
@@ -144,13 +148,13 @@ def _parse_oxc_spec(
drop = bool(column.get("drop") is True)
return OxcLocalCallableValidatorSpec(
- name=name,
- drop=drop,
- target_columns=target_columns,
- batch_size=batch_size,
- code_lang=code_lang,
- validation_mode=validation_mode,
- code_shape=code_shape,
+ name = name,
+ drop = drop,
+ target_columns = target_columns,
+ batch_size = batch_size,
+ code_lang = code_lang,
+ validation_mode = validation_mode,
+ code_shape = code_shape,
)
@@ -172,18 +176,16 @@ def _parse_oxc_validation_marker(fn_name: str) -> tuple[str, str, str]:
return "javascript", "syntax", "auto"
code_lang = parts[0] if parts[0] in _OXC_LANG_TO_NODE_LANG else "javascript"
mode = parts[1] if parts[1] in _OXC_VALIDATION_MODES else "syntax"
- code_shape = parts[2] if len(parts) >= 3 and parts[2] in _OXC_CODE_SHAPES else "auto"
+ code_shape = (
+ parts[2] if len(parts) >= 3 and parts[2] in _OXC_CODE_SHAPES else "auto"
+ )
return code_lang, mode, code_shape
-@lru_cache(maxsize=8)
+@lru_cache(maxsize = 8)
def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape: str):
node_lang = _OXC_LANG_TO_NODE_LANG.get(lang, "js")
- mode = (
- validation_mode
- if validation_mode in _OXC_VALIDATION_MODES
- else "syntax"
- )
+ mode = validation_mode if validation_mode in _OXC_VALIDATION_MODES else "syntax"
normalized_code_shape = code_shape if code_shape in _OXC_CODE_SHAPES else "auto"
def _validator(df):
@@ -197,14 +199,17 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape:
code_values = (
["" for _ in range(row_count)]
if not code_column
- else ["" if value is None else str(value) for value in df[code_column].tolist()]
+ else [
+ "" if value is None else str(value)
+ for value in df[code_column].tolist()
+ ]
)
results = _run_oxc_batch(
- node_lang=node_lang,
- validation_mode=mode,
- code_shape=normalized_code_shape,
- code_values=code_values,
+ node_lang = node_lang,
+ validation_mode = mode,
+ code_shape = normalized_code_shape,
+ code_values = code_values,
)
if len(results) != row_count:
results = _fallback_results(
@@ -213,9 +218,7 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape:
)
return pd.DataFrame(results)
- _validator.__name__ = (
- f"{OXC_VALIDATION_FN_MARKER}_{node_lang}_{mode.replace('+', '_')}_{normalized_code_shape}"
- )
+ _validator.__name__ = f"{OXC_VALIDATION_FN_MARKER}_{node_lang}_{mode.replace('+', '_')}_{normalized_code_shape}"
return _validator
@@ -247,12 +250,12 @@ def _run_oxc_batch(
env["TEMP"] = tmp_dir_str
proc = subprocess.run(
["node", str(_OXC_RUNNER_PATH)],
- cwd=str(_OXC_TOOL_DIR),
- input=json.dumps(payload),
- text=True,
- capture_output=True,
- check=False,
- env=env,
+ cwd = str(_OXC_TOOL_DIR),
+ input = json.dumps(payload),
+ text = True,
+ capture_output = True,
+ check = False,
+ env = env,
)
except (OSError, ValueError) as exc:
logger.warning("OXC subprocess launch failed: %s", exc)
@@ -298,13 +301,21 @@ def _run_oxc_batch(
warning_count_raw = item.get("warning_count")
out.append(
{
- "is_valid": bool(is_valid_raw) if isinstance(is_valid_raw, bool) else False,
- "error_count": int(error_count_raw) if isinstance(error_count_raw, int) else 0,
+ "is_valid": bool(is_valid_raw)
+ if isinstance(is_valid_raw, bool)
+ else False,
+ "error_count": int(error_count_raw)
+ if isinstance(error_count_raw, int)
+ else 0,
"error_message": str(message_raw or ""),
- "severity": str(severity_raw) if isinstance(severity_raw, str) else None,
+ "severity": str(severity_raw)
+ if isinstance(severity_raw, str)
+ else None,
"code": str(code_raw) if isinstance(code_raw, str) else None,
"labels": labels_raw if isinstance(labels_raw, list) else [],
- "codeframe": str(codeframe_raw) if isinstance(codeframe_raw, str) else None,
+ "codeframe": str(codeframe_raw)
+ if isinstance(codeframe_raw, str)
+ else None,
"warning_count": int(warning_count_raw)
if isinstance(warning_count_raw, int)
else 0,
diff --git a/studio/backend/core/data_recipe/oxc-validator/validate.mjs b/studio/backend/core/data_recipe/oxc-validator/validate.mjs
index 7d2f206ce0..ad61fb5a9e 100644
--- a/studio/backend/core/data_recipe/oxc-validator/validate.mjs
+++ b/studio/backend/core/data_recipe/oxc-validator/validate.mjs
@@ -1,3 +1,6 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
import { spawnSync } from "node:child_process";
import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py
index e90bb01c38..1c86ac42d9 100644
--- a/studio/backend/core/data_recipe/service.py
+++ b/studio/backend/core/data_recipe/service.py
@@ -14,6 +14,7 @@ from .local_callable_validators import (
register_oxc_local_callable_validators,
split_oxc_local_callable_validators,
)
+
_IMAGE_CONTEXT_PATCHED = False
@@ -21,7 +22,9 @@ def _encode_bytes_to_base64(value: bytes | bytearray) -> str:
return base64.b64encode(bytes(value)).decode("utf-8")
-def _load_image_file_to_base64(path_value: str, *, base_path: str | None = None) -> str | None:
+def _load_image_file_to_base64(
+ path_value: str, *, base_path: str | None = None
+) -> str | None:
try:
path = Path(path_value)
candidates: list[Path] = []
@@ -53,7 +56,7 @@ def _pil_image_to_base64(value: Any) -> str | None:
image_format = str(getattr(value, "format", "") or "").upper()
if image_format not in {"PNG", "JPEG", "JPG", "WEBP", "GIF"}:
image_format = "PNG"
- value.save(buffer, format=image_format)
+ value.save(buffer, format = image_format)
return _encode_bytes_to_base64(buffer.getvalue())
@@ -93,7 +96,7 @@ def _normalize_image_context_value(value: Any, *, base_path: str | None = None)
path_value = value.get("path")
if isinstance(path_value, str) and path_value.strip():
- if as_base64 := _load_image_file_to_base64(path_value, base_path=base_path):
+ if as_base64 := _load_image_file_to_base64(path_value, base_path = base_path):
return as_base64
return path_value
@@ -116,8 +119,10 @@ def _apply_data_designer_image_context_patch() -> None:
original_auto_resolve = ImageContext._auto_resolve_context_value
- def _patched_auto_resolve(self: Any, context_value: Any, base_path: str | None) -> Any:
- normalized = _normalize_image_context_value(context_value, base_path=base_path)
+ def _patched_auto_resolve(
+ self: Any, context_value: Any, base_path: str | None
+ ) -> Any:
+ normalized = _normalize_image_context_value(context_value, base_path = base_path)
return original_auto_resolve(self, normalized, base_path)
ImageContext._auto_resolve_context_value = _patched_auto_resolve
@@ -137,12 +142,12 @@ def build_model_providers(recipe: dict[str, Any]):
api_key = os.getenv(api_key_env)
providers.append(
ModelProvider(
- name=provider["name"],
- endpoint=provider["endpoint"],
- provider_type=provider.get("provider_type", "openai"),
- api_key=api_key,
- extra_headers=provider.get("extra_headers"),
- extra_body=provider.get("extra_body"),
+ name = provider["name"],
+ endpoint = provider["endpoint"],
+ provider_type = provider.get("provider_type", "openai"),
+ api_key = api_key,
+ extra_headers = provider.get("extra_headers"),
+ extra_body = provider.get("extra_body"),
)
)
@@ -170,10 +175,10 @@ def build_mcp_providers(
args = []
providers.append(
LocalStdioMCPProvider(
- name=str(provider.get("name", "")),
- command=str(provider.get("command", "")),
- args=[str(value) for value in args],
- env={str(key): str(value) for key, value in env.items()},
+ name = str(provider.get("name", "")),
+ command = str(provider.get("command", "")),
+ args = [str(value) for value in args],
+ env = {str(key): str(value) for key, value in env.items()},
)
)
continue
@@ -185,10 +190,10 @@ def build_mcp_providers(
api_key = os.getenv(str(api_key_env))
providers.append(
MCPProvider(
- name=str(provider.get("name", "")),
- endpoint=str(provider.get("endpoint", "")),
- provider_type=str(provider_type),
- api_key=str(api_key) if api_key else None,
+ name = str(provider.get("name", "")),
+ endpoint = str(provider.get("endpoint", "")),
+ provider_type = str(provider_type),
+ api_key = str(api_key) if api_key else None,
)
)
return providers
@@ -209,8 +214,8 @@ def build_config_builder(recipe: dict[str, Any]):
)
builder = DataDesignerConfigBuilder.from_config({"data_designer": recipe_core})
register_oxc_local_callable_validators(
- builder=builder,
- specs=oxc_local_callable_specs,
+ builder = builder,
+ specs = oxc_local_callable_specs,
)
# DataDesignerConfigBuilder.from_config currently skips processors.
@@ -223,7 +228,7 @@ def build_config_builder(recipe: dict[str, Any]):
continue
kwargs = {k: v for k, v in processor.items() if k != "processor_type"}
builder.add_processor(
- processor_type=ProcessorType(processor_type_raw),
+ processor_type = ProcessorType(processor_type_raw),
**kwargs,
)
@@ -239,9 +244,9 @@ def create_data_designer(
from data_designer.interface.data_designer import DataDesigner
return DataDesigner(
- artifact_path=artifact_path,
- model_providers=build_model_providers(recipe),
- mcp_providers=build_mcp_providers(recipe),
+ artifact_path = artifact_path,
+ model_providers = build_model_providers(recipe),
+ mcp_providers = build_mcp_providers(recipe),
)
@@ -257,11 +262,11 @@ def preview_recipe(
) -> tuple[list[dict[str, Any]], dict[str, Any] | None, dict[str, Any] | None]:
builder = build_config_builder(recipe)
designer = create_data_designer(recipe)
- results = designer.preview(builder, num_records=num_records)
+ results = designer.preview(builder, num_records = num_records)
dataset: list[dict[str, Any]] = []
if results.dataset is not None:
- raw_rows = results.dataset.to_dict(orient="records")
+ raw_rows = results.dataset.to_dict(orient = "records")
dataset = [to_jsonable(row) for row in raw_rows]
artifacts = (
@@ -272,7 +277,7 @@ def preview_recipe(
analysis = (
None
if results.analysis is None
- else to_jsonable(results.analysis.model_dump(mode="json"))
+ else to_jsonable(results.analysis.model_dump(mode = "json"))
)
return dataset, artifacts, analysis
diff --git a/studio/backend/core/export/__init__.py b/studio/backend/core/export/__init__.py
index 3021d15a4d..16fba368ef 100644
--- a/studio/backend/core/export/__init__.py
+++ b/studio/backend/core/export/__init__.py
@@ -8,13 +8,14 @@ The default get_export_backend() returns an ExportOrchestrator that
delegates to a subprocess. The original ExportBackend runs inside
the subprocess and can be imported directly from .export when needed.
"""
+
from .orchestrator import ExportOrchestrator, get_export_backend
# Expose ExportOrchestrator as ExportBackend for backward compat
ExportBackend = ExportOrchestrator
__all__ = [
- 'ExportBackend',
- 'ExportOrchestrator',
- 'get_export_backend',
+ "ExportBackend",
+ "ExportOrchestrator",
+ "get_export_backend",
]
diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py
index 2057f2a0ff..a6f67c5bae 100644
--- a/studio/backend/core/export/export.py
+++ b/studio/backend/core/export/export.py
@@ -5,6 +5,7 @@
"""
Export backend - handles model exporting in various formats
"""
+
import glob
import json
import structlog
@@ -50,7 +51,7 @@ def _apply_wsl_sudo_patch():
try:
import unsloth_zoo.llama_cpp as llama_cpp_module
- def _wsl_do_we_need_sudo(system_type="debian"):
+ def _wsl_do_we_need_sudo(system_type = "debian"):
logger.info(
"WSL detected β skipping sudo check "
"(build deps pre-installed by setup.sh)"
@@ -59,16 +60,14 @@ def _apply_wsl_sudo_patch():
llama_cpp_module.do_we_need_sudo = _wsl_do_we_need_sudo
logger.info(
- "Applied WSL sudo patch to "
- "unsloth_zoo.llama_cpp.do_we_need_sudo"
+ "Applied WSL sudo patch to " "unsloth_zoo.llama_cpp.do_we_need_sudo"
)
except Exception as e:
logger.warning(f"Could not apply WSL sudo patch: {e}")
# Model card template
-MODEL_CARD = \
-"""---
+MODEL_CARD = """---
base_model: {base_model}
tags:
- text-generation-inference
@@ -92,6 +91,7 @@ This {model_type} model was trained 2x faster with [Unsloth](https://github.com/
[
](https://github.com/unslothai/unsloth)
"""
+
class ExportBackend:
"""Handles model export operations"""
@@ -130,7 +130,9 @@ class ExportBackend:
logger.error(f"Error during memory cleanup: {e}")
return False
- def scan_checkpoints(self, outputs_dir: str = str(outputs_root())) -> List[Tuple[str, List[Tuple[str, str]]]]:
+ def scan_checkpoints(
+ self, outputs_dir: str = str(outputs_root())
+ ) -> List[Tuple[str, List[Tuple[str, str]]]]:
"""
Scan outputs folder for training runs and their checkpoints.
@@ -138,13 +140,16 @@ class ExportBackend:
List of tuples: [(model_name, [(display_name, checkpoint_path), ...]), ...]
"""
from utils.models.checkpoints import scan_checkpoints
- return scan_checkpoints(outputs_dir=outputs_dir)
- def load_checkpoint(self,
- checkpoint_path: str,
- max_seq_length: int = 2048,
- load_in_4bit: bool = True,
- trust_remote_code: bool = False) -> Tuple[bool, str]:
+ return scan_checkpoints(outputs_dir = outputs_dir)
+
+ def load_checkpoint(
+ self,
+ checkpoint_path: str,
+ max_seq_length: int = 2048,
+ load_in_4bit: bool = True,
+ trust_remote_code: bool = False,
+ ) -> Tuple[bool, str]:
"""
Load a checkpoint for export.
@@ -174,81 +179,85 @@ class ExportBackend:
self.is_vision = not self._audio_type and is_vision_model(model_id)
# Load model based on type
- if self._audio_type == 'csm':
+ if self._audio_type == "csm":
from unsloth import FastModel
from transformers import CsmForConditionalGeneration
+
logger.info("Loading as CSM audio model...")
model, tokenizer = FastModel.from_pretrained(
- model_name=checkpoint_path,
- max_seq_length=max_seq_length,
- dtype=None,
- auto_model=CsmForConditionalGeneration,
- load_in_4bit=False,
- trust_remote_code=trust_remote_code,
+ model_name = checkpoint_path,
+ max_seq_length = max_seq_length,
+ dtype = None,
+ auto_model = CsmForConditionalGeneration,
+ load_in_4bit = False,
+ trust_remote_code = trust_remote_code,
)
- elif self._audio_type == 'whisper':
+ elif self._audio_type == "whisper":
from unsloth import FastModel
from transformers import WhisperForConditionalGeneration
+
logger.info("Loading as Whisper audio model...")
model, tokenizer = FastModel.from_pretrained(
- model_name=checkpoint_path,
- dtype=None,
- load_in_4bit=False,
- auto_model=WhisperForConditionalGeneration,
- trust_remote_code=trust_remote_code,
+ model_name = checkpoint_path,
+ dtype = None,
+ load_in_4bit = False,
+ auto_model = WhisperForConditionalGeneration,
+ trust_remote_code = trust_remote_code,
)
- elif self._audio_type == 'snac':
+ elif self._audio_type == "snac":
logger.info("Loading as SNAC (Orpheus) audio model...")
model, tokenizer = FastLanguageModel.from_pretrained(
- model_name=checkpoint_path,
- max_seq_length=max_seq_length,
- dtype=None,
- load_in_4bit=load_in_4bit,
- trust_remote_code=trust_remote_code,
+ model_name = checkpoint_path,
+ max_seq_length = max_seq_length,
+ dtype = None,
+ load_in_4bit = load_in_4bit,
+ trust_remote_code = trust_remote_code,
)
- elif self._audio_type == 'bicodec':
+ elif self._audio_type == "bicodec":
from unsloth import FastModel
+
logger.info("Loading as BiCodec (Spark-TTS) audio model...")
model, tokenizer = FastModel.from_pretrained(
- model_name=checkpoint_path,
- max_seq_length=max_seq_length,
- dtype=torch.float32,
- load_in_4bit=False,
- trust_remote_code=trust_remote_code,
+ model_name = checkpoint_path,
+ max_seq_length = max_seq_length,
+ dtype = torch.float32,
+ load_in_4bit = False,
+ trust_remote_code = trust_remote_code,
)
- elif self._audio_type == 'dac':
+ elif self._audio_type == "dac":
from unsloth import FastModel
+
logger.info("Loading as DAC (OuteTTS) audio model...")
model, tokenizer = FastModel.from_pretrained(
- model_name=checkpoint_path,
- max_seq_length=max_seq_length,
- load_in_4bit=False,
- trust_remote_code=trust_remote_code,
+ model_name = checkpoint_path,
+ max_seq_length = max_seq_length,
+ load_in_4bit = False,
+ trust_remote_code = trust_remote_code,
)
elif self.is_vision:
logger.info("Loading as vision model...")
model, processor = FastVisionModel.from_pretrained(
- model_name=checkpoint_path,
- max_seq_length=max_seq_length,
- dtype=None,
- load_in_4bit=load_in_4bit,
- trust_remote_code=trust_remote_code,
+ model_name = checkpoint_path,
+ max_seq_length = max_seq_length,
+ dtype = None,
+ load_in_4bit = load_in_4bit,
+ trust_remote_code = trust_remote_code,
)
tokenizer = processor # For vision models, processor acts as tokenizer
else:
logger.info("Loading as text model...")
model, tokenizer = FastLanguageModel.from_pretrained(
- model_name=checkpoint_path,
- max_seq_length=max_seq_length,
- dtype=None,
- load_in_4bit=load_in_4bit,
- trust_remote_code=trust_remote_code,
+ model_name = checkpoint_path,
+ max_seq_length = max_seq_length,
+ dtype = None,
+ load_in_4bit = load_in_4bit,
+ trust_remote_code = trust_remote_code,
)
# Check if PEFT model
@@ -273,28 +282,35 @@ class ExportBackend:
except Exception as e:
logger.error(f"Error loading checkpoint: {e}")
import traceback
+
logger.error(traceback.format_exc())
return False, f"Failed to load checkpoint: {str(e)}"
def _write_export_metadata(self, save_directory: str):
"""Write export_metadata.json with base model info for Chat page discovery."""
try:
- base_model = get_base_model_from_lora(self.current_checkpoint) if self.current_checkpoint else None
+ base_model = (
+ get_base_model_from_lora(self.current_checkpoint)
+ if self.current_checkpoint
+ else None
+ )
metadata = {"base_model": base_model}
metadata_path = os.path.join(save_directory, "export_metadata.json")
with open(metadata_path, "w") as f:
- json.dump(metadata, f, indent=2)
+ json.dump(metadata, f, indent = 2)
logger.info(f"Wrote export metadata to {metadata_path}")
except Exception as e:
logger.warning(f"Could not write export metadata: {e}")
- def export_merged_model(self,
- save_directory: str,
- format_type: str = "16-bit (FP16)",
- push_to_hub: bool = False,
- repo_id: Optional[str] = None,
- hf_token: Optional[str] = None,
- private: bool = False) -> Tuple[bool, str]:
+ def export_merged_model(
+ self,
+ save_directory: str,
+ format_type: str = "16-bit (FP16)",
+ push_to_hub: bool = False,
+ repo_id: Optional[str] = None,
+ hf_token: Optional[str] = None,
+ private: bool = False,
+ ) -> Tuple[bool, str]:
"""
Export merged model (for PEFT models).
@@ -319,7 +335,7 @@ class ExportBackend:
# Determine save method
if format_type == "4-bit (FP4)":
save_method = "merged_4bit_forced"
- elif self._audio_type == 'whisper':
+ elif self._audio_type == "whisper":
# Whisper uses save_method=None for local 16-bit merged save
save_method = None
else: # 16-bit (FP16)
@@ -332,9 +348,7 @@ class ExportBackend:
ensure_dir(Path(save_directory))
self.current_model.save_pretrained_merged(
- save_directory,
- self.current_tokenizer,
- save_method=save_method
+ save_directory, self.current_tokenizer, save_method = save_method
)
# Write export metadata so the Chat page can identify the base model
@@ -344,18 +358,23 @@ class ExportBackend:
# Push to hub if requested
if push_to_hub:
if not repo_id or not hf_token:
- return False, "Repository ID and Hugging Face token required for Hub upload"
+ return (
+ False,
+ "Repository ID and Hugging Face token required for Hub upload",
+ )
logger.info(f"Pushing merged model to Hub: {repo_id}")
# Whisper uses save_method=None for local but "merged_16bit" for hub push
- hub_save_method = save_method if save_method is not None else "merged_16bit"
+ hub_save_method = (
+ save_method if save_method is not None else "merged_16bit"
+ )
self.current_model.push_to_hub_merged(
repo_id,
self.current_tokenizer,
- save_method=hub_save_method,
- token=hf_token,
- private=private
+ save_method = hub_save_method,
+ token = hf_token,
+ private = private,
)
logger.info(f"Model pushed successfully to {repo_id}")
@@ -364,16 +383,19 @@ class ExportBackend:
except Exception as e:
logger.error(f"Error exporting merged model: {e}")
import traceback
+
logger.error(traceback.format_exc())
return False, f"Export failed: {str(e)}"
- def export_base_model(self,
- save_directory: str,
- push_to_hub: bool = False,
- repo_id: Optional[str] = None,
- hf_token: Optional[str] = None,
- private: bool = False,
- base_model_id: Optional[str] = None) -> Tuple[bool, str]:
+ def export_base_model(
+ self,
+ save_directory: str,
+ push_to_hub: bool = False,
+ repo_id: Optional[str] = None,
+ hf_token: Optional[str] = None,
+ private: bool = False,
+ base_model_id: Optional[str] = None,
+ ) -> Tuple[bool, str]:
"""
Export base model (for non-PEFT models).
@@ -384,7 +406,10 @@ class ExportBackend:
return False, "No model loaded. Please select a checkpoint first."
if self.is_peft:
- return False, "This is a PEFT model. Use 'Merged Model' export type instead."
+ return (
+ False,
+ "This is a PEFT model. Use 'Merged Model' export type instead.",
+ )
try:
# Save locally if requested
@@ -403,40 +428,47 @@ class ExportBackend:
# Push to hub if requested
if push_to_hub:
if not repo_id or not hf_token:
- return False, "Repository ID and Hugging Face token required for Hub upload"
+ return (
+ False,
+ "Repository ID and Hugging Face token required for Hub upload",
+ )
logger.info(f"Pushing base model to Hub: {repo_id}")
# Get base model name from request or model config
- base_model = base_model_id or self.current_model.config._name_or_path or "unknown"
+ base_model = (
+ base_model_id
+ or self.current_model.config._name_or_path
+ or "unknown"
+ )
# Create repo
- hf_api = HfApi(token=hf_token)
+ hf_api = HfApi(token = hf_token)
repo_id = PushToHubMixin._create_repo(
PushToHubMixin,
- repo_id=repo_id,
- private=private,
- token=hf_token,
+ repo_id = repo_id,
+ private = private,
+ token = hf_token,
)
username = repo_id.split("/")[0]
# Create and push model card
content = MODEL_CARD.format(
- username=username,
- base_model=base_model,
- model_type=self.current_model.config.model_type,
- method="",
- extra="unsloth",
+ username = username,
+ base_model = base_model,
+ model_type = self.current_model.config.model_type,
+ method = "",
+ extra = "unsloth",
)
card = ModelCard(content)
- card.push_to_hub(repo_id, token=hf_token, commit_message="Unsloth Model Card")
+ card.push_to_hub(
+ repo_id, token = hf_token, commit_message = "Unsloth Model Card"
+ )
# Upload model files
if save_directory:
hf_api.upload_folder(
- folder_path=save_directory,
- repo_id=repo_id,
- repo_type="model"
+ folder_path = save_directory, repo_id = repo_id, repo_type = "model"
)
logger.info(f"Model pushed successfully to {repo_id}")
else:
@@ -447,16 +479,18 @@ class ExportBackend:
except Exception as e:
logger.error(f"Error exporting base model: {e}")
import traceback
+
logger.error(traceback.format_exc())
return False, f"Export failed: {str(e)}"
-
- def export_gguf(self,
- save_directory: str,
- quantization_method: str = "Q4_K_M",
- push_to_hub: bool = False,
- repo_id: Optional[str] = None,
- hf_token: Optional[str] = None) -> Tuple[bool, str]:
+ def export_gguf(
+ self,
+ save_directory: str,
+ quantization_method: str = "Q4_K_M",
+ push_to_hub: bool = False,
+ repo_id: Optional[str] = None,
+ hf_token: Optional[str] = None,
+ ) -> Tuple[bool, str]:
"""
Export model in GGUF format.
@@ -505,17 +539,21 @@ class ExportBackend:
self.current_model.save_pretrained_gguf(
model_save_path,
self.current_tokenizer,
- quantization_method=quant_method
+ quantization_method = quant_method,
)
# Relocate GGUF artifacts into the export directory.
# convert_to_gguf writes .gguf files to cwd (repo root)
# because --outfile is a relative path like "model.Q4_K_M.gguf".
- new_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs
+ new_ggufs = (
+ set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs
+ )
for src in sorted(new_ggufs):
dest = os.path.join(abs_save_dir, os.path.basename(src))
shutil.move(src, dest)
- logger.info(f"Relocated GGUF: {os.path.basename(src)} β {abs_save_dir}/")
+ logger.info(
+ f"Relocated GGUF: {os.path.basename(src)} β {abs_save_dir}/"
+ )
# Flatten any .gguf files from subdirectories into abs_save_dir.
# save_pretrained_gguf may create subdirs (e.g. model_gguf/)
@@ -528,7 +566,7 @@ class ExportBackend:
shutil.move(str(src), dest)
logger.info(f"Relocated GGUF: {src.name} β {abs_save_dir}/")
# Clean up the subdirectory (intermediate HF files, etc.)
- shutil.rmtree(str(sub), ignore_errors=True)
+ shutil.rmtree(str(sub), ignore_errors = True)
logger.info(f"Cleaned up subdirectory: {sub.name}")
# Write export metadata so the Chat page can identify the base model
@@ -546,15 +584,18 @@ class ExportBackend:
# Push to hub if requested
if push_to_hub:
if not repo_id or not hf_token:
- return False, "Repository ID and Hugging Face token required for Hub upload"
+ return (
+ False,
+ "Repository ID and Hugging Face token required for Hub upload",
+ )
logger.info(f"Pushing GGUF model to Hub: {repo_id}")
self.current_model.push_to_hub_gguf(
repo_id,
self.current_tokenizer,
- quantization_method=quant_method,
- token=hf_token
+ quantization_method = quant_method,
+ token = hf_token,
)
logger.info(f"GGUF model pushed successfully to {repo_id}")
@@ -563,15 +604,18 @@ class ExportBackend:
except Exception as e:
logger.error(f"Error exporting GGUF model: {e}")
import traceback
+
logger.error(traceback.format_exc())
return False, f"GGUF export failed: {str(e)}"
- def export_lora_adapter(self,
- save_directory: str,
- push_to_hub: bool = False,
- repo_id: Optional[str] = None,
- hf_token: Optional[str] = None,
- private: bool = False) -> Tuple[bool, str]:
+ def export_lora_adapter(
+ self,
+ save_directory: str,
+ push_to_hub: bool = False,
+ repo_id: Optional[str] = None,
+ hf_token: Optional[str] = None,
+ private: bool = False,
+ ) -> Tuple[bool, str]:
"""
Export LoRA adapter only (not merged).
@@ -598,19 +642,16 @@ class ExportBackend:
# Push to hub if requested
if push_to_hub:
if not repo_id or not hf_token:
- return False, "Repository ID and Hugging Face token required for Hub upload"
+ return (
+ False,
+ "Repository ID and Hugging Face token required for Hub upload",
+ )
logger.info(f"Pushing LoRA adapter to Hub: {repo_id}")
- self.current_model.push_to_hub(
- repo_id,
- token=hf_token,
- private=private
- )
+ self.current_model.push_to_hub(repo_id, token = hf_token, private = private)
self.current_tokenizer.push_to_hub(
- repo_id,
- token=hf_token,
- private=private
+ repo_id, token = hf_token, private = private
)
logger.info(f"Adapter pushed successfully to {repo_id}")
@@ -619,6 +660,7 @@ class ExportBackend:
except Exception as e:
logger.error(f"Error exporting LoRA adapter: {e}")
import traceback
+
logger.error(traceback.format_exc())
return False, f"Adapter export failed: {str(e)}"
@@ -626,6 +668,7 @@ class ExportBackend:
# Global export backend instance
_export_backend = None
+
def get_export_backend() -> ExportBackend:
"""Get or create the global export backend instance"""
global _export_backend
diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py
index 49b89cc3b1..a9fbe659b3 100644
--- a/studio/backend/core/export/orchestrator.py
+++ b/studio/backend/core/export/orchestrator.py
@@ -13,6 +13,7 @@ the old subprocess is killed and a new one is spawned with the correct version.
Pattern follows core/inference/orchestrator.py.
"""
+
import atexit
import structlog
from loggers import get_logger
@@ -65,13 +66,13 @@ class ExportOrchestrator:
self._resp_queue = _CTX.Queue()
self._proc = _CTX.Process(
- target=run_export_process,
- kwargs={
+ target = run_export_process,
+ kwargs = {
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,
"config": config,
},
- daemon=True,
+ daemon = True,
)
self._proc.start()
logger.info("Export subprocess started (pid=%s)", self._proc.pid)
@@ -93,7 +94,7 @@ class ExportOrchestrator:
# 3. Wait for graceful shutdown
try:
- self._proc.join(timeout=timeout)
+ self._proc.join(timeout = timeout)
except Exception:
pass
@@ -102,14 +103,14 @@ class ExportOrchestrator:
logger.warning("Export subprocess did not exit gracefully, terminating")
try:
self._proc.terminate()
- self._proc.join(timeout=5)
+ self._proc.join(timeout = 5)
except Exception:
pass
if self._proc is not None and self._proc.is_alive():
logger.warning("Subprocess still alive after terminate, killing")
try:
self._proc.kill()
- self._proc.join(timeout=3)
+ self._proc.join(timeout = 3)
except Exception:
pass
@@ -120,7 +121,7 @@ class ExportOrchestrator:
def _cleanup(self):
"""atexit handler."""
- self._shutdown_subprocess(timeout=5.0)
+ self._shutdown_subprocess(timeout = 5.0)
def _ensure_subprocess_alive(self) -> bool:
"""Check if subprocess is alive."""
@@ -144,15 +145,13 @@ class ExportOrchestrator:
if self._resp_queue is None:
return None
try:
- return self._resp_queue.get(timeout=timeout)
+ return self._resp_queue.get(timeout = timeout)
except queue.Empty:
return None
except (EOFError, OSError, ValueError):
return None
- def _wait_response(
- self, expected_type: str, timeout: float = 3600.0
- ) -> dict:
+ def _wait_response(self, expected_type: str, timeout: float = 3600.0) -> dict:
"""Block until a response of the expected type arrives.
Export operations can take a very long time β GGUF conversion for
@@ -163,7 +162,7 @@ class ExportOrchestrator:
while time.monotonic() < deadline:
remaining = max(0.1, deadline - time.monotonic())
- resp = self._read_resp(timeout=min(remaining, 2.0))
+ resp = self._read_resp(timeout = min(remaining, 2.0))
if resp is None:
# Check subprocess health
@@ -187,7 +186,8 @@ class ExportOrchestrator:
# Other response types during wait β skip
logger.debug(
"Skipping response type '%s' while waiting for '%s'",
- rtype, expected_type,
+ rtype,
+ expected_type,
)
raise RuntimeError(
@@ -233,15 +233,15 @@ class ExportOrchestrator:
if self._ensure_subprocess_alive():
self._shutdown_subprocess()
elif self._proc is not None:
- self._shutdown_subprocess(timeout=2)
+ self._shutdown_subprocess(timeout = 2)
logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path)
self._spawn_subprocess(sub_config)
try:
- resp = self._wait_response("loaded", timeout=300)
+ resp = self._wait_response("loaded", timeout = 300)
except RuntimeError as exc:
- self._shutdown_subprocess(timeout=5)
+ self._shutdown_subprocess(timeout = 5)
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
@@ -271,14 +271,17 @@ class ExportOrchestrator:
private: bool = False,
) -> Tuple[bool, str]:
"""Export merged PEFT model."""
- return self._run_export("merged", {
- "save_directory": save_directory,
- "format_type": format_type,
- "push_to_hub": push_to_hub,
- "repo_id": repo_id,
- "hf_token": hf_token,
- "private": private,
- })
+ return self._run_export(
+ "merged",
+ {
+ "save_directory": save_directory,
+ "format_type": format_type,
+ "push_to_hub": push_to_hub,
+ "repo_id": repo_id,
+ "hf_token": hf_token,
+ "private": private,
+ },
+ )
def export_base_model(
self,
@@ -290,14 +293,17 @@ class ExportOrchestrator:
base_model_id: Optional[str] = None,
) -> Tuple[bool, str]:
"""Export base model (non-PEFT)."""
- return self._run_export("base", {
- "save_directory": save_directory,
- "push_to_hub": push_to_hub,
- "repo_id": repo_id,
- "hf_token": hf_token,
- "private": private,
- "base_model_id": base_model_id,
- })
+ return self._run_export(
+ "base",
+ {
+ "save_directory": save_directory,
+ "push_to_hub": push_to_hub,
+ "repo_id": repo_id,
+ "hf_token": hf_token,
+ "private": private,
+ "base_model_id": base_model_id,
+ },
+ )
def export_gguf(
self,
@@ -308,13 +314,16 @@ class ExportOrchestrator:
hf_token: Optional[str] = None,
) -> Tuple[bool, str]:
"""Export model in GGUF format."""
- return self._run_export("gguf", {
- "save_directory": save_directory,
- "quantization_method": quantization_method,
- "push_to_hub": push_to_hub,
- "repo_id": repo_id,
- "hf_token": hf_token,
- })
+ return self._run_export(
+ "gguf",
+ {
+ "save_directory": save_directory,
+ "quantization_method": quantization_method,
+ "push_to_hub": push_to_hub,
+ "repo_id": repo_id,
+ "hf_token": hf_token,
+ },
+ )
def export_lora_adapter(
self,
@@ -325,13 +334,16 @@ class ExportOrchestrator:
private: bool = False,
) -> Tuple[bool, str]:
"""Export LoRA adapter only."""
- return self._run_export("lora", {
- "save_directory": save_directory,
- "push_to_hub": push_to_hub,
- "repo_id": repo_id,
- "hf_token": hf_token,
- "private": private,
- })
+ return self._run_export(
+ "lora",
+ {
+ "save_directory": save_directory,
+ "push_to_hub": push_to_hub,
+ "repo_id": repo_id,
+ "hf_token": hf_token,
+ "private": private,
+ },
+ )
def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str]:
"""Send an export command to the subprocess and wait for result."""
@@ -344,7 +356,7 @@ class ExportOrchestrator:
self._send_cmd(cmd)
resp = self._wait_response(
f"export_{export_type}_done",
- timeout=3600, # GGUF for 30B+ models can take 30+ min
+ timeout = 3600, # GGUF for 30B+ models can take 30+ min
)
return resp.get("success", False), resp.get("message", "")
except RuntimeError as exc:
@@ -361,7 +373,7 @@ class ExportOrchestrator:
try:
self._send_cmd({"type": "cleanup"})
- resp = self._wait_response("cleanup_done", timeout=30)
+ resp = self._wait_response("cleanup_done", timeout = 30)
success = resp.get("success", False)
except RuntimeError:
success = False
@@ -379,7 +391,8 @@ class ExportOrchestrator:
) -> List[Tuple[str, list]]:
"""Scan for checkpoints β no ML imports needed, runs locally."""
from utils.models.checkpoints import scan_checkpoints
- return scan_checkpoints(outputs_dir=outputs_dir)
+
+ return scan_checkpoints(outputs_dir = outputs_dir)
# ========== GLOBAL INSTANCE ==========
diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py
index 197b9fa4a8..1128a5014a 100644
--- a/studio/backend/core/export/worker.py
+++ b/studio/backend/core/export/worker.py
@@ -14,6 +14,7 @@ shutdown) via mp.Queue.
Pattern follows core/inference/worker.py and core/training/worker.py.
"""
+
from __future__ import annotations
import structlog
@@ -43,25 +44,45 @@ def _activate_transformers_version(model_name: str) -> None:
resolved = _resolve_base_model(model_name)
if needs_transformers_5(resolved):
- venv_t5 = os.path.join(os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5")
+ venv_t5 = os.path.join(
+ os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5"
+ )
if os.path.isdir(venv_t5):
sys.path.insert(0, venv_t5)
logger.info("Activated transformers 5.x from %s", venv_t5)
else:
-
# Fallback: pip install at runtime (slower, ~10-15s)
logger.warning(".venv_t5 not found at %s β installing at runtime", venv_t5)
import subprocess as sp
- os.makedirs(venv_t5, exist_ok=True)
+
+ os.makedirs(venv_t5, exist_ok = True)
r1 = sp.run(
- [sys.executable, "-m", "pip", "install", "--target", venv_t5,
- "--no-deps", "transformers==5.2.0"],
- stdout=sp.PIPE, stderr=sp.STDOUT,
+ [
+ sys.executable,
+ "-m",
+ "pip",
+ "install",
+ "--target",
+ venv_t5,
+ "--no-deps",
+ "transformers==5.2.0",
+ ],
+ stdout = sp.PIPE,
+ stderr = sp.STDOUT,
)
r2 = sp.run(
- [sys.executable, "-m", "pip", "install", "--target", venv_t5,
- "--no-deps", "huggingface_hub==1.3.0"],
- stdout=sp.PIPE, stderr=sp.STDOUT,
+ [
+ sys.executable,
+ "-m",
+ "pip",
+ "install",
+ "--target",
+ venv_t5,
+ "--no-deps",
+ "huggingface_hub==1.3.0",
+ ],
+ stdout = sp.PIPE,
+ stderr = sp.STDOUT,
)
if r1.returncode != 0 or r2.returncode != 0:
raise RuntimeError(
@@ -92,37 +113,46 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
trust_remote_code = cmd.get("trust_remote_code", False)
try:
- _send_response(resp_queue, {
- "type": "status",
- "message": f"Loading checkpoint: {checkpoint_path}",
- "ts": time.time(),
- })
-
- success, message = backend.load_checkpoint(
- checkpoint_path=checkpoint_path,
- max_seq_length=max_seq_length,
- load_in_4bit=load_in_4bit,
- trust_remote_code=trust_remote_code,
+ _send_response(
+ resp_queue,
+ {
+ "type": "status",
+ "message": f"Loading checkpoint: {checkpoint_path}",
+ "ts": time.time(),
+ },
)
- _send_response(resp_queue, {
- "type": "loaded",
- "success": success,
- "message": message,
- "checkpoint": checkpoint_path if success else None,
- "is_vision": backend.is_vision if success else False,
- "is_peft": backend.is_peft if success else False,
- "ts": time.time(),
- })
+ success, message = backend.load_checkpoint(
+ checkpoint_path = checkpoint_path,
+ max_seq_length = max_seq_length,
+ load_in_4bit = load_in_4bit,
+ trust_remote_code = trust_remote_code,
+ )
+
+ _send_response(
+ resp_queue,
+ {
+ "type": "loaded",
+ "success": success,
+ "message": message,
+ "checkpoint": checkpoint_path if success else None,
+ "is_vision": backend.is_vision if success else False,
+ "is_peft": backend.is_peft if success else False,
+ "ts": time.time(),
+ },
+ )
except Exception as exc:
- _send_response(resp_queue, {
- "type": "loaded",
- "success": False,
- "message": str(exc),
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "loaded",
+ "success": False,
+ "message": str(exc),
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
@@ -133,74 +163,86 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
try:
if export_type == "merged":
success, message = backend.export_merged_model(
- save_directory=cmd.get("save_directory", ""),
- format_type=cmd.get("format_type", "16-bit (FP16)"),
- push_to_hub=cmd.get("push_to_hub", False),
- repo_id=cmd.get("repo_id"),
- hf_token=cmd.get("hf_token"),
- private=cmd.get("private", False),
+ save_directory = cmd.get("save_directory", ""),
+ format_type = cmd.get("format_type", "16-bit (FP16)"),
+ push_to_hub = cmd.get("push_to_hub", False),
+ repo_id = cmd.get("repo_id"),
+ hf_token = cmd.get("hf_token"),
+ private = cmd.get("private", False),
)
elif export_type == "base":
success, message = backend.export_base_model(
- save_directory=cmd.get("save_directory", ""),
- push_to_hub=cmd.get("push_to_hub", False),
- repo_id=cmd.get("repo_id"),
- hf_token=cmd.get("hf_token"),
- private=cmd.get("private", False),
- base_model_id=cmd.get("base_model_id"),
+ save_directory = cmd.get("save_directory", ""),
+ push_to_hub = cmd.get("push_to_hub", False),
+ repo_id = cmd.get("repo_id"),
+ hf_token = cmd.get("hf_token"),
+ private = cmd.get("private", False),
+ base_model_id = cmd.get("base_model_id"),
)
elif export_type == "gguf":
success, message = backend.export_gguf(
- save_directory=cmd.get("save_directory", ""),
- quantization_method=cmd.get("quantization_method", "Q4_K_M"),
- push_to_hub=cmd.get("push_to_hub", False),
- repo_id=cmd.get("repo_id"),
- hf_token=cmd.get("hf_token"),
+ save_directory = cmd.get("save_directory", ""),
+ quantization_method = cmd.get("quantization_method", "Q4_K_M"),
+ push_to_hub = cmd.get("push_to_hub", False),
+ repo_id = cmd.get("repo_id"),
+ hf_token = cmd.get("hf_token"),
)
elif export_type == "lora":
success, message = backend.export_lora_adapter(
- save_directory=cmd.get("save_directory", ""),
- push_to_hub=cmd.get("push_to_hub", False),
- repo_id=cmd.get("repo_id"),
- hf_token=cmd.get("hf_token"),
- private=cmd.get("private", False),
+ save_directory = cmd.get("save_directory", ""),
+ push_to_hub = cmd.get("push_to_hub", False),
+ repo_id = cmd.get("repo_id"),
+ hf_token = cmd.get("hf_token"),
+ private = cmd.get("private", False),
)
else:
success, message = False, f"Unknown export type: {export_type}"
- _send_response(resp_queue, {
- "type": response_type,
- "success": success,
- "message": message,
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": response_type,
+ "success": success,
+ "message": message,
+ "ts": time.time(),
+ },
+ )
except Exception as exc:
- _send_response(resp_queue, {
- "type": response_type,
- "success": False,
- "message": str(exc),
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": response_type,
+ "success": False,
+ "message": str(exc),
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
def _handle_cleanup(backend, resp_queue: Any) -> None:
"""Handle a cleanup command."""
try:
success = backend.cleanup_memory()
- _send_response(resp_queue, {
- "type": "cleanup_done",
- "success": success,
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "cleanup_done",
+ "success": success,
+ "ts": time.time(),
+ },
+ )
except Exception as exc:
- _send_response(resp_queue, {
- "type": "cleanup_done",
- "success": False,
- "message": str(exc),
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "cleanup_done",
+ "success": False,
+ "message": str(exc),
+ "ts": time.time(),
+ },
+ )
def run_export_process(
@@ -219,16 +261,19 @@ def run_export_process(
import queue as _queue
os.environ["TOKENIZERS_PARALLELISM"] = "false"
- os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
+ os.environ["PYTHONWARNINGS"] = (
+ "ignore" # Suppress warnings at C-level before imports
+ )
import warnings
from loggers.config import LogConfig
+
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
LogConfig.setup_logging(
- service_name="unsloth-studio-export-worker",
- env=os.getenv("ENVIRONMENT_TYPE", "production"),
+ service_name = "unsloth-studio-export-worker",
+ env = os.getenv("ENVIRONMENT_TYPE", "production"),
)
checkpoint_path = config["checkpoint_path"]
@@ -237,18 +282,22 @@ def run_export_process(
try:
_activate_transformers_version(checkpoint_path)
except Exception as exc:
- _send_response(resp_queue, {
- "type": "error",
- "error": f"Failed to activate transformers version: {exc}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "error",
+ "error": f"Failed to activate transformers version: {exc}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
return
# ββ 1b. On Windows, check Triton availability (must be before import torch) ββ
if sys.platform == "win32":
try:
import triton # noqa: F401
+
logger.info("Triton available β torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
@@ -259,11 +308,14 @@ def run_export_process(
# ββ 2. Import ML libraries (fresh in this clean process) ββ
try:
- _send_response(resp_queue, {
- "type": "status",
- "message": "Importing ML libraries...",
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "status",
+ "message": "Importing ML libraries...",
+ "ts": time.time(),
+ },
+ )
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
@@ -272,15 +324,21 @@ def run_export_process(
from core.export.export import ExportBackend
import transformers
- logger.info("Export subprocess loaded transformers %s", transformers.__version__)
+
+ logger.info(
+ "Export subprocess loaded transformers %s", transformers.__version__
+ )
except Exception as exc:
- _send_response(resp_queue, {
- "type": "error",
- "error": f"Failed to import ML libraries: {exc}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "error",
+ "error": f"Failed to import ML libraries: {exc}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
return
# ββ 3. Create export backend and load initial checkpoint ββ
@@ -290,12 +348,15 @@ def run_export_process(
_handle_load(backend, config, resp_queue)
except Exception as exc:
- _send_response(resp_queue, {
- "type": "error",
- "error": f"Failed to initialize export backend: {exc}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "error",
+ "error": f"Failed to initialize export backend: {exc}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
return
# ββ 4. Command loop β process commands until shutdown ββ
@@ -303,7 +364,7 @@ def run_export_process(
while True:
try:
- cmd = cmd_queue.get(timeout=1.0)
+ cmd = cmd_queue.get(timeout = 1.0)
except _queue.Empty:
continue
except (EOFError, OSError):
@@ -329,13 +390,16 @@ def run_export_process(
_handle_cleanup(backend, resp_queue)
elif cmd_type == "status":
- _send_response(resp_queue, {
- "type": "status_response",
- "checkpoint": backend.current_checkpoint,
- "is_vision": backend.is_vision,
- "is_peft": backend.is_peft,
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "status_response",
+ "checkpoint": backend.current_checkpoint,
+ "is_vision": backend.is_vision,
+ "is_peft": backend.is_peft,
+ "ts": time.time(),
+ },
+ )
elif cmd_type == "shutdown":
logger.info("Shutdown command received, cleaning up and exiting")
@@ -343,25 +407,36 @@ def run_export_process(
backend.cleanup_memory()
except Exception:
pass
- _send_response(resp_queue, {
- "type": "shutdown_ack",
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "shutdown_ack",
+ "ts": time.time(),
+ },
+ )
return
else:
logger.warning("Unknown command type: %s", cmd_type)
- _send_response(resp_queue, {
- "type": "error",
- "error": f"Unknown command type: {cmd_type}",
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "error",
+ "error": f"Unknown command type: {cmd_type}",
+ "ts": time.time(),
+ },
+ )
except Exception as exc:
- logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info=True)
- _send_response(resp_queue, {
- "type": "error",
- "error": f"Command '{cmd_type}' failed: {exc}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ logger.error(
+ "Error handling command '%s': %s", cmd_type, exc, exc_info = True
+ )
+ _send_response(
+ resp_queue,
+ {
+ "type": "error",
+ "error": f"Command '{cmd_type}' failed: {exc}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py
index 97f5a03f79..35318f6357 100644
--- a/studio/backend/core/inference/__init__.py
+++ b/studio/backend/core/inference/__init__.py
@@ -8,6 +8,7 @@ The default get_inference_backend() returns an InferenceOrchestrator that
delegates to a subprocess. The original InferenceBackend runs inside
the subprocess and can be imported directly from .inference when needed.
"""
+
from .orchestrator import InferenceOrchestrator, get_inference_backend
from .llama_cpp import LlamaCppBackend
@@ -15,8 +16,8 @@ from .llama_cpp import LlamaCppBackend
InferenceBackend = InferenceOrchestrator
__all__ = [
- 'InferenceBackend',
- 'InferenceOrchestrator',
- 'get_inference_backend',
- 'LlamaCppBackend',
+ "InferenceBackend",
+ "InferenceOrchestrator",
+ "get_inference_backend",
+ "LlamaCppBackend",
]
diff --git a/studio/backend/core/inference/audio_codecs.py b/studio/backend/core/inference/audio_codecs.py
index b22dd1deab..3a418d921d 100644
--- a/studio/backend/core/inference/audio_codecs.py
+++ b/studio/backend/core/inference/audio_codecs.py
@@ -5,6 +5,7 @@
Audio codec loading and decoding for TTS inference.
Supports: SNAC (Orpheus), CSM (Sesame), BiCodec (Spark), DAC (OuteTTS)
"""
+
import io
import re
import wave
@@ -45,7 +46,12 @@ class AudioCodecManager:
self._bicodec_repo_path = None
self._dac_audio_codec = None
- def load_codec(self, audio_type: str, device: str = "cuda", model_repo_path: Optional[str] = None) -> None:
+ def load_codec(
+ self,
+ audio_type: str,
+ device: str = "cuda",
+ model_repo_path: Optional[str] = None,
+ ) -> None:
"""Load the appropriate codec for the given audio type."""
if audio_type == "snac":
self._load_snac(device)
@@ -64,7 +70,10 @@ class AudioCodecManager:
if self._snac_model is not None:
return
from snac import SNAC
- self._snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
+
+ self._snac_model = (
+ SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
+ )
logger.info("Loaded SNAC codec (24kHz)")
def _load_bicodec(self, device: str, model_repo_path: Optional[str] = None) -> None:
@@ -76,13 +85,22 @@ class AudioCodecManager:
# Clone SparkAudio/Spark-TTS GitHub repo for the sparktts Python package
# (same approach as training β the HF model repos don't contain the package)
- spark_code_dir = os.path.join(os.path.dirname(model_repo_path or "."), "Spark-TTS")
+ spark_code_dir = os.path.join(
+ os.path.dirname(model_repo_path or "."), "Spark-TTS"
+ )
sparktts_pkg = os.path.join(spark_code_dir, "sparktts")
if not os.path.isdir(sparktts_pkg):
logger.info(f"Cloning SparkAudio/Spark-TTS to {spark_code_dir}...")
subprocess.run(
- ["git", "clone", "--depth", "1", "https://github.com/SparkAudio/Spark-TTS", spark_code_dir],
- check=True,
+ [
+ "git",
+ "clone",
+ "--depth",
+ "1",
+ "https://github.com/SparkAudio/Spark-TTS",
+ spark_code_dir,
+ ],
+ check = True,
)
if spark_code_dir not in sys.path:
@@ -112,8 +130,15 @@ class AudioCodecManager:
if not os.path.isdir(outetts_pkg):
logger.info(f"Cloning edwko/OuteTTS to {outetts_code_dir}...")
subprocess.run(
- ["git", "clone", "--depth", "1", "https://github.com/edwko/OuteTTS", outetts_code_dir],
- check=True,
+ [
+ "git",
+ "clone",
+ "--depth",
+ "1",
+ "https://github.com/edwko/OuteTTS",
+ outetts_code_dir,
+ ],
+ check = True,
)
# Remove files that pull in heavy / incompatible dependencies
# (matches notebook: gguf_model.py is under models/, others under outetts/)
@@ -134,17 +159,19 @@ class AudioCodecManager:
from outetts.models.config import ModelConfig as OuteTTSModelConfig
dummy_config = OuteTTSModelConfig(
- tokenizer_path="OuteAI/Llama-OuteTTS-1.0-1B",
- device=device,
- audio_codec_path=None,
+ tokenizer_path = "OuteAI/Llama-OuteTTS-1.0-1B",
+ device = device,
+ audio_codec_path = None,
)
- processor = AudioProcessor(config=dummy_config)
+ processor = AudioProcessor(config = dummy_config)
self._dac_audio_codec = processor.audio_codec
logger.info("Loaded DAC audio codec")
# ββ Decoders βββββββββββββββββββββββββββββββββββββββββββββββββ
- def decode_snac(self, generated_ids: torch.Tensor, device: str) -> Tuple[bytes, int]:
+ def decode_snac(
+ self, generated_ids: torch.Tensor, device: str
+ ) -> Tuple[bytes, int]:
"""
Decode SNAC tokens (Orpheus) into WAV bytes.
@@ -155,12 +182,14 @@ class AudioCodecManager:
Returns (wav_bytes, 24000).
"""
# Find START_OF_SPEECH token (128257)
- token_indices = (generated_ids == 128257).nonzero(as_tuple=True)
+ token_indices = (generated_ids == 128257).nonzero(as_tuple = True)
if len(token_indices[1]) > 0:
- cropped = generated_ids[:, token_indices[1][-1] + 1:]
+ cropped = generated_ids[:, token_indices[1][-1] + 1 :]
else:
# Gracefully fall back to using entire output if marker not found
- logger.warning("No START_OF_SPEECH token (128257) found β using full generated output")
+ logger.warning(
+ "No START_OF_SPEECH token (128257) found β using full generated output"
+ )
cropped = generated_ids
row = cropped[0]
@@ -213,14 +242,20 @@ class AudioCodecManager:
semantic_matches = re.findall(r"<\|bicodec_semantic_(\d+)\|>", generated_text)
global_matches = re.findall(r"<\|bicodec_global_(\d+)\|>", generated_text)
- logger.info(f"BiCodec decode: {len(global_matches)} global tokens, {len(semantic_matches)} semantic tokens")
+ logger.info(
+ f"BiCodec decode: {len(global_matches)} global tokens, {len(semantic_matches)} semantic tokens"
+ )
if len(global_matches) < 10:
- logger.info(f"BiCodec generated text (first 500 chars): {generated_text[:500]}")
+ logger.info(
+ f"BiCodec generated text (first 500 chars): {generated_text[:500]}"
+ )
if not semantic_matches:
raise ValueError("No bicodec_semantic tokens found in generated output")
- semantic_ids = torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0)
+ semantic_ids = (
+ torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0)
+ )
# Speaker encoder expects exactly 32 global tokens (token_num=32 in BiCodec config).
# Pad with zeros or truncate to 32.
@@ -260,7 +295,7 @@ class AudioCodecManager:
c1 = c1[:t]
c2 = c2[:t]
- codes = torch.tensor([[c1, c2]], dtype=torch.int64).to(device)
+ codes = torch.tensor([[c1, c2]], dtype = torch.int64).to(device)
with torch.no_grad():
audio = self._dac_audio_codec.decode(codes)
diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py
index a2ebbc0f5c..ff02816953 100644
--- a/studio/backend/core/inference/inference.py
+++ b/studio/backend/core/inference/inference.py
@@ -4,6 +4,7 @@
"""
Core inference backend - streamlined
"""
+
from unsloth import FastLanguageModel, FastVisionModel
from unsloth.chat_templates import get_chat_template
from transformers import TextStreamer
@@ -24,9 +25,9 @@ import structlog
from loggers import get_logger
-
logger = get_logger(__name__)
+
class InferenceBackend:
"""Unified inference backend supporting text, vision, and LoRA models"""
@@ -52,6 +53,7 @@ class InferenceBackend:
# concurrent compare-mode requests race on the GPU. The lock is
# acquired by the *background generation thread*, not the event-loop.
import threading
+
self._generation_lock = threading.Lock()
self._model_state_lock = threading.Lock()
@@ -62,13 +64,15 @@ class InferenceBackend:
# API supports -1 as "disable top-k"; transformers expects 0 to disable.
return 0 if top_k < 0 else top_k
- def load_model(self,
- config: ModelConfig,
- max_seq_length: int = 2048,
- dtype = None,
- load_in_4bit: bool = True,
- hf_token: Optional[str] = None,
- trust_remote_code: bool = False) -> bool:
+ def load_model(
+ self,
+ config: ModelConfig,
+ max_seq_length: int = 2048,
+ dtype = None,
+ load_in_4bit: bool = True,
+ hf_token: Optional[str] = None,
+ trust_remote_code: bool = False,
+ ) -> bool:
"""
Load any model: base, LoRA adapter, text, or vision.
"""
@@ -104,18 +108,21 @@ class InferenceBackend:
if config.is_audio:
audio_type = config.audio_type
adapter_info = " (LoRA adapter)" if config.is_lora else ""
- logger.info(f"Loading audio ({audio_type}) model{adapter_info}: {model_name}")
+ logger.info(
+ f"Loading audio ({audio_type}) model{adapter_info}: {model_name}"
+ )
log_gpu_memory(f"Before loading {model_name}")
if audio_type == "csm":
from unsloth import FastModel
from transformers import CsmForConditionalGeneration
+
model, processor = FastModel.from_pretrained(
config.path,
- auto_model=CsmForConditionalGeneration,
- load_in_4bit=False,
- token=hf_token if hf_token and hf_token.strip() else None,
- trust_remote_code=trust_remote_code,
+ auto_model = CsmForConditionalGeneration,
+ load_in_4bit = False,
+ token = hf_token if hf_token and hf_token.strip() else None,
+ trust_remote_code = trust_remote_code,
)
FastModel.for_inference(model)
self.models[model_name]["model"] = model
@@ -135,34 +142,42 @@ class InferenceBackend:
else:
# base_model is an HF ID β download it
from huggingface_hub import snapshot_download
+
local_dir = base_path.split("/")[-1]
- repo_path = snapshot_download(base_path, local_dir=local_dir)
+ repo_path = snapshot_download(
+ base_path, local_dir = local_dir
+ )
abs_repo_path = os.path.abspath(repo_path)
- logger.info(f"Spark-TTS LoRA: loading adapter from {config.path}, BiCodec from {abs_repo_path}")
+ logger.info(
+ f"Spark-TTS LoRA: loading adapter from {config.path}, BiCodec from {abs_repo_path}"
+ )
model, tokenizer = FastModel.from_pretrained(
config.path,
- dtype=torch.float32,
- load_in_4bit=False,
- token=hf_token if hf_token and hf_token.strip() else None,
- trust_remote_code=trust_remote_code,
+ dtype = torch.float32,
+ load_in_4bit = False,
+ token = hf_token if hf_token and hf_token.strip() else None,
+ trust_remote_code = trust_remote_code,
)
else:
# Base model: download full HF repo, then load from /LLM subfolder
from huggingface_hub import snapshot_download
+
hf_repo = config.path
local_dir = hf_repo.split("/")[-1]
- repo_path = snapshot_download(hf_repo, local_dir=local_dir)
+ repo_path = snapshot_download(hf_repo, local_dir = local_dir)
abs_repo_path = os.path.abspath(repo_path)
llm_path = os.path.join(abs_repo_path, "LLM")
- logger.info(f"Spark-TTS: downloaded repo to {repo_path}, loading LLM from {llm_path}")
+ logger.info(
+ f"Spark-TTS: downloaded repo to {repo_path}, loading LLM from {llm_path}"
+ )
model, tokenizer = FastModel.from_pretrained(
llm_path,
- dtype=torch.float32,
- load_in_4bit=False,
- token=hf_token if hf_token and hf_token.strip() else None,
- trust_remote_code=trust_remote_code,
+ dtype = torch.float32,
+ load_in_4bit = False,
+ token = hf_token if hf_token and hf_token.strip() else None,
+ trust_remote_code = trust_remote_code,
)
FastModel.for_inference(model)
@@ -172,12 +187,13 @@ class InferenceBackend:
elif audio_type == "dac":
# OuteTTS uses FastModel (not FastLanguageModel)
from unsloth import FastModel
+
model, tokenizer = FastModel.from_pretrained(
config.path,
- max_seq_length=max_seq_length,
- load_in_4bit=False,
- token=hf_token if hf_token and hf_token.strip() else None,
- trust_remote_code=trust_remote_code,
+ max_seq_length = max_seq_length,
+ load_in_4bit = False,
+ token = hf_token if hf_token and hf_token.strip() else None,
+ trust_remote_code = trust_remote_code,
)
FastModel.for_inference(model)
self.models[model_name]["model"] = model
@@ -186,28 +202,30 @@ class InferenceBackend:
# Whisper ASR β uses FastModel with WhisperForConditionalGeneration
from unsloth import FastModel
from transformers import WhisperForConditionalGeneration
+
model, tokenizer = FastModel.from_pretrained(
config.path,
- auto_model=WhisperForConditionalGeneration,
- whisper_language="English",
- whisper_task="transcribe",
- load_in_4bit=False,
- token=hf_token if hf_token and hf_token.strip() else None,
- trust_remote_code=trust_remote_code,
+ auto_model = WhisperForConditionalGeneration,
+ whisper_language = "English",
+ whisper_task = "transcribe",
+ load_in_4bit = False,
+ token = hf_token if hf_token and hf_token.strip() else None,
+ trust_remote_code = trust_remote_code,
)
FastModel.for_inference(model)
model.eval()
# Create ASR pipeline (per notebook)
from transformers import pipeline as hf_pipeline
+
whisper_pipe = hf_pipeline(
"automatic-speech-recognition",
- model=model,
- tokenizer=tokenizer.tokenizer,
- feature_extractor=tokenizer.feature_extractor,
- processor=tokenizer,
- return_language=True,
- torch_dtype=torch.float16,
+ model = model,
+ tokenizer = tokenizer.tokenizer,
+ feature_extractor = tokenizer.feature_extractor,
+ processor = tokenizer,
+ return_language = True,
+ torch_dtype = torch.float16,
)
self.models[model_name]["model"] = model
self.models[model_name]["tokenizer"] = tokenizer
@@ -215,11 +233,11 @@ class InferenceBackend:
else:
# SNAC (Orpheus) uses FastLanguageModel
model, tokenizer = FastLanguageModel.from_pretrained(
- model_name=config.path,
- max_seq_length=max_seq_length,
- load_in_4bit=False,
- token=hf_token if hf_token and hf_token.strip() else None,
- trust_remote_code=trust_remote_code,
+ model_name = config.path,
+ max_seq_length = max_seq_length,
+ load_in_4bit = False,
+ token = hf_token if hf_token and hf_token.strip() else None,
+ trust_remote_code = trust_remote_code,
)
FastLanguageModel.for_inference(model)
self.models[model_name]["model"] = model
@@ -229,7 +247,9 @@ class InferenceBackend:
# (Whisper is ASR, audio_vlm is audio input β neither needs a codec)
if audio_type not in ("whisper", "audio_vlm"):
model_repo_path = self.models[model_name].get("model_repo_path")
- self._audio_codec_manager.load_codec(audio_type, self.device, model_repo_path=model_repo_path)
+ self._audio_codec_manager.load_codec(
+ audio_type, self.device, model_repo_path = model_repo_path
+ )
self.active_model_name = model_name
self.loading_models.discard(model_name)
@@ -238,7 +258,9 @@ class InferenceBackend:
return True
model_type = "vision" if config.is_vision else "text"
- adapter_info = " (LoRA adapter)" if self.models[model_name]["is_lora"] else ""
+ adapter_info = (
+ " (LoRA adapter)" if self.models[model_name]["is_lora"] else ""
+ )
logger.info(f"Loading {model_type} model{adapter_info}: {model_name}")
log_gpu_memory(f"Before loading {model_name}")
@@ -246,12 +268,12 @@ class InferenceBackend:
if config.is_vision:
# Vision model (or vision LoRA adapter)
model, processor = FastVisionModel.from_pretrained(
- model_name=config.path, # Can be base model OR LoRA adapter path
- max_seq_length=max_seq_length,
- dtype=dtype,
- load_in_4bit=load_in_4bit,
- token=hf_token if hf_token and hf_token.strip() else None,
- trust_remote_code=trust_remote_code,
+ model_name = config.path, # Can be base model OR LoRA adapter path
+ max_seq_length = max_seq_length,
+ dtype = dtype,
+ load_in_4bit = load_in_4bit,
+ token = hf_token if hf_token and hf_token.strip() else None,
+ trust_remote_code = trust_remote_code,
)
# Apply inference optimization
@@ -261,10 +283,16 @@ class InferenceBackend:
# instead of a proper Processor for some models (e.g. Gemma-3).
# In that case, load the real processor from the base model.
from transformers import ProcessorMixin
- if not (isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")):
+
+ if not (
+ isinstance(processor, ProcessorMixin)
+ or hasattr(processor, "image_processor")
+ ):
# For LoRA adapters, use the base model. For local merged exports,
# read export_metadata.json to find the original base model.
- processor_source = config.base_model if config.is_lora else config.identifier
+ processor_source = (
+ config.base_model if config.is_lora else config.identifier
+ )
if not config.is_lora and config.is_local:
_meta_path = Path(config.path) / "export_metadata.json"
try:
@@ -279,12 +307,15 @@ class InferenceBackend:
f"for '{model_name}' β loading proper processor from '{processor_source}'"
)
from transformers import AutoProcessor
+
processor = AutoProcessor.from_pretrained(
processor_source,
- token=hf_token if hf_token and hf_token.strip() else None,
- trust_remote_code=trust_remote_code,
+ token = hf_token if hf_token and hf_token.strip() else None,
+ trust_remote_code = trust_remote_code,
+ )
+ logger.info(
+ f"Loaded {type(processor).__name__} from {processor_source}"
)
- logger.info(f"Loaded {type(processor).__name__} from {processor_source}")
self.models[model_name]["model"] = model
self.models[model_name]["tokenizer"] = processor
@@ -293,12 +324,12 @@ class InferenceBackend:
else:
# Text model (or text LoRA adapter)
model, tokenizer = FastLanguageModel.from_pretrained(
- model_name=config.path, # Can be base model OR LoRA adapter path
- max_seq_length=max_seq_length,
- dtype=dtype,
- load_in_4bit=load_in_4bit,
- token=hf_token if hf_token and hf_token.strip() else None,
- trust_remote_code=trust_remote_code,
+ model_name = config.path, # Can be base model OR LoRA adapter path
+ max_seq_length = max_seq_length,
+ dtype = dtype,
+ load_in_4bit = load_in_4bit,
+ token = hf_token if hf_token and hf_token.strip() else None,
+ trust_remote_code = trust_remote_code,
)
# Apply inference optimization
@@ -351,6 +382,7 @@ class InferenceBackend:
# Remove stale compiled cache so the next model gets a fresh one
from utils.cache_cleanup import clear_unsloth_compiled_cache
+
clear_unsloth_compiled_cache()
logger.info(f"Model '{model_name}' successfully unloaded.")
@@ -359,7 +391,9 @@ class InferenceBackend:
logger.error(f"Error while unloading model '{model_name}': {e}")
return False
else:
- logger.warning(f"Attempted to unload model '{model_name}', but it was not found in the registry.")
+ logger.warning(
+ f"Attempted to unload model '{model_name}', but it was not found in the registry."
+ )
return True
def revert_to_base_model(self, base_model_name: str) -> bool:
@@ -384,7 +418,7 @@ class InferenceBackend:
# After model.unload(), the base model may still carry a peft_config
# attribute. Removing it ensures PeftModel.from_pretrained() gets
# a clean base model without "multiple adapters" warnings.
- if hasattr(model, 'peft_config'):
+ if hasattr(model, "peft_config"):
del model.peft_config
logger.info(f"Model '{base_model_name}' reverted to clean base state.")
@@ -393,12 +427,18 @@ class InferenceBackend:
except Exception as e:
logger.error(f"Failed to revert model to base state: {e}")
import traceback
+
logger.error(traceback.format_exc())
return False
- def load_for_eval(self, lora_path: str, max_seq_length: int = 2048,
- dtype = None, load_in_4bit: bool = True,
- hf_token: Optional[str] = None) -> Tuple[bool, Optional[str], Optional[str]]:
+ def load_for_eval(
+ self,
+ lora_path: str,
+ max_seq_length: int = 2048,
+ dtype = None,
+ load_in_4bit: bool = True,
+ hf_token: Optional[str] = None,
+ ) -> Tuple[bool, Optional[str], Optional[str]]:
"""
Final Corrected Version:
Ensures the base model and the specified adapter are loaded.
@@ -406,6 +446,7 @@ class InferenceBackend:
"""
try:
from utils.models import ModelConfig
+
lora_config = ModelConfig.from_lora_path(lora_path, hf_token)
if not lora_config:
return False, None, None
@@ -413,10 +454,16 @@ class InferenceBackend:
base_model_name = lora_config.base_model
# 1. Load the base model if it's not already in memory
- if base_model_name not in self.models or not self.models[base_model_name].get("model"):
+ if base_model_name not in self.models or not self.models[
+ base_model_name
+ ].get("model"):
logger.info(f"Base model '{base_model_name}' not loaded, loading now.")
- base_config = ModelConfig.from_ui_selection(base_model_name, None, is_lora=False)
- if not self.load_model(base_config, max_seq_length, dtype, load_in_4bit, hf_token):
+ base_config = ModelConfig.from_ui_selection(
+ base_model_name, None, is_lora = False
+ )
+ if not self.load_model(
+ base_config, max_seq_length, dtype, load_in_4bit, hf_token
+ ):
return False, None, None
self.active_model_name = base_model_name
@@ -427,9 +474,9 @@ class InferenceBackend:
# 3. Call our robust load_adapter function to ensure this specific adapter is loaded.
# It will only load from disk if the model doesn't already have it.
adapter_success = self.load_adapter(
- base_model_name=base_model_name,
- adapter_path=lora_path,
- adapter_name=adapter_name
+ base_model_name = base_model_name,
+ adapter_path = lora_path,
+ adapter_name = adapter_name,
)
if not adapter_success:
return False, base_model_name, None
@@ -440,10 +487,13 @@ class InferenceBackend:
except Exception as e:
logger.error(f"Error during load_for_eval: {e}")
import traceback
+
logger.error(traceback.format_exc())
return False, None, None
- def load_adapter(self, base_model_name: str, adapter_path: str, adapter_name: str) -> bool:
+ def load_adapter(
+ self, base_model_name: str, adapter_path: str, adapter_name: str
+ ) -> bool:
"""
Loads an adapter onto the model ONLY if it's not already attached.
"""
@@ -451,20 +501,26 @@ class InferenceBackend:
# Check if this adapter name is already part of the model's config. This is the most reliable check.
if hasattr(model, "peft_config") and adapter_name in model.peft_config:
- logger.info(f"Adapter '{adapter_name}' is already attached to the model. Skipping load.")
+ logger.info(
+ f"Adapter '{adapter_name}' is already attached to the model. Skipping load."
+ )
return True
try:
- logger.info(f"Loading new adapter '{adapter_name}' from '{adapter_path}' onto {base_model_name}")
- model.load_adapter(adapter_path, adapter_name=adapter_name)
+ logger.info(
+ f"Loading new adapter '{adapter_name}' from '{adapter_path}' onto {base_model_name}"
+ )
+ model.load_adapter(adapter_path, adapter_name = adapter_name)
# Update our internal registry ONLY after a successful load.
if "loaded_adapters" not in self.models[base_model_name]:
self.models[base_model_name]["loaded_adapters"] = {}
self.models[base_model_name]["loaded_adapters"][adapter_name] = adapter_path
- total_adapters = len(getattr(model, 'peft_config', {}))
- logger.info(f"Adapter '{adapter_name}' loaded successfully. (Total unique adapters on model: {total_adapters})")
+ total_adapters = len(getattr(model, "peft_config", {}))
+ logger.info(
+ f"Adapter '{adapter_name}' loaded successfully. (Total unique adapters on model: {total_adapters})"
+ )
return True
except Exception as e:
logger.error(f"Failed to load adapter '{adapter_name}': {e}")
@@ -513,15 +569,21 @@ class InferenceBackend:
if use_adapter is False:
# Disable LoRA layers β base model output
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
- logger.info(f"Compare mode: disabling adapters on '{base}' for base model generation")
+ logger.info(
+ f"Compare mode: disabling adapters on '{base}' for base model generation"
+ )
model.base_model.disable_adapter_layers()
else:
- logger.info(f"Compare mode: model '{base}' is not a PeftModel, already base")
+ logger.info(
+ f"Compare mode: model '{base}' is not a PeftModel, already base"
+ )
elif use_adapter is True:
# Re-enable LoRA layers β adapter output
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
- logger.info(f"Compare mode: enabling adapters on '{base}' for LoRA generation")
+ logger.info(
+ f"Compare mode: enabling adapters on '{base}' for LoRA generation"
+ )
model.base_model.enable_adapter_layers()
else:
logger.warning("use_adapter=true but model is not a PeftModel")
@@ -529,16 +591,20 @@ class InferenceBackend:
elif isinstance(use_adapter, str):
# Enable adapters and set the specific one active
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
- logger.info(f"Compare mode: enabling adapter '{use_adapter}' on '{base}'")
+ logger.info(
+ f"Compare mode: enabling adapter '{use_adapter}' on '{base}'"
+ )
model.base_model.enable_adapter_layers()
self.set_active_adapter(base, use_adapter)
else:
- logger.warning(f"use_adapter='{use_adapter}' but model is not a PeftModel")
+ logger.warning(
+ f"use_adapter='{use_adapter}' but model is not a PeftModel"
+ )
def generate_with_adapter_control(
self,
use_adapter: Optional[Union[bool, str]] = None,
- cancel_event=None,
+ cancel_event = None,
**gen_kwargs,
) -> Generator[str, None, None]:
"""
@@ -554,49 +620,53 @@ class InferenceBackend:
**gen_kwargs: Forwarded to generate_chat_response.
"""
yield from self._generate_chat_response_inner(
- cancel_event=cancel_event, _adapter_state=use_adapter, **gen_kwargs
+ cancel_event = cancel_event, _adapter_state = use_adapter, **gen_kwargs
)
- def generate_chat_response(self,
- messages: list,
- system_prompt: str,
- image=None,
- temperature: float = 0.7,
- top_p: float = 0.9,
- top_k: int = 40,
- min_p: float = 0.0,
- max_new_tokens: int = 256,
- repetition_penalty: float = 1.1,
- cancel_event=None) -> Generator[str, None, None]:
+ def generate_chat_response(
+ self,
+ messages: list,
+ system_prompt: str,
+ image = None,
+ temperature: float = 0.7,
+ top_p: float = 0.9,
+ top_k: int = 40,
+ min_p: float = 0.0,
+ max_new_tokens: int = 256,
+ repetition_penalty: float = 1.1,
+ cancel_event = None,
+ ) -> Generator[str, None, None]:
"""
Generate response for text or vision models.
The generation lock is acquired by the background generation thread.
"""
yield from self._generate_chat_response_inner(
- messages=messages,
- system_prompt=system_prompt,
- image=image,
- temperature=temperature,
- top_p=top_p,
- top_k=top_k,
- min_p=min_p,
- max_new_tokens=max_new_tokens,
- repetition_penalty=repetition_penalty,
- cancel_event=cancel_event,
+ messages = messages,
+ system_prompt = system_prompt,
+ image = image,
+ temperature = temperature,
+ top_p = top_p,
+ top_k = top_k,
+ min_p = min_p,
+ max_new_tokens = max_new_tokens,
+ repetition_penalty = repetition_penalty,
+ cancel_event = cancel_event,
)
- def _generate_chat_response_inner(self,
- messages: list,
- system_prompt: str = "",
- image=None,
- temperature: float = 0.7,
- top_p: float = 0.9,
- top_k: int = 40,
- min_p: float = 0.0,
- max_new_tokens: int = 256,
- repetition_penalty: float = 1.1,
- cancel_event=None,
- _adapter_state=None) -> Generator[str, None, None]:
+ def _generate_chat_response_inner(
+ self,
+ messages: list,
+ system_prompt: str = "",
+ image = None,
+ temperature: float = 0.7,
+ top_p: float = 0.9,
+ top_k: int = 40,
+ min_p: float = 0.0,
+ max_new_tokens: int = 256,
+ repetition_penalty: float = 1.1,
+ cancel_event = None,
+ _adapter_state = None,
+ ) -> Generator[str, None, None]:
"""
Inner generation logic. Called by both generate_chat_response
and generate_with_adapter_control.
@@ -621,16 +691,24 @@ class InferenceBackend:
# FastVisionModel may return a raw tokenizer (e.g. GemmaTokenizerFast)
# instead of a proper ProcessorMixin for some models (e.g. Gemma-3).
from transformers import ProcessorMixin
+
processor = model_info.get("processor")
- has_image_processing = (
- processor is not None
- and (isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor"))
+ has_image_processing = processor is not None and (
+ isinstance(processor, ProcessorMixin)
+ or hasattr(processor, "image_processor")
)
if has_image_processing:
yield from self._generate_vision_response(
- messages, system_prompt, image,
- temperature, top_p, top_k, min_p, max_new_tokens, repetition_penalty,
- cancel_event=cancel_event,
+ messages,
+ system_prompt,
+ image,
+ temperature,
+ top_p,
+ top_k,
+ min_p,
+ max_new_tokens,
+ repetition_penalty,
+ cancel_event = cancel_event,
)
return
else:
@@ -645,31 +723,36 @@ class InferenceBackend:
# Step 1: Apply get_chat_template if model is in mapper
try:
- from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, get_tokenizer_chat_template
+ from utils.datasets import (
+ MODEL_TO_TEMPLATE_MAPPER,
+ get_tokenizer_chat_template,
+ )
model_name_lower = self.active_model_name.lower()
# Check if model has a registered template
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
- logger.info(f"Applying chat template '{template_name}' for {self.active_model_name}")
+ logger.info(
+ f"Applying chat template '{template_name}' for {self.active_model_name}"
+ )
# This modifies the tokenizer with the correct template
tokenizer = get_chat_template(
tokenizer,
- chat_template=template_name,
+ chat_template = template_name,
)
else:
- logger.info(f"No registered template for {self.active_model_name}, using tokenizer default")
+ logger.info(
+ f"No registered template for {self.active_model_name}, using tokenizer default"
+ )
except Exception as e:
logger.warning(f"Could not apply get_chat_template: {e}")
# Step 2: Format with tokenizer.apply_chat_template()
try:
formatted_prompt = tokenizer.apply_chat_template(
- messages,
- tokenize=False,
- add_generation_prompt=True
+ messages, tokenize = False, add_generation_prompt = True
)
logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...")
except Exception as e:
@@ -679,14 +762,30 @@ class InferenceBackend:
# Step 3: Generate
yield from self.generate_stream(
- formatted_prompt, temperature, top_p, top_k, min_p, max_new_tokens, repetition_penalty,
- cancel_event=cancel_event,
- _adapter_state=_adapter_state,
+ formatted_prompt,
+ temperature,
+ top_p,
+ top_k,
+ min_p,
+ max_new_tokens,
+ repetition_penalty,
+ cancel_event = cancel_event,
+ _adapter_state = _adapter_state,
)
- def _generate_vision_response(self, messages, system_prompt, image,
- temperature, top_p, top_k, min_p, max_new_tokens,
- repetition_penalty, cancel_event=None) -> Generator[str, None, None]:
+ def _generate_vision_response(
+ self,
+ messages,
+ system_prompt,
+ image,
+ temperature,
+ top_p,
+ top_k,
+ min_p,
+ max_new_tokens,
+ repetition_penalty,
+ cancel_event = None,
+ ) -> Generator[str, None, None]:
"""Handle vision model generation with true token-by-token streaming."""
model_info = self.models[self.active_model_name]
model = model_info["model"]
@@ -699,8 +798,9 @@ class InferenceBackend:
user_message = ""
if messages and messages[-1]["role"] == "user":
import re
+
user_message = messages[-1]["content"]
- user_message = re.sub(r']*>', '', user_message).strip()
+ user_message = re.sub(r"]*>", "", user_message).strip()
if not user_message:
user_message = "Describe this image." if image else "Hello"
@@ -712,22 +812,26 @@ class InferenceBackend:
"role": "user",
"content": [
{"type": "image"},
- {"type": "text", "text": user_message}
+ {"type": "text", "text": user_message},
],
}
]
- input_text = processor.apply_chat_template(vision_messages, add_generation_prompt=True, tokenize=False)
+ input_text = processor.apply_chat_template(
+ vision_messages, add_generation_prompt = True, tokenize = False
+ )
inputs = processor(
image,
input_text,
- add_special_tokens=False,
- return_tensors="pt",
+ add_special_tokens = False,
+ return_tensors = "pt",
).to(self.device)
else:
# Text-only for vision model
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
- inputs = raw_tokenizer(formatted_prompt, return_tensors="pt").to(self.device)
+ inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(
+ self.device
+ )
# Stream with TextIteratorStreamer + background thread
try:
@@ -736,21 +840,21 @@ class InferenceBackend:
streamer = TextIteratorStreamer(
raw_tokenizer,
- skip_prompt=True,
- skip_special_tokens=True,
- timeout=0.2,
+ skip_prompt = True,
+ skip_special_tokens = True,
+ timeout = 0.2,
)
generation_kwargs = dict(
**inputs,
- streamer=streamer,
- max_new_tokens=max_new_tokens,
- use_cache=True,
- do_sample=temperature > 0,
- temperature=temperature,
- top_p=top_p,
- top_k=top_k,
- min_p=min_p,
+ streamer = streamer,
+ max_new_tokens = max_new_tokens,
+ use_cache = True,
+ do_sample = temperature > 0,
+ temperature = temperature,
+ top_p = top_p,
+ top_k = top_k,
+ min_p = min_p,
)
err: dict[str, str] = {}
@@ -768,11 +872,12 @@ class InferenceBackend:
except Exception:
pass
- thread = threading.Thread(target=generate_fn)
+ thread = threading.Thread(target = generate_fn)
thread.start()
output = ""
from queue import Empty
+
try:
while True:
if cancel_event is not None and cancel_event.is_set():
@@ -792,9 +897,11 @@ class InferenceBackend:
finally:
if cancel_event is not None:
cancel_event.set()
- thread.join(timeout=10)
+ thread.join(timeout = 10)
if thread.is_alive():
- logger.warning("Vision generation thread did not exit after cancel/join timeout")
+ logger.warning(
+ "Vision generation thread did not exit after cancel/join timeout"
+ )
if err.get("msg"):
yield f"Error: {err['msg']}"
@@ -803,10 +910,19 @@ class InferenceBackend:
logger.error(f"Vision generation error: {e}")
yield f"Error: {str(e)}"
- def generate_audio_input_response(self, messages, system_prompt, audio_array,
- temperature, top_p, top_k, min_p,
- max_new_tokens, repetition_penalty,
- cancel_event=None) -> Generator[str, None, None]:
+ def generate_audio_input_response(
+ self,
+ messages,
+ system_prompt,
+ audio_array,
+ temperature,
+ top_p,
+ top_k,
+ min_p,
+ max_new_tokens,
+ repetition_penalty,
+ cancel_event = None,
+ ) -> Generator[str, None, None]:
"""Handle audio input (ASR) generation β accepts audio numpy array, streams text output.
Uses processor.apply_chat_template with audio embedded in messages (Gemma 3n pattern).
@@ -846,11 +962,11 @@ class InferenceBackend:
# apply_chat_template handles audio embedding + tokenization in one step
inputs = processor.apply_chat_template(
audio_messages,
- add_generation_prompt=True,
- tokenize=True,
- return_dict=True,
- return_tensors="pt",
- truncation=False,
+ add_generation_prompt = True,
+ tokenize = True,
+ return_dict = True,
+ return_tensors = "pt",
+ truncation = False,
).to(self.device)
try:
@@ -859,18 +975,18 @@ class InferenceBackend:
streamer = TextIteratorStreamer(
raw_tokenizer,
- skip_prompt=True,
- skip_special_tokens=True,
- timeout=0.2,
+ skip_prompt = True,
+ skip_special_tokens = True,
+ timeout = 0.2,
)
# Notebook uses do_sample=False for ASR (greedy decoding for accuracy)
generation_kwargs = dict(
**inputs,
- streamer=streamer,
- max_new_tokens=max_new_tokens,
- use_cache=True,
- do_sample=False,
+ streamer = streamer,
+ max_new_tokens = max_new_tokens,
+ use_cache = True,
+ do_sample = False,
)
err: dict[str, str] = {}
@@ -888,7 +1004,7 @@ class InferenceBackend:
except Exception:
pass
- thread = threading.Thread(target=generate_fn)
+ thread = threading.Thread(target = generate_fn)
thread.start()
output = ""
@@ -910,9 +1026,11 @@ class InferenceBackend:
finally:
if cancel_event is not None:
cancel_event.set()
- thread.join(timeout=10)
+ thread.join(timeout = 10)
if thread.is_alive():
- logger.warning("Audio input generation thread did not exit after cancel/join timeout")
+ logger.warning(
+ "Audio input generation thread did not exit after cancel/join timeout"
+ )
if err.get("msg"):
yield f"Error: {err['msg']}"
@@ -921,7 +1039,9 @@ class InferenceBackend:
logger.error(f"Audio input generation error: {e}")
yield f"Error: {str(e)}"
- def generate_whisper_response(self, audio_array, cancel_event=None) -> Generator[str, None, None]:
+ def generate_whisper_response(
+ self, audio_array, cancel_event = None
+ ) -> Generator[str, None, None]:
"""Whisper ASR β takes audio numpy array, yields transcribed text.
Uses the pre-built transformers pipeline (created during model loading).
@@ -943,16 +1063,18 @@ class InferenceBackend:
logger.error(f"Whisper ASR error: {e}")
yield f"Error: {str(e)}"
- def generate_stream(self,
- prompt: str,
- temperature: float = 0.7,
- top_p: float = 0.9,
- top_k: int = 40,
- min_p: float = 0.0,
- max_new_tokens: int = 256,
- repetition_penalty: float = 1.1,
- cancel_event=None,
- _adapter_state=None) -> Generator[str, None, None]:
+ def generate_stream(
+ self,
+ prompt: str,
+ temperature: float = 0.7,
+ top_p: float = 0.9,
+ top_k: int = 40,
+ min_p: float = 0.0,
+ max_new_tokens: int = 256,
+ repetition_penalty: float = 1.1,
+ cancel_event = None,
+ _adapter_state = None,
+ ) -> Generator[str, None, None]:
"""Generate streaming text response (text models only).
_adapter_state: if not None, the background thread toggles adapters
@@ -971,30 +1093,32 @@ class InferenceBackend:
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
try:
- inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+ inputs = tokenizer(prompt, return_tensors = "pt").to(model.device)
from transformers import TextIteratorStreamer
import threading
streamer = TextIteratorStreamer(
tokenizer,
- skip_prompt=True,
- skip_special_tokens=True,
- timeout=0.2,
+ skip_prompt = True,
+ skip_special_tokens = True,
+ timeout = 0.2,
)
generation_kwargs = dict(
**inputs,
- streamer=streamer,
- max_new_tokens=max_new_tokens,
- temperature=temperature,
- top_p=top_p,
- top_k=top_k,
- min_p=min_p,
- repetition_penalty=repetition_penalty,
- do_sample=temperature > 0,
- eos_token_id=tokenizer.eos_token_id,
- pad_token_id=tokenizer.eos_token_id if tokenizer.pad_token_id is None else tokenizer.pad_token_id,
+ streamer = streamer,
+ max_new_tokens = max_new_tokens,
+ temperature = temperature,
+ top_p = top_p,
+ top_k = top_k,
+ min_p = min_p,
+ repetition_penalty = repetition_penalty,
+ do_sample = temperature > 0,
+ eos_token_id = tokenizer.eos_token_id,
+ pad_token_id = tokenizer.eos_token_id
+ if tokenizer.pad_token_id is None
+ else tokenizer.pad_token_id,
)
if cancel_event is not None:
from transformers.generation.stopping_criteria import (
@@ -1029,11 +1153,12 @@ class InferenceBackend:
pass
err: dict[str, str] = {}
- thread = threading.Thread(target=generate_fn)
+ thread = threading.Thread(target = generate_fn)
thread.start()
output = ""
from queue import Empty
+
try:
while True:
if cancel_event is not None and cancel_event.is_set():
@@ -1053,9 +1178,11 @@ class InferenceBackend:
finally:
if cancel_event is not None:
cancel_event.set()
- thread.join(timeout=10)
+ thread.join(timeout = 10)
if thread.is_alive():
- logger.warning("Generation thread did not exit after cancel/join timeout")
+ logger.warning(
+ "Generation thread did not exit after cancel/join timeout"
+ )
if err.get("msg"):
yield f"Error: {err['msg']}"
@@ -1100,83 +1227,139 @@ class InferenceBackend:
self._apply_adapter_state(use_adapter)
if audio_type == "snac":
- return self._generate_snac(model, tokenizer, text, temperature, top_p, max_new_tokens, repetition_penalty)
+ return self._generate_snac(
+ model,
+ tokenizer,
+ text,
+ temperature,
+ top_p,
+ max_new_tokens,
+ repetition_penalty,
+ )
elif audio_type == "csm":
processor = model_info.get("processor", tokenizer)
return self._generate_csm(model, processor, text, max_new_tokens)
elif audio_type == "bicodec":
- return self._generate_bicodec(model, tokenizer, text, temperature, top_k, max_new_tokens)
+ return self._generate_bicodec(
+ model, tokenizer, text, temperature, top_k, max_new_tokens
+ )
elif audio_type == "dac":
- return self._generate_dac(model, tokenizer, text, temperature, top_k, top_p, min_p, max_new_tokens, repetition_penalty)
+ return self._generate_dac(
+ model,
+ tokenizer,
+ text,
+ temperature,
+ top_k,
+ top_p,
+ min_p,
+ max_new_tokens,
+ repetition_penalty,
+ )
else:
raise RuntimeError(f"Unknown audio_type: {audio_type}")
- def _generate_snac(self, model, tokenizer, text, temperature, top_p, max_new_tokens, repetition_penalty):
+ def _generate_snac(
+ self,
+ model,
+ tokenizer,
+ text,
+ temperature,
+ top_p,
+ max_new_tokens,
+ repetition_penalty,
+ ):
"""Generate audio using SNAC codec (Orpheus)."""
device = model.device
- start_token = torch.tensor([[128259]], device=device) # START_OF_HUMAN
- end_tokens = torch.tensor([[128009, 128260]], device=device) # EOT, END_OF_HUMAN
- text_ids = tokenizer(text, return_tensors="pt").input_ids.to(device)
- input_ids = torch.cat([start_token, text_ids, end_tokens], dim=1)
+ start_token = torch.tensor([[128259]], device = device) # START_OF_HUMAN
+ end_tokens = torch.tensor(
+ [[128009, 128260]], device = device
+ ) # EOT, END_OF_HUMAN
+ text_ids = tokenizer(text, return_tensors = "pt").input_ids.to(device)
+ input_ids = torch.cat([start_token, text_ids, end_tokens], dim = 1)
attention_mask = torch.ones_like(input_ids)
generated = model.generate(
- input_ids=input_ids,
- attention_mask=attention_mask,
- max_new_tokens=max_new_tokens,
- do_sample=True,
- temperature=temperature,
- top_p=top_p,
- repetition_penalty=repetition_penalty,
- eos_token_id=128258, # END_OF_SPEECH
- use_cache=True,
+ input_ids = input_ids,
+ attention_mask = attention_mask,
+ max_new_tokens = max_new_tokens,
+ do_sample = True,
+ temperature = temperature,
+ top_p = top_p,
+ repetition_penalty = repetition_penalty,
+ eos_token_id = 128258, # END_OF_SPEECH
+ use_cache = True,
)
return self._audio_codec_manager.decode_snac(generated, str(device))
def _generate_csm(self, model, processor, text, max_new_tokens):
"""Generate audio using CSM (Sesame)."""
speaker_id = 0
- inputs = processor(f"[{speaker_id}]{text}", add_special_tokens=True, return_tensors="pt").to(model.device)
- audio_values = model.generate(**inputs, max_new_tokens=max_new_tokens, output_audio=True)
+ inputs = processor(
+ f"[{speaker_id}]{text}", add_special_tokens = True, return_tensors = "pt"
+ ).to(model.device)
+ audio_values = model.generate(
+ **inputs, max_new_tokens = max_new_tokens, output_audio = True
+ )
return self._audio_codec_manager.decode_csm(audio_values)
- def _generate_bicodec(self, model, tokenizer, text, temperature, top_k, max_new_tokens):
+ def _generate_bicodec(
+ self, model, tokenizer, text, temperature, top_k, max_new_tokens
+ ):
"""Generate audio using BiCodec (Spark-TTS)."""
- prompt = "<|task_tts|><|start_content|>" + text + "<|end_content|><|start_global_token|>"
- inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
+ prompt = (
+ "<|task_tts|><|start_content|>"
+ + text
+ + "<|end_content|><|start_global_token|>"
+ )
+ inputs = tokenizer([prompt], return_tensors = "pt").to(model.device)
generated = model.generate(
**inputs,
- max_new_tokens=max_new_tokens,
- do_sample=True,
- temperature=temperature,
- top_k=top_k,
- eos_token_id=tokenizer.eos_token_id,
- pad_token_id=tokenizer.pad_token_id,
+ max_new_tokens = max_new_tokens,
+ do_sample = True,
+ temperature = temperature,
+ top_k = top_k,
+ eos_token_id = tokenizer.eos_token_id,
+ pad_token_id = tokenizer.pad_token_id,
)
- new_tokens = generated[:, inputs.input_ids.shape[1]:]
- decoded_text = tokenizer.batch_decode(new_tokens, skip_special_tokens=False)[0]
+ new_tokens = generated[:, inputs.input_ids.shape[1] :]
+ decoded_text = tokenizer.batch_decode(new_tokens, skip_special_tokens = False)[0]
return self._audio_codec_manager.decode_bicodec(decoded_text, str(model.device))
- def _generate_dac(self, model, tokenizer, text, temperature, top_k, top_p, min_p, max_new_tokens, repetition_penalty):
+ def _generate_dac(
+ self,
+ model,
+ tokenizer,
+ text,
+ temperature,
+ top_k,
+ top_p,
+ min_p,
+ max_new_tokens,
+ repetition_penalty,
+ ):
"""Generate audio using DAC (OuteTTS). Follows Oute_TTS_(1B).ipynb exactly."""
# Monkey-patch RepetitionPenaltyLogitsProcessor with a 64-token penalty
# window (same as the OuteTTS notebook) to avoid degenerate repetition.
self._patch_repetition_penalty_processor()
- prompt = "<|im_start|>\n<|text_start|>" + text + "<|text_end|>\n<|audio_start|><|global_features_start|>\n"
+ prompt = (
+ "<|im_start|>\n<|text_start|>"
+ + text
+ + "<|text_end|>\n<|audio_start|><|global_features_start|>\n"
+ )
with torch.inference_mode():
- with torch.amp.autocast('cuda', dtype=model.dtype):
- inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
+ with torch.amp.autocast("cuda", dtype = model.dtype):
+ inputs = tokenizer([prompt], return_tensors = "pt").to(model.device)
generated = model.generate(
**inputs,
- temperature=temperature,
- top_k=top_k,
- top_p=top_p,
- min_p=min_p,
- repetition_penalty=repetition_penalty,
- max_new_tokens=max_new_tokens,
+ temperature = temperature,
+ top_k = top_k,
+ top_p = top_p,
+ min_p = min_p,
+ repetition_penalty = repetition_penalty,
+ max_new_tokens = max_new_tokens,
)
- decoded_text = tokenizer.batch_decode(generated, skip_special_tokens=False)[0]
+ decoded_text = tokenizer.batch_decode(generated, skip_special_tokens = False)[0]
return self._audio_codec_manager.decode_dac(decoded_text, str(model.device))
_repetition_penalty_patched = False
@@ -1199,11 +1382,15 @@ class InferenceBackend:
def __init__(self, penalty: float):
self.penalty_last_n = 64
if not isinstance(penalty, float) or penalty <= 0:
- raise ValueError(f"`penalty` has to be a positive float, but is {penalty}")
+ raise ValueError(
+ f"`penalty` has to be a positive float, but is {penalty}"
+ )
self.penalty = penalty
@torch.no_grad()
- def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+ def __call__(
+ self, input_ids: torch.LongTensor, scores: torch.FloatTensor
+ ) -> torch.FloatTensor:
if self.penalty_last_n == 0 or self.penalty == 1.0:
return scores
batch_size, seq_len = input_ids.shape
@@ -1217,11 +1404,17 @@ class InferenceBackend:
if token_id >= vocab_size:
continue
logit = scores[b, token_id]
- scores[b, token_id] = logit * self.penalty if logit <= 0 else logit / self.penalty
+ scores[b, token_id] = (
+ logit * self.penalty if logit <= 0 else logit / self.penalty
+ )
return scores
- generation_utils.RepetitionPenaltyLogitsProcessor = RepetitionPenaltyLogitsProcessorPatch
- logger.info("Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS")
+ generation_utils.RepetitionPenaltyLogitsProcessor = (
+ RepetitionPenaltyLogitsProcessorPatch
+ )
+ logger.info(
+ "Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS"
+ )
def format_chat_prompt(self, messages: list, system_prompt: str = None) -> str:
if not self.active_model_name or self.active_model_name not in self.models:
@@ -1232,7 +1425,9 @@ class InferenceBackend:
logger.error("Tokenizer not loaded for active model")
return ""
- chat_template_info = self.models[self.active_model_name].get("chat_template_info", {})
+ chat_template_info = self.models[self.active_model_name].get(
+ "chat_template_info", {}
+ )
tokenizer = self.models[self.active_model_name]["tokenizer"]
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
@@ -1249,12 +1444,15 @@ class InferenceBackend:
if role in ["system", "user", "assistant"] and content.strip():
if role == last_role:
- logger.debug(f"Skipping consecutive {role} message to maintain alternation")
+ logger.debug(
+ f"Skipping consecutive {role} message to maintain alternation"
+ )
continue
if role == "user":
import re
- clean_content = re.sub(r'<[^>]+>', '', content).strip()
+
+ clean_content = re.sub(r"<[^>]+>", "", content).strip()
if clean_content:
chat_messages.append({"role": role, "content": clean_content})
last_role = role
@@ -1265,7 +1463,9 @@ class InferenceBackend:
continue
if chat_messages and chat_messages[-1]["role"] == "assistant":
- logger.debug("Removing final assistant message to ensure proper alternation")
+ logger.debug(
+ "Removing final assistant message to ensure proper alternation"
+ )
chat_messages.pop()
logger.info(f"Sending {len(chat_messages)} messages to tokenizer:")
@@ -1274,31 +1474,44 @@ class InferenceBackend:
try:
formatted_prompt = tokenizer.apply_chat_template(
- chat_messages,
- tokenize=False,
- add_generation_prompt=True
+ chat_messages, tokenize = False, add_generation_prompt = True
)
logger.info(f"Successfully applied tokenizer's native chat template")
return formatted_prompt
except Exception as e:
error_msg = str(e).lower()
- if "chat_template is not set" in error_msg or "no template argument" in error_msg:
- logger.info(f"Base model detected - no built-in chat template available, using fallback formatting")
+ if (
+ "chat_template is not set" in error_msg
+ or "no template argument" in error_msg
+ ):
+ logger.info(
+ f"Base model detected - no built-in chat template available, using fallback formatting"
+ )
else:
logger.warning(f"Failed to apply tokenizer chat template: {e}")
- logger.debug(f"""Failed with messages: {[f"{m['role']}: {m['content'][:30]}..." for m in chat_messages]}""")
+ logger.debug(
+ f"""Failed with messages: {[f"{m['role']}: {m['content'][:30]}..." for m in chat_messages]}"""
+ )
if chat_template_info.get("has_template", False):
- logger.info("Falling back to manual template formatting based on detected patterns")
+ logger.info(
+ "Falling back to manual template formatting based on detected patterns"
+ )
template_type = chat_template_info.get("format_type", "generic")
- manual_prompt = self._format_chat_manual(chat_messages, template_type, chat_template_info.get("special_tokens", {}))
+ manual_prompt = self._format_chat_manual(
+ chat_messages,
+ template_type,
+ chat_template_info.get("special_tokens", {}),
+ )
logger.info(f"Manual template result: {manual_prompt[:200]}...")
return manual_prompt
else:
logger.info("Using generic chat formatting for base model")
return self._format_generic_template(chat_messages, {})
- def _format_chat_manual(self, messages: list, template_type: str, special_tokens: dict) -> str:
+ def _format_chat_manual(
+ self, messages: list, template_type: str, special_tokens: dict
+ ) -> str:
"""
Manual chat formatting fallback for when tokenizer template fails
@@ -1329,7 +1542,9 @@ class InferenceBackend:
for msg in messages:
role = msg["role"]
content = msg["content"]
- formatted += f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>"
+ formatted += (
+ f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>"
+ )
formatted += "<|start_header_id|>assistant<|end_header_id|>\n\n"
return formatted
@@ -1358,7 +1573,10 @@ class InferenceBackend:
formatted += f"[INST] {user_content} [/INST]"
- if i + 1 < len(conversation) and conversation[i + 1]["role"] == "assistant":
+ if (
+ i + 1 < len(conversation)
+ and conversation[i + 1]["role"] == "assistant"
+ ):
formatted += f" {conversation[i + 1]['content']}"
i += 2
else:
@@ -1435,11 +1653,11 @@ class InferenceBackend:
try:
# This is a common pattern for Unsloth/Hugging Face models
- if hasattr(model, 'past_key_values'):
+ if hasattr(model, "past_key_values"):
model.past_key_values = None
- if hasattr(model, 'generation_config'):
- if hasattr(model.generation_config, 'past_key_values'):
- model.generation_config.past_key_values = None
+ if hasattr(model, "generation_config"):
+ if hasattr(model.generation_config, "past_key_values"):
+ model.generation_config.past_key_values = None
logger.debug(f"Reset generation state for model: {model_name}")
except Exception as e:
@@ -1456,6 +1674,7 @@ class InferenceBackend:
logger.debug("Cleared GPU cache")
import gc
+
gc.collect()
logger.info("Performed comprehensive generation state reset")
@@ -1468,8 +1687,9 @@ class InferenceBackend:
return None
if img.size[0] > max_size or img.size[1] > max_size:
from PIL import Image
- ratio = min(max_size/img.size[0], max_size/img.size[1])
- new_size = (int(img.size[0]*ratio), int(img.size[1]*ratio))
+
+ ratio = min(max_size / img.size[0], max_size / img.size[1])
+ new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
return img.resize(new_size, Image.Resampling.LANCZOS)
return img
@@ -1483,7 +1703,9 @@ class InferenceBackend:
return text.strip()
def _load_chat_template_info(self, model_name: str):
- if model_name not in self.models or not self.models[model_name].get("tokenizer"):
+ if model_name not in self.models or not self.models[model_name].get(
+ "tokenizer"
+ ):
return
tokenizer = self.models[model_name]["tokenizer"]
@@ -1497,29 +1719,43 @@ class InferenceBackend:
try:
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER
- #Try exact match first
+
+ # Try exact match first
model_name_lower = model_name.lower()
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
- chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
- logger.info(f"Detected template '{chat_template_info['template_name']}' for {model_name} from mapper")
+ chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[
+ model_name_lower
+ ]
+ logger.info(
+ f"Detected template '{chat_template_info['template_name']}' for {model_name} from mapper"
+ )
else:
# Try partial match (for variants like model_name-bnb-4bit)
for key in MODEL_TO_TEMPLATE_MAPPER:
if key in model_name_lower or model_name_lower in key:
- chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[key]
- logger.info(f"Detected template '{chat_template_info['template_name']}' for {model_name} (partial match)")
+ chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[
+ key
+ ]
+ logger.info(
+ f"Detected template '{chat_template_info['template_name']}' for {model_name} (partial match)"
+ )
break
except Exception as e:
- logger.warning(f"Could not detect template from mapper for {model_name}: {e}")
+ logger.warning(
+ f"Could not detect template from mapper for {model_name}: {e}"
+ )
try:
- if hasattr(tokenizer, 'chat_template') and tokenizer.chat_template:
+ if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
chat_template_info["has_template"] = True
chat_template_info["template"] = tokenizer.chat_template
template_str = tokenizer.chat_template.lower()
- if "start_header_id" in template_str and "end_header_id" in template_str:
+ if (
+ "start_header_id" in template_str
+ and "end_header_id" in template_str
+ ):
chat_template_info["format_type"] = "llama3"
elif "[inst]" in template_str and "[/inst]" in template_str:
chat_template_info["format_type"] = "mistral"
@@ -1530,21 +1766,25 @@ class InferenceBackend:
else:
chat_template_info["format_type"] = "custom"
- logger.info(f"Loaded chat template for {model_name} (detected as {chat_template_info['format_type']} format)")
+ logger.info(
+ f"Loaded chat template for {model_name} (detected as {chat_template_info['format_type']} format)"
+ )
logger.debug(f"Template preview: {tokenizer.chat_template[:200]}...")
special_tokens = {}
- if hasattr(tokenizer, 'bos_token') and tokenizer.bos_token:
+ if hasattr(tokenizer, "bos_token") and tokenizer.bos_token:
special_tokens["bos_token"] = tokenizer.bos_token
- if hasattr(tokenizer, 'eos_token') and tokenizer.eos_token:
+ if hasattr(tokenizer, "eos_token") and tokenizer.eos_token:
special_tokens["eos_token"] = tokenizer.eos_token
- if hasattr(tokenizer, 'pad_token') and tokenizer.pad_token:
+ if hasattr(tokenizer, "pad_token") and tokenizer.pad_token:
special_tokens["pad_token"] = tokenizer.pad_token
chat_template_info["special_tokens"] = special_tokens
else:
- logger.info(f"No chat template found for {model_name}, will use generic formatting")
+ logger.info(
+ f"No chat template found for {model_name}, will use generic formatting"
+ )
except Exception as e:
logger.error(f"Error loading chat template info for {model_name}: {e}")
@@ -1552,10 +1792,13 @@ class InferenceBackend:
self.models[model_name]["chat_template_info"] = chat_template_info
if chat_template_info["has_template"]:
- logger.info(f"Chat template loaded for {model_name}: {chat_template_info['format_type']} format")
+ logger.info(
+ f"Chat template loaded for {model_name}: {chat_template_info['format_type']} format"
+ )
else:
- logger.info(f"No built-in chat template for {model_name}, will use generic formatting")
-
+ logger.info(
+ f"No built-in chat template for {model_name}, will use generic formatting"
+ )
def get_current_model(self) -> Optional[str]:
"""Get currently active model name"""
@@ -1569,11 +1812,13 @@ class InferenceBackend:
"""Get name of currently loading model"""
return next(iter(self.loading_models)) if self.loading_models else None
- def load_model_simple(self,
- model_path: str,
- hf_token: Optional[str] = None,
- max_seq_length: int = 2048,
- load_in_4bit: bool = True) -> bool:
+ def load_model_simple(
+ self,
+ model_path: str,
+ hf_token: Optional[str] = None,
+ max_seq_length: int = 2048,
+ load_in_4bit: bool = True,
+ ) -> bool:
"""
Simple model loading wrapper for chat interface.
Accepts model path as string and handles ModelConfig creation internally.
@@ -1591,17 +1836,17 @@ class InferenceBackend:
# Create config from string path
config = ModelConfig.from_ui_selection(
model_path,
- lora_path=None, # No LoRA for chat
- is_lora=False
+ lora_path = None, # No LoRA for chat
+ is_lora = False,
)
# Call existing load_model with config
return self.load_model(
- config=config,
- max_seq_length=max_seq_length,
- dtype=None, # Auto-detect
- load_in_4bit=load_in_4bit,
- hf_token=hf_token
+ config = config,
+ max_seq_length = max_seq_length,
+ dtype = None, # Auto-detect
+ load_in_4bit = load_in_4bit,
+ hf_token = hf_token,
)
except Exception as e:
@@ -1609,9 +1854,9 @@ class InferenceBackend:
return False
-
# Global inference backend instance
inference_backend = InferenceBackend()
+
def get_inference_backend() -> InferenceBackend:
return inference_backend
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 13973b17f8..21a25d51b2 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -7,6 +7,7 @@ llama-server inference backend for GGUF models.
Manages a llama-server subprocess and proxies chat completions
through its OpenAI-compatible /v1/chat/completions endpoint.
"""
+
import atexit
import json
import structlog
@@ -146,7 +147,9 @@ class LlamaCppBackend:
if build_path.is_file():
return str(build_path)
if sys.platform == "win32":
- win_path = project_root / "llama.cpp" / "build" / "bin" / "Release" / binary_name
+ win_path = (
+ project_root / "llama.cpp" / "build" / "bin" / "Release" / binary_name
+ )
if win_path.is_file():
return str(win_path)
@@ -168,7 +171,7 @@ class LlamaCppBackend:
def _find_free_port() -> int:
"""Find an available TCP port."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
- s.bind(("", 0))
+ s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
# ββ Stdout drain (prevents pipe deadlock on Windows) βββββββββ
@@ -258,15 +261,19 @@ class LlamaCppBackend:
try:
import re
from huggingface_hub import list_repo_files
- files = list_repo_files(hf_repo, token=hf_token)
+
+ files = list_repo_files(hf_repo, token = hf_token)
variant_lower = hf_variant.lower()
# Use word-boundary matching so "Q8_0" doesn't also
# match "IQ8_0" or other superset variant names.
boundary = re.compile(
- r'(? bool:
"""Check if subprocess is alive."""
@@ -173,15 +176,13 @@ class InferenceOrchestrator:
if self._resp_queue is None:
return None
try:
- return self._resp_queue.get(timeout=timeout)
+ return self._resp_queue.get(timeout = timeout)
except queue.Empty:
return None
except (EOFError, OSError, ValueError):
return None
- def _wait_response(
- self, expected_type: str, timeout: float = 120.0
- ) -> dict:
+ def _wait_response(self, expected_type: str, timeout: float = 120.0) -> dict:
"""Block until a response of the expected type arrives.
Also handles 'status' and 'error' events during the wait.
@@ -192,7 +193,7 @@ class InferenceOrchestrator:
while time.monotonic() < deadline:
remaining = max(0.1, deadline - time.monotonic())
- resp = self._read_resp(timeout=min(remaining, 1.0))
+ resp = self._read_resp(timeout = min(remaining, 1.0))
if resp is None:
# Check subprocess health
@@ -214,7 +215,11 @@ class InferenceOrchestrator:
continue
# Other response types during wait β skip
- logger.debug("Skipping response type '%s' while waiting for '%s'", rtype, expected_type)
+ logger.debug(
+ "Skipping response type '%s' while waiting for '%s'",
+ rtype,
+ expected_type,
+ )
raise RuntimeError(
f"Timeout waiting for '{expected_type}' response after {timeout}s"
@@ -241,7 +246,7 @@ class InferenceOrchestrator:
"""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
- resp = self._read_resp(timeout=min(0.5, deadline - time.monotonic()))
+ resp = self._read_resp(timeout = min(0.5, deadline - time.monotonic()))
if resp is None:
if not self._ensure_subprocess_alive():
return
@@ -259,7 +264,7 @@ class InferenceOrchestrator:
self,
config, # ModelConfig
max_seq_length: int = 2048,
- dtype=None,
+ dtype = None,
load_in_4bit: bool = True,
hf_token: Optional[str] = None,
trust_remote_code: bool = False,
@@ -298,14 +303,15 @@ class InferenceOrchestrator:
elif self._proc is not None:
# Dead subprocess β clean up
- self._shutdown_subprocess(timeout=2)
+ self._shutdown_subprocess(timeout = 2)
logger.info(
"Spawning fresh inference subprocess for '%s' (transformers %s.x)",
- model_name, needed_major,
+ model_name,
+ needed_major,
)
self._spawn_subprocess(sub_config)
- resp = self._wait_response("loaded", timeout=180)
+ resp = self._wait_response("loaded", timeout = 180)
# Update local state from response
if resp.get("success"):
@@ -346,11 +352,13 @@ class InferenceOrchestrator:
return True
try:
- self._send_cmd({
- "type": "unload",
- "model_name": model_name,
- })
- resp = self._wait_response("unloaded", timeout=30)
+ self._send_cmd(
+ {
+ "type": "unload",
+ "model_name": model_name,
+ }
+ )
+ resp = self._wait_response("unloaded", timeout = 30)
# Update local state
self.models.pop(model_name, None)
@@ -372,40 +380,40 @@ class InferenceOrchestrator:
self,
messages: list,
system_prompt: str = "",
- image=None,
+ image = None,
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_new_tokens: int = 256,
repetition_penalty: float = 1.1,
- cancel_event=None,
+ cancel_event = None,
) -> Generator[str, None, None]:
"""Generate response, streaming tokens from subprocess."""
yield from self._generate_inner(
- messages=messages,
- system_prompt=system_prompt,
- image=image,
- temperature=temperature,
- top_p=top_p,
- top_k=top_k,
- min_p=min_p,
- max_new_tokens=max_new_tokens,
- repetition_penalty=repetition_penalty,
- cancel_event=cancel_event,
- use_adapter=None,
+ messages = messages,
+ system_prompt = system_prompt,
+ image = image,
+ temperature = temperature,
+ top_p = top_p,
+ top_k = top_k,
+ min_p = min_p,
+ max_new_tokens = max_new_tokens,
+ repetition_penalty = repetition_penalty,
+ cancel_event = cancel_event,
+ use_adapter = None,
)
def generate_with_adapter_control(
self,
use_adapter: Optional[Union[bool, str]] = None,
- cancel_event=None,
+ cancel_event = None,
**gen_kwargs,
) -> Generator[str, None, None]:
"""Generate with adapter control, streaming tokens from subprocess."""
yield from self._generate_inner(
- use_adapter=use_adapter,
- cancel_event=cancel_event,
+ use_adapter = use_adapter,
+ cancel_event = cancel_event,
**gen_kwargs,
)
@@ -413,15 +421,15 @@ class InferenceOrchestrator:
self,
messages: list = None,
system_prompt: str = "",
- image=None,
+ image = None,
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_new_tokens: int = 256,
repetition_penalty: float = 1.1,
- cancel_event=None,
- use_adapter=None,
+ cancel_event = None,
+ use_adapter = None,
) -> Generator[str, None, None]:
"""Inner generation logic β sends command to subprocess, yields tokens.
@@ -442,32 +450,32 @@ class InferenceOrchestrator:
# can consume and drop each other's token events.
with self._gen_lock:
yield from self._generate_locked(
- messages=messages,
- system_prompt=system_prompt,
- image=image,
- temperature=temperature,
- top_p=top_p,
- top_k=top_k,
- min_p=min_p,
- max_new_tokens=max_new_tokens,
- repetition_penalty=repetition_penalty,
- cancel_event=cancel_event,
- use_adapter=use_adapter,
+ messages = messages,
+ system_prompt = system_prompt,
+ image = image,
+ temperature = temperature,
+ top_p = top_p,
+ top_k = top_k,
+ min_p = min_p,
+ max_new_tokens = max_new_tokens,
+ repetition_penalty = repetition_penalty,
+ cancel_event = cancel_event,
+ use_adapter = use_adapter,
)
def _generate_locked(
self,
messages: list = None,
system_prompt: str = "",
- image=None,
+ image = None,
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_new_tokens: int = 256,
repetition_penalty: float = 1.1,
- cancel_event=None,
- use_adapter=None,
+ cancel_event = None,
+ use_adapter = None,
) -> Generator[str, None, None]:
"""Actual generation logic β must be called under _gen_lock."""
request_id = str(uuid.uuid4())
@@ -503,7 +511,7 @@ class InferenceOrchestrator:
# Yield tokens from response queue β we are the only reader
# because _gen_lock is held.
while True:
- resp = self._read_resp(timeout=30.0)
+ resp = self._read_resp(timeout = 30.0)
if resp is None:
# Check subprocess health
@@ -531,7 +539,7 @@ class InferenceOrchestrator:
# Wait for the subprocess to acknowledge cancellation
# (gen_done/gen_error) so stale events don't leak into
# the next generation request.
- self._drain_until_gen_done(timeout=5.0)
+ self._drain_until_gen_done(timeout = 5.0)
return
yield resp.get("text", "")
@@ -577,6 +585,7 @@ class InferenceOrchestrator:
raise RuntimeError("No active model")
import uuid
+
request_id = str(uuid.uuid4())
cmd = {
@@ -599,11 +608,13 @@ class InferenceOrchestrator:
deadline = time.monotonic() + 120.0
while time.monotonic() < deadline:
remaining = max(0.1, deadline - time.monotonic())
- resp = self._read_resp(timeout=min(remaining, 1.0))
+ resp = self._read_resp(timeout = min(remaining, 1.0))
if resp is None:
if not self._ensure_subprocess_alive():
- raise RuntimeError("Inference subprocess crashed during audio generation")
+ raise RuntimeError(
+ "Inference subprocess crashed during audio generation"
+ )
continue
rtype = resp.get("type", "")
@@ -627,15 +638,15 @@ class InferenceOrchestrator:
def generate_whisper_response(
self,
audio_array,
- cancel_event=None,
+ cancel_event = None,
) -> Generator[str, None, None]:
"""Whisper ASR β sends audio to subprocess, yields text."""
yield from self._generate_audio_input_inner(
- audio_array=audio_array,
- audio_type="whisper",
- messages=[],
- system_prompt="",
- cancel_event=cancel_event,
+ audio_array = audio_array,
+ audio_type = "whisper",
+ messages = [],
+ system_prompt = "",
+ cancel_event = cancel_event,
)
def generate_audio_input_response(
@@ -649,21 +660,21 @@ class InferenceOrchestrator:
min_p: float = 0.0,
max_new_tokens: int = 512,
repetition_penalty: float = 1.1,
- cancel_event=None,
+ cancel_event = None,
) -> Generator[str, None, None]:
"""Audio input generation (e.g. Gemma 3n) β streams text tokens."""
yield from self._generate_audio_input_inner(
- audio_array=audio_array,
- audio_type=None, # worker will use generate_audio_input_response
- messages=messages,
- system_prompt=system_prompt,
- temperature=temperature,
- top_p=top_p,
- top_k=top_k,
- min_p=min_p,
- max_new_tokens=max_new_tokens,
- repetition_penalty=repetition_penalty,
- cancel_event=cancel_event,
+ audio_array = audio_array,
+ audio_type = None, # worker will use generate_audio_input_response
+ messages = messages,
+ system_prompt = system_prompt,
+ temperature = temperature,
+ top_p = top_p,
+ top_k = top_k,
+ min_p = min_p,
+ max_new_tokens = max_new_tokens,
+ repetition_penalty = repetition_penalty,
+ cancel_event = cancel_event,
)
def _generate_audio_input_inner(
@@ -678,7 +689,7 @@ class InferenceOrchestrator:
min_p: float = 0.0,
max_new_tokens: int = 512,
repetition_penalty: float = 1.1,
- cancel_event=None,
+ cancel_event = None,
) -> Generator[str, None, None]:
"""Shared inner logic for audio input generation (Whisper + ASR)."""
if not self._ensure_subprocess_alive():
@@ -690,10 +701,15 @@ class InferenceOrchestrator:
with self._gen_lock:
import uuid
+
request_id = str(uuid.uuid4())
# Convert numpy array to list for mp.Queue serialization
- audio_data = audio_array.tolist() if hasattr(audio_array, 'tolist') else list(audio_array)
+ audio_data = (
+ audio_array.tolist()
+ if hasattr(audio_array, "tolist")
+ else list(audio_array)
+ )
cmd = {
"type": "generate_audio_input",
@@ -718,7 +734,7 @@ class InferenceOrchestrator:
# Yield tokens β same pattern as _generate_locked
while True:
- resp = self._read_resp(timeout=30.0)
+ resp = self._read_resp(timeout = 30.0)
if resp is None:
if not self._ensure_subprocess_alive():
@@ -738,7 +754,7 @@ class InferenceOrchestrator:
if rtype == "token":
if cancel_event is not None and cancel_event.is_set():
self._cancel_generation()
- self._drain_until_gen_done(timeout=5.0)
+ self._drain_until_gen_done(timeout = 5.0)
return
yield resp.get("text", "")
@@ -761,6 +777,7 @@ class InferenceOrchestrator:
return None
if img.size[0] > max_size or img.size[1] > max_size:
from PIL import Image
+
ratio = min(max_size / img.size[0], max_size / img.size[1])
new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
return img.resize(new_size, Image.Resampling.LANCZOS)
@@ -770,7 +787,7 @@ class InferenceOrchestrator:
def _pil_to_base64(img) -> str:
"""Convert a PIL Image to base64 string for IPC."""
buf = BytesIO()
- img.save(buf, format="PNG")
+ img.save(buf, format = "PNG")
return base64.b64encode(buf.getvalue()).decode("ascii")
def get_current_model(self) -> Optional[str]:
diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py
index 3b878b348f..b266f5e0d4 100644
--- a/studio/backend/core/inference/worker.py
+++ b/studio/backend/core/inference/worker.py
@@ -13,6 +13,7 @@ The subprocess stays alive while a model is loaded, accepting commands
Pattern follows core/training/worker.py.
"""
+
from __future__ import annotations
import base64
@@ -45,7 +46,9 @@ def _activate_transformers_version(model_name: str) -> None:
resolved = _resolve_base_model(model_name)
if needs_transformers_5(resolved):
- venv_t5 = os.path.join(os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5")
+ venv_t5 = os.path.join(
+ os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5"
+ )
if os.path.isdir(venv_t5):
sys.path.insert(0, venv_t5)
logger.info("Activated transformers 5.x from %s", venv_t5)
@@ -53,16 +56,35 @@ def _activate_transformers_version(model_name: str) -> None:
# Fallback: pip install at runtime (slower, ~10-15s)
logger.warning(".venv_t5 not found at %s β installing at runtime", venv_t5)
import subprocess as sp
- os.makedirs(venv_t5, exist_ok=True)
+
+ os.makedirs(venv_t5, exist_ok = True)
r1 = sp.run(
- [sys.executable, "-m", "pip", "install", "--target", venv_t5,
- "--no-deps", "transformers==5.2.0"],
- stdout=sp.PIPE, stderr=sp.STDOUT,
+ [
+ sys.executable,
+ "-m",
+ "pip",
+ "install",
+ "--target",
+ venv_t5,
+ "--no-deps",
+ "transformers==5.2.0",
+ ],
+ stdout = sp.PIPE,
+ stderr = sp.STDOUT,
)
r2 = sp.run(
- [sys.executable, "-m", "pip", "install", "--target", venv_t5,
- "--no-deps", "huggingface_hub==1.3.0"],
- stdout=sp.PIPE, stderr=sp.STDOUT,
+ [
+ sys.executable,
+ "-m",
+ "pip",
+ "install",
+ "--target",
+ venv_t5,
+ "--no-deps",
+ "huggingface_hub==1.3.0",
+ ],
+ stdout = sp.PIPE,
+ stderr = sp.STDOUT,
)
if r1.returncode != 0 or r2.returncode != 0:
raise RuntimeError(
@@ -80,6 +102,7 @@ def _activate_transformers_version(model_name: str) -> None:
def _decode_image(image_base64: str):
"""Decode base64 string to PIL.Image."""
from PIL import Image
+
image_data = base64.b64decode(image_base64)
return Image.open(BytesIO(image_data))
@@ -90,6 +113,7 @@ def _resize_image(img, max_size: int = 800):
return None
if img.size[0] > max_size or img.size[1] > max_size:
from PIL import Image
+
ratio = min(max_size / img.size[0], max_size / img.size[1])
new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
return img.resize(new_size, Image.Resampling.LANCZOS)
@@ -114,9 +138,9 @@ def _build_model_config(config: dict):
gguf_variant = config.get("gguf_variant")
mc = ModelConfig.from_identifier(
- model_id=model_name,
- hf_token=hf_token,
- gguf_variant=gguf_variant,
+ model_id = model_name,
+ hf_token = hf_token,
+ gguf_variant = gguf_variant,
)
if not mc:
raise ValueError(f"Invalid model identifier: {model_name}")
@@ -136,6 +160,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
if mc.is_lora and mc.path:
import json
from pathlib import Path
+
adapter_cfg_path = Path(mc.path) / "adapter_config.json"
if adapter_cfg_path.exists():
try:
@@ -143,24 +168,34 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
adapter_cfg = json.load(f)
training_method = adapter_cfg.get("unsloth_training_method")
if training_method == "lora" and load_in_4bit:
- logger.info("adapter_config.json says lora β setting load_in_4bit=False")
+ logger.info(
+ "adapter_config.json says lora β setting load_in_4bit=False"
+ )
load_in_4bit = False
elif training_method == "qlora" and not load_in_4bit:
- logger.info("adapter_config.json says qlora β setting load_in_4bit=True")
+ logger.info(
+ "adapter_config.json says qlora β setting load_in_4bit=True"
+ )
load_in_4bit = True
elif not training_method:
- if mc.base_model and "-bnb-4bit" not in mc.base_model.lower() and load_in_4bit:
- logger.info("No training method, base model has no -bnb-4bit β setting load_in_4bit=False")
+ if (
+ mc.base_model
+ and "-bnb-4bit" not in mc.base_model.lower()
+ and load_in_4bit
+ ):
+ logger.info(
+ "No training method, base model has no -bnb-4bit β setting load_in_4bit=False"
+ )
load_in_4bit = False
except Exception as e:
logger.warning("Could not read adapter_config.json: %s", e)
success = backend.load_model(
- config=mc,
- max_seq_length=config.get("max_seq_length", 2048),
- load_in_4bit=load_in_4bit,
- hf_token=hf_token,
- trust_remote_code=config.get("trust_remote_code", False),
+ config = mc,
+ max_seq_length = config.get("max_seq_length", 2048),
+ load_in_4bit = load_in_4bit,
+ hf_token = hf_token,
+ trust_remote_code = config.get("trust_remote_code", False),
)
if success:
@@ -175,28 +210,37 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"audio_type": getattr(mc, "audio_type", None),
"has_audio_input": getattr(mc, "has_audio_input", False),
}
- _send_response(resp_queue, {
- "type": "loaded",
- "success": True,
- "model_info": model_info,
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "loaded",
+ "success": True,
+ "model_info": model_info,
+ "ts": time.time(),
+ },
+ )
else:
- _send_response(resp_queue, {
- "type": "loaded",
- "success": False,
- "error": "Failed to load model",
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "loaded",
+ "success": False,
+ "error": "Failed to load model",
+ "ts": time.time(),
+ },
+ )
except Exception as exc:
- _send_response(resp_queue, {
- "type": "loaded",
- "success": False,
- "error": str(exc),
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "loaded",
+ "success": False,
+ "error": str(exc),
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
def _handle_generate(
@@ -240,7 +284,7 @@ def _handle_generate(
use_adapter = cmd.get("use_adapter")
if use_adapter is not None:
generator = backend.generate_with_adapter_control(
- use_adapter=use_adapter,
+ use_adapter = use_adapter,
**gen_kwargs,
)
else:
@@ -254,29 +298,38 @@ def _handle_generate(
logger.info("Generation cancelled for request %s", request_id)
break
- _send_response(resp_queue, {
- "type": "token",
- "request_id": request_id,
- "text": cumulative_text,
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "token",
+ "request_id": request_id,
+ "text": cumulative_text,
+ "ts": time.time(),
+ },
+ )
- _send_response(resp_queue, {
- "type": "gen_done",
- "request_id": request_id,
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "gen_done",
+ "request_id": request_id,
+ "ts": time.time(),
+ },
+ )
logger.info("Finished text generation for request_id=%s", request_id)
except Exception as exc:
- logger.error("Generation error: %s", exc, exc_info=True)
- _send_response(resp_queue, {
- "type": "gen_error",
- "request_id": request_id,
- "error": str(exc),
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ logger.error("Generation error: %s", exc, exc_info = True)
+ _send_response(
+ resp_queue,
+ {
+ "type": "gen_error",
+ "request_id": request_id,
+ "error": str(exc),
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
def _handle_generate_audio(
@@ -289,35 +342,41 @@ def _handle_generate_audio(
try:
logger.info("Starting audio generation for request_id=%s", request_id)
wav_bytes, sample_rate = backend.generate_audio_response(
- text=cmd["text"],
- temperature=cmd.get("temperature", 0.6),
- top_p=cmd.get("top_p", 0.95),
- top_k=cmd.get("top_k", 50),
- min_p=cmd.get("min_p", 0.0),
- max_new_tokens=cmd.get("max_new_tokens", 2048),
- repetition_penalty=cmd.get("repetition_penalty", 1.1),
- use_adapter=cmd.get("use_adapter"),
+ text = cmd["text"],
+ temperature = cmd.get("temperature", 0.6),
+ top_p = cmd.get("top_p", 0.95),
+ top_k = cmd.get("top_k", 50),
+ min_p = cmd.get("min_p", 0.0),
+ max_new_tokens = cmd.get("max_new_tokens", 2048),
+ repetition_penalty = cmd.get("repetition_penalty", 1.1),
+ use_adapter = cmd.get("use_adapter"),
)
# Send WAV bytes as base64 (bytes can't go through mp.Queue directly)
- _send_response(resp_queue, {
- "type": "audio_done",
- "request_id": request_id,
- "wav_base64": base64.b64encode(wav_bytes).decode("ascii"),
- "sample_rate": sample_rate,
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "audio_done",
+ "request_id": request_id,
+ "wav_base64": base64.b64encode(wav_bytes).decode("ascii"),
+ "sample_rate": sample_rate,
+ "ts": time.time(),
+ },
+ )
logger.info("Finished audio generation for request_id=%s", request_id)
except Exception as exc:
- logger.error("Audio generation error: %s", exc, exc_info=True)
- _send_response(resp_queue, {
- "type": "audio_error",
- "request_id": request_id,
- "error": str(exc),
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ logger.error("Audio generation error: %s", exc, exc_info = True)
+ _send_response(
+ resp_queue,
+ {
+ "type": "audio_error",
+ "request_id": request_id,
+ "error": str(exc),
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
def _handle_generate_audio_input(
@@ -333,59 +392,70 @@ def _handle_generate_audio_input(
import numpy as np
# Decode audio array from list (numpy arrays can't go through mp.Queue)
- audio_array = np.array(cmd["audio_data"], dtype=np.float32)
+ audio_array = np.array(cmd["audio_data"], dtype = np.float32)
audio_type = cmd.get("audio_type")
if audio_type == "whisper":
generator = backend.generate_whisper_response(
- audio_array=audio_array,
- cancel_event=cancel_event,
+ audio_array = audio_array,
+ cancel_event = cancel_event,
)
else:
generator = backend.generate_audio_input_response(
- messages=cmd.get("messages", []),
- system_prompt=cmd.get("system_prompt", ""),
- audio_array=audio_array,
- temperature=cmd.get("temperature", 0.7),
- top_p=cmd.get("top_p", 0.9),
- top_k=cmd.get("top_k", 40),
- min_p=cmd.get("min_p", 0.0),
- max_new_tokens=cmd.get("max_new_tokens", 512),
- repetition_penalty=cmd.get("repetition_penalty", 1.1),
- cancel_event=cancel_event,
+ messages = cmd.get("messages", []),
+ system_prompt = cmd.get("system_prompt", ""),
+ audio_array = audio_array,
+ temperature = cmd.get("temperature", 0.7),
+ top_p = cmd.get("top_p", 0.9),
+ top_k = cmd.get("top_k", 40),
+ min_p = cmd.get("min_p", 0.0),
+ max_new_tokens = cmd.get("max_new_tokens", 512),
+ repetition_penalty = cmd.get("repetition_penalty", 1.1),
+ cancel_event = cancel_event,
)
logger.info("Starting audio input generation for request_id=%s", request_id)
for text_chunk in generator:
if cancel_event.is_set():
- logger.info("Audio input generation cancelled for request %s", request_id)
+ logger.info(
+ "Audio input generation cancelled for request %s", request_id
+ )
break
- _send_response(resp_queue, {
- "type": "token",
- "request_id": request_id,
- "text": text_chunk,
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "token",
+ "request_id": request_id,
+ "text": text_chunk,
+ "ts": time.time(),
+ },
+ )
- _send_response(resp_queue, {
- "type": "gen_done",
- "request_id": request_id,
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "gen_done",
+ "request_id": request_id,
+ "ts": time.time(),
+ },
+ )
logger.info("Finished audio input generation for request_id=%s", request_id)
except Exception as exc:
- logger.error("Audio input generation error: %s", exc, exc_info=True)
- _send_response(resp_queue, {
- "type": "gen_error",
- "request_id": request_id,
- "error": str(exc),
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ logger.error("Audio input generation error: %s", exc, exc_info = True)
+ _send_response(
+ resp_queue,
+ {
+ "type": "gen_error",
+ "request_id": request_id,
+ "error": str(exc),
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
@@ -397,19 +467,25 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
elif backend.active_model_name:
backend.unload_model(backend.active_model_name)
- _send_response(resp_queue, {
- "type": "unloaded",
- "model_name": model_name,
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "unloaded",
+ "model_name": model_name,
+ "ts": time.time(),
+ },
+ )
except Exception as exc:
logger.error("Unload error: %s", exc)
- _send_response(resp_queue, {
- "type": "unloaded",
- "model_name": model_name,
- "error": str(exc),
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "unloaded",
+ "model_name": model_name,
+ "error": str(exc),
+ "ts": time.time(),
+ },
+ )
def run_inference_process(
@@ -428,16 +504,19 @@ def run_inference_process(
config: Initial configuration dict with model info.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
- os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
+ os.environ["PYTHONWARNINGS"] = (
+ "ignore" # Suppress warnings at C-level before imports
+ )
import warnings
from loggers.config import LogConfig
+
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
LogConfig.setup_logging(
- service_name="unsloth-studio-inference-worker",
- env=os.getenv("ENVIRONMENT_TYPE", "production"),
+ service_name = "unsloth-studio-inference-worker",
+ env = os.getenv("ENVIRONMENT_TYPE", "production"),
)
model_name = config["model_name"]
@@ -446,18 +525,22 @@ def run_inference_process(
try:
_activate_transformers_version(model_name)
except Exception as exc:
- _send_response(resp_queue, {
- "type": "error",
- "error": f"Failed to activate transformers version: {exc}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "error",
+ "error": f"Failed to activate transformers version: {exc}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
return
# ββ 1b. On Windows, check Triton availability (must be before import torch) ββ
if sys.platform == "win32":
try:
import triton # noqa: F401
+
logger.info("Triton available β torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
@@ -468,11 +551,14 @@ def run_inference_process(
# ββ 2. Import ML libraries (fresh in this clean process) ββ
try:
- _send_response(resp_queue, {
- "type": "status",
- "message": "Importing ML libraries...",
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "status",
+ "message": "Importing ML libraries...",
+ "ts": time.time(),
+ },
+ )
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
@@ -481,36 +567,46 @@ def run_inference_process(
from core.inference.inference import InferenceBackend
import transformers
+
logger.info("Subprocess loaded transformers %s", transformers.__version__)
except Exception as exc:
- _send_response(resp_queue, {
- "type": "error",
- "error": f"Failed to import ML libraries: {exc}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "error",
+ "error": f"Failed to import ML libraries: {exc}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
return
# ββ 3. Create inference backend and load initial model ββ
try:
backend = InferenceBackend()
- _send_response(resp_queue, {
- "type": "status",
- "message": "Loading model...",
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "status",
+ "message": "Loading model...",
+ "ts": time.time(),
+ },
+ )
_handle_load(backend, config, resp_queue)
except Exception as exc:
- _send_response(resp_queue, {
- "type": "error",
- "error": f"Failed to initialize inference backend: {exc}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "error",
+ "error": f"Failed to initialize inference backend: {exc}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
return
# ββ 4. Command loop β process commands until shutdown ββ
@@ -520,7 +616,7 @@ def run_inference_process(
while True:
try:
- cmd = cmd_queue.get(timeout=1.0)
+ cmd = cmd_queue.get(timeout = 1.0)
except _queue.Empty:
continue
except (EOFError, OSError):
@@ -564,26 +660,32 @@ def run_inference_process(
elif cmd_type == "reset":
cancel_event.set()
backend.reset_generation_state()
- _send_response(resp_queue, {
- "type": "reset_ack",
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "reset_ack",
+ "ts": time.time(),
+ },
+ )
elif cmd_type == "status":
# Return current status
- _send_response(resp_queue, {
- "type": "status_response",
- "active_model": backend.active_model_name,
- "models": {
- name: {
- "is_vision": info.get("is_vision", False),
- "is_lora": info.get("is_lora", False),
- }
- for name, info in backend.models.items()
+ _send_response(
+ resp_queue,
+ {
+ "type": "status_response",
+ "active_model": backend.active_model_name,
+ "models": {
+ name: {
+ "is_vision": info.get("is_vision", False),
+ "is_lora": info.get("is_lora", False),
+ }
+ for name, info in backend.models.items()
+ },
+ "loading": list(backend.loading_models),
+ "ts": time.time(),
},
- "loading": list(backend.loading_models),
- "ts": time.time(),
- })
+ )
elif cmd_type == "shutdown":
logger.info("Shutdown command received, exiting")
@@ -593,25 +695,36 @@ def run_inference_process(
backend.unload_model(model_name)
except Exception:
pass
- _send_response(resp_queue, {
- "type": "shutdown_ack",
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "shutdown_ack",
+ "ts": time.time(),
+ },
+ )
return
else:
logger.warning("Unknown command type: %s", cmd_type)
- _send_response(resp_queue, {
- "type": "error",
- "error": f"Unknown command type: {cmd_type}",
- "ts": time.time(),
- })
+ _send_response(
+ resp_queue,
+ {
+ "type": "error",
+ "error": f"Unknown command type: {cmd_type}",
+ "ts": time.time(),
+ },
+ )
except Exception as exc:
- logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info=True)
- _send_response(resp_queue, {
- "type": "error",
- "error": f"Command '{cmd_type}' failed: {exc}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ logger.error(
+ "Error handling command '%s': %s", cmd_type, exc, exc_info = True
+ )
+ _send_response(
+ resp_queue,
+ {
+ "type": "error",
+ "error": f"Command '{cmd_type}' failed: {exc}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ },
+ )
diff --git a/studio/backend/core/training/__init__.py b/studio/backend/core/training/__init__.py
index 3b92aa704d..6e19805eb0 100644
--- a/studio/backend/core/training/__init__.py
+++ b/studio/backend/core/training/__init__.py
@@ -4,10 +4,11 @@
"""
Training submodule - Training backends and trainer classes
"""
+
from .training import TrainingBackend, TrainingProgress, get_training_backend
__all__ = [
- 'TrainingProgress',
- 'TrainingBackend',
- 'get_training_backend',
+ "TrainingProgress",
+ "TrainingBackend",
+ "get_training_backend",
]
diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index f933b5582e..cee84a8809 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -5,13 +5,16 @@
Unsloth Training Backend
Integrates Unsloth training capabilities with the FastAPI backend
"""
+
import os
import sys
+
# Prevent tokenizer parallelism deadlocks when datasets uses multiprocessing fork
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import torch
from utils.hardware import clear_gpu_cache, safe_num_proc
+
torch._dynamo.config.recompile_limit = 64
from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
from unsloth.chat_templates import get_chat_template
@@ -31,13 +34,17 @@ from datasets import Dataset, load_dataset
from utils.models import is_vision_model, detect_audio_type
from utils.datasets import format_and_template_dataset
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER
-from utils.paths import ensure_dir, resolve_dataset_path, resolve_output_dir, resolve_tensorboard_dir
+from utils.paths import (
+ ensure_dir,
+ resolve_dataset_path,
+ resolve_output_dir,
+ resolve_tensorboard_dir,
+)
from trl import SFTTrainer, SFTConfig
logger = get_logger(__name__)
-
def _build_report_targets(training_args) -> list[str] | str:
report_to: list[str] = []
if training_args.get("enable_wandb", False):
@@ -50,6 +57,7 @@ def _build_report_targets(training_args) -> list[str] | str:
@dataclass
class TrainingProgress:
"""Training progress tracking"""
+
epoch: float = 0
step: int = 0
total_steps: int = 0
@@ -65,6 +73,7 @@ class TrainingProgress:
num_tokens: Optional[int] = None
eval_loss: Optional[float] = None
+
class UnslothTrainer:
"""
Unsloth Training Backend
@@ -85,10 +94,16 @@ class UnslothTrainer:
# Model state tracking
self.is_vlm = False
self.is_audio = False
- self.is_audio_vlm = False # Multimodal model (e.g. Gemma 3N) trained on audio data
+ self.is_audio_vlm = (
+ False # Multimodal model (e.g. Gemma 3N) trained on audio data
+ )
self._audio_type = None # 'csm', 'whisper', 'snac', 'bicodec', 'dac'
- self._cuda_audio_used = False # Set once after audio CUDA preprocessing; never cleared
- self._spark_tts_repo_dir = None # Path to downloaded Spark-TTS repo (for BiCodecTokenizer)
+ self._cuda_audio_used = (
+ False # Set once after audio CUDA preprocessing; never cleared
+ )
+ self._spark_tts_repo_dir = (
+ None # Path to downloaded Spark-TTS repo (for BiCodecTokenizer)
+ )
self.model_name = None
# Training metrics tracking
@@ -102,9 +117,9 @@ class UnslothTrainer:
# Store training context for later transfer
self.training_context = {
- 'base_model_name': None,
- 'output_dir': None,
- 'is_lora': True, # Default to LoRA
+ "base_model_name": None,
+ "output_dir": None,
+ "is_lora": True, # Default to LoRA
}
def pre_detect_and_load_tokenizer(
@@ -136,7 +151,7 @@ class UnslothTrainer:
# --- Detect audio type (reads config.json only, no VRAM) ---
self._audio_type = detect_audio_type(model_name, hf_token)
- if self._audio_type == 'audio_vlm':
+ if self._audio_type == "audio_vlm":
self.is_audio = False
self.is_audio_vlm = is_dataset_audio
self._audio_type = None
@@ -153,21 +168,30 @@ class UnslothTrainer:
logger.info(
"pre_detect: audio_type=%s, is_audio=%s, is_audio_vlm=%s, is_vlm=%s",
- self._audio_type, self.is_audio, self.is_audio_vlm, self.is_vlm,
+ self._audio_type,
+ self.is_audio,
+ self.is_audio_vlm,
+ self.is_vlm,
)
# --- Load lightweight tokenizer/processor (CPU only, no VRAM) ---
# Whisper needs AutoProcessor (has feature_extractor + tokenizer).
# All others work with AutoTokenizer (CSM loads its own processor inline).
- if self._audio_type == 'whisper':
+ if self._audio_type == "whisper":
from transformers import AutoProcessor
+
self.tokenizer = AutoProcessor.from_pretrained(
- model_name, trust_remote_code=trust_remote_code, token=hf_token,
+ model_name,
+ trust_remote_code = trust_remote_code,
+ token = hf_token,
)
else:
from transformers import AutoTokenizer
+
self.tokenizer = AutoTokenizer.from_pretrained(
- model_name, trust_remote_code=trust_remote_code, token=hf_token,
+ model_name,
+ trust_remote_code = trust_remote_code,
+ token = hf_token,
)
logger.info("Pre-loaded tokenizer for %s", model_name)
@@ -193,15 +217,16 @@ class UnslothTrainer:
def _create_progress_callback(self):
"""Create a TrainerCallback for progress tracking. Reused by all training branches."""
from transformers import TrainerCallback
+
trainer_ref = self
class _ProgressCallback(TrainerCallback):
- def on_log(self, args, state, control, logs=None, **kwargs):
+ def on_log(self, args, state, control, logs = None, **kwargs):
if not logs:
return
- loss_value = logs.get('loss', logs.get('train_loss', 0.0))
+ loss_value = logs.get("loss", logs.get("train_loss", 0.0))
current_step = state.global_step
- grad_norm = logs.get('grad_norm', None)
+ grad_norm = logs.get("grad_norm", None)
elapsed_seconds = None
if trainer_ref.training_start_time is not None:
@@ -213,25 +238,27 @@ class UnslothTrainer:
if total_steps > 0:
steps_remaining = total_steps - current_step
if steps_remaining > 0:
- eta_seconds = (elapsed_seconds / current_step) * steps_remaining
+ eta_seconds = (
+ elapsed_seconds / current_step
+ ) * steps_remaining
num_tokens = getattr(state, "num_input_tokens_seen", None)
trainer_ref._update_progress(
- step=current_step,
- epoch=round(state.epoch, 2) if state.epoch else 0,
- loss=loss_value,
- learning_rate=logs.get('learning_rate', 0.0),
- elapsed_seconds=elapsed_seconds,
- eta_seconds=eta_seconds,
- grad_norm=grad_norm,
- num_tokens=num_tokens,
- eval_loss=logs.get('eval_loss', None),
- status_message="",
+ step = current_step,
+ epoch = round(state.epoch, 2) if state.epoch else 0,
+ loss = loss_value,
+ learning_rate = logs.get("learning_rate", 0.0),
+ elapsed_seconds = elapsed_seconds,
+ eta_seconds = eta_seconds,
+ grad_norm = grad_norm,
+ num_tokens = num_tokens,
+ eval_loss = logs.get("eval_loss", None),
+ status_message = "",
)
def on_epoch_end(self, args, state, control, **kwargs):
- trainer_ref._update_progress(epoch=state.epoch, step=state.global_step)
+ trainer_ref._update_progress(epoch = state.epoch, step = state.global_step)
def on_step_end(self, args, state, control, **kwargs):
if trainer_ref.should_stop:
@@ -241,29 +268,35 @@ class UnslothTrainer:
return _ProgressCallback()
- def _calculate_total_steps(self, num_samples, batch_size, grad_accum, num_epochs, max_steps):
+ def _calculate_total_steps(
+ self, num_samples, batch_size, grad_accum, num_epochs, max_steps
+ ):
"""Calculate total training steps from dataset size and training params."""
if max_steps and max_steps > 0:
return max_steps
len_dataloader = math.ceil(num_samples / batch_size)
- steps_per_epoch = max(len_dataloader // grad_accum + int(len_dataloader % grad_accum > 0), 1)
+ steps_per_epoch = max(
+ len_dataloader // grad_accum + int(len_dataloader % grad_accum > 0), 1
+ )
return steps_per_epoch * num_epochs
- def _build_audio_training_args(self, training_args, output_dir, *, extra_args=None):
+ def _build_audio_training_args(self, training_args, output_dir, *, extra_args = None):
"""Build training args dict for audio branches.
Constructs the common config (batch size, lr, warmup, fp16/bf16, etc.)
and applies per-branch overrides via extra_args.
"""
- batch_size = training_args.get('batch_size', 2)
- gradient_accumulation_steps = training_args.get('gradient_accumulation_steps', 4)
- warmup_steps_val = training_args.get('warmup_steps', 5)
- max_steps_val = training_args.get('max_steps', 0)
- learning_rate = training_args.get('learning_rate', 2e-4)
- weight_decay = training_args.get('weight_decay', 0.001)
- lr_scheduler_type = training_args.get('lr_scheduler_type', 'linear')
- random_seed = training_args.get('random_seed', 3407)
- optim_value = training_args.get('optim', 'adamw_8bit')
+ batch_size = training_args.get("batch_size", 2)
+ gradient_accumulation_steps = training_args.get(
+ "gradient_accumulation_steps", 4
+ )
+ warmup_steps_val = training_args.get("warmup_steps", 5)
+ max_steps_val = training_args.get("max_steps", 0)
+ learning_rate = training_args.get("learning_rate", 2e-4)
+ weight_decay = training_args.get("weight_decay", 0.001)
+ lr_scheduler_type = training_args.get("lr_scheduler_type", "linear")
+ random_seed = training_args.get("random_seed", 3407)
+ optim_value = training_args.get("optim", "adamw_8bit")
config = {
"per_device_train_batch_size": batch_size,
@@ -290,10 +323,10 @@ class UnslothTrainer:
if max_steps_val and max_steps_val > 0:
config["max_steps"] = max_steps_val
else:
- config["num_train_epochs"] = training_args.get('num_epochs', 3)
+ config["num_train_epochs"] = training_args.get("num_epochs", 3)
# save_steps
- save_steps_val = training_args.get('save_steps', 0)
+ save_steps_val = training_args.get("save_steps", 0)
if save_steps_val and save_steps_val > 0:
config["save_steps"] = save_steps_val
config["save_strategy"] = "steps"
@@ -304,7 +337,7 @@ class UnslothTrainer:
return config
- def _finalize_training(self, output_dir, label=""):
+ def _finalize_training(self, output_dir, label = ""):
"""Save model after training and update progress. Used by all training branches."""
if self.should_stop and self.save_on_stop:
self.trainer.save_model()
@@ -313,13 +346,15 @@ class UnslothTrainer:
msg = f"{label} training stopped" if label else "Training stopped"
logger.info(f"\n{msg}. Model saved to {output_dir}\n")
self._update_progress(
- is_training=False,
- status_message=f"Training stopped. Model saved to {output_dir}",
+ is_training = False,
+ status_message = f"Training stopped. Model saved to {output_dir}",
)
elif self.should_stop:
msg = f"{label} training cancelled" if label else "Training cancelled"
logger.info(f"\n{msg}.\n")
- self._update_progress(is_training=False, status_message="Training cancelled.")
+ self._update_progress(
+ is_training = False, status_message = "Training cancelled."
+ )
else:
self.trainer.save_model()
self.tokenizer.save_pretrained(output_dir)
@@ -327,9 +362,9 @@ class UnslothTrainer:
msg = f"{label} training completed" if label else "Training completed"
logger.info(f"\n{msg}! Model saved to {output_dir}\n")
self._update_progress(
- is_training=False,
- is_completed=True,
- status_message=f"Training completed! Model saved to {output_dir}",
+ is_training = False,
+ is_completed = True,
+ status_message = f"Training completed! Model saved to {output_dir}",
)
def _cleanup_audio_artifacts(self):
@@ -349,7 +384,9 @@ class UnslothTrainer:
]
# Spark-TTS path is relative to the downloaded repo
if self._spark_tts_repo_dir:
- spark_code_dir = os.path.join(os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS")
+ spark_code_dir = os.path.join(
+ os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS"
+ )
audio_paths.append(spark_code_dir)
removed_paths = []
@@ -359,14 +396,16 @@ class UnslothTrainer:
removed_paths.append(path)
# Remove stale audio modules from sys.modules
- prefixes = ('snac', 'whisper', 'sparktts', 'outetts')
+ prefixes = ("snac", "whisper", "sparktts", "outetts")
removed_modules = [key for key in _sys.modules if key.startswith(prefixes)]
for key in removed_modules:
del _sys.modules[key]
if removed_paths or removed_modules:
- logger.info(f"Cleaned up audio artifacts: {len(removed_paths)} paths, "
- f"{len(removed_modules)} modules\n")
+ logger.info(
+ f"Cleaned up audio artifacts: {len(removed_paths)} paths, "
+ f"{len(removed_modules)} modules\n"
+ )
def _resolve_audio_columns(self, dataset, custom_format_mapping: dict = None):
"""Resolve audio, text, and speaker columns from user mapping or hardcoded fallback.
@@ -389,11 +428,22 @@ class UnslothTrainer:
speaker_col = col
# Use mapping if both required columns exist in the dataset
if audio_col and audio_col in cols and text_col and text_col in cols:
- return {"audio_col": audio_col, "text_col": text_col, "speaker_col": speaker_col}
+ return {
+ "audio_col": audio_col,
+ "text_col": text_col,
+ "speaker_col": speaker_col,
+ }
# Hardcoded fallback (existing behavior)
audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None)
- text_col = next((c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), None)
+ text_col = next(
+ (
+ c
+ for c in cols
+ if c.lower() in ("text", "sentence", "transcript", "transcription")
+ ),
+ None,
+ )
speaker_col = None
if "source" in cols:
@@ -401,20 +451,27 @@ class UnslothTrainer:
elif "speaker_id" in cols:
speaker_col = "speaker_id"
- return {"audio_col": audio_col, "text_col": text_col, "speaker_col": speaker_col}
+ return {
+ "audio_col": audio_col,
+ "text_col": text_col,
+ "speaker_col": speaker_col,
+ }
-
- def load_model(self,
- model_name: str,
- max_seq_length: int = 2048,
- load_in_4bit: bool = True,
- hf_token: Optional[str] = None,
- is_dataset_image: bool = False,
- is_dataset_audio: bool = False,
- trust_remote_code: bool = False) -> bool:
+ def load_model(
+ self,
+ model_name: str,
+ max_seq_length: int = 2048,
+ load_in_4bit: bool = True,
+ hf_token: Optional[str] = None,
+ is_dataset_image: bool = False,
+ is_dataset_audio: bool = False,
+ trust_remote_code: bool = False,
+ ) -> bool:
"""Load model for training (supports both text and vision models)"""
self.load_in_4bit = load_in_4bit # Store for training_meta.json
- self.trust_remote_code = trust_remote_code # For AutoProcessor etc. used during training
+ self.trust_remote_code = (
+ trust_remote_code # For AutoProcessor etc. used during training
+ )
try:
if self.model is not None:
del self.model
@@ -439,9 +496,10 @@ class UnslothTrainer:
# restores original class definitions so Unsloth can re-compile cleanly.
import sys as _sys
import importlib
+
for _key, _mod in list(_sys.modules.items()):
- if 'transformers.models.' in _key and '.modeling_' in _key:
- if hasattr(_mod, '__UNSLOTH_PATCHED__'):
+ if "transformers.models." in _key and ".modeling_" in _key:
+ if hasattr(_mod, "__UNSLOTH_PATCHED__"):
try:
importlib.reload(_mod)
except Exception:
@@ -449,13 +507,16 @@ class UnslothTrainer:
# Remove stale compiled cache so the new model gets a fresh one
from utils.cache_cleanup import clear_unsloth_compiled_cache
+
clear_unsloth_compiled_cache()
# Detect audio model type dynamically (config.json + tokenizer)
self._audio_type = detect_audio_type(model_name, hf_token)
# audio_vlm is detected as an audio_type now, handle it separately
- if self._audio_type == 'audio_vlm':
+ if self._audio_type == "audio_vlm":
self.is_audio = False
- self.is_audio_vlm = is_dataset_audio # Only use audio VLM path if dataset has audio
+ self.is_audio_vlm = (
+ is_dataset_audio # Only use audio VLM path if dataset has audio
+ )
self._audio_type = None
else:
self.is_audio = self._audio_type is not None
@@ -470,25 +531,33 @@ class UnslothTrainer:
self.model_name = model_name
self.max_seq_length = max_seq_length
- logger.info(f"Audio type: {self._audio_type}, is_audio: {self.is_audio}, is_audio_vlm: {self.is_audio_vlm}")
- logger.info(f"Dataset has images: {is_dataset_image}, audio: {is_dataset_audio}")
+ logger.info(
+ f"Audio type: {self._audio_type}, is_audio: {self.is_audio}, is_audio_vlm: {self.is_audio_vlm}"
+ )
+ logger.info(
+ f"Dataset has images: {is_dataset_image}, audio: {is_dataset_audio}"
+ )
logger.info(f"Using VLM path: {self.is_vlm}")
# Reset training state for new run
self._update_progress(
- is_training=True,
- is_completed=False,
- error=None,
- step=0,
- loss=0.0,
- epoch=0
+ is_training = True,
+ is_completed = False,
+ error = None,
+ step = 0,
+ loss = 0.0,
+ epoch = 0,
)
# Update UI immediately with loading message
- model_display = model_name.split('/')[-1] if '/' in model_name else model_name
- model_type_label = 'audio' if self.is_audio else ('vision' if self.is_vlm else 'text')
+ model_display = (
+ model_name.split("/")[-1] if "/" in model_name else model_name
+ )
+ model_type_label = (
+ "audio" if self.is_audio else ("vision" if self.is_vlm else "text")
+ )
self._update_progress(
- status_message=f"Loading {model_type_label} model... {model_display}"
+ status_message = f"Loading {model_type_label} model... {model_display}"
)
logger.info(f"\nLoading {model_type_label} model: {model_name}")
@@ -499,10 +568,11 @@ class UnslothTrainer:
# Proactive gated-model check: verify access BEFORE from_pretrained.
# Catches ALL gated/private models (text, vision, audio) globally.
- if '/' in model_name: # Only check HF repo IDs, not local paths
+ if "/" in model_name: # Only check HF repo IDs, not local paths
try:
from huggingface_hub import model_info as hf_model_info
- info = hf_model_info(model_name, token=hf_token or None)
+
+ info = hf_model_info(model_name, token = hf_token or None)
# model_info succeeds even for gated repos (metadata is public),
# but info.gated tells us if files require acceptance/token.
if info.gated and not hf_token:
@@ -510,49 +580,57 @@ class UnslothTrainer:
f"Access denied for '{model_name}'. This model is gated. "
f"Please add a Hugging Face token with access and try again."
)
- logger.error(f"Model '{model_name}' is gated (gated={info.gated}) and no HF token provided")
- self._update_progress(error=friendly, is_training=False)
+ logger.error(
+ f"Model '{model_name}' is gated (gated={info.gated}) and no HF token provided"
+ )
+ self._update_progress(error = friendly, is_training = False)
return False
except Exception as gate_err:
- from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
+ from huggingface_hub.utils import (
+ GatedRepoError,
+ RepositoryNotFoundError,
+ )
+
if isinstance(gate_err, (GatedRepoError, RepositoryNotFoundError)):
friendly = (
f"Access denied for '{model_name}'. This model is gated or private. "
f"Please add a Hugging Face token with access and try again."
)
logger.error(f"Gated model check failed: {gate_err}")
- self._update_progress(error=friendly, is_training=False)
+ self._update_progress(error = friendly, is_training = False)
return False
# Branch based on model type
- if self._audio_type == 'csm':
+ if self._audio_type == "csm":
# CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False
from unsloth import FastModel
from transformers import CsmForConditionalGeneration
+
self.model, self.tokenizer = FastModel.from_pretrained(
- model_name=model_name,
- max_seq_length=max_seq_length,
- dtype=None,
- auto_model=CsmForConditionalGeneration,
- load_in_4bit=False,
- token=hf_token,
- trust_remote_code=trust_remote_code,
+ model_name = model_name,
+ max_seq_length = max_seq_length,
+ dtype = None,
+ auto_model = CsmForConditionalGeneration,
+ load_in_4bit = False,
+ token = hf_token,
+ trust_remote_code = trust_remote_code,
)
logger.info("Loaded CSM audio model")
- elif self._audio_type == 'whisper':
+ elif self._audio_type == "whisper":
# Whisper: FastModel + auto_model=WhisperForConditionalGeneration + load_in_4bit=False
from unsloth import FastModel
from transformers import WhisperForConditionalGeneration
+
self.model, self.tokenizer = FastModel.from_pretrained(
- model_name=model_name,
- dtype=None,
- load_in_4bit=False,
- auto_model=WhisperForConditionalGeneration,
- whisper_language="English",
- whisper_task="transcribe",
- token=hf_token,
- trust_remote_code=trust_remote_code,
+ model_name = model_name,
+ dtype = None,
+ load_in_4bit = False,
+ auto_model = WhisperForConditionalGeneration,
+ whisper_language = "English",
+ whisper_task = "transcribe",
+ token = hf_token,
+ trust_remote_code = trust_remote_code,
)
# Configure generation settings (notebook lines 100-105)
self.model.generation_config.language = "<|en|>"
@@ -561,19 +639,21 @@ class UnslothTrainer:
self.model.generation_config.forced_decoder_ids = None
logger.info("Loaded Whisper audio model (FastModel)")
- elif self._audio_type == 'snac':
+ elif self._audio_type == "snac":
# Orpheus: language model with audio codec tokens
self.model, self.tokenizer = FastLanguageModel.from_pretrained(
- model_name=model_name,
- max_seq_length=max_seq_length,
- dtype=None,
- load_in_4bit=load_in_4bit,
- token=hf_token,
- trust_remote_code=trust_remote_code,
+ model_name = model_name,
+ max_seq_length = max_seq_length,
+ dtype = None,
+ load_in_4bit = load_in_4bit,
+ token = hf_token,
+ trust_remote_code = trust_remote_code,
+ )
+ logger.info(
+ f"Loaded {self._audio_type} audio model (FastLanguageModel)"
)
- logger.info(f"Loaded {self._audio_type} audio model (FastLanguageModel)")
- elif self._audio_type == 'bicodec':
+ elif self._audio_type == "bicodec":
# Spark-TTS: download full repo (contains sparktts package + BiCodec weights),
# then load only the LLM subfolder with FastModel.
# model_name may be:
@@ -593,29 +673,32 @@ class UnslothTrainer:
local_dir = model_name.split("/")[-1]
llm_path = f"{local_dir}/LLM"
- repo_path = snapshot_download(hf_repo, local_dir=local_dir)
- self._spark_tts_repo_dir = os.path.abspath(repo_path) # Absolute path for sys.path
+ repo_path = snapshot_download(hf_repo, local_dir = local_dir)
+ self._spark_tts_repo_dir = os.path.abspath(
+ repo_path
+ ) # Absolute path for sys.path
llm_path = os.path.join(self._spark_tts_repo_dir, "LLM")
self.model, self.tokenizer = FastModel.from_pretrained(
- model_name=llm_path,
- max_seq_length=max_seq_length,
- dtype=torch.float32, # Spark-TTS requires float32
- load_in_4bit=False,
- token=hf_token,
- trust_remote_code=trust_remote_code,
+ model_name = llm_path,
+ max_seq_length = max_seq_length,
+ dtype = torch.float32, # Spark-TTS requires float32
+ load_in_4bit = False,
+ token = hf_token,
+ trust_remote_code = trust_remote_code,
)
logger.info("Loaded Spark-TTS (bicodec) model")
- elif self._audio_type == 'dac':
+ elif self._audio_type == "dac":
# OuteTTS: uses FastModel (not FastLanguageModel) with load_in_4bit=False
from unsloth import FastModel
+
self.model, self.tokenizer = FastModel.from_pretrained(
model_name,
- max_seq_length=max_seq_length,
- load_in_4bit=False,
- token=hf_token,
- trust_remote_code=trust_remote_code,
+ max_seq_length = max_seq_length,
+ load_in_4bit = False,
+ token = hf_token,
+ trust_remote_code = trust_remote_code,
)
logger.info("Loaded OuteTTS (dac) model (FastModel)")
@@ -623,57 +706,71 @@ class UnslothTrainer:
# Audio VLM: multimodal model trained on audio (e.g. Gemma 3N)
# Uses FastModel (general loader) β returns (model, processor)
from unsloth import FastModel
+
self.model, self.tokenizer = FastModel.from_pretrained(
- model_name=model_name,
- max_seq_length=max_seq_length,
- dtype=None,
- load_in_4bit=load_in_4bit,
- token=hf_token,
- trust_remote_code=trust_remote_code,
+ model_name = model_name,
+ max_seq_length = max_seq_length,
+ dtype = None,
+ load_in_4bit = load_in_4bit,
+ token = hf_token,
+ trust_remote_code = trust_remote_code,
)
logger.info("Loaded audio VLM model (FastModel)")
elif self.is_vlm:
# Load vision model - returns (model, tokenizer)
self.model, self.tokenizer = FastVisionModel.from_pretrained(
- model_name=model_name,
- max_seq_length=max_seq_length,
- dtype=None, # Auto-detect
- load_in_4bit=load_in_4bit,
- token=hf_token,
- trust_remote_code=trust_remote_code,
+ model_name = model_name,
+ max_seq_length = max_seq_length,
+ dtype = None, # Auto-detect
+ load_in_4bit = load_in_4bit,
+ token = hf_token,
+ trust_remote_code = trust_remote_code,
)
logger.info("Loaded vision model")
# Diagnostic: check if FastVisionModel returned a real Processor or a raw tokenizer
from transformers import ProcessorMixin
+
tok = self.tokenizer
- has_image_proc = isinstance(tok, ProcessorMixin) or hasattr(tok, "image_processor")
- logger.info(f"\n[VLM Diagnostic] FastVisionModel returned: {type(tok).__name__}")
- logger.info(f"[VLM Diagnostic] Is ProcessorMixin: {isinstance(tok, ProcessorMixin)}")
- logger.info(f"[VLM Diagnostic] Has image_processor: {hasattr(tok, 'image_processor')}")
- logger.info(f"[VLM Diagnostic] Usable as vision processor: {has_image_proc}\n")
+ has_image_proc = isinstance(tok, ProcessorMixin) or hasattr(
+ tok, "image_processor"
+ )
+ logger.info(
+ f"\n[VLM Diagnostic] FastVisionModel returned: {type(tok).__name__}"
+ )
+ logger.info(
+ f"[VLM Diagnostic] Is ProcessorMixin: {isinstance(tok, ProcessorMixin)}"
+ )
+ logger.info(
+ f"[VLM Diagnostic] Has image_processor: {hasattr(tok, 'image_processor')}"
+ )
+ logger.info(
+ f"[VLM Diagnostic] Usable as vision processor: {has_image_proc}\n"
+ )
else:
# Load text model - returns (model, tokenizer)
self.model, self.tokenizer = FastLanguageModel.from_pretrained(
- model_name=model_name,
- max_seq_length=max_seq_length,
- dtype=None, # Auto-detect
- load_in_4bit=load_in_4bit,
- token=hf_token,
- trust_remote_code=trust_remote_code,
+ model_name = model_name,
+ max_seq_length = max_seq_length,
+ dtype = None, # Auto-detect
+ load_in_4bit = load_in_4bit,
+ token = hf_token,
+ trust_remote_code = trust_remote_code,
)
logger.info("Loaded text model")
if self.should_stop:
return False
- self._update_progress(status_message="Model loaded successfully")
+ self._update_progress(status_message = "Model loaded successfully")
logger.info("Model loaded successfully")
return True
except OSError as e:
- if "could not get source code" in str(e) and not getattr(self, '_source_code_retried', False):
+ if "could not get source code" in str(e) and not getattr(
+ self, "_source_code_retried", False
+ ):
# Unsloth's patching can leave stale state that makes
# inspect.getsource() fail when switching model families
# (e.g. gemma3 β gemma3n). The load always succeeds on a
@@ -681,48 +778,77 @@ class UnslothTrainer:
# imports clean up the stale state as a side effect.
self._source_code_retried = True
logger.info(f"\n'could not get source code' β retrying once...\n")
- return self.load_model(model_name, max_seq_length, load_in_4bit, hf_token,
- is_dataset_image, is_dataset_audio, trust_remote_code)
+ return self.load_model(
+ model_name,
+ max_seq_length,
+ load_in_4bit,
+ hf_token,
+ is_dataset_image,
+ is_dataset_audio,
+ trust_remote_code,
+ )
error_msg = str(e)
error_lower = error_msg.lower()
- if any(k in error_lower for k in ("gated repo", "access to it at", "401", "403", "unauthorized", "forbidden")):
+ if any(
+ k in error_lower
+ for k in (
+ "gated repo",
+ "access to it at",
+ "401",
+ "403",
+ "unauthorized",
+ "forbidden",
+ )
+ ):
error_msg = (
f"Access denied for '{model_name}'. This model is gated or private. "
f"Please add a Hugging Face token with access and try again."
)
logger.error(f"Error loading model: {e}")
- self._update_progress(error=error_msg, is_training=False)
+ self._update_progress(error = error_msg, is_training = False)
return False
except Exception as e:
error_msg = str(e)
# Catch gated/auth errors and surface a friendly message
error_lower = error_msg.lower()
- if any(k in error_lower for k in ("gated repo", "access to it at", "401", "403", "unauthorized", "forbidden")):
+ if any(
+ k in error_lower
+ for k in (
+ "gated repo",
+ "access to it at",
+ "401",
+ "403",
+ "unauthorized",
+ "forbidden",
+ )
+ ):
error_msg = (
f"Access denied for '{model_name}'. This model is gated or private. "
f"Please add a Hugging Face token with access and try again."
)
logger.error(f"Error loading model: {e}")
- self._update_progress(error=error_msg, is_training=False)
+ self._update_progress(error = error_msg, is_training = False)
return False
finally:
self._source_code_retried = False
- def prepare_model_for_training(self,
- use_lora: bool = True,
- # Vision-specific LoRA parameters (only used if is_vlm=True)
- finetune_vision_layers: bool = True,
- finetune_language_layers: bool = True,
- finetune_attention_modules: bool = True,
- finetune_mlp_modules: bool = True,
- # Standard LoRA parameters
- target_modules: list = None,
- lora_r: int = 16,
- lora_alpha: int = 16,
- lora_dropout: float = 0.0,
- use_gradient_checkpointing: str = "unsloth",
- use_rslora: bool = False,
- use_loftq: bool = False) -> bool:
+ def prepare_model_for_training(
+ self,
+ use_lora: bool = True,
+ # Vision-specific LoRA parameters (only used if is_vlm=True)
+ finetune_vision_layers: bool = True,
+ finetune_language_layers: bool = True,
+ finetune_attention_modules: bool = True,
+ finetune_mlp_modules: bool = True,
+ # Standard LoRA parameters
+ target_modules: list = None,
+ lora_r: int = 16,
+ lora_alpha: int = 16,
+ lora_dropout: float = 0.0,
+ use_gradient_checkpointing: str = "unsloth",
+ use_rslora: bool = False,
+ use_loftq: bool = False,
+ ) -> bool:
"""
Prepare model for training (with optional LoRA).
"""
@@ -730,10 +856,11 @@ class UnslothTrainer:
if self.model is None:
raise ValueError("Model not loaded. Call load_model() first.")
-
# Full finetuning mode - skip PEFT entirely
if not use_lora:
- self._update_progress(status_message="Full finetuning mode - no LoRA adapters")
+ self._update_progress(
+ status_message = "Full finetuning mode - no LoRA adapters"
+ )
logger.info("Full finetuning mode - training all parameters\n")
return True
@@ -744,15 +871,27 @@ class UnslothTrainer:
target_modules = "all-linear"
else:
target_modules = [m for m in target_modules if m != "all-linear"]
- elif target_modules is None or (isinstance(target_modules, list) and len(target_modules) == 0):
- target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
- "gate_proj", "up_proj", "down_proj"]
+ elif target_modules is None or (
+ isinstance(target_modules, list) and len(target_modules) == 0
+ ):
+ target_modules = [
+ "q_proj",
+ "k_proj",
+ "v_proj",
+ "o_proj",
+ "gate_proj",
+ "up_proj",
+ "down_proj",
+ ]
# Validate and normalize gradient_checkpointing
# Must be one of: True, False, or "unsloth"
if isinstance(use_gradient_checkpointing, str):
use_gradient_checkpointing = use_gradient_checkpointing.strip().lower()
- if use_gradient_checkpointing == "" or use_gradient_checkpointing == "unsloth":
+ if (
+ use_gradient_checkpointing == ""
+ or use_gradient_checkpointing == "unsloth"
+ ):
use_gradient_checkpointing = "unsloth"
elif use_gradient_checkpointing in ("true", "1", "yes"):
use_gradient_checkpointing = True
@@ -760,102 +899,122 @@ class UnslothTrainer:
use_gradient_checkpointing = False
else:
# Invalid value, default to "unsloth"
- logger.warning(f"Invalid gradient_checkpointing value: {use_gradient_checkpointing}, defaulting to 'unsloth'")
+ logger.warning(
+ f"Invalid gradient_checkpointing value: {use_gradient_checkpointing}, defaulting to 'unsloth'"
+ )
use_gradient_checkpointing = "unsloth"
elif use_gradient_checkpointing not in (True, False, "unsloth"):
# Invalid type or value, default to "unsloth"
- logger.warning(f"Invalid gradient_checkpointing type/value: {use_gradient_checkpointing}, defaulting to 'unsloth'")
+ logger.warning(
+ f"Invalid gradient_checkpointing type/value: {use_gradient_checkpointing}, defaulting to 'unsloth'"
+ )
use_gradient_checkpointing = "unsloth"
# Verify model is loaded
if self.model is None:
error_msg = "Model is None - model was not loaded properly"
logger.error(error_msg)
- self._update_progress(error=error_msg)
+ self._update_progress(error = error_msg)
return False
# Check if model has the expected attributes
- if not hasattr(self.model, 'config'):
+ if not hasattr(self.model, "config"):
error_msg = "Model does not have config attribute - model may not be loaded correctly"
logger.error(error_msg)
- self._update_progress(error=error_msg)
+ self._update_progress(error = error_msg)
return False
- logger.info(f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n")
- logger.info(f"Gradient checkpointing: {use_gradient_checkpointing} (type: {type(use_gradient_checkpointing).__name__})\n")
+ logger.info(
+ f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n"
+ )
+ logger.info(
+ f"Gradient checkpointing: {use_gradient_checkpointing} (type: {type(use_gradient_checkpointing).__name__})\n"
+ )
# Branch based on model type: audio, audio_vlm, vision, or text
- if self._audio_type in ('csm', 'bicodec', 'dac') or self.is_audio_vlm:
+ if self._audio_type in ("csm", "bicodec", "dac") or self.is_audio_vlm:
# Models using FastModel.get_peft_model (codec audio + audio VLM)
from unsloth import FastModel
- label = self._audio_type or 'audio_vlm'
+
+ label = self._audio_type or "audio_vlm"
logger.info(f"{label} LoRA configuration:")
logger.info(f" - Target modules: {target_modules}")
if self.is_audio_vlm:
logger.info(f" - Finetune vision layers: {finetune_vision_layers}")
- logger.info(f" - Finetune language layers: {finetune_language_layers}")
- logger.info(f" - Finetune attention modules: {finetune_attention_modules}")
+ logger.info(
+ f" - Finetune language layers: {finetune_language_layers}"
+ )
+ logger.info(
+ f" - Finetune attention modules: {finetune_attention_modules}"
+ )
logger.info(f" - Finetune MLP modules: {finetune_mlp_modules}")
logger.info()
peft_kwargs = dict(
- r=lora_r,
- target_modules=target_modules,
- lora_alpha=lora_alpha,
- lora_dropout=lora_dropout,
- bias="none",
- use_gradient_checkpointing=use_gradient_checkpointing,
- random_state=3407,
- use_rslora=use_rslora,
- loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None,
+ r = lora_r,
+ target_modules = target_modules,
+ lora_alpha = lora_alpha,
+ lora_dropout = lora_dropout,
+ bias = "none",
+ use_gradient_checkpointing = use_gradient_checkpointing,
+ random_state = 3407,
+ use_rslora = use_rslora,
+ loftq_config = {"loftq_bits": 4, "loftq_iter": 1}
+ if use_loftq
+ else None,
)
# Audio VLM models support VLM-style layer selection
if self.is_audio_vlm:
peft_kwargs.update(
- finetune_vision_layers=finetune_vision_layers,
- finetune_language_layers=finetune_language_layers,
- finetune_attention_modules=finetune_attention_modules,
- finetune_mlp_modules=finetune_mlp_modules,
+ finetune_vision_layers = finetune_vision_layers,
+ finetune_language_layers = finetune_language_layers,
+ finetune_attention_modules = finetune_attention_modules,
+ finetune_mlp_modules = finetune_mlp_modules,
)
self.model = FastModel.get_peft_model(self.model, **peft_kwargs)
- elif self._audio_type == 'whisper':
+ elif self._audio_type == "whisper":
# Phase 2: Whisper uses FastModel.get_peft_model with task_type=None
from unsloth import FastModel
+
logger.info(f"Audio model (whisper) LoRA configuration:")
logger.info(f" - Target modules: {target_modules}\n")
self.model = FastModel.get_peft_model(
self.model,
- r=lora_r,
- target_modules=target_modules,
- lora_alpha=lora_alpha,
- lora_dropout=lora_dropout,
- bias="none",
- use_gradient_checkpointing=use_gradient_checkpointing,
- random_state=3407,
- use_rslora=use_rslora,
- loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None,
- task_type=None,
+ r = lora_r,
+ target_modules = target_modules,
+ lora_alpha = lora_alpha,
+ lora_dropout = lora_dropout,
+ bias = "none",
+ use_gradient_checkpointing = use_gradient_checkpointing,
+ random_state = 3407,
+ use_rslora = use_rslora,
+ loftq_config = {"loftq_bits": 4, "loftq_iter": 1}
+ if use_loftq
+ else None,
+ task_type = None,
)
- elif self._audio_type == 'snac':
+ elif self._audio_type == "snac":
# Orpheus uses FastLanguageModel.get_peft_model
logger.info(f"Audio model ({self._audio_type}) LoRA configuration:")
logger.info(f" - Target modules: {target_modules}\n")
self.model = FastLanguageModel.get_peft_model(
self.model,
- r=lora_r,
- target_modules=target_modules,
- lora_alpha=lora_alpha,
- lora_dropout=lora_dropout,
- bias="none",
- use_gradient_checkpointing=use_gradient_checkpointing,
- random_state=3407,
- use_rslora=use_rslora,
- loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None,
+ r = lora_r,
+ target_modules = target_modules,
+ lora_alpha = lora_alpha,
+ lora_dropout = lora_dropout,
+ bias = "none",
+ use_gradient_checkpointing = use_gradient_checkpointing,
+ random_state = 3407,
+ use_rslora = use_rslora,
+ loftq_config = {"loftq_bits": 4, "loftq_iter": 1}
+ if use_loftq
+ else None,
)
elif self.is_vlm:
@@ -863,24 +1022,28 @@ class UnslothTrainer:
logger.info(f"Vision model LoRA configuration:")
logger.info(f" - Finetune vision layers: {finetune_vision_layers}")
logger.info(f" - Finetune language layers: {finetune_language_layers}")
- logger.info(f" - Finetune attention modules: {finetune_attention_modules}")
+ logger.info(
+ f" - Finetune attention modules: {finetune_attention_modules}"
+ )
logger.info(f" - Finetune MLP modules: {finetune_mlp_modules}\n")
self.model = FastVisionModel.get_peft_model(
self.model,
- finetune_vision_layers=finetune_vision_layers,
- finetune_language_layers=finetune_language_layers,
- finetune_attention_modules=finetune_attention_modules,
- finetune_mlp_modules=finetune_mlp_modules,
- r=lora_r,
- target_modules=target_modules,
- lora_alpha=lora_alpha,
- lora_dropout=lora_dropout,
- bias="none",
- use_gradient_checkpointing=use_gradient_checkpointing,
- random_state=3407,
- use_rslora=use_rslora,
- loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None,
+ finetune_vision_layers = finetune_vision_layers,
+ finetune_language_layers = finetune_language_layers,
+ finetune_attention_modules = finetune_attention_modules,
+ finetune_mlp_modules = finetune_mlp_modules,
+ r = lora_r,
+ target_modules = target_modules,
+ lora_alpha = lora_alpha,
+ lora_dropout = lora_dropout,
+ bias = "none",
+ use_gradient_checkpointing = use_gradient_checkpointing,
+ random_state = 3407,
+ use_rslora = use_rslora,
+ loftq_config = {"loftq_bits": 4, "loftq_iter": 1}
+ if use_loftq
+ else None,
)
else:
# Text model LoRA
@@ -889,15 +1052,17 @@ class UnslothTrainer:
self.model = FastLanguageModel.get_peft_model(
self.model,
- r=lora_r,
- target_modules=target_modules,
- lora_alpha=lora_alpha,
- lora_dropout=lora_dropout,
- bias="none",
- use_gradient_checkpointing=use_gradient_checkpointing,
- random_state=3407,
- use_rslora=use_rslora,
- loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None,
+ r = lora_r,
+ target_modules = target_modules,
+ lora_alpha = lora_alpha,
+ lora_dropout = lora_dropout,
+ bias = "none",
+ use_gradient_checkpointing = use_gradient_checkpointing,
+ random_state = 3407,
+ use_rslora = use_rslora,
+ loftq_config = {"loftq_bits": 4, "loftq_iter": 1}
+ if use_loftq
+ else None,
)
# Check if stopped during LoRA preparation
@@ -905,20 +1070,25 @@ class UnslothTrainer:
logger.info("Stopped during LoRA configuration\n")
return False
- self._update_progress(status_message="LoRA adapters configured")
+ self._update_progress(status_message = "LoRA adapters configured")
logger.info("LoRA adapters configured successfully\n")
return True
except Exception as e:
import traceback
import sys
- error_details = f"{type(e).__name__}: {str(e)}" if str(e) else f"{type(e).__name__} (no message)"
+
+ error_details = (
+ f"{type(e).__name__}: {str(e)}"
+ if str(e)
+ else f"{type(e).__name__} (no message)"
+ )
full_traceback = traceback.format_exc()
logger.error(f"Error preparing model: {error_details}")
logger.error(f"Full traceback:\n{full_traceback}")
logger.info(f"\n[ERROR] Error preparing model: {error_details}")
logger.info(f"[ERROR] Full traceback:\n{full_traceback}")
- self._update_progress(error=error_details)
+ self._update_progress(error = error_details)
return False
def _apply_csm_forward_fix(self):
@@ -946,29 +1116,44 @@ class UnslothTrainer:
# Keys that the depth decoder and its sub-layers actually understand
_TRANSFORMERS_KWARGS = {
- 'num_items_in_batch', 'output_hidden_states', 'output_attentions',
- 'output_router_logits', 'cu_seq_lens_q', 'cu_seq_lens_k',
- 'max_length_q', 'max_length_k',
+ "num_items_in_batch",
+ "output_hidden_states",
+ "output_attentions",
+ "output_router_logits",
+ "cu_seq_lens_q",
+ "cu_seq_lens_k",
+ "max_length_q",
+ "max_length_k",
}
def _fixed_csm_forward(
self,
- input_ids=None, input_values=None, attention_mask=None,
- input_values_cutoffs=None, position_ids=None, past_key_values=None,
- inputs_embeds=None, labels=None, use_cache=None,
- cache_position=None, logits_to_keep=0, **kwargs,
+ input_ids = None,
+ input_values = None,
+ attention_mask = None,
+ input_values_cutoffs = None,
+ position_ids = None,
+ past_key_values = None,
+ inputs_embeds = None,
+ labels = None,
+ use_cache = None,
+ cache_position = None,
+ logits_to_keep = 0,
+ **kwargs,
):
# Strip non-standard kwargs injected by Unsloth/PEFT (causal_mask,
# num_logits_to_keep, task_ids, return_dict, etc.)
- output_attentions = kwargs.pop('output_attentions', None)
- output_hidden_states = kwargs.pop('output_hidden_states', None)
- kwargs.pop('return_dict', None)
- kwargs.pop('causal_mask', None)
- kwargs.pop('num_logits_to_keep', None)
- kwargs.pop('task_ids', None)
+ output_attentions = kwargs.pop("output_attentions", None)
+ output_hidden_states = kwargs.pop("output_hidden_states", None)
+ kwargs.pop("return_dict", None)
+ kwargs.pop("causal_mask", None)
+ kwargs.pop("num_logits_to_keep", None)
+ kwargs.pop("task_ids", None)
# Only keep recognized TransformersKwargs
- clean_kwargs = {k: v for k, v in kwargs.items() if k in _TRANSFORMERS_KWARGS}
+ clean_kwargs = {
+ k: v for k, v in kwargs.items() if k in _TRANSFORMERS_KWARGS
+ }
if input_ids is not None and input_ids.ndim == 2:
merged = self._merge_input_ids_with_input_values(
@@ -979,18 +1164,22 @@ class UnslothTrainer:
input_ids = None
backbone_outputs = self.backbone_model(
- input_ids=input_ids, attention_mask=attention_mask,
- position_ids=position_ids, past_key_values=past_key_values,
- inputs_embeds=inputs_embeds, use_cache=use_cache,
- cache_position=cache_position,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
+ input_ids = input_ids,
+ attention_mask = attention_mask,
+ position_ids = position_ids,
+ past_key_values = past_key_values,
+ inputs_embeds = inputs_embeds,
+ use_cache = use_cache,
+ cache_position = cache_position,
+ output_attentions = output_attentions,
+ output_hidden_states = output_hidden_states,
**clean_kwargs,
)
backbone_hidden_states = backbone_outputs[0]
slice_indices = (
- slice(-logits_to_keep, None) if isinstance(logits_to_keep, int)
+ slice(-logits_to_keep, None)
+ if isinstance(logits_to_keep, int)
else logits_to_keep
)
backbone_logits = self.lm_head(backbone_hidden_states[:, slice_indices, :])
@@ -1002,17 +1191,21 @@ class UnslothTrainer:
if labels is not None:
backbone_labels = labels[:, :, 0]
backbone_loss = self.loss_function(
- logits=backbone_logits, labels=backbone_labels,
- vocab_size=self.config.vocab_size, **clean_kwargs,
+ logits = backbone_logits,
+ labels = backbone_labels,
+ vocab_size = self.config.vocab_size,
+ **clean_kwargs,
)
- train_mask = ~(labels[:, :, 1:] == -100).all(dim=-1)
- depth_decoder_input_ids = labels[train_mask][..., :self.config.num_codebooks - 1]
+ train_mask = ~(labels[:, :, 1:] == -100).all(dim = -1)
+ depth_decoder_input_ids = labels[train_mask][
+ ..., : self.config.num_codebooks - 1
+ ]
depth_decoder_input_ids = nn.functional.pad(
- depth_decoder_input_ids, (1, 0), value=0
+ depth_decoder_input_ids, (1, 0), value = 0
)
- train_idxs = train_mask.nonzero(as_tuple=True)
+ train_idxs = train_mask.nonzero(as_tuple = True)
backbone_last_hidden_states = backbone_hidden_states[
train_idxs[0], train_idxs[1] - 1, :
]
@@ -1021,18 +1214,19 @@ class UnslothTrainer:
# Build clean kwargs for depth decoder
dd_kwargs = clean_kwargs.copy()
# Scale num_items_in_batch for depth decoder (31 codebooks)
- if 'num_items_in_batch' in dd_kwargs:
- dd_kwargs['num_items_in_batch'] = (
- dd_kwargs['num_items_in_batch'] * (self.config.num_codebooks - 1)
- )
+ if "num_items_in_batch" in dd_kwargs:
+ dd_kwargs["num_items_in_batch"] = dd_kwargs[
+ "num_items_in_batch"
+ ] * (self.config.num_codebooks - 1)
depth_decoder_outputs = self.depth_decoder(
- input_ids=depth_decoder_input_ids,
- backbone_last_hidden_state=backbone_last_hidden_states,
- use_cache=False, return_dict=True,
- labels=depth_decoder_labels,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
+ input_ids = depth_decoder_input_ids,
+ backbone_last_hidden_state = backbone_last_hidden_states,
+ use_cache = False,
+ return_dict = True,
+ labels = depth_decoder_labels,
+ output_attentions = output_attentions,
+ output_hidden_states = output_hidden_states,
**dd_kwargs,
)
@@ -1049,21 +1243,27 @@ class UnslothTrainer:
loss = backbone_loss + depth_decoder_loss
return CsmOutputWithPast(
- loss=loss, backbone_loss=backbone_loss,
- depth_decoder_loss=depth_decoder_loss, logits=backbone_logits,
- past_key_values=backbone_outputs.past_key_values,
- hidden_states=backbone_outputs.hidden_states,
- attentions=backbone_outputs.attentions,
- depth_decoder_logits=(
+ loss = loss,
+ backbone_loss = backbone_loss,
+ depth_decoder_loss = depth_decoder_loss,
+ logits = backbone_logits,
+ past_key_values = backbone_outputs.past_key_values,
+ hidden_states = backbone_outputs.hidden_states,
+ attentions = backbone_outputs.attentions,
+ depth_decoder_logits = (
depth_decoder_outputs.logits if depth_decoder_outputs else None
),
- depth_decoder_past_key_values=(
- depth_decoder_outputs.past_key_values if depth_decoder_outputs else None
+ depth_decoder_past_key_values = (
+ depth_decoder_outputs.past_key_values
+ if depth_decoder_outputs
+ else None
),
- depth_decoder_hidden_states=(
- depth_decoder_outputs.hidden_states if depth_decoder_outputs else None
+ depth_decoder_hidden_states = (
+ depth_decoder_outputs.hidden_states
+ if depth_decoder_outputs
+ else None
),
- depth_decoder_attentions=(
+ depth_decoder_attentions = (
depth_decoder_outputs.attentions if depth_decoder_outputs else None
),
)
@@ -1075,7 +1275,7 @@ class UnslothTrainer:
CsmForConditionalGeneration.forward = _fixed_csm_forward
logger.info("Applied CSM forward fix (class + instance level)\n")
- def _preprocess_csm_dataset(self, dataset, custom_format_mapping=None):
+ def _preprocess_csm_dataset(self, dataset, custom_format_mapping = None):
"""Preprocess dataset for CSM TTS training (exact notebook copy)."""
from transformers import AutoProcessor
from datasets import Audio
@@ -1083,13 +1283,13 @@ class UnslothTrainer:
processor = AutoProcessor.from_pretrained(
self.model_name,
- trust_remote_code=getattr(self, "trust_remote_code", False),
+ trust_remote_code = getattr(self, "trust_remote_code", False),
)
# Strip pad_to_multiple_of from tokenizer init_kwargs β fine-tuned models
# (e.g. keanteng/sesame-csm-elise) save it in tokenizer_config.json, and
# _merge_kwargs leaks it into audio_kwargs where EncodecFeatureExtractor rejects it.
- processor.tokenizer.init_kwargs.pop('pad_to_multiple_of', None)
+ processor.tokenizer.init_kwargs.pop("pad_to_multiple_of", None)
# Resolve columns from user mapping or hardcoded fallback
resolved = self._resolve_audio_columns(dataset, custom_format_mapping)
@@ -1098,21 +1298,35 @@ class UnslothTrainer:
speaker_key = resolved["speaker_col"]
if audio_col is None:
- raise ValueError(f"No audio column found in dataset. Columns: {dataset.column_names}")
+ raise ValueError(
+ f"No audio column found in dataset. Columns: {dataset.column_names}"
+ )
if text_col is None:
- raise ValueError(f"No text column found in dataset. Columns: {dataset.column_names}")
+ raise ValueError(
+ f"No text column found in dataset. Columns: {dataset.column_names}"
+ )
if speaker_key is None:
- logger.info("No speaker found, adding default 'source' of 0 for all examples\n")
+ logger.info(
+ "No speaker found, adding default 'source' of 0 for all examples\n"
+ )
dataset = dataset.add_column("source", ["0"] * len(dataset))
speaker_key = "source"
- logger.info(f"CSM preprocessing: audio_col='{audio_col}', text_col='{text_col}', speaker_key='{speaker_key}'\n")
+ logger.info(
+ f"CSM preprocessing: audio_col='{audio_col}', text_col='{text_col}', speaker_key='{speaker_key}'\n"
+ )
- dataset = dataset.cast_column(audio_col, Audio(sampling_rate=24000))
+ dataset = dataset.cast_column(audio_col, Audio(sampling_rate = 24000))
- required_keys = ["input_ids", "attention_mask", "labels", "input_values", "input_values_cutoffs"]
+ required_keys = [
+ "input_ids",
+ "attention_mask",
+ "labels",
+ "input_values",
+ "input_values_cutoffs",
+ ]
- self._update_progress(status_message="Preprocessing CSM dataset...")
+ self._update_progress(status_message = "Preprocessing CSM dataset...")
processed_examples = []
skipped = 0
for idx in range(len(dataset)):
@@ -1122,31 +1336,33 @@ class UnslothTrainer:
example = dataset[idx]
try:
- conversation = [{
- "role": str(example[speaker_key]),
- "content": [
- {"type": "text", "text": example.get(text_col, "")},
- {"type": "audio", "path": example[audio_col]["array"]},
- ],
- }]
+ conversation = [
+ {
+ "role": str(example[speaker_key]),
+ "content": [
+ {"type": "text", "text": example.get(text_col, "")},
+ {"type": "audio", "path": example[audio_col]["array"]},
+ ],
+ }
+ ]
# NOTE: pad_to_multiple_of intentionally omitted from text_kwargs β
# CsmProcessor._merge_kwargs leaks it to EncodecFeatureExtractor which rejects it.
model_inputs = processor.apply_chat_template(
conversation,
- tokenize=True,
- return_dict=True,
- output_labels=True,
- text_kwargs={
+ tokenize = True,
+ return_dict = True,
+ output_labels = True,
+ text_kwargs = {
"padding": "max_length",
"max_length": 256,
"padding_side": "right",
},
- audio_kwargs={
+ audio_kwargs = {
"sampling_rate": 24_000,
"max_length": 240001,
"padding": "max_length",
},
- common_kwargs={"return_tensors": "pt"},
+ common_kwargs = {"return_tensors": "pt"},
)
out = {}
@@ -1168,7 +1384,7 @@ class UnslothTrainer:
if (idx + 1) % 100 == 0:
self._update_progress(
- status_message=f"Preprocessing CSM... {idx + 1}/{len(dataset)}"
+ status_message = f"Preprocessing CSM... {idx + 1}/{len(dataset)}"
)
if not processed_examples:
@@ -1177,11 +1393,13 @@ class UnslothTrainer:
)
result_dataset = Dataset.from_list(processed_examples)
- logger.info(f"CSM preprocessing complete: {len(result_dataset)} examples "
- f"({skipped} skipped)\n")
+ logger.info(
+ f"CSM preprocessing complete: {len(result_dataset)} examples "
+ f"({skipped} skipped)\n"
+ )
return result_dataset
- def _format_audio_vlm_dataset(self, dataset, custom_format_mapping=None):
+ def _format_audio_vlm_dataset(self, dataset, custom_format_mapping = None):
"""Format dataset as audio chat messages for multimodal models (e.g. Gemma 3N).
Expects columns: audio (Audio), text (str).
@@ -1201,7 +1419,7 @@ class UnslothTrainer:
self._audio_vlm_audio_col = audio_col
# Cast audio to 16kHz (standard for speech models)
- dataset = dataset.cast_column(audio_col, Audio(sampling_rate=16000))
+ dataset = dataset.cast_column(audio_col, Audio(sampling_rate = 16000))
def format_messages(samples):
formatted = {"messages": []}
@@ -1209,26 +1427,35 @@ class UnslothTrainer:
audio = samples[audio_col][idx]["array"]
label = str(samples[text_col][idx])
message = [
- {"role": "system", "content": [
- {"type": "text", "text": "You are an assistant that transcribes speech accurately."}
- ]},
- {"role": "user", "content": [
- {"type": "audio", "audio": audio},
- {"type": "text", "text": "Please transcribe this audio."}
- ]},
- {"role": "assistant", "content": [
- {"type": "text", "text": label}
- ]},
+ {
+ "role": "system",
+ "content": [
+ {
+ "type": "text",
+ "text": "You are an assistant that transcribes speech accurately.",
+ }
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ {"type": "audio", "audio": audio},
+ {"type": "text", "text": "Please transcribe this audio."},
+ ],
+ },
+ {"role": "assistant", "content": [{"type": "text", "text": label}]},
]
formatted["messages"].append(message)
return formatted
- self._update_progress(status_message="Formatting audio VLM dataset...")
- dataset = dataset.map(format_messages, batched=True, batch_size=4, num_proc=safe_num_proc(4))
+ self._update_progress(status_message = "Formatting audio VLM dataset...")
+ dataset = dataset.map(
+ format_messages, batched = True, batch_size = 4, num_proc = safe_num_proc(4)
+ )
logger.info(f"Audio VLM dataset formatted: {len(dataset)} examples\n")
return dataset
- def _preprocess_snac_dataset(self, dataset, custom_format_mapping=None):
+ def _preprocess_snac_dataset(self, dataset, custom_format_mapping = None):
"""Preprocess dataset for Orpheus TTS training with SNAC codec.
Mirrors Orpheus_(3B)-TTS.ipynb: encode audio with SNAC (24kHz, 3 hierarchical
@@ -1246,13 +1473,13 @@ class UnslothTrainer:
# Orpheus special token IDs (hardcoded in tokenizer vocabulary)
START_OF_HUMAN = 128259
- END_OF_HUMAN = 128260
- START_OF_AI = 128261
- END_OF_AI = 128262
+ END_OF_HUMAN = 128260
+ START_OF_AI = 128261
+ END_OF_AI = 128262
START_OF_SPEECH = 128257
- END_OF_SPEECH = 128258
- END_OF_TEXT = 128009
- AUDIO_OFFSET = 128266
+ END_OF_SPEECH = 128258
+ END_OF_TEXT = 128009
+ AUDIO_OFFSET = 128266
resolved = self._resolve_audio_columns(dataset, custom_format_mapping)
audio_col = resolved["audio_col"]
@@ -1266,25 +1493,37 @@ class UnslothTrainer:
# Cast audio column so datasets 4.x AudioDecoder objects are decoded to dicts
from datasets import Audio
- dataset = dataset.cast_column(audio_col, Audio(sampling_rate=SNAC_SAMPLE_RATE))
+
+ dataset = dataset.cast_column(audio_col, Audio(sampling_rate = SNAC_SAMPLE_RATE))
# Get dataset sample rate from first example (after cast, always SNAC_SAMPLE_RATE)
first_audio = dataset[0][audio_col]
- ds_sample_rate = first_audio.get("sampling_rate", SNAC_SAMPLE_RATE) if isinstance(first_audio, dict) else SNAC_SAMPLE_RATE
+ ds_sample_rate = (
+ first_audio.get("sampling_rate", SNAC_SAMPLE_RATE)
+ if isinstance(first_audio, dict)
+ else SNAC_SAMPLE_RATE
+ )
# Load SNAC codec model
- self._update_progress(status_message="Loading SNAC codec model...")
+ self._update_progress(status_message = "Loading SNAC codec model...")
logger.info("Loading SNAC codec model...\n")
from snac import SNAC
+
snac_model = SNAC.from_pretrained(SNAC_MODEL_NAME)
snac_model = snac_model.to(device).eval()
# Resample transform (created once)
- resample_transform = T.Resample(orig_freq=ds_sample_rate, new_freq=SNAC_SAMPLE_RATE) if ds_sample_rate != SNAC_SAMPLE_RATE else None
+ resample_transform = (
+ T.Resample(orig_freq = ds_sample_rate, new_freq = SNAC_SAMPLE_RATE)
+ if ds_sample_rate != SNAC_SAMPLE_RATE
+ else None
+ )
- self._update_progress(status_message="Encoding audio with SNAC...")
- logger.info(f"SNAC preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
- f"has_source={has_source}, ds_sample_rate={ds_sample_rate}\n")
+ self._update_progress(status_message = "Encoding audio with SNAC...")
+ logger.info(
+ f"SNAC preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
+ f"has_source={has_source}, ds_sample_rate={ds_sample_rate}\n"
+ )
processed_examples = []
skipped = 0
@@ -1306,7 +1545,11 @@ class UnslothTrainer:
continue
# --- Encode audio with SNAC (notebook lines 122-142) ---
- waveform = torch.from_numpy(audio_data["array"]).unsqueeze(0).to(dtype=torch.float32)
+ waveform = (
+ torch.from_numpy(audio_data["array"])
+ .unsqueeze(0)
+ .to(dtype = torch.float32)
+ )
if resample_transform is not None:
waveform = resample_transform(waveform)
@@ -1318,12 +1561,22 @@ class UnslothTrainer:
all_codes = []
for i in range(codes[0].shape[1]):
all_codes.append(codes[0][0][i].item() + AUDIO_OFFSET)
- all_codes.append(codes[1][0][2*i].item() + AUDIO_OFFSET + 4096)
- all_codes.append(codes[2][0][4*i].item() + AUDIO_OFFSET + (2*4096))
- all_codes.append(codes[2][0][(4*i)+1].item() + AUDIO_OFFSET + (3*4096))
- all_codes.append(codes[1][0][(2*i)+1].item() + AUDIO_OFFSET + (4*4096))
- all_codes.append(codes[2][0][(4*i)+2].item() + AUDIO_OFFSET + (5*4096))
- all_codes.append(codes[2][0][(4*i)+3].item() + AUDIO_OFFSET + (6*4096))
+ all_codes.append(codes[1][0][2 * i].item() + AUDIO_OFFSET + 4096)
+ all_codes.append(
+ codes[2][0][4 * i].item() + AUDIO_OFFSET + (2 * 4096)
+ )
+ all_codes.append(
+ codes[2][0][(4 * i) + 1].item() + AUDIO_OFFSET + (3 * 4096)
+ )
+ all_codes.append(
+ codes[1][0][(2 * i) + 1].item() + AUDIO_OFFSET + (4 * 4096)
+ )
+ all_codes.append(
+ codes[2][0][(4 * i) + 2].item() + AUDIO_OFFSET + (5 * 4096)
+ )
+ all_codes.append(
+ codes[2][0][(4 * i) + 3].item() + AUDIO_OFFSET + (6 * 4096)
+ )
if len(all_codes) == 0:
skipped += 1
@@ -1333,12 +1586,16 @@ class UnslothTrainer:
deduped = all_codes[:7]
for i in range(7, len(all_codes), 7):
if all_codes[i] != deduped[-7]:
- deduped.extend(all_codes[i:i+7])
+ deduped.extend(all_codes[i : i + 7])
all_codes = deduped
# --- Build text tokens (notebook lines 217-224) ---
- text_prompt = f"{example[speaker_col]}: {text}" if has_source and example.get(speaker_col) else text
- text_ids = tokenizer.encode(text_prompt, add_special_tokens=True)
+ text_prompt = (
+ f"{example[speaker_col]}: {text}"
+ if has_source and example.get(speaker_col)
+ else text
+ )
+ text_ids = tokenizer.encode(text_prompt, add_special_tokens = True)
text_ids.append(END_OF_TEXT)
# --- Build full input_ids (notebook lines 225-234) ---
@@ -1360,11 +1617,13 @@ class UnslothTrainer:
labels = list(input_ids)
attention_mask = [1] * len(input_ids)
- processed_examples.append({
- "input_ids": input_ids,
- "labels": labels,
- "attention_mask": attention_mask,
- })
+ processed_examples.append(
+ {
+ "input_ids": input_ids,
+ "labels": labels,
+ "attention_mask": attention_mask,
+ }
+ )
except Exception as e:
logger.warning(f"Error processing SNAC example {idx}: {e}")
@@ -1374,7 +1633,7 @@ class UnslothTrainer:
# Progress update every 100 examples
if (idx + 1) % 100 == 0:
self._update_progress(
- status_message=f"Encoding audio... {idx + 1}/{len(dataset)}"
+ status_message = f"Encoding audio... {idx + 1}/{len(dataset)}"
)
# Free SNAC model from GPU
@@ -1382,6 +1641,7 @@ class UnslothTrainer:
snac_model.to("cpu")
del snac_model
import gc
+
gc.collect()
torch.cuda.empty_cache()
self._cuda_audio_used = True
@@ -1392,11 +1652,13 @@ class UnslothTrainer:
)
result_dataset = Dataset.from_list(processed_examples)
- logger.info(f"SNAC preprocessing complete: {len(result_dataset)} examples "
- f"({skipped} skipped)\n")
+ logger.info(
+ f"SNAC preprocessing complete: {len(result_dataset)} examples "
+ f"({skipped} skipped)\n"
+ )
return result_dataset
- def _preprocess_bicodec_dataset(self, dataset, custom_format_mapping=None):
+ def _preprocess_bicodec_dataset(self, dataset, custom_format_mapping = None):
"""Preprocess dataset for Spark-TTS training with BiCodec tokenizer.
Mirrors Spark_TTS_(0_5B).ipynb: encode audio with BiCodec (semantic + global tokens),
@@ -1413,14 +1675,23 @@ class UnslothTrainer:
# The sparktts Python package lives in the SparkAudio/Spark-TTS GitHub repo,
# NOT in the unsloth/Spark-TTS-0.5B HF model repo. Clone it if needed.
- spark_code_dir = os.path.join(os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS")
+ spark_code_dir = os.path.join(
+ os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS"
+ )
sparktts_pkg = os.path.join(spark_code_dir, "sparktts")
if not os.path.isdir(sparktts_pkg):
- self._update_progress(status_message="Cloning Spark-TTS code repo...")
+ self._update_progress(status_message = "Cloning Spark-TTS code repo...")
logger.info(f"Cloning SparkAudio/Spark-TTS to {spark_code_dir}...\n")
subprocess.run(
- ["git", "clone", "--depth", "1", "https://github.com/SparkAudio/Spark-TTS", spark_code_dir],
- check=True,
+ [
+ "git",
+ "clone",
+ "--depth",
+ "1",
+ "https://github.com/SparkAudio/Spark-TTS",
+ spark_code_dir,
+ ],
+ check = True,
)
if spark_code_dir not in sys.path:
@@ -1443,18 +1714,21 @@ class UnslothTrainer:
# Cast audio column so datasets 4.x AudioDecoder objects are decoded to dicts.
# Don't resample here β BiCodec's target_sr may differ; the loop handles resampling.
from datasets import Audio
+
dataset = dataset.cast_column(audio_col, Audio())
# Load BiCodec tokenizer
- self._update_progress(status_message="Loading BiCodec tokenizer...")
+ self._update_progress(status_message = "Loading BiCodec tokenizer...")
logger.info("Loading BiCodec tokenizer...\n")
audio_tokenizer = BiCodecTokenizer(self._spark_tts_repo_dir, device)
- target_sr = audio_tokenizer.config['sample_rate']
+ target_sr = audio_tokenizer.config["sample_rate"]
- self._update_progress(status_message="Encoding audio with BiCodec...")
- logger.info(f"BiCodec preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
- f"has_source={has_source}, target_sr={target_sr}\n")
+ self._update_progress(status_message = "Encoding audio with BiCodec...")
+ logger.info(
+ f"BiCodec preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
+ f"has_source={has_source}, target_sr={target_sr}\n"
+ )
def extract_wav2vec2_features(wavs: torch.Tensor) -> torch.Tensor:
"""Extract wav2vec2 features (average of layers 11, 14, 16)."""
@@ -1464,11 +1738,13 @@ class UnslothTrainer:
processed = audio_tokenizer.processor(
wav_np,
- sampling_rate=16000,
- return_tensors="pt",
- padding=True,
+ sampling_rate = 16000,
+ return_tensors = "pt",
+ padding = True,
+ )
+ input_values = processed.input_values.to(
+ audio_tokenizer.feature_extractor.device
)
- input_values = processed.input_values.to(audio_tokenizer.feature_extractor.device)
model_output = audio_tokenizer.feature_extractor(input_values)
if model_output.hidden_states is None:
@@ -1505,7 +1781,7 @@ class UnslothTrainer:
# Resample if needed
if sampling_rate != target_sr:
- resampler = T.Resample(orig_freq=sampling_rate, new_freq=target_sr)
+ resampler = T.Resample(orig_freq = sampling_rate, new_freq = target_sr)
audio_tensor_temp = torch.from_numpy(audio_array).float()
audio_array = resampler(audio_tensor_temp).numpy()
@@ -1517,8 +1793,12 @@ class UnslothTrainer:
ref_wav_np = audio_tokenizer.get_ref_clip(audio_array)
# Prepare tensors
- audio_tensor = torch.from_numpy(audio_array).unsqueeze(0).float().to(device)
- ref_wav_tensor = torch.from_numpy(ref_wav_np).unsqueeze(0).float().to(device)
+ audio_tensor = (
+ torch.from_numpy(audio_array).unsqueeze(0).float().to(device)
+ )
+ ref_wav_tensor = (
+ torch.from_numpy(ref_wav_np).unsqueeze(0).float().to(device)
+ )
# Extract wav2vec2 features
feat = extract_wav2vec2_features(audio_tensor)
@@ -1530,31 +1810,45 @@ class UnslothTrainer:
}
# BiCodec tokenize
- semantic_token_ids, global_token_ids = audio_tokenizer.model.tokenize(batch)
+ semantic_token_ids, global_token_ids = audio_tokenizer.model.tokenize(
+ batch
+ )
global_tokens = "".join(
- [f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze().cpu().numpy()]
+ [
+ f"<|bicodec_global_{i}|>"
+ for i in global_token_ids.squeeze().cpu().numpy()
+ ]
)
semantic_tokens = "".join(
- [f"<|bicodec_semantic_{i}|>" for i in semantic_token_ids.squeeze().cpu().numpy()]
+ [
+ f"<|bicodec_semantic_{i}|>"
+ for i in semantic_token_ids.squeeze().cpu().numpy()
+ ]
)
# Format text with source prefix if available
- text_content = f"{example[speaker_col]}: {text}" if has_source and example.get(speaker_col) else text
+ text_content = (
+ f"{example[speaker_col]}: {text}"
+ if has_source and example.get(speaker_col)
+ else text
+ )
- formatted = "".join([
- "<|task_tts|>",
- "<|start_content|>",
- text_content,
- "<|end_content|>",
- "<|start_global_token|>",
- global_tokens,
- "<|end_global_token|>",
- "<|start_semantic_token|>",
- semantic_tokens,
- "<|end_semantic_token|>",
- "<|im_end|>",
- ])
+ formatted = "".join(
+ [
+ "<|task_tts|>",
+ "<|start_content|>",
+ text_content,
+ "<|end_content|>",
+ "<|start_global_token|>",
+ global_tokens,
+ "<|end_global_token|>",
+ "<|start_semantic_token|>",
+ semantic_tokens,
+ "<|end_semantic_token|>",
+ "<|im_end|>",
+ ]
+ )
processed_examples.append({"text": formatted})
@@ -1566,7 +1860,7 @@ class UnslothTrainer:
# Progress update every 100 examples
if (idx + 1) % 100 == 0:
self._update_progress(
- status_message=f"Encoding audio with BiCodec... {idx + 1}/{len(dataset)}"
+ status_message = f"Encoding audio with BiCodec... {idx + 1}/{len(dataset)}"
)
# Free BiCodec model from GPU
@@ -1575,6 +1869,7 @@ class UnslothTrainer:
audio_tokenizer.feature_extractor.cpu()
del audio_tokenizer
import gc
+
gc.collect()
torch.cuda.empty_cache()
self._cuda_audio_used = True
@@ -1585,15 +1880,17 @@ class UnslothTrainer:
)
result_dataset = Dataset.from_list(processed_examples)
- logger.info(f"BiCodec preprocessing complete: {len(result_dataset)} examples "
- f"({skipped} skipped)\n")
+ logger.info(
+ f"BiCodec preprocessing complete: {len(result_dataset)} examples "
+ f"({skipped} skipped)\n"
+ )
# Debug: show first example text (truncated)
sample = result_dataset[0]["text"]
logger.info(f"Sample text (first 200 chars): {sample[:200]}...\n")
logger.info(f"Sample text length: {len(sample)} chars\n")
return result_dataset
- def _preprocess_dac_dataset(self, dataset, custom_format_mapping=None):
+ def _preprocess_dac_dataset(self, dataset, custom_format_mapping = None):
"""Preprocess dataset for OuteTTS training with DAC codec.
Mirrors Oute_TTS_(1B).ipynb DataCreationV3: uses Whisper for word timings,
@@ -1613,15 +1910,23 @@ class UnslothTrainer:
# Clone OuteTTS repo (same as audio_codecs._load_dac)
import subprocess
+
base_dir = os.path.dirname(os.path.abspath(__file__))
outetts_code_dir = os.path.join(base_dir, "inference", "OuteTTS")
outetts_pkg = os.path.join(outetts_code_dir, "outetts")
if not os.path.isdir(outetts_pkg):
- self._update_progress(status_message="Cloning OuteTTS code repo...")
+ self._update_progress(status_message = "Cloning OuteTTS code repo...")
logger.info(f"Cloning edwko/OuteTTS to {outetts_code_dir}...\n")
subprocess.run(
- ["git", "clone", "--depth", "1", "https://github.com/edwko/OuteTTS", outetts_code_dir],
- check=True,
+ [
+ "git",
+ "clone",
+ "--depth",
+ "1",
+ "https://github.com/edwko/OuteTTS",
+ outetts_code_dir,
+ ],
+ check = True,
)
for fpath in [
os.path.join(outetts_pkg, "models", "gguf_model.py"),
@@ -1651,29 +1956,35 @@ class UnslothTrainer:
# Cast audio to 24kHz (notebook: dataset.cast_column("audio", Audio(sampling_rate=24000)))
from datasets import Audio
- dataset = dataset.cast_column(audio_col, Audio(sampling_rate=24000))
+
+ dataset = dataset.cast_column(audio_col, Audio(sampling_rate = 24000))
logger.info("Cast audio column to 24kHz\n")
# Load Whisper for word timings
- self._update_progress(status_message="Loading Whisper model for word timings...")
+ self._update_progress(
+ status_message = "Loading Whisper model for word timings..."
+ )
logger.info("Loading Whisper model for word timings...\n")
import whisper
- whisper_model = whisper.load_model("turbo", device=device)
+
+ whisper_model = whisper.load_model("turbo", device = device)
# Load OuteTTS AudioProcessor + PromptProcessor
- self._update_progress(status_message="Loading OuteTTS AudioProcessor...")
+ self._update_progress(status_message = "Loading OuteTTS AudioProcessor...")
logger.info("Loading OuteTTS AudioProcessor...\n")
model_tokenizer_path = "OuteAI/Llama-OuteTTS-1.0-1B"
dummy_config = OuteTTSModelConfig(
- tokenizer_path=model_tokenizer_path,
- device=device,
- audio_codec_path=None,
+ tokenizer_path = model_tokenizer_path,
+ device = device,
+ audio_codec_path = None,
)
- audio_processor = AudioProcessor(config=dummy_config)
+ audio_processor = AudioProcessor(config = dummy_config)
prompt_processor = PromptProcessor(model_tokenizer_path)
- self._update_progress(status_message="Preprocessing audio with OuteTTS...")
- logger.info(f"DAC preprocessing: audio_col='{audio_col}', text_col='{text_col}'\n")
+ self._update_progress(status_message = "Preprocessing audio with OuteTTS...")
+ logger.info(
+ f"DAC preprocessing: audio_col='{audio_col}', text_col='{text_col}'\n"
+ )
processed_examples = []
skipped = 0
@@ -1694,28 +2005,30 @@ class UnslothTrainer:
skipped += 1
continue
- audio_array = np.array(audio_data["array"], dtype=np.float32)
+ audio_array = np.array(audio_data["array"], dtype = np.float32)
sampling_rate = audio_data.get("sampling_rate", 24000)
# Convert to WAV bytes (Whisper needs a file path)
buf = io.BytesIO()
- sf.write(buf, audio_array, sampling_rate, format="WAV", subtype="FLOAT")
+ sf.write(buf, audio_array, sampling_rate, format = "WAV", subtype = "FLOAT")
buf.seek(0)
audio_bytes = buf.getvalue()
# 1. Get word timings from Whisper
with tempfile.NamedTemporaryFile(
- suffix=".wav",
- delete=False,
- dir=str(ensure_dir(tmp_root())),
+ suffix = ".wav",
+ delete = False,
+ dir = str(ensure_dir(tmp_root())),
) as tmp:
tmp.write(audio_bytes)
tmp.flush()
tmp_path = tmp.name
try:
- whisper_result = whisper_model.transcribe(tmp_path, word_timestamps=True)
+ whisper_result = whisper_model.transcribe(
+ tmp_path, word_timestamps = True
+ )
finally:
- Path(tmp_path).unlink(missing_ok=True)
+ Path(tmp_path).unlink(missing_ok = True)
normalized_transcript = text_normalizations(text)
words_with_timings = []
@@ -1724,11 +2037,13 @@ class UnslothTrainer:
for word_info in segment.get("words", []):
cleaned = word_info["word"].strip()
if cleaned:
- words_with_timings.append({
- "word": cleaned,
- "start": float(word_info["start"]),
- "end": float(word_info["end"]),
- })
+ words_with_timings.append(
+ {
+ "word": cleaned,
+ "start": float(word_info["start"]),
+ "end": float(word_info["end"]),
+ }
+ )
if not words_with_timings:
skipped += 1
@@ -1757,16 +2072,17 @@ class UnslothTrainer:
if (idx + 1) % 100 == 0:
self._update_progress(
- status_message=f"Preprocessing audio with OuteTTS... {idx + 1}/{len(dataset)}"
+ status_message = f"Preprocessing audio with OuteTTS... {idx + 1}/{len(dataset)}"
)
# Free Whisper from GPU (notebook: data_processor.whisper_model.to('cpu'))
logger.info("Moving Whisper model to CPU...\n")
- whisper_model.to('cpu')
+ whisper_model.to("cpu")
del whisper_model
del audio_processor
del prompt_processor
import gc
+
gc.collect()
torch.cuda.empty_cache()
self._cuda_audio_used = True
@@ -1777,13 +2093,17 @@ class UnslothTrainer:
)
result_dataset = HFDataset.from_list(processed_examples)
- logger.info(f"DAC preprocessing complete: {len(result_dataset)} examples "
- f"({skipped} skipped)\n")
+ logger.info(
+ f"DAC preprocessing complete: {len(result_dataset)} examples "
+ f"({skipped} skipped)\n"
+ )
sample = result_dataset[0]["text"]
logger.info(f"Sample text (first 200 chars): {sample[:200]}...\n")
return result_dataset
- def _preprocess_whisper_dataset(self, dataset, eval_split=None, custom_format_mapping=None):
+ def _preprocess_whisper_dataset(
+ self, dataset, eval_split = None, custom_format_mapping = None
+ ):
"""Preprocess dataset for Whisper speech-to-text training.
Mirrors Whisper.ipynb: extract audio features with Whisper's feature
@@ -1803,20 +2123,24 @@ class UnslothTrainer:
)
# Cast audio to 16kHz (Whisper's expected sample rate)
- dataset = dataset.cast_column(audio_col, Audio(sampling_rate=WHISPER_SAMPLE_RATE))
+ dataset = dataset.cast_column(
+ audio_col, Audio(sampling_rate = WHISPER_SAMPLE_RATE)
+ )
# Train/eval split (notebook does dataset.train_test_split)
eval_dataset_raw = None
if eval_split:
- splits = dataset.train_test_split(test_size=0.06, seed=42)
+ splits = dataset.train_test_split(test_size = 0.06, seed = 42)
dataset = splits["train"]
eval_dataset_raw = splits["test"]
- self._update_progress(status_message="Processing audio for Whisper...")
- logger.info(f"Whisper preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
- f"samples={len(dataset)}\n")
+ self._update_progress(status_message = "Processing audio for Whisper...")
+ logger.info(
+ f"Whisper preprocessing: audio_col='{audio_col}', text_col='{text_col}', "
+ f"samples={len(dataset)}\n"
+ )
- def process_split(ds, split_name="train"):
+ def process_split(ds, split_name = "train"):
processed = []
skipped = 0
for idx in range(len(ds)):
@@ -1828,53 +2152,67 @@ class UnslothTrainer:
try:
audio_data = example.get(audio_col)
text = example.get(text_col)
- if audio_data is None or audio_data.get("array") is None or not text:
+ if (
+ audio_data is None
+ or audio_data.get("array") is None
+ or not text
+ ):
skipped += 1
continue
# Extract audio features (notebook line 112-115)
features = self.tokenizer.feature_extractor(
- audio_data["array"], sampling_rate=audio_data["sampling_rate"]
+ audio_data["array"], sampling_rate = audio_data["sampling_rate"]
)
# Tokenize text (notebook line 116)
tokenized_text = self.tokenizer.tokenizer(text)
- processed.append({
- "input_features": features.input_features[0],
- "labels": tokenized_text.input_ids,
- })
+ processed.append(
+ {
+ "input_features": features.input_features[0],
+ "labels": tokenized_text.input_ids,
+ }
+ )
except Exception as e:
- logger.warning(f"Error processing Whisper {split_name} example {idx}: {e}")
+ logger.warning(
+ f"Error processing Whisper {split_name} example {idx}: {e}"
+ )
skipped += 1
continue
if (idx + 1) % 100 == 0:
self._update_progress(
- status_message=f"Processing {split_name} audio... {idx + 1}/{len(ds)}"
+ status_message = f"Processing {split_name} audio... {idx + 1}/{len(ds)}"
)
- logger.info(f"Whisper {split_name} preprocessing: {len(processed)} examples ({skipped} skipped)\n")
+ logger.info(
+ f"Whisper {split_name} preprocessing: {len(processed)} examples ({skipped} skipped)\n"
+ )
return processed
train_data = process_split(dataset, "train")
- eval_data = process_split(eval_dataset_raw, "eval") if eval_dataset_raw else None
+ eval_data = (
+ process_split(eval_dataset_raw, "eval") if eval_dataset_raw else None
+ )
if not train_data:
raise ValueError("No valid examples after Whisper preprocessing")
return (train_data, eval_data)
- def load_and_format_dataset(self,
- dataset_source: str,
- format_type: str = "auto",
- local_datasets: list = None,
- custom_format_mapping: dict = None,
- subset: str = None,
- train_split: str = "train",
- eval_split: str = None,
- eval_steps: float = 0.00,
- dataset_slice_start: int = None,
- dataset_slice_end: int = None) -> Optional[tuple]:
+ def load_and_format_dataset(
+ self,
+ dataset_source: str,
+ format_type: str = "auto",
+ local_datasets: list = None,
+ custom_format_mapping: dict = None,
+ subset: str = None,
+ train_split: str = "train",
+ eval_split: str = None,
+ eval_steps: float = 0.00,
+ dataset_slice_start: int = None,
+ dataset_slice_end: int = None,
+ ) -> Optional[tuple]:
"""
Load and prepare dataset for training.
@@ -1888,7 +2226,9 @@ class UnslothTrainer:
try:
dataset = None
eval_dataset = None
- has_separate_eval_source = False # True if eval comes from a separate HF split
+ has_separate_eval_source = (
+ False # True if eval comes from a separate HF split
+ )
eval_enabled = eval_steps is not None and eval_steps > 0
if local_datasets:
@@ -1919,35 +2259,41 @@ class UnslothTrainer:
continue
# Fall through to single-file detection for dirs with json/csv
candidates: list[Path] = []
- for ext in ('.json', '.jsonl', '.csv', '.parquet'):
+ for ext in (".json", ".jsonl", ".csv", ".parquet"):
candidates.extend(sorted(file_path_obj.glob(f"*{ext}")))
if candidates:
all_files.extend(str(c) for c in candidates)
continue
- raise ValueError(f"No supported data files in directory: {file_path_obj}")
+ raise ValueError(
+ f"No supported data files in directory: {file_path_obj}"
+ )
else:
all_files.append(str(file_path_obj))
if all_files:
# Determine loader type from the first file extension
first_ext = Path(all_files[0]).suffix.lower()
- if first_ext in ('.json', '.jsonl'):
- loader = 'json'
- elif first_ext == '.csv':
- loader = 'csv'
- elif first_ext == '.parquet':
- loader = 'parquet'
+ if first_ext in (".json", ".jsonl"):
+ loader = "json"
+ elif first_ext == ".csv":
+ loader = "csv"
+ elif first_ext == ".parquet":
+ loader = "parquet"
else:
- raise ValueError(f"Unsupported local dataset format: {all_files[0]}")
+ raise ValueError(
+ f"Unsupported local dataset format: {all_files[0]}"
+ )
- dataset = load_dataset(loader, data_files=all_files, split='train')
+ dataset = load_dataset(loader, data_files = all_files, split = "train")
# Check if stopped during dataset loading
if self.should_stop:
logger.info("Stopped during dataset loading\n")
return None
- self._update_progress(status_message=f"Loaded {len(dataset)} samples from local files")
+ self._update_progress(
+ status_message = f"Loaded {len(dataset)} samples from local files"
+ )
logger.info(f"Loaded {len(dataset)} samples from local files\n")
logger.info(f"[DEBUG] Dataset cache_files: {dataset.cache_files}\n")
@@ -1959,9 +2305,11 @@ class UnslothTrainer:
load_kwargs["name"] = subset
_slice_start = dataset_slice_start or 0
- if (dataset_slice_end is not None
- and dataset_slice_end >= 0
- and dataset_slice_end >= _slice_start):
+ if (
+ dataset_slice_end is not None
+ and dataset_slice_end >= 0
+ and dataset_slice_end >= _slice_start
+ ):
# Manual slice β stream only the rows we need instead of
# downloading the entire dataset.
rows_to_stream = dataset_slice_end + 1
@@ -1970,14 +2318,14 @@ class UnslothTrainer:
f"(start={dataset_slice_start}, end={dataset_slice_end}), "
f"streaming {rows_to_stream} rows\n"
)
- stream = load_dataset(**load_kwargs, streaming=True)
+ stream = load_dataset(**load_kwargs, streaming = True)
dataset = Dataset.from_list(list(stream.take(rows_to_stream)))
logger.info(
f"[dataset-slice] Downloaded {len(dataset)} rows "
f"(requested {rows_to_stream})\n"
)
self._update_progress(
- status_message=f"Streamed {len(dataset)} rows from HuggingFace"
+ status_message = f"Streamed {len(dataset)} rows from HuggingFace"
)
else:
dataset = load_dataset(**load_kwargs)
@@ -1987,8 +2335,12 @@ class UnslothTrainer:
logger.info("Stopped during dataset loading\n")
return None
- self._update_progress(status_message=f"Loaded dataset from HuggingFace: {dataset_source}")
- logger.info(f"Loaded dataset from Hugging Face: {dataset_source} ({len(dataset)} rows)\n")
+ self._update_progress(
+ status_message = f"Loaded dataset from HuggingFace: {dataset_source}"
+ )
+ logger.info(
+ f"Loaded dataset from Hugging Face: {dataset_source} ({len(dataset)} rows)\n"
+ )
# Resolve eval split from a separate HF split (explicit or auto-detected)
if eval_enabled:
@@ -2001,20 +2353,26 @@ class UnslothTrainer:
eval_load_kwargs["name"] = subset
eval_dataset = load_dataset(**eval_load_kwargs)
has_separate_eval_source = True
- logger.info(f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n")
+ logger.info(
+ f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n"
+ )
elif eval_split and eval_split == effective_train:
# Same split as training β will do 80/20 split after formatting
- logger.info(f"Eval split '{eval_split}' is the same as train split β will split 80/20\n")
+ logger.info(
+ f"Eval split '{eval_split}' is the same as train split β will split 80/20\n"
+ )
else:
# Auto-detect eval split from HF (returns a separate dataset, or None)
eval_dataset = self._auto_detect_eval_split_from_hf(
- dataset_source=dataset_source,
- subset=subset,
+ dataset_source = dataset_source,
+ subset = subset,
)
if eval_dataset is not None:
has_separate_eval_source = True
else:
- logger.info("Eval disabled (eval_steps <= 0), skipping eval split detection\n")
+ logger.info(
+ "Eval disabled (eval_steps <= 0), skipping eval split detection\n"
+ )
if dataset is None:
raise ValueError("No dataset provided")
@@ -2023,13 +2381,21 @@ class UnslothTrainer:
if dataset_slice_start is not None or dataset_slice_end is not None:
total_rows = len(dataset)
start = dataset_slice_start if dataset_slice_start is not None else 0
- end = dataset_slice_end if dataset_slice_end is not None else total_rows - 1
+ end = (
+ dataset_slice_end
+ if dataset_slice_end is not None
+ else total_rows - 1
+ )
# Clamp to valid range
start = max(0, min(start, total_rows - 1))
end = max(start, min(end, total_rows - 1))
dataset = dataset.select(range(start, end + 1))
- logger.info(f"Sliced dataset to rows [{start}, {end}]: {len(dataset)} of {total_rows} rows\n")
- self._update_progress(status_message=f"Sliced dataset to {len(dataset)} rows (indices {start}-{end})")
+ logger.info(
+ f"Sliced dataset to rows [{start}, {end}]: {len(dataset)} of {total_rows} rows\n"
+ )
+ self._update_progress(
+ status_message = f"Sliced dataset to {len(dataset)} rows (indices {start}-{end})"
+ )
# Check if stopped before applying template
if self.should_stop:
@@ -2037,30 +2403,38 @@ class UnslothTrainer:
return None
# ========== AUDIO MODELS: custom preprocessing ==========
- if self._audio_type == 'csm':
+ if self._audio_type == "csm":
processed = self._preprocess_csm_dataset(dataset, custom_format_mapping)
return (processed, None)
- elif self._audio_type == 'whisper':
+ elif self._audio_type == "whisper":
train_data, eval_data = self._preprocess_whisper_dataset(
- dataset, eval_split=eval_split, custom_format_mapping=custom_format_mapping
+ dataset,
+ eval_split = eval_split,
+ custom_format_mapping = custom_format_mapping,
)
return (train_data, eval_data)
- elif self._audio_type == 'snac':
- processed = self._preprocess_snac_dataset(dataset, custom_format_mapping)
+ elif self._audio_type == "snac":
+ processed = self._preprocess_snac_dataset(
+ dataset, custom_format_mapping
+ )
return (processed, None)
- elif self._audio_type == 'bicodec':
- processed = self._preprocess_bicodec_dataset(dataset, custom_format_mapping)
+ elif self._audio_type == "bicodec":
+ processed = self._preprocess_bicodec_dataset(
+ dataset, custom_format_mapping
+ )
return ({"dataset": processed, "final_format": "audio_bicodec"}, None)
- elif self._audio_type == 'dac':
+ elif self._audio_type == "dac":
processed = self._preprocess_dac_dataset(dataset, custom_format_mapping)
return ({"dataset": processed, "final_format": "audio_dac"}, None)
elif self.is_audio_vlm:
- formatted = self._format_audio_vlm_dataset(dataset, custom_format_mapping)
+ formatted = self._format_audio_vlm_dataset(
+ dataset, custom_format_mapping
+ )
return (formatted, None)
# ========== FORMAT FIRST ==========
@@ -2068,13 +2442,13 @@ class UnslothTrainer:
dataset_info = format_and_template_dataset(
dataset,
- model_name=self.model_name,
- tokenizer=self.tokenizer,
- is_vlm=self.is_vlm,
- format_type=format_type,
- dataset_name=dataset_source,
- custom_format_mapping=custom_format_mapping,
- progress_callback=self._update_progress,
+ model_name = self.model_name,
+ tokenizer = self.tokenizer,
+ is_vlm = self.is_vlm,
+ format_type = format_type,
+ dataset_name = dataset_source,
+ custom_format_mapping = custom_format_mapping,
+ progress_callback = self._update_progress,
)
# Check if stopped during formatting
@@ -2087,10 +2461,12 @@ class UnslothTrainer:
errors = dataset_info.get("errors", [])
error_msg = "; ".join(errors) if errors else "Dataset formatting failed"
logger.error(f"Dataset conversion failed: {error_msg}")
- self._update_progress(error=error_msg)
+ self._update_progress(error = error_msg)
return None
- self._update_progress(status_message=f"Dataset formatted and ready for training")
+ self._update_progress(
+ status_message = f"Dataset formatted and ready for training"
+ )
logger.info(f"Dataset formatted successfully\n")
# ========== THEN SPLIT ==========
@@ -2099,12 +2475,12 @@ class UnslothTrainer:
logger.info(f"Formatting eval dataset ({len(eval_dataset)} rows)...\n")
eval_info = format_and_template_dataset(
eval_dataset,
- model_name=self.model_name,
- tokenizer=self.tokenizer,
- is_vlm=self.is_vlm,
- format_type=format_type,
- dataset_name=dataset_source,
- custom_format_mapping=custom_format_mapping,
+ model_name = self.model_name,
+ tokenizer = self.tokenizer,
+ is_vlm = self.is_vlm,
+ format_type = format_type,
+ dataset_name = dataset_source,
+ custom_format_mapping = custom_format_mapping,
)
eval_dataset = eval_info["dataset"]
logger.info(f"Eval dataset formatted successfully\n")
@@ -2120,14 +2496,16 @@ class UnslothTrainer:
except Exception as e:
logger.error(f"Error loading dataset: {e}")
- self._update_progress(error=str(e))
+ self._update_progress(error = str(e))
return None
- def _auto_detect_eval_split_from_hf(self, dataset_source: str,
- subset: str) -> Optional[Dataset]:
+ def _auto_detect_eval_split_from_hf(
+ self, dataset_source: str, subset: str
+ ) -> Optional[Dataset]:
"""Auto-detect an eval split from HF dataset (separate named split only)."""
try:
from datasets import get_dataset_split_names
+
load_kwargs = {"path": dataset_source}
if subset:
load_kwargs["config_name"] = subset
@@ -2142,10 +2520,14 @@ class UnslothTrainer:
eval_load_kwargs["name"] = subset
candidate_ds = load_dataset(**eval_load_kwargs)
if len(candidate_ds) >= 16:
- logger.info(f"Auto-detected eval split '{candidate}' with {len(candidate_ds)} rows\n")
+ logger.info(
+ f"Auto-detected eval split '{candidate}' with {len(candidate_ds)} rows\n"
+ )
return candidate_ds
else:
- logger.info(f"Found eval split '{candidate}' but only {len(candidate_ds)} rows (< 16), skipping\n")
+ logger.info(
+ f"Found eval split '{candidate}' but only {len(candidate_ds)} rows (< 16), skipping\n"
+ )
except Exception as e:
logger.warning(f"Could not check dataset splits: {e}")
@@ -2172,42 +2554,45 @@ class UnslothTrainer:
eval_size = min(eval_size, n // 2)
logger.info(f"Auto-splitting: {eval_size} rows for eval from {n} total\n")
- split_result = dataset.train_test_split(test_size=eval_size, seed=3407)
- logger.info(f"Split complete: {len(split_result['train'])} train, {len(split_result['test'])} eval\n")
- return (split_result['train'], split_result['test'])
+ split_result = dataset.train_test_split(test_size = eval_size, seed = 3407)
+ logger.info(
+ f"Split complete: {len(split_result['train'])} train, {len(split_result['test'])} eval\n"
+ )
+ return (split_result["train"], split_result["test"])
- def start_training(self,
- dataset: Dataset,
- eval_dataset: Dataset = None,
- eval_steps: float = 0.00,
- output_dir: str | None = None,
- num_epochs: int = 3,
- learning_rate: float = 5e-5,
- batch_size: int = 2,
- gradient_accumulation_steps: int = 4,
- warmup_steps: int = None,
- warmup_ratio: float = None,
- max_steps: int = 0,
- save_steps: int = 0,
- weight_decay: float = 0.01,
- random_seed: int = 3407,
- packing: bool = False,
- train_on_completions: bool = False,
- enable_wandb: bool = False,
- wandb_project: str = "unsloth-training",
- wandb_token: str = None,
- enable_tensorboard: bool = False,
- tensorboard_dir: str | None = None,
- **kwargs) -> bool:
+ def start_training(
+ self,
+ dataset: Dataset,
+ eval_dataset: Dataset = None,
+ eval_steps: float = 0.00,
+ output_dir: str | None = None,
+ num_epochs: int = 3,
+ learning_rate: float = 5e-5,
+ batch_size: int = 2,
+ gradient_accumulation_steps: int = 4,
+ warmup_steps: int = None,
+ warmup_ratio: float = None,
+ max_steps: int = 0,
+ save_steps: int = 0,
+ weight_decay: float = 0.01,
+ random_seed: int = 3407,
+ packing: bool = False,
+ train_on_completions: bool = False,
+ enable_wandb: bool = False,
+ wandb_project: str = "unsloth-training",
+ wandb_token: str = None,
+ enable_tensorboard: bool = False,
+ tensorboard_dir: str | None = None,
+ **kwargs,
+ ) -> bool:
"""Start training in a separate thread"""
if self.is_training:
logger.warning("Training already in progress")
return False
-
if self.model is None or self.tokenizer is None:
- self._update_progress(error="Model not loaded")
+ self._update_progress(error = "Model not loaded")
return False
# Pre-import heavy transformers modules on the main thread.
@@ -2220,7 +2605,8 @@ class UnslothTrainer:
TrainingArguments as _TrainingArguments,
TrainerCallback as _TrainerCallback,
)
- if self._audio_type == 'whisper':
+
+ if self._audio_type == "whisper":
from transformers import ( # noqa: F401
Seq2SeqTrainer as _Seq2SeqTrainer,
Seq2SeqTrainingArguments as _Seq2SeqTrainingArguments,
@@ -2228,31 +2614,31 @@ class UnslothTrainer:
# Start training in separate thread
self.training_thread = threading.Thread(
- target=self._train_worker,
- args=(dataset,),
- kwargs={
- 'output_dir': output_dir,
- 'num_epochs': num_epochs,
- 'learning_rate': learning_rate,
- 'batch_size': batch_size,
- 'gradient_accumulation_steps': gradient_accumulation_steps,
- 'warmup_steps': warmup_steps,
- 'warmup_ratio': warmup_ratio,
- 'max_steps': max_steps,
- 'save_steps': save_steps,
- 'weight_decay': weight_decay,
- 'random_seed': random_seed,
- 'packing': packing,
- 'train_on_completions': train_on_completions,
- 'enable_wandb': enable_wandb,
- 'wandb_project': wandb_project,
- 'wandb_token': wandb_token,
- 'enable_tensorboard': enable_tensorboard,
- 'tensorboard_dir': tensorboard_dir,
- 'eval_dataset': eval_dataset,
- 'eval_steps': eval_steps,
- **kwargs
- }
+ target = self._train_worker,
+ args = (dataset,),
+ kwargs = {
+ "output_dir": output_dir,
+ "num_epochs": num_epochs,
+ "learning_rate": learning_rate,
+ "batch_size": batch_size,
+ "gradient_accumulation_steps": gradient_accumulation_steps,
+ "warmup_steps": warmup_steps,
+ "warmup_ratio": warmup_ratio,
+ "max_steps": max_steps,
+ "save_steps": save_steps,
+ "weight_decay": weight_decay,
+ "random_seed": random_seed,
+ "packing": packing,
+ "train_on_completions": train_on_completions,
+ "enable_wandb": enable_wandb,
+ "wandb_project": wandb_project,
+ "wandb_token": wandb_token,
+ "enable_tensorboard": enable_tensorboard,
+ "tensorboard_dir": tensorboard_dir,
+ "eval_dataset": eval_dataset,
+ "eval_steps": eval_steps,
+ **kwargs,
+ },
)
self.should_stop = False
@@ -2269,100 +2655,130 @@ class UnslothTrainer:
"""Worker function for training (runs in separate thread)"""
try:
# Store training parameters for metrics calculation
- self.batch_size = training_args.get('batch_size', 2)
- self.max_seq_length = training_args.get('max_seq_length', 2048)
- self.gradient_accumulation_steps = training_args.get('gradient_accumulation_steps', 4)
-
+ self.batch_size = training_args.get("batch_size", 2)
+ self.max_seq_length = training_args.get("max_seq_length", 2048)
+ self.gradient_accumulation_steps = training_args.get(
+ "gradient_accumulation_steps", 4
+ )
+
# Set training start time
self.training_start_time = time.time()
-
- self._update_progress(is_training=True, error=None)
+
+ self._update_progress(is_training = True, error = None)
# Setup logging
- if training_args.get('enable_wandb', False) and training_args.get('wandb_token'):
- os.environ["WANDB_API_KEY"] = training_args['wandb_token']
+ if training_args.get("enable_wandb", False) and training_args.get(
+ "wandb_token"
+ ):
+ os.environ["WANDB_API_KEY"] = training_args["wandb_token"]
import wandb
- wandb.init(project=training_args.get('wandb_project', 'unsloth-training'))
+
+ wandb.init(
+ project = training_args.get("wandb_project", "unsloth-training")
+ )
# Create output directory
output_dir = str(resolve_output_dir(training_args.get("output_dir")))
ensure_dir(Path(output_dir))
# ========== AUDIO TRAINER BRANCH ==========
- if self._audio_type == 'csm':
+ if self._audio_type == "csm":
# CSM uses plain HF Trainer (NOT SFTTrainer)
# Needs remove_unused_columns=False for depth decoder (input_values + cutoffs)
from transformers import Trainer as HFTrainer, TrainingArguments
+
self._apply_csm_forward_fix()
- config = self._build_audio_training_args(training_args, output_dir, extra_args={
- "remove_unused_columns": False,
- })
+ config = self._build_audio_training_args(
+ training_args,
+ output_dir,
+ extra_args = {
+ "remove_unused_columns": False,
+ },
+ )
self.trainer = HFTrainer(
- model=self.model, train_dataset=dataset,
- args=TrainingArguments(**config),
+ model = self.model,
+ train_dataset = dataset,
+ args = TrainingArguments(**config),
)
self.trainer.add_callback(self._create_progress_callback())
- batch_size = training_args.get('batch_size', 2)
+ batch_size = training_args.get("batch_size", 2)
total = self._calculate_total_steps(
- len(dataset), batch_size,
- training_args.get('gradient_accumulation_steps', 4),
- training_args.get('num_epochs', 3),
- training_args.get('max_steps', 0),
+ len(dataset),
+ batch_size,
+ training_args.get("gradient_accumulation_steps", 4),
+ training_args.get("num_epochs", 3),
+ training_args.get("max_steps", 0),
+ )
+ self._update_progress(
+ total_steps = total, status_message = "Starting CSM training..."
)
- self._update_progress(total_steps=total, status_message="Starting CSM training...")
logger.info(f"CSM training config: {config}\n")
self.trainer.train()
self._finalize_training(output_dir, "CSM")
return
- elif self._audio_type == 'snac':
+ elif self._audio_type == "snac":
# Orpheus: language model with SNAC codec tokens β plain HF Trainer
# DataCollatorForSeq2Seq dynamically pads variable-length sequences per batch
# (text + audio codes vary in length) and pads labels with -100.
- from transformers import Trainer as HFTrainer, TrainingArguments, DataCollatorForSeq2Seq
+ from transformers import (
+ Trainer as HFTrainer,
+ TrainingArguments,
+ DataCollatorForSeq2Seq,
+ )
config = self._build_audio_training_args(training_args, output_dir)
self.trainer = HFTrainer(
- model=self.model, train_dataset=dataset,
- args=TrainingArguments(**config),
- data_collator=DataCollatorForSeq2Seq(
- tokenizer=self.tokenizer, padding=True, pad_to_multiple_of=8,
+ model = self.model,
+ train_dataset = dataset,
+ args = TrainingArguments(**config),
+ data_collator = DataCollatorForSeq2Seq(
+ tokenizer = self.tokenizer,
+ padding = True,
+ pad_to_multiple_of = 8,
),
)
self.trainer.add_callback(self._create_progress_callback())
- batch_size = training_args.get('batch_size', 2)
+ batch_size = training_args.get("batch_size", 2)
total = self._calculate_total_steps(
- len(dataset), batch_size,
- training_args.get('gradient_accumulation_steps', 4),
- training_args.get('num_epochs', 3),
- training_args.get('max_steps', 0),
+ len(dataset),
+ batch_size,
+ training_args.get("gradient_accumulation_steps", 4),
+ training_args.get("num_epochs", 3),
+ training_args.get("max_steps", 0),
+ )
+ self._update_progress(
+ total_steps = total, status_message = "Starting SNAC training..."
)
- self._update_progress(total_steps=total, status_message="Starting SNAC training...")
logger.info(f"SNAC training config: {config}\n")
self.trainer.train()
self._finalize_training(output_dir, "SNAC")
return
- elif self._audio_type == 'whisper':
+ elif self._audio_type == "whisper":
# Whisper: Seq2SeqTrainer with custom speech collator
from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments
from utils.datasets import DataCollatorSpeechSeq2SeqWithPadding
- eval_dataset = training_args.get('eval_dataset', None)
+ eval_dataset = training_args.get("eval_dataset", None)
extra = {"remove_unused_columns": False, "label_names": ["labels"]}
if eval_dataset:
extra["eval_strategy"] = "steps"
- extra["eval_steps"] = training_args.get('eval_steps', 5)
+ extra["eval_steps"] = training_args.get("eval_steps", 5)
- config = self._build_audio_training_args(training_args, output_dir, extra_args=extra)
+ config = self._build_audio_training_args(
+ training_args, output_dir, extra_args = extra
+ )
trainer_kwargs = {
"model": self.model,
"train_dataset": dataset,
- "data_collator": DataCollatorSpeechSeq2SeqWithPadding(processor=self.tokenizer),
+ "data_collator": DataCollatorSpeechSeq2SeqWithPadding(
+ processor = self.tokenizer
+ ),
"processing_class": self.tokenizer.feature_extractor,
"args": Seq2SeqTrainingArguments(**config),
}
@@ -2372,27 +2788,37 @@ class UnslothTrainer:
self.trainer = Seq2SeqTrainer(**trainer_kwargs)
self.trainer.add_callback(self._create_progress_callback())
- batch_size = training_args.get('batch_size', 2)
+ batch_size = training_args.get("batch_size", 2)
total = self._calculate_total_steps(
- len(dataset), batch_size,
- training_args.get('gradient_accumulation_steps', 4),
- training_args.get('num_epochs', 3),
- training_args.get('max_steps', 0),
+ len(dataset),
+ batch_size,
+ training_args.get("gradient_accumulation_steps", 4),
+ training_args.get("num_epochs", 3),
+ training_args.get("max_steps", 0),
+ )
+ self._update_progress(
+ total_steps = total, status_message = "Starting Whisper training..."
)
- self._update_progress(total_steps=total, status_message="Starting Whisper training...")
logger.info(f"Whisper training config: {config}\n")
self.trainer.train()
self._finalize_training(output_dir, "Whisper")
return
- elif self._audio_type is not None and self._audio_type not in ('bicodec', 'dac'):
+ elif self._audio_type is not None and self._audio_type not in (
+ "bicodec",
+ "dac",
+ ):
# bicodec/dac use the standard SFTTrainer text path below
- raise NotImplementedError(f"Audio training for '{self._audio_type}' not yet implemented")
+ raise NotImplementedError(
+ f"Audio training for '{self._audio_type}' not yet implemented"
+ )
# ========== DATA COLLATOR SELECTION ==========
# Detect special model types
model_name_lower = self.model_name.lower()
- is_deepseek_ocr = "deepseek" in model_name_lower and "ocr" in model_name_lower
+ is_deepseek_ocr = (
+ "deepseek" in model_name_lower and "ocr" in model_name_lower
+ )
logger.info("Configuring data collator...\n")
@@ -2409,7 +2835,7 @@ class UnslothTrainer:
"snapshot_download('unsloth/DeepSeek-OCR', local_dir='deepseek_ocr')"
)
logger.error(error_msg)
- self._update_progress(error=error_msg, is_training=False)
+ self._update_progress(error = error_msg, is_training = False)
return
try:
@@ -2418,19 +2844,21 @@ class UnslothTrainer:
logger.info("Configuring DeepSeek OCR data collator...\n")
FastVisionModel.for_training(self.model)
data_collator = DeepSeekOCRDataCollator(
- tokenizer=self.tokenizer,
- model=self.model,
- image_size=640,
- base_size=1024,
- crop_mode=True,
- train_on_responses_only=training_args.get('train_on_completions', False),
+ tokenizer = self.tokenizer,
+ model = self.model,
+ image_size = 640,
+ base_size = 1024,
+ crop_mode = True,
+ train_on_responses_only = training_args.get(
+ "train_on_completions", False
+ ),
)
logger.info("DeepSeek OCR data collator configured successfully\n")
except Exception as e:
logger.error(f"Failed to configure DeepSeek OCR collator: {e}")
error_msg = f"Error configuring DeepSeek OCR: {str(e)}"
- self._update_progress(error=error_msg, is_training=False)
+ self._update_progress(error = error_msg, is_training = False)
return
elif self.is_audio_vlm:
@@ -2439,26 +2867,33 @@ class UnslothTrainer:
logger.info("Configuring audio VLM data collator...\n")
processor = self.tokenizer # FastModel returns processor as tokenizer
- audio_col_name = getattr(self, '_audio_vlm_audio_col', 'audio')
+ audio_col_name = getattr(self, "_audio_vlm_audio_col", "audio")
def audio_vlm_collate_fn(examples):
texts = []
audios = []
for example in examples:
text = processor.apply_chat_template(
- example["messages"], tokenize=False, add_generation_prompt=False
+ example["messages"],
+ tokenize = False,
+ add_generation_prompt = False,
).strip()
texts.append(text)
audios.append(example[audio_col_name]["array"])
batch = processor(
- text=texts, audio=audios, return_tensors="pt", padding=True
+ text = texts, audio = audios, return_tensors = "pt", padding = True
)
# Labels = input_ids with special tokens masked
labels = batch["input_ids"].clone()
labels[labels == processor.tokenizer.pad_token_id] = -100
- for attr in ('audio_token_id', 'image_token_id', 'boi_token_id', 'eoi_token_id'):
+ for attr in (
+ "audio_token_id",
+ "image_token_id",
+ "boi_token_id",
+ "eoi_token_id",
+ ):
token_id = getattr(processor.tokenizer, attr, None)
if token_id is not None:
labels[labels == token_id] = -100
@@ -2479,41 +2914,52 @@ class UnslothTrainer:
# ========== TRAINING CONFIGURATION ==========
# Handle warmup_steps vs warmup_ratio
- warmup_steps_val = training_args.get('warmup_steps', None)
- warmup_ratio_val = training_args.get('warmup_ratio', None)
-
- lr_value = training_args.get('learning_rate', 2e-4)
- logger.info(f"[DEBUG] learning_rate from training_args: {lr_value} (type: {type(lr_value).__name__})\n")
+ warmup_steps_val = training_args.get("warmup_steps", None)
+ warmup_ratio_val = training_args.get("warmup_ratio", None)
+
+ lr_value = training_args.get("learning_rate", 2e-4)
+ logger.info(
+ f"[DEBUG] learning_rate from training_args: {lr_value} (type: {type(lr_value).__name__})\n"
+ )
config_args = {
- "per_device_train_batch_size": training_args.get('batch_size', 2),
- "gradient_accumulation_steps": training_args.get('gradient_accumulation_steps', 4),
- "num_train_epochs": training_args.get('num_epochs', 3), # Default to epochs
+ "per_device_train_batch_size": training_args.get("batch_size", 2),
+ "gradient_accumulation_steps": training_args.get(
+ "gradient_accumulation_steps", 4
+ ),
+ "num_train_epochs": training_args.get(
+ "num_epochs", 3
+ ), # Default to epochs
"learning_rate": lr_value,
"fp16": not is_bfloat16_supported(),
"bf16": is_bfloat16_supported(),
"logging_steps": 1,
- "weight_decay": training_args.get('weight_decay', 0.01),
- "seed": training_args.get('random_seed', 3407),
+ "weight_decay": training_args.get("weight_decay", 0.01),
+ "seed": training_args.get("random_seed", 3407),
"output_dir": output_dir,
"report_to": _build_report_targets(training_args),
"include_num_input_tokens_seen": True, # Enable token counting
- "dataset_num_proc": 1 if (self.is_audio or self.is_audio_vlm or self._cuda_audio_used) else safe_num_proc(max(1, os.cpu_count() // 4)),
- "max_seq_length": training_args.get('max_seq_length', 2048),
+ "dataset_num_proc": 1
+ if (self.is_audio or self.is_audio_vlm or self._cuda_audio_used)
+ else safe_num_proc(max(1, os.cpu_count() // 4)),
+ "max_seq_length": training_args.get("max_seq_length", 2048),
}
if training_args.get("enable_tensorboard", False):
config_args["logging_dir"] = str(
resolve_tensorboard_dir(training_args.get("tensorboard_dir"))
)
- logger.info(f"[DEBUG] dataset_num_proc={config_args['dataset_num_proc']} (is_audio={self.is_audio}, is_audio_vlm={self.is_audio_vlm}, _cuda_audio_used={self._cuda_audio_used})")
+ logger.info(
+ f"[DEBUG] dataset_num_proc={config_args['dataset_num_proc']} (is_audio={self.is_audio}, is_audio_vlm={self.is_audio_vlm}, _cuda_audio_used={self._cuda_audio_used})"
+ )
# On Windows with transformers 5.x, disable DataLoader multiprocessing
# to avoid issues with modified sys.path (.venv_t5) in spawned workers.
if sys.platform == "win32":
import transformers as _tf
+
if _tf.__version__.startswith("5."):
config_args["dataloader_num_workers"] = 0
-
+
# Add warmup parameter - use warmup_ratio if provided, otherwise warmup_steps
if warmup_ratio_val is not None:
config_args["warmup_ratio"] = warmup_ratio_val
@@ -2527,13 +2973,13 @@ class UnslothTrainer:
logger.info(f"Using default warmup_steps: 5\n")
# Add save_steps if specified
- save_steps_val = training_args.get('save_steps', 0)
+ save_steps_val = training_args.get("save_steps", 0)
if save_steps_val and save_steps_val > 0:
config_args["save_steps"] = save_steps_val
config_args["save_strategy"] = "steps"
# If max_steps is specified, use it instead of epochs
- max_steps_val = training_args.get('max_steps', 0)
+ max_steps_val = training_args.get("max_steps", 0)
if max_steps_val and max_steps_val > 0:
del config_args["num_train_epochs"] # Remove epochs
config_args["max_steps"] = max_steps_val # Use steps instead
@@ -2542,62 +2988,74 @@ class UnslothTrainer:
logger.info(f"Training for {config_args['num_train_epochs']} epochs\n")
# ========== EVAL CONFIGURATION ==========
- eval_dataset = training_args.get('eval_dataset', None)
- eval_steps_val = training_args.get('eval_steps', 0.00)
+ eval_dataset = training_args.get("eval_dataset", None)
+ eval_steps_val = training_args.get("eval_steps", 0.00)
if eval_dataset is not None:
if eval_steps_val > 0:
config_args["eval_strategy"] = "steps"
config_args["eval_steps"] = eval_steps_val
- logger.info(f"β
Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n")
+ logger.info(
+ f"β
Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n"
+ )
logger.info(f"Eval dataset: {len(eval_dataset)} rows\n")
else:
- logger.info(f"β οΈ Eval dataset provided but eval_steps={eval_steps_val} (disabled)\n")
+ logger.info(
+ f"β οΈ Eval dataset provided but eval_steps={eval_steps_val} (disabled)\n"
+ )
logger.info("To enable evaluation, set eval_steps > 0.0\n")
else:
logger.info("No eval dataset β evaluation disabled\n")
# Add model-specific parameters
# Use optim and lr_scheduler_type from training_args if provided, otherwise use defaults
- optim_value = training_args.get('optim', "adamw_8bit")
- lr_scheduler_type_value = training_args.get('lr_scheduler_type', "linear")
-
+ optim_value = training_args.get("optim", "adamw_8bit")
+ lr_scheduler_type_value = training_args.get("lr_scheduler_type", "linear")
+
if self.is_vlm or self.is_audio_vlm:
# Vision / audio VLM config (both need skip_prepare_dataset + remove_unused_columns)
label = "audio VLM" if self.is_audio_vlm else "vision"
logger.info(f"Configuring {label} model training parameters\n")
# Use provided values or defaults for vision models
- optim_value = training_args.get('optim', "adamw_torch_fused")
- lr_scheduler_type_value = training_args.get('lr_scheduler_type', "cosine")
- config_args.update({
- "optim": optim_value,
- "lr_scheduler_type": lr_scheduler_type_value,
- "gradient_checkpointing": True,
- "gradient_checkpointing_kwargs": {"use_reentrant": False},
- "max_grad_norm": 0.3,
- "remove_unused_columns": False,
- "dataset_text_field": "",
- "dataset_kwargs": {"skip_prepare_dataset": True},
- "max_length": training_args.get('max_seq_length', 2048),
- })
+ optim_value = training_args.get("optim", "adamw_torch_fused")
+ lr_scheduler_type_value = training_args.get(
+ "lr_scheduler_type", "cosine"
+ )
+ config_args.update(
+ {
+ "optim": optim_value,
+ "lr_scheduler_type": lr_scheduler_type_value,
+ "gradient_checkpointing": True,
+ "gradient_checkpointing_kwargs": {"use_reentrant": False},
+ "max_grad_norm": 0.3,
+ "remove_unused_columns": False,
+ "dataset_text_field": "",
+ "dataset_kwargs": {"skip_prepare_dataset": True},
+ "max_length": training_args.get("max_seq_length", 2048),
+ }
+ )
else:
logger.info("Configuring text model training parameters\n")
- config_args.update({
- "optim": optim_value,
- "lr_scheduler_type": lr_scheduler_type_value,
- "dataset_text_field": "text",
- })
+ config_args.update(
+ {
+ "optim": optim_value,
+ "lr_scheduler_type": lr_scheduler_type_value,
+ "dataset_text_field": "text",
+ }
+ )
# Only add packing for text models (not DeepSeek OCR which is VLM)
if not is_deepseek_ocr:
- packing_enabled = training_args.get('packing', False)
+ packing_enabled = training_args.get("packing", False)
config_args["packing"] = packing_enabled
- logger.info(f"Sequence packing: {'enabled' if packing_enabled else 'disabled'}\n")
+ logger.info(
+ f"Sequence packing: {'enabled' if packing_enabled else 'disabled'}\n"
+ )
# Audio codec overrides β BiCodec/DAC use the text SFTTrainer path
- if self._audio_type == 'bicodec':
+ if self._audio_type == "bicodec":
config_args["packing"] = False
logger.info("Applied BiCodec overrides: packing=False\n")
- elif self._audio_type == 'dac':
+ elif self._audio_type == "dac":
config_args["packing"] = False
logger.info("Applied DAC overrides: packing=False\n")
@@ -2608,8 +3066,14 @@ class UnslothTrainer:
if self.is_audio_vlm:
# Audio VLM (e.g. Gemma 3N + audio): raw Dataset from _format_audio_vlm_dataset
# Notebook uses processing_class=processor.tokenizer (text tokenizer only)
- train_dataset = dataset if isinstance(dataset, Dataset) else dataset['dataset']
- processing_class = self.tokenizer.tokenizer if hasattr(self.tokenizer, 'tokenizer') else self.tokenizer
+ train_dataset = (
+ dataset if isinstance(dataset, Dataset) else dataset["dataset"]
+ )
+ processing_class = (
+ self.tokenizer.tokenizer
+ if hasattr(self.tokenizer, "tokenizer")
+ else self.tokenizer
+ )
trainer_kwargs = {
"model": self.model,
"train_dataset": train_dataset,
@@ -2622,7 +3086,9 @@ class UnslothTrainer:
self.trainer = SFTTrainer(**trainer_kwargs)
elif self.is_vlm:
# Image VLM: dataset is dict wrapper from format_and_template_dataset
- train_dataset = dataset['dataset'] if isinstance(dataset, dict) else dataset
+ train_dataset = (
+ dataset["dataset"] if isinstance(dataset, dict) else dataset
+ )
trainer_kwargs = {
"model": self.model,
"train_dataset": train_dataset,
@@ -2640,15 +3106,20 @@ class UnslothTrainer:
# ProcessorMixin β sets _is_vlm=True β skips _prepare_dataset entirely,
# and the 'text' column never gets tokenized to 'input_ids'.
from transformers import ProcessorMixin
+
sft_tokenizer = self.tokenizer
- if isinstance(self.tokenizer, ProcessorMixin) and hasattr(self.tokenizer, 'tokenizer'):
- logger.info(f" β οΈ Unwrapping Processor β raw tokenizer for text-only SFTTrainer")
+ if isinstance(self.tokenizer, ProcessorMixin) and hasattr(
+ self.tokenizer, "tokenizer"
+ ):
+ logger.info(
+ f" β οΈ Unwrapping Processor β raw tokenizer for text-only SFTTrainer"
+ )
sft_tokenizer = self.tokenizer.tokenizer
trainer_kwargs = {
"model": self.model,
"tokenizer": sft_tokenizer,
- "train_dataset": dataset['dataset'],
+ "train_dataset": dataset["dataset"],
"data_collator": data_collator,
"args": SFTConfig(**config_args),
}
@@ -2665,11 +3136,18 @@ class UnslothTrainer:
# Determine if we should train on responses only
instruction_part = None
response_part = None
- train_on_responses_enabled = training_args.get('train_on_completions', False)
+ train_on_responses_enabled = training_args.get(
+ "train_on_completions", False
+ )
# DeepSeek OCR handles this internally in its collator, so skip
# Audio VLM handles label masking in its collator, so skip
- if train_on_responses_enabled and not self.is_audio_vlm and not self.is_audio and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'):
+ if (
+ train_on_responses_enabled
+ and not self.is_audio_vlm
+ and not self.is_audio
+ and not (is_deepseek_ocr or dataset["final_format"].lower() == "alpaca")
+ ):
try:
logger.info("Configuring train on responses only...\n")
@@ -2681,16 +3159,26 @@ class UnslothTrainer:
logger.info(f"Detected template: {template_name}\n")
if template_name in TEMPLATE_TO_RESPONSES_MAPPER:
- instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["instruction"]
- response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"]
+ instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[
+ template_name
+ ]["instruction"]
+ response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name][
+ "response"
+ ]
- logger.info(f"Instruction marker: {instruction_part[:50]}...\n")
+ logger.info(
+ f"Instruction marker: {instruction_part[:50]}...\n"
+ )
logger.info(f"Response marker: {response_part[:50]}...\n")
else:
- logger.info(f"No response mapping found for template: {template_name}\n")
+ logger.info(
+ f"No response mapping found for template: {template_name}\n"
+ )
train_on_responses_enabled = False
else:
- logger.info(f"No template mapping found for model: {self.model_name}\n")
+ logger.info(
+ f"No template mapping found for model: {self.model_name}\n"
+ )
train_on_responses_enabled = False
except Exception as e:
@@ -2698,15 +3186,22 @@ class UnslothTrainer:
train_on_responses_enabled = False
# Apply train on responses only if we have valid parts
- if train_on_responses_enabled and instruction_part and response_part and not self.is_audio_vlm and not self.is_audio and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'):
+ if (
+ train_on_responses_enabled
+ and instruction_part
+ and response_part
+ and not self.is_audio_vlm
+ and not self.is_audio
+ and not (is_deepseek_ocr or dataset["final_format"].lower() == "alpaca")
+ ):
try:
from unsloth.chat_templates import train_on_responses_only
self.trainer = train_on_responses_only(
self.trainer,
- instruction_part=instruction_part,
- response_part=response_part,
- num_proc=config_args["dataset_num_proc"],
+ instruction_part = instruction_part,
+ response_part = response_part,
+ num_proc = config_args["dataset_num_proc"],
)
logger.info("Train on responses only configured successfully\n")
@@ -2719,10 +3214,14 @@ class UnslothTrainer:
filtered_len = len(self.trainer.train_dataset)
original_len = len(dataset["dataset"])
dropped = original_len - filtered_len
- drop_pct = round(100 * dropped / original_len, 1) if original_len > 0 else 0
+ drop_pct = (
+ round(100 * dropped / original_len, 1)
+ if original_len > 0
+ else 0
+ )
if filtered_len == 0 or drop_pct > 30:
- max_seq = training_args.get('max_seq_length', 2048)
+ max_seq = training_args.get("max_seq_length", 2048)
error_msg = (
f"{dropped}/{original_len} samples ({drop_pct}%) "
f"were dropped after applying 'train on responses "
@@ -2733,7 +3232,7 @@ class UnslothTrainer:
f"or disabling 'Train on completions'."
)
logger.error(error_msg)
- self._update_progress(error=error_msg, is_training=False)
+ self._update_progress(error = error_msg, is_training = False)
return
if dropped > 0:
@@ -2747,12 +3246,23 @@ class UnslothTrainer:
# [DEBUG] Decode first sample AFTER train_on_completions applied
try:
_row = self.trainer.train_dataset[0]
- _space = self.tokenizer(" ", add_special_tokens=False).input_ids[0]
- print("[DEBUG] === After train_on_completions ===", flush=True)
- print(f"[DEBUG] input_ids decoded:\n{self.tokenizer.decode(_row['input_ids'])}\n", flush=True)
- print(f"[DEBUG] labels decoded (-100 β space):\n{self.tokenizer.decode([_space if x == -100 else x for x in _row['labels']])}\n", flush=True)
+ _space = self.tokenizer(
+ " ", add_special_tokens = False
+ ).input_ids[0]
+ print("[DEBUG] === After train_on_completions ===", flush = True)
+ print(
+ f"[DEBUG] input_ids decoded:\n{self.tokenizer.decode(_row['input_ids'])}\n",
+ flush = True,
+ )
+ print(
+ f"[DEBUG] labels decoded (-100 β space):\n{self.tokenizer.decode([_space if x == -100 else x for x in _row['labels']])}\n",
+ flush = True,
+ )
except Exception as _dbg_e:
- print(f"[DEBUG] Could not decode post-completions sample: {_dbg_e}", flush=True)
+ print(
+ f"[DEBUG] Could not decode post-completions sample: {_dbg_e}",
+ flush = True,
+ )
except Exception as e:
logger.warning(f"Failed to apply train on responses only: {e}")
@@ -2766,18 +3276,21 @@ class UnslothTrainer:
# ========== PROGRESS TRACKING ==========
self.trainer.add_callback(self._create_progress_callback())
- num_samples = len(dataset['dataset'] if isinstance(dataset, dict) else dataset)
- batch_size = training_args.get('batch_size', 2)
- total_steps = self._calculate_total_steps(
- num_samples, batch_size,
- training_args.get('gradient_accumulation_steps', 4),
- training_args.get('num_epochs', 3),
- training_args.get('max_steps', 0),
+ num_samples = len(
+ dataset["dataset"] if isinstance(dataset, dict) else dataset
)
- self._update_progress(total_steps=total_steps)
+ batch_size = training_args.get("batch_size", 2)
+ total_steps = self._calculate_total_steps(
+ num_samples,
+ batch_size,
+ training_args.get("gradient_accumulation_steps", 4),
+ training_args.get("num_epochs", 3),
+ training_args.get("max_steps", 0),
+ )
+ self._update_progress(total_steps = total_steps)
# ========== START TRAINING ==========
- self._update_progress(status_message="Starting training...")
+ self._update_progress(status_message = "Starting training...")
logger.info("Starting training...\n")
self.trainer.train()
@@ -2786,9 +3299,10 @@ class UnslothTrainer:
except Exception as e:
import traceback
+
logger.error(f"Training error: {e}")
logger.error(f"Full traceback:\n{traceback.format_exc()}")
- self._update_progress(is_training=False, error=str(e))
+ self._update_progress(is_training = False, error = str(e))
finally:
self.is_training = False
@@ -2815,10 +3329,12 @@ class UnslothTrainer:
method = "lora"
config["unsloth_training_method"] = method
- logger.info(f"Patching adapter_config.json with unsloth_training_method='{method}'")
+ logger.info(
+ f"Patching adapter_config.json with unsloth_training_method='{method}'"
+ )
with open(config_path, "w") as f:
- json.dump(config, f, indent=2)
+ json.dump(config, f, indent = 2)
except Exception as e:
logger.warning(f"Failed to patch adapter_config.json: {e}")
@@ -2833,7 +3349,7 @@ class UnslothTrainer:
if save
else "Cancelling training..."
)
- self._update_progress(status_message=stop_msg)
+ self._update_progress(status_message = stop_msg)
# If trainer exists, try to stop it gracefully
if self.trainer:
@@ -2872,13 +3388,16 @@ def _ensure_deepseek_ocr_installed():
try:
# Try importing to see if already available
from deepseek_ocr.modeling_deepseekocr import format_messages
+
logger.info("DeepSeek OCR module already available")
return True
except ImportError:
pass
try:
- logger.info("DeepSeek OCR module not found. Auto-installing from HuggingFace...")
+ logger.info(
+ "DeepSeek OCR module not found. Auto-installing from HuggingFace..."
+ )
logger.info("\n Downloading DeepSeek OCR module from HuggingFace...\n")
from huggingface_hub import snapshot_download
@@ -2893,9 +3412,7 @@ def _ensure_deepseek_ocr_installed():
local_dir = os.path.join(parent_dir, "deepseek_ocr")
snapshot_download(
- "unsloth/DeepSeek-OCR",
- local_dir=local_dir,
- local_dir_use_symlinks=False
+ "unsloth/DeepSeek-OCR", local_dir = local_dir, local_dir_use_symlinks = False
)
# Add to sys.path if not already there
@@ -2914,9 +3431,11 @@ def _ensure_deepseek_ocr_installed():
logger.info(f"\nβ Failed to install DeepSeek OCR module: {e}\n")
return False
+
# Global trainer instance
_trainer_instance = None
+
def get_trainer() -> UnslothTrainer:
"""Get global trainer instance"""
global _trainer_instance
diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index 649e48b96f..96a92af942 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -13,6 +13,7 @@ worker's mp.Queue, and exposes the same API surface to routes/training.py.
Pattern follows core/data_recipe/jobs/manager.py.
"""
+
import math
import multiprocessing as mp
import queue
@@ -39,6 +40,7 @@ PLOT_HEIGHT = 3.5
class TrainingProgress:
"""Mirror of trainer.TrainingProgress β kept here so the parent process
never needs to import the heavy ML modules."""
+
epoch: float = 0
step: int = 0
total_steps: int = 0
@@ -109,7 +111,7 @@ class TrainingBackend:
# Join prior pump thread to prevent it from consuming events
# from the new job's queue (it reads self._event_queue dynamically).
if self._pump_thread is not None and self._pump_thread.is_alive():
- self._pump_thread.join(timeout=5.0)
+ self._pump_thread.join(timeout = 5.0)
if self._pump_thread.is_alive():
logger.warning("Previous pump thread did not exit within 5s")
self._pump_thread = None
@@ -117,7 +119,9 @@ class TrainingBackend:
# Reset state
self._should_stop = False
self._cancel_requested = False
- self._progress = TrainingProgress(is_training=True, status_message="Initializing training...")
+ self._progress = TrainingProgress(
+ is_training = True, status_message = "Initializing training..."
+ )
self.loss_history.clear()
self.lr_history.clear()
self.step_history.clear()
@@ -172,7 +176,9 @@ class TrainingBackend:
"train_on_completions": kwargs.get("train_on_completions", False),
"finetune_vision_layers": kwargs.get("finetune_vision_layers", True),
"finetune_language_layers": kwargs.get("finetune_language_layers", True),
- "finetune_attention_modules": kwargs.get("finetune_attention_modules", True),
+ "finetune_attention_modules": kwargs.get(
+ "finetune_attention_modules", True
+ ),
"finetune_mlp_modules": kwargs.get("finetune_mlp_modules", True),
"enable_wandb": kwargs.get("enable_wandb", False),
"wandb_token": kwargs.get("wandb_token"),
@@ -193,19 +199,19 @@ class TrainingBackend:
self._stop_queue = _CTX.Queue()
self._proc = _CTX.Process(
- target=run_training_process,
- kwargs={
+ target = run_training_process,
+ kwargs = {
"event_queue": self._event_queue,
"stop_queue": self._stop_queue,
"config": config,
},
- daemon=True,
+ daemon = True,
)
self._proc.start()
logger.info("Training subprocess started (pid=%s)", self._proc.pid)
# Start event pump thread
- self._pump_thread = threading.Thread(target=self._pump_loop, daemon=True)
+ self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread.start()
return True
@@ -224,7 +230,8 @@ class TrainingBackend:
# Update progress immediately for responsive UI
self._progress.status_message = (
"Stopping training and saving checkpoint..."
- if save else "Cancelling training..."
+ if save
+ else "Cancelling training..."
)
return True
@@ -232,15 +239,17 @@ class TrainingBackend:
"""Force-kill the training subprocess so state can be reset immediately."""
with self._lock:
if self._proc is not None and self._proc.is_alive():
- logger.info("Force-terminating training subprocess (pid=%s)", self._proc.pid)
+ logger.info(
+ "Force-terminating training subprocess (pid=%s)", self._proc.pid
+ )
self._proc.terminate()
proc = self._proc
if proc is not None:
- proc.join(timeout=5.0)
+ proc.join(timeout = 5.0)
if proc.is_alive():
proc.kill()
- proc.join(timeout=2.0)
+ proc.join(timeout = 2.0)
def is_training_active(self) -> bool:
"""Check if training is currently active."""
@@ -262,9 +271,29 @@ class TrainingBackend:
# Check status message for activity indicators
status_lower = (p.status_message or "").lower()
- if any(k in status_lower for k in ["cancelled", "canceled", "stopped", "completed", "ready to train"]):
+ if any(
+ k in status_lower
+ for k in [
+ "cancelled",
+ "canceled",
+ "stopped",
+ "completed",
+ "ready to train",
+ ]
+ ):
return False
- if any(k in status_lower for k in ["loading", "preparing", "training", "configuring", "tokenizing", "starting", "importing"]):
+ if any(
+ k in status_lower
+ for k in [
+ "loading",
+ "preparing",
+ "training",
+ "configuring",
+ "tokenizing",
+ "starting",
+ "importing",
+ ]
+ ):
return True
return False
@@ -282,7 +311,7 @@ class TrainingBackend:
def refresh_plot_for_theme(self, theme: str) -> Optional[plt.Figure]:
"""Refresh plot with new theme."""
- if theme and isinstance(theme, str) and theme in ['light', 'dark']:
+ if theme and isinstance(theme, str) and theme in ["light", "dark"]:
self.current_theme = theme
if self.loss_history:
with self._lock:
@@ -296,6 +325,7 @@ class TrainingBackend:
class _TrainerShim:
"""Minimal shim so routes that access backend.trainer.* still work."""
+
def __init__(self, backend: "TrainingBackend"):
self._backend = backend
self.should_stop = False
@@ -333,7 +363,7 @@ class TrainingBackend:
return
# Try to read an event
- event = self._read_queue(self._event_queue, timeout_sec=0.25)
+ event = self._read_queue(self._event_queue, timeout_sec = 0.25)
if event is not None:
self._handle_event(event)
continue
@@ -354,7 +384,10 @@ class TrainingBackend:
self._progress.status_message = "Training stopped."
else:
self._progress.is_training = False
- self._progress.error = self._progress.error or "Training process exited unexpectedly"
+ self._progress.error = (
+ self._progress.error
+ or "Training process exited unexpectedly"
+ )
return
def _handle_event(self, event: dict) -> None:
@@ -366,8 +399,12 @@ class TrainingBackend:
self._progress.step = event.get("step", self._progress.step)
self._progress.epoch = event.get("epoch", self._progress.epoch)
self._progress.loss = event.get("loss", self._progress.loss)
- self._progress.learning_rate = event.get("learning_rate", self._progress.learning_rate)
- self._progress.total_steps = event.get("total_steps", self._progress.total_steps)
+ self._progress.learning_rate = event.get(
+ "learning_rate", self._progress.learning_rate
+ )
+ self._progress.total_steps = event.get(
+ "total_steps", self._progress.total_steps
+ )
self._progress.elapsed_seconds = event.get("elapsed_seconds")
self._progress.eta_seconds = event.get("eta_seconds")
self._progress.grad_norm = event.get("grad_norm")
@@ -428,7 +465,7 @@ class TrainingBackend:
@staticmethod
def _read_queue(q: Any, timeout_sec: float) -> Optional[dict]:
try:
- return q.get(timeout=timeout_sec)
+ return q.get(timeout = timeout_sec)
except queue.Empty:
return None
except (EOFError, OSError, ValueError):
@@ -449,28 +486,30 @@ class TrainingBackend:
# Plot generation (unchanged from original)
# ------------------------------------------------------------------
- def _create_loss_plot(self, progress: TrainingProgress, theme: str = "light") -> plt.Figure:
+ def _create_loss_plot(
+ self, progress: TrainingProgress, theme: str = "light"
+ ) -> plt.Figure:
"""Create training loss plot with theme-aware styling."""
- plt.close('all')
+ plt.close("all")
LIGHT_STYLE = {
"facecolor": "#ffffff",
"grid_color": "#d1d5db",
"line": "#16b88a",
"text": "#1f2937",
- "empty_text": "#6b7280"
+ "empty_text": "#6b7280",
}
DARK_STYLE = {
"facecolor": "#292929",
"grid_color": "#404040",
"line": "#4ade80",
"text": "#e5e7eb",
- "empty_text": "#9ca3af"
+ "empty_text": "#9ca3af",
}
style = LIGHT_STYLE if theme == "light" else DARK_STYLE
- fig, ax = plt.subplots(figsize=(PLOT_WIDTH, PLOT_HEIGHT))
+ fig, ax = plt.subplots(figsize = (PLOT_WIDTH, PLOT_HEIGHT))
fig.patch.set_facecolor(style["facecolor"])
ax.set_facecolor(style["facecolor"])
@@ -478,8 +517,15 @@ class TrainingBackend:
steps = self.step_history
losses = self.loss_history
scatter_color = "#60a5fa"
- ax.scatter(steps, losses, s=16, alpha=0.6, color=scatter_color,
- linewidths=0, label="Training Loss (raw)")
+ ax.scatter(
+ steps,
+ losses,
+ s = 16,
+ alpha = 0.6,
+ color = scatter_color,
+ linewidths = 0,
+ label = "Training Loss (raw)",
+ )
MA_WINDOW = 20
window = min(MA_WINDOW, len(losses))
@@ -495,15 +541,21 @@ class TrainingBackend:
denom = i - start + 1
ma.append((cumsum[i + 1] - cumsum[start]) / denom)
- ax.plot(steps, ma, color=style["line"], linewidth=2.5, alpha=0.95,
- label=f"Moving Avg ({ma[-1]:.4f})")
+ ax.plot(
+ steps,
+ ma,
+ color = style["line"],
+ linewidth = 2.5,
+ alpha = 0.95,
+ label = f"Moving Avg ({ma[-1]:.4f})",
+ )
- leg = ax.legend(frameon=False, fontsize=9)
+ leg = ax.legend(frameon = False, fontsize = 9)
for t in leg.get_texts():
t.set_color(style["text"])
- ax.set_xlabel('Steps', fontsize=10, color=style["text"])
- ax.set_ylabel('Loss', fontsize=10, color=style["text"])
+ ax.set_xlabel("Steps", fontsize = 10, color = style["text"])
+ ax.set_ylabel("Loss", fontsize = 10, color = style["text"])
if progress.error:
title = f"Error: {progress.error}"
@@ -516,17 +568,31 @@ class TrainingBackend:
else:
title = "Training Loss"
- ax.set_title(title, fontsize=11, fontweight='bold', pad=10, color=style["text"])
- ax.grid(True, alpha=0.4, linestyle='--', color=style["grid_color"])
- ax.tick_params(colors=style["text"], which='both')
- ax.spines['top'].set_visible(False)
- ax.spines['right'].set_visible(False)
- ax.spines['bottom'].set_color(style["text"])
- ax.spines['left'].set_color(style["text"])
+ ax.set_title(
+ title, fontsize = 11, fontweight = "bold", pad = 10, color = style["text"]
+ )
+ ax.grid(True, alpha = 0.4, linestyle = "--", color = style["grid_color"])
+ ax.tick_params(colors = style["text"], which = "both")
+ ax.spines["top"].set_visible(False)
+ ax.spines["right"].set_visible(False)
+ ax.spines["bottom"].set_color(style["text"])
+ ax.spines["left"].set_color(style["text"])
else:
- display_msg = progress.status_message if progress.status_message else 'Waiting for training data...'
- ax.text(0.5, 0.5, display_msg, ha='center', va='center', fontsize=16,
- color=style["empty_text"], transform=ax.transAxes)
+ display_msg = (
+ progress.status_message
+ if progress.status_message
+ else "Waiting for training data..."
+ )
+ ax.text(
+ 0.5,
+ 0.5,
+ display_msg,
+ ha = "center",
+ va = "center",
+ fontsize = 16,
+ color = style["empty_text"],
+ transform = ax.transAxes,
+ )
ax.set_xticks([])
ax.set_yticks([])
for spine in ax.spines.values():
@@ -544,7 +610,8 @@ class TrainingBackend:
"""
logger.info(
"_transfer_to_inference_backend: subprocess training β "
- "model must be loaded from disk (output_dir=%s)", self._output_dir
+ "model must be loaded from disk (output_dir=%s)",
+ self._output_dir,
)
return False
diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index 0eedeb5135..b5cf984bf9 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -10,6 +10,7 @@ solving the transformers version-switching problem completely.
Pattern follows core/data_recipe/jobs/worker.py.
"""
+
from __future__ import annotations
import structlog
@@ -39,7 +40,9 @@ def _activate_transformers_version(model_name: str) -> None:
resolved = _resolve_base_model(model_name)
if needs_transformers_5(resolved):
- venv_t5 = os.path.join(os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5")
+ venv_t5 = os.path.join(
+ os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5"
+ )
if os.path.isdir(venv_t5):
sys.path.insert(0, venv_t5)
logger.info("Activated transformers 5.x from %s", venv_t5)
@@ -47,16 +50,35 @@ def _activate_transformers_version(model_name: str) -> None:
# Fallback: pip install at runtime (slower, ~10-15s)
logger.warning(".venv_t5 not found at %s β installing at runtime", venv_t5)
import subprocess as sp
- os.makedirs(venv_t5, exist_ok=True)
+
+ os.makedirs(venv_t5, exist_ok = True)
r1 = sp.run(
- [sys.executable, "-m", "pip", "install", "--target", venv_t5,
- "--no-deps", "transformers==5.2.0"],
- stdout=sp.PIPE, stderr=sp.STDOUT,
+ [
+ sys.executable,
+ "-m",
+ "pip",
+ "install",
+ "--target",
+ venv_t5,
+ "--no-deps",
+ "transformers==5.2.0",
+ ],
+ stdout = sp.PIPE,
+ stderr = sp.STDOUT,
)
r2 = sp.run(
- [sys.executable, "-m", "pip", "install", "--target", venv_t5,
- "--no-deps", "huggingface_hub==1.3.0"],
- stdout=sp.PIPE, stderr=sp.STDOUT,
+ [
+ sys.executable,
+ "-m",
+ "pip",
+ "install",
+ "--target",
+ venv_t5,
+ "--no-deps",
+ "huggingface_hub==1.3.0",
+ ],
+ stdout = sp.PIPE,
+ stderr = sp.STDOUT,
)
if r1.returncode != 0 or r2.returncode != 0:
raise RuntimeError(
@@ -85,16 +107,19 @@ def run_training_process(
config: Training configuration dict with all parameters.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
- os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
+ os.environ["PYTHONWARNINGS"] = (
+ "ignore" # Suppress warnings at C-level before imports
+ )
import warnings
from loggers.config import LogConfig
+
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
-
+
LogConfig.setup_logging(
- service_name="unsloth-studio-training-worker",
- env=os.getenv("ENVIRONMENT_TYPE", "production"),
+ service_name = "unsloth-studio-training-worker",
+ env = os.getenv("ENVIRONMENT_TYPE", "production"),
)
model_name = config["model_name"]
@@ -103,18 +128,21 @@ def run_training_process(
try:
_activate_transformers_version(model_name)
except Exception as exc:
- event_queue.put({
- "type": "error",
- "error": f"Failed to activate transformers version: {exc}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": f"Failed to activate transformers version: {exc}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ }
+ )
return
# ββ 1b. On Windows, check Triton availability (must be before import torch) ββ
if sys.platform == "win32":
try:
import triton # noqa: F401
+
logger.info("Triton available β torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
@@ -132,17 +160,25 @@ def run_training_process(
sys.path.insert(0, backend_path)
from core.training.trainer import UnslothTrainer, TrainingProgress
- from utils.paths import ensure_dir, resolve_output_dir, resolve_tensorboard_dir, datasets_root
+ from utils.paths import (
+ ensure_dir,
+ resolve_output_dir,
+ resolve_tensorboard_dir,
+ datasets_root,
+ )
import transformers
+
logger.info("Subprocess loaded transformers %s", transformers.__version__)
except Exception as exc:
- event_queue.put({
- "type": "error",
- "error": f"Failed to import ML libraries: {exc}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": f"Failed to import ML libraries: {exc}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ }
+ )
return
# ββ 2b. EMBEDDING MODEL FAST-PATH ββ
@@ -153,12 +189,14 @@ def run_training_process(
try:
_run_embedding_training(event_queue, stop_queue, config)
except Exception as exc:
- event_queue.put({
- "type": "error",
- "error": str(exc),
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": str(exc),
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ }
+ )
return
# ββ 3. Create a fresh trainer instance ββ
@@ -169,21 +207,23 @@ def run_training_process(
has_train_loss = progress.step >= 0 and progress.loss > 0
has_eval_loss = progress.eval_loss is not None
if has_train_loss or has_eval_loss:
- event_queue.put({
- "type": "progress",
- "step": progress.step,
- "epoch": progress.epoch,
- "loss": progress.loss,
- "learning_rate": progress.learning_rate,
- "total_steps": progress.total_steps,
- "elapsed_seconds": progress.elapsed_seconds,
- "eta_seconds": progress.eta_seconds,
- "grad_norm": progress.grad_norm,
- "num_tokens": progress.num_tokens,
- "eval_loss": progress.eval_loss,
- "status_message": progress.status_message,
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "progress",
+ "step": progress.step,
+ "epoch": progress.epoch,
+ "loss": progress.loss,
+ "learning_rate": progress.learning_rate,
+ "total_steps": progress.total_steps,
+ "elapsed_seconds": progress.elapsed_seconds,
+ "eta_seconds": progress.eta_seconds,
+ "grad_norm": progress.grad_norm,
+ "num_tokens": progress.num_tokens,
+ "eval_loss": progress.eval_loss,
+ "status_message": progress.status_message,
+ "ts": time.time(),
+ }
+ )
if progress.status_message:
_send_status(event_queue, progress.status_message)
@@ -196,7 +236,7 @@ def run_training_process(
def _poll_stop():
while True:
try:
- msg = stop_queue.get(timeout=1.0)
+ msg = stop_queue.get(timeout = 1.0)
if msg and msg.get("type") == "stop":
save = msg.get("save", True)
trainer.should_stop = True
@@ -208,7 +248,7 @@ def run_training_process(
except (EOFError, OSError):
return
- stop_thread = threading.Thread(target=_poll_stop, daemon=True)
+ stop_thread = threading.Thread(target = _poll_stop, daemon = True)
stop_thread.start()
# ββ 4. Execute the training pipeline ββ
@@ -222,12 +262,12 @@ def run_training_process(
# ββ 4a. Lightweight detection + tokenizer (no VRAM) ββ
_send_status(event_queue, "Detecting model type...")
trainer.pre_detect_and_load_tokenizer(
- model_name=model_name,
- max_seq_length=config["max_seq_length"],
- hf_token=hf_token,
- is_dataset_image=config.get("is_dataset_image", False),
- is_dataset_audio=config.get("is_dataset_audio", False),
- trust_remote_code=config.get("trust_remote_code", False),
+ model_name = model_name,
+ max_seq_length = config["max_seq_length"],
+ hf_token = hf_token,
+ is_dataset_image = config.get("is_dataset_image", False),
+ is_dataset_audio = config.get("is_dataset_audio", False),
+ trust_remote_code = config.get("trust_remote_code", False),
)
if trainer.should_stop:
event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
@@ -237,16 +277,16 @@ def run_training_process(
_send_status(event_queue, "Loading and formatting dataset...")
hf_dataset = config.get("hf_dataset", "")
dataset_result = trainer.load_and_format_dataset(
- dataset_source=hf_dataset if hf_dataset and hf_dataset.strip() else None,
- format_type=config.get("format_type", ""),
- local_datasets=config.get("local_datasets") or None,
- custom_format_mapping=config.get("custom_format_mapping"),
- subset=config.get("subset"),
- train_split=config.get("train_split", "train"),
- eval_split=config.get("eval_split"),
- eval_steps=config.get("eval_steps", 0.00),
- dataset_slice_start=config.get("dataset_slice_start"),
- dataset_slice_end=config.get("dataset_slice_end"),
+ dataset_source = hf_dataset if hf_dataset and hf_dataset.strip() else None,
+ format_type = config.get("format_type", ""),
+ local_datasets = config.get("local_datasets") or None,
+ custom_format_mapping = config.get("custom_format_mapping"),
+ subset = config.get("subset"),
+ train_split = config.get("train_split", "train"),
+ eval_split = config.get("eval_split"),
+ eval_steps = config.get("eval_steps", 0.00),
+ dataset_slice_start = config.get("dataset_slice_start"),
+ dataset_slice_end = config.get("dataset_slice_end"),
)
if isinstance(dataset_result, tuple):
@@ -260,13 +300,19 @@ def run_training_process(
# or a raw Dataset for audio paths
try:
ds = dataset["dataset"] if isinstance(dataset, dict) else dataset
- print(f"\n[DEBUG] Dataset loaded BEFORE model. type={type(ds).__name__}, len={len(ds)}", flush=True)
- print(f"[DEBUG] Columns: {ds.column_names}", flush=True)
+ print(
+ f"\n[DEBUG] Dataset loaded BEFORE model. type={type(ds).__name__}, len={len(ds)}",
+ flush = True,
+ )
+ print(f"[DEBUG] Columns: {ds.column_names}", flush = True)
sample = ds[0]
preview = {k: str(v)[:300] for k, v in sample.items()}
- print(f"[DEBUG] First sample: {preview}\n", flush=True)
+ print(f"[DEBUG] First sample: {preview}\n", flush = True)
except Exception as e:
- print(f"[DEBUG] Could not preview first sample: {type(e).__name__}: {e}", flush=True)
+ print(
+ f"[DEBUG] Could not preview first sample: {type(e).__name__}: {e}",
+ flush = True,
+ )
# Disable eval if eval_steps <= 0
eval_steps = config.get("eval_steps", 0.00)
@@ -276,88 +322,114 @@ def run_training_process(
# Tell the parent process that eval is configured so the frontend
# shows "Waiting for first evaluation step..." instead of "not configured"
if eval_dataset is not None:
- event_queue.put({
- "type": "eval_configured",
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "eval_configured",
+ "ts": time.time(),
+ }
+ )
if dataset is None or trainer.should_stop:
if trainer.should_stop:
- event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
+ event_queue.put(
+ {"type": "complete", "output_dir": None, "ts": time.time()}
+ )
else:
- event_queue.put({
- "type": "error",
- "error": trainer.training_progress.error or "Failed to load dataset",
- "stack": "", "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": trainer.training_progress.error
+ or "Failed to load dataset",
+ "stack": "",
+ "ts": time.time(),
+ }
+ )
return
# ββ 4c. Load training model (uses VRAM β dataset already formatted) ββ
_send_status(event_queue, "Loading model...")
success = trainer.load_model(
- model_name=model_name,
- max_seq_length=config["max_seq_length"],
- load_in_4bit=config["load_in_4bit"],
- hf_token=hf_token,
- is_dataset_image=config.get("is_dataset_image", False),
- is_dataset_audio=config.get("is_dataset_audio", False),
- trust_remote_code=config.get("trust_remote_code", False),
+ model_name = model_name,
+ max_seq_length = config["max_seq_length"],
+ load_in_4bit = config["load_in_4bit"],
+ hf_token = hf_token,
+ is_dataset_image = config.get("is_dataset_image", False),
+ is_dataset_audio = config.get("is_dataset_audio", False),
+ trust_remote_code = config.get("trust_remote_code", False),
)
if not success or trainer.should_stop:
if trainer.should_stop:
- event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
+ event_queue.put(
+ {"type": "complete", "output_dir": None, "ts": time.time()}
+ )
else:
error_msg = trainer.training_progress.error or "Failed to load model"
- event_queue.put({
- "type": "error",
- "error": error_msg,
- "stack": "", "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": error_msg,
+ "stack": "",
+ "ts": time.time(),
+ }
+ )
return
# ββ 4d. Prepare model (LoRA or full finetuning) ββ
training_type = config.get("training_type", "LoRA/QLoRA")
- use_lora = (training_type == "LoRA/QLoRA")
+ use_lora = training_type == "LoRA/QLoRA"
if use_lora:
_send_status(event_queue, "Configuring LoRA adapters...")
success = trainer.prepare_model_for_training(
- use_lora=True,
- finetune_vision_layers=config.get("finetune_vision_layers", True),
- finetune_language_layers=config.get("finetune_language_layers", True),
- finetune_attention_modules=config.get("finetune_attention_modules", True),
- finetune_mlp_modules=config.get("finetune_mlp_modules", True),
- target_modules=config.get("target_modules"),
- lora_r=config.get("lora_r", 16),
- lora_alpha=config.get("lora_alpha", 16),
- lora_dropout=config.get("lora_dropout", 0.0),
- use_gradient_checkpointing=config.get("gradient_checkpointing", "unsloth"),
- use_rslora=config.get("use_rslora", False),
- use_loftq=config.get("use_loftq", False),
+ use_lora = True,
+ finetune_vision_layers = config.get("finetune_vision_layers", True),
+ finetune_language_layers = config.get("finetune_language_layers", True),
+ finetune_attention_modules = config.get(
+ "finetune_attention_modules", True
+ ),
+ finetune_mlp_modules = config.get("finetune_mlp_modules", True),
+ target_modules = config.get("target_modules"),
+ lora_r = config.get("lora_r", 16),
+ lora_alpha = config.get("lora_alpha", 16),
+ lora_dropout = config.get("lora_dropout", 0.0),
+ use_gradient_checkpointing = config.get(
+ "gradient_checkpointing", "unsloth"
+ ),
+ use_rslora = config.get("use_rslora", False),
+ use_loftq = config.get("use_loftq", False),
)
else:
_send_status(event_queue, "Preparing model for full finetuning...")
- success = trainer.prepare_model_for_training(use_lora=False)
+ success = trainer.prepare_model_for_training(use_lora = False)
if not success or trainer.should_stop:
if trainer.should_stop:
- event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
+ event_queue.put(
+ {"type": "complete", "output_dir": None, "ts": time.time()}
+ )
else:
- event_queue.put({
- "type": "error",
- "error": trainer.training_progress.error or "Failed to prepare model",
- "stack": "", "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": trainer.training_progress.error
+ or "Failed to prepare model",
+ "stack": "",
+ "ts": time.time(),
+ }
+ )
return
# Convert learning rate
try:
lr_value = float(config.get("learning_rate", "2e-4"))
except ValueError:
- event_queue.put({
- "type": "error",
- "error": f"Invalid learning rate: {config.get('learning_rate')}",
- "stack": "", "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": f"Invalid learning rate: {config.get('learning_rate')}",
+ "stack": "",
+ "ts": time.time(),
+ }
+ )
return
# Generate output dir
@@ -379,64 +451,72 @@ def run_training_process(
trainer._train_worker(
dataset,
- output_dir=output_dir,
- num_epochs=config.get("num_epochs", 3),
- learning_rate=lr_value,
- batch_size=config.get("batch_size", 2),
- gradient_accumulation_steps=config.get("gradient_accumulation_steps", 4),
- warmup_steps=config.get("warmup_steps"),
- warmup_ratio=config.get("warmup_ratio"),
- max_steps=max_steps if max_steps and max_steps > 0 else 0,
- save_steps=save_steps if save_steps and save_steps > 0 else 0,
- weight_decay=config.get("weight_decay", 0.01),
- random_seed=config.get("random_seed", 3407),
- packing=config.get("packing", False),
- train_on_completions=config.get("train_on_completions", False),
- enable_wandb=config.get("enable_wandb", False),
- wandb_project=config.get("wandb_project", "unsloth-training"),
- wandb_token=config.get("wandb_token"),
- enable_tensorboard=config.get("enable_tensorboard", False),
- tensorboard_dir=tensorboard_dir,
- eval_dataset=eval_dataset,
- eval_steps=eval_steps,
- max_seq_length=config.get("max_seq_length", 2048),
- optim=config.get("optim", "adamw_8bit"),
- lr_scheduler_type=config.get("lr_scheduler_type", "linear"),
+ output_dir = output_dir,
+ num_epochs = config.get("num_epochs", 3),
+ learning_rate = lr_value,
+ batch_size = config.get("batch_size", 2),
+ gradient_accumulation_steps = config.get("gradient_accumulation_steps", 4),
+ warmup_steps = config.get("warmup_steps"),
+ warmup_ratio = config.get("warmup_ratio"),
+ max_steps = max_steps if max_steps and max_steps > 0 else 0,
+ save_steps = save_steps if save_steps and save_steps > 0 else 0,
+ weight_decay = config.get("weight_decay", 0.01),
+ random_seed = config.get("random_seed", 3407),
+ packing = config.get("packing", False),
+ train_on_completions = config.get("train_on_completions", False),
+ enable_wandb = config.get("enable_wandb", False),
+ wandb_project = config.get("wandb_project", "unsloth-training"),
+ wandb_token = config.get("wandb_token"),
+ enable_tensorboard = config.get("enable_tensorboard", False),
+ tensorboard_dir = tensorboard_dir,
+ eval_dataset = eval_dataset,
+ eval_steps = eval_steps,
+ max_seq_length = config.get("max_seq_length", 2048),
+ optim = config.get("optim", "adamw_8bit"),
+ lr_scheduler_type = config.get("lr_scheduler_type", "linear"),
)
# Check final state
progress = trainer.get_training_progress()
if progress.error:
- event_queue.put({
- "type": "error",
- "error": progress.error,
- "stack": "",
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": progress.error,
+ "stack": "",
+ "ts": time.time(),
+ }
+ )
else:
- event_queue.put({
- "type": "complete",
- "output_dir": output_dir,
- "status_message": progress.status_message or "Training completed",
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "complete",
+ "output_dir": output_dir,
+ "status_message": progress.status_message or "Training completed",
+ "ts": time.time(),
+ }
+ )
except Exception as exc:
- event_queue.put({
- "type": "error",
- "error": str(exc),
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": str(exc),
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ }
+ )
def _send_status(event_queue: Any, message: str) -> None:
"""Send a status update to the parent process."""
- event_queue.put({
- "type": "status",
- "message": message,
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "status",
+ "message": message,
+ "ts": time.time(),
+ }
+ )
def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> None:
@@ -469,14 +549,17 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
from sentence_transformers.training_args import BatchSamplers
from datasets import load_dataset, Dataset
from transformers import TrainerCallback
+ from utils.paths import datasets_root, resolve_output_dir
except ImportError as e:
- event_queue.put({
- "type": "error",
- "error": f"Failed to import embedding libraries: {e}. "
- "Ensure 'sentence_transformers' and 'unsloth' are installed.",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": f"Failed to import embedding libraries: {e}. "
+ "Ensure 'sentence_transformers' and 'unsloth' are installed.",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ }
+ )
return
# ββ Stop signal handling ββ
@@ -487,18 +570,21 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
nonlocal _should_stop, _save_on_stop
while True:
try:
- msg = stop_queue.get(timeout=1.0)
+ msg = stop_queue.get(timeout = 1.0)
if msg and msg.get("type") == "stop":
_save_on_stop = msg.get("save", True)
_should_stop = True
- logger.info("Embedding training: stop signal received (save=%s)", _save_on_stop)
+ logger.info(
+ "Embedding training: stop signal received (save=%s)",
+ _save_on_stop,
+ )
return
except _queue.Empty:
continue
except (EOFError, OSError):
return
- stop_thread = threading.Thread(target=_poll_stop, daemon=True)
+ stop_thread = threading.Thread(target = _poll_stop, daemon = True)
stop_thread.start()
# ββ 2. Load model ββ
@@ -508,21 +594,23 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
hf_token = hf_token if hf_token and hf_token.strip() else None
max_seq_length = config.get("max_seq_length", 512)
training_type = config.get("training_type", "LoRA/QLoRA")
- use_lora = (training_type == "LoRA/QLoRA")
+ use_lora = training_type == "LoRA/QLoRA"
model = FastSentenceTransformer.from_pretrained(
- model_name=model_name,
- max_seq_length=max_seq_length,
- full_finetuning=not use_lora,
- token=hf_token,
+ model_name = model_name,
+ max_seq_length = max_seq_length,
+ full_finetuning = not use_lora,
+ token = hf_token,
)
except Exception as e:
- event_queue.put({
- "type": "error",
- "error": f"Failed to load embedding model '{model_name}': {e}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": f"Failed to load embedding model '{model_name}': {e}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ }
+ )
return
if _should_stop:
@@ -540,24 +628,29 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
model = FastSentenceTransformer.get_peft_model(
model,
- r=config.get("lora_r", 32),
- target_modules=config.get("target_modules") or ["q_proj", "k_proj", "v_proj", "o_proj"],
- lora_alpha=config.get("lora_alpha", 64),
- lora_dropout=config.get("lora_dropout", 0.0),
- bias="none",
- use_gradient_checkpointing=gradient_checkpointing,
- random_state=config.get("random_seed", 3407),
- use_rslora=config.get("use_rslora", False),
- loftq_config={"loftq_bits": 4, "loftq_iter": 1} if config.get("use_loftq") else None,
- task_type="FEATURE_EXTRACTION",
+ r = config.get("lora_r", 32),
+ target_modules = config.get("target_modules")
+ or ["q_proj", "k_proj", "v_proj", "o_proj"],
+ lora_alpha = config.get("lora_alpha", 64),
+ lora_dropout = config.get("lora_dropout", 0.0),
+ bias = "none",
+ use_gradient_checkpointing = gradient_checkpointing,
+ random_state = config.get("random_seed", 3407),
+ use_rslora = config.get("use_rslora", False),
+ loftq_config = {"loftq_bits": 4, "loftq_iter": 1}
+ if config.get("use_loftq")
+ else None,
+ task_type = "FEATURE_EXTRACTION",
)
except Exception as e:
- event_queue.put({
- "type": "error",
- "error": f"Failed to configure LoRA for embedding model: {e}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": f"Failed to configure LoRA for embedding model: {e}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ }
+ )
return
if _should_stop:
@@ -578,16 +671,21 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
dataset = load_dataset(
hf_dataset.strip(),
subset,
- split=train_split,
- token=hf_token,
+ split = train_split,
+ token = hf_token,
)
elif local_datasets:
# Load from local file(s) β mirrors the non-embedding pipeline's
# directory handling so recipe outputs (parquet-files/) work.
all_files: list[str] = []
for dataset_file in local_datasets:
- file_path = dataset_file if os.path.isabs(dataset_file) else os.path.join(
- str(datasets_root()), dataset_file,
+ file_path = (
+ dataset_file
+ if os.path.isabs(dataset_file)
+ else os.path.join(
+ str(datasets_root()),
+ dataset_file,
+ )
)
if os.path.isdir(file_path):
file_path_obj = Path(file_path)
@@ -606,7 +704,9 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
if candidates:
all_files.extend(str(c) for c in candidates)
continue
- raise ValueError(f"No supported data files in directory: {file_path_obj}")
+ raise ValueError(
+ f"No supported data files in directory: {file_path_obj}"
+ )
else:
all_files.append(file_path)
@@ -619,14 +719,19 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
elif first_ext == ".parquet":
loader = "parquet"
else:
- raise ValueError(f"Unsupported local dataset format: {all_files[0]}")
- dataset = load_dataset(loader, data_files=all_files, split="train")
+ raise ValueError(
+ f"Unsupported local dataset format: {all_files[0]}"
+ )
+ dataset = load_dataset(loader, data_files = all_files, split = "train")
else:
- event_queue.put({
- "type": "error",
- "error": "No dataset specified for embedding training.",
- "stack": "", "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": "No dataset specified for embedding training.",
+ "stack": "",
+ "ts": time.time(),
+ }
+ )
return
# Apply dataset slicing if specified
@@ -639,12 +744,14 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
logger.info(f"Embedding dataset loaded: {len(dataset)} samples")
except Exception as e:
- event_queue.put({
- "type": "error",
- "error": f"Failed to load dataset: {e}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": f"Failed to load dataset: {e}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ }
+ )
return
if _should_stop:
@@ -659,16 +766,21 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
try:
lr_value = float(config.get("learning_rate", "2e-4"))
except ValueError:
- event_queue.put({
- "type": "error",
- "error": f"Invalid learning rate: {config.get('learning_rate')}",
- "stack": "", "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": f"Invalid learning rate: {config.get('learning_rate')}",
+ "stack": "",
+ "ts": time.time(),
+ }
+ )
return
output_dir = config.get("output_dir")
if not output_dir:
- output_dir = f"./outputs/{model_name.replace('/', '_')}_{int(time.time())}"
+ output_dir = str(
+ resolve_output_dir(f"{model_name.replace('/', '_')}_{int(time.time())}")
+ )
num_epochs = config.get("num_epochs", 2)
batch_size = config.get("batch_size", 256)
@@ -728,7 +840,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
class _EmbeddingProgressCallback(TrainerCallback):
"""Sends training progress events to the parent process via event_queue."""
- def on_log(self, args, state, control, logs=None, **kwargs):
+ def on_log(self, args, state, control, logs = None, **kwargs):
if not logs:
return
loss_value = logs.get("loss", logs.get("train_loss", 0.0))
@@ -741,21 +853,23 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
if remaining > 0:
eta = (elapsed / current_step) * remaining
- event_queue.put({
- "type": "progress",
- "step": current_step,
- "epoch": round(state.epoch, 2) if state.epoch else 0,
- "loss": loss_value,
- "learning_rate": logs.get("learning_rate", 0.0),
- "total_steps": total_steps,
- "elapsed_seconds": elapsed,
- "eta_seconds": eta,
- "grad_norm": logs.get("grad_norm"),
- "num_tokens": getattr(state, "num_input_tokens_seen", None),
- "eval_loss": logs.get("eval_loss"),
- "status_message": "",
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "progress",
+ "step": current_step,
+ "epoch": round(state.epoch, 2) if state.epoch else 0,
+ "loss": loss_value,
+ "learning_rate": logs.get("learning_rate", 0.0),
+ "total_steps": total_steps,
+ "elapsed_seconds": elapsed,
+ "eta_seconds": eta,
+ "grad_norm": logs.get("grad_norm"),
+ "num_tokens": getattr(state, "num_input_tokens_seen", None),
+ "eval_loss": logs.get("eval_loss"),
+ "status_message": "",
+ "ts": time.time(),
+ }
+ )
def on_step_end(self, args, state, control, **kwargs):
if _should_stop:
@@ -767,31 +881,35 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
_send_status(event_queue, "Starting embedding training...")
try:
trainer = SentenceTransformerTrainer(
- model=model,
- train_dataset=dataset,
- loss=loss,
- args=args,
- callbacks=[_EmbeddingProgressCallback()],
+ model = model,
+ train_dataset = dataset,
+ loss = loss,
+ args = args,
+ callbacks = [_EmbeddingProgressCallback()],
)
trainer.train()
except Exception as e:
- event_queue.put({
- "type": "error",
- "error": f"Embedding training failed: {e}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": f"Embedding training failed: {e}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ }
+ )
return
# ββ 10. Save model ββ
if _should_stop and not _save_on_stop:
- event_queue.put({
- "type": "complete",
- "output_dir": None,
- "status_message": "Training cancelled",
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "complete",
+ "output_dir": None,
+ "status_message": "Training cancelled",
+ "ts": time.time(),
+ }
+ )
return
_send_status(event_queue, "Saving model...")
@@ -801,18 +919,22 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
logger.info("Embedding model saved to %s", output_dir)
except Exception as e:
logger.error("Failed to save embedding model: %s", e)
- event_queue.put({
- "type": "error",
- "error": f"Training completed but failed to save: {e}",
- "stack": traceback.format_exc(limit=20),
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "error",
+ "error": f"Training completed but failed to save: {e}",
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ }
+ )
return
# ββ 11. Done ββ
- event_queue.put({
- "type": "complete",
- "output_dir": output_dir,
- "status_message": "Embedding training completed",
- "ts": time.time(),
- })
+ event_queue.put(
+ {
+ "type": "complete",
+ "output_dir": output_dir,
+ "status_message": "Embedding training completed",
+ "ts": time.time(),
+ }
+ )
diff --git a/studio/backend/loggers/__init__.py b/studio/backend/loggers/__init__.py
index 81ae4948b3..721dde4688 100644
--- a/studio/backend/loggers/__init__.py
+++ b/studio/backend/loggers/__init__.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
from .handlers import get_logger
__all__ = ["get_logger"]
diff --git a/studio/backend/loggers/config.py b/studio/backend/loggers/config.py
index d6247a139a..0d32a64657 100644
--- a/studio/backend/loggers/config.py
+++ b/studio/backend/loggers/config.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
"""Logging configuration for structured logging with structlog.
This module provides centralized logging configuration with environment-specific
@@ -19,6 +22,7 @@ from typing import Optional
import structlog
+
class LogConfig:
"""Structured logging configuration for the application.
@@ -41,9 +45,9 @@ class LogConfig:
log_level = getattr(logging, log_level_name, logging.INFO)
structlog.configure(
- processors=[
+ processors = [
# Reorder processors to control field order
- structlog.processors.TimeStamper(fmt="iso"), # timestamp first
+ structlog.processors.TimeStamper(fmt = "iso"), # timestamp first
structlog.processors.add_log_level, # level second
structlog.contextvars.merge_contextvars,
# Custom processor to flatten the extra field
@@ -59,14 +63,14 @@ class LogConfig:
},
},
(
- structlog.processors.JSONRenderer(sort_keys=False) # Preserve order
+ structlog.processors.JSONRenderer(sort_keys = False) # Preserve order
if env == "production"
else structlog.dev.ConsoleRenderer()
),
],
- wrapper_class=structlog.make_filtering_bound_logger(log_level),
- logger_factory=structlog.PrintLoggerFactory(file=sys.stdout),
- cache_logger_on_first_use=True,
+ wrapper_class = structlog.make_filtering_bound_logger(log_level),
+ logger_factory = structlog.PrintLoggerFactory(file = sys.stdout),
+ cache_logger_on_first_use = True,
)
return structlog.get_logger(service_name)
diff --git a/studio/backend/loggers/handlers.py b/studio/backend/loggers/handlers.py
index 08f692c256..3add92ea1e 100644
--- a/studio/backend/loggers/handlers.py
+++ b/studio/backend/loggers/handlers.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
"""Logging handlers and middleware for structured logging.
This module provides FastAPI middleware and structlog processors for:
@@ -36,21 +39,23 @@ class LoggingMiddleware(BaseHTTPMiddleware):
"/api/train/status",
"/api/train/metrics",
"/api/train/hardware",
- "/api/system"
+ "/api/system",
}
is_excluded = (
request.url.path in EXCLUDED_PATHS
or request.url.path.startswith("/assets/")
- or request.url.path.endswith((".png", ".jpg", ".jpeg", ".ico", ".woff", ".woff2", ".ttf"))
+ or request.url.path.endswith(
+ (".png", ".jpg", ".jpeg", ".ico", ".woff", ".woff2", ".ttf")
+ )
)
if not is_excluded:
logger.info(
"request_completed",
- method=request.method,
- path=request.url.path,
- status_code=response.status_code,
- process_time_ms=round(process_time, 2),
+ method = request.method,
+ path = request.url.path,
+ status_code = response.status_code,
+ process_time_ms = round(process_time, 2),
)
return response
@@ -58,10 +63,10 @@ class LoggingMiddleware(BaseHTTPMiddleware):
except Exception as e:
logger.error(
"request_failed",
- path=request.url.path,
- method=request.method,
- error=str(e),
- exc_info=True,
+ path = request.url.path,
+ method = request.method,
+ error = str(e),
+ exc_info = True,
)
raise
diff --git a/studio/backend/main.py b/studio/backend/main.py
index f0ba4be4f4..46615bb5f0 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -4,9 +4,11 @@
"""
Main FastAPI application for Unsloth UI Backend
"""
+
import os
+
# Suppress annoying C-level dependency warnings globally
-os.environ["PYTHONWARNINGS"] = "ignore"
+os.environ["PYTHONWARNINGS"] = "ignore"
import secrets
import shutil
@@ -54,7 +56,7 @@ async def lifespan(app: FastAPI):
# Version switching now uses .venv_t5/ (pre-installed by setup.sh).
overlay_dir = Path(__file__).resolve().parent.parent.parent / ".venv_overlay"
if overlay_dir.is_dir():
- shutil.rmtree(overlay_dir, ignore_errors=True)
+ shutil.rmtree(overlay_dir, ignore_errors = True)
# Detect hardware first β sets DEVICE global used everywhere
detect_hardware()
@@ -62,12 +64,14 @@ async def lifespan(app: FastAPI):
# Disable flex attention on Blackwell+ GPUs (sm_120 and above)
if get_device() == DeviceType.CUDA:
import torch
+
props = torch.cuda.get_device_properties(0)
sm_version = props.major * 10 + props.minor
if sm_version >= 120:
os.environ["UNSLOTH_ENABLE_FLEX_ATTENTION"] = "0"
import structlog
from loggers import get_logger
+
get_logger(__name__).info(
f"GPU sm_{sm_version} detected β setting UNSLOTH_FLEX_ATTENTION=0"
)
@@ -75,13 +79,16 @@ async def lifespan(app: FastAPI):
# Pre-cache the helper GGUF model for LLM-assisted dataset detection.
# Runs in a background thread so it doesn't block server startup.
import threading
+
def _precache():
try:
from utils.datasets.llm_assist import precache_helper_gguf
+
precache_helper_gguf()
except Exception:
pass # non-critical
- threading.Thread(target=_precache, daemon=True).start()
+
+ threading.Thread(target = _precache, daemon = True).start()
if not storage.is_initialized():
setup_token = secrets.token_urlsafe(32)
@@ -100,10 +107,10 @@ async def lifespan(app: FastAPI):
# Create FastAPI app
app = FastAPI(
- title="Unsloth UI Backend",
- version="1.0.0",
- description="Backend API for Unsloth UI - Training and Model Management",
- lifespan=lifespan,
+ title = "Unsloth UI Backend",
+ version = "1.0.0",
+ description = "Backend API for Unsloth UI - Training and Model Management",
+ lifespan = lifespan,
)
# Initialize structured logging
@@ -111,8 +118,8 @@ from loggers.config import LogConfig
from loggers.handlers import LoggingMiddleware
logger = LogConfig.setup_logging(
- service_name="unsloth-studio-backend",
- env=os.getenv("ENVIRONMENT_TYPE", "production")
+ service_name = "unsloth-studio-backend",
+ env = os.getenv("ENVIRONMENT_TYPE", "production"),
)
app.add_middleware(LoggingMiddleware)
@@ -120,38 +127,39 @@ app.add_middleware(LoggingMiddleware)
# CORS middleware
app.add_middleware(
CORSMiddleware,
- allow_origins=["*"], # In production, specify allowed origins
- allow_credentials=True,
- allow_methods=["*"],
- allow_headers=["*"],
+ allow_origins = ["*"], # In production, specify allowed origins
+ allow_credentials = True,
+ allow_methods = ["*"],
+ allow_headers = ["*"],
)
# ============ Register API Routes ============
# Register routers
-app.include_router(auth_router, prefix="/api/auth", tags=["auth"])
-app.include_router(training_router, prefix="/api/train", tags=["training"])
-app.include_router(models_router, prefix="/api/models", tags=["models"])
-app.include_router(inference_router, prefix="/api/inference", tags=["inference"])
+app.include_router(auth_router, prefix = "/api/auth", tags = ["auth"])
+app.include_router(training_router, prefix = "/api/train", tags = ["training"])
+app.include_router(models_router, prefix = "/api/models", tags = ["models"])
+app.include_router(inference_router, prefix = "/api/inference", tags = ["inference"])
# OpenAI-compatible endpoints: mount the same inference router at /v1
# so external tools (Open WebUI, SillyTavern, etc.) can use the
# standard /v1/chat/completions path.
-app.include_router(inference_router, prefix="/v1", tags=["openai-compat"])
-app.include_router(datasets_router, prefix="/api/datasets", tags=["datasets"])
-app.include_router(data_recipe_router, prefix="/api/data-recipe", tags=["data-recipe"])
-app.include_router(export_router, prefix="/api/export", tags=["export"])
+app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"])
+app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"])
+app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"])
+app.include_router(export_router, prefix = "/api/export", tags = ["export"])
# ============ Health and System Endpoints ============
+
@app.get("/api/health")
async def health_check():
"""Health check endpoint"""
return {
"status": "healthy",
"timestamp": datetime.now().isoformat(),
- "service": "Unsloth UI Backend"
+ "service": "Unsloth UI Backend",
}
@@ -167,11 +175,13 @@ async def get_system_info():
gpu_info = {"available": mem_info.get("available", False), "devices": []}
if mem_info.get("available"):
- gpu_info["devices"].append({
- "index": mem_info.get("device", 0),
- "name": mem_info.get("device_name", "Unknown"),
- "memory_total_gb": round(mem_info.get("total_gb", 0), 2),
- })
+ gpu_info["devices"].append(
+ {
+ "index": mem_info.get("device", 0),
+ "name": mem_info.get("device_name", "Unknown"),
+ "memory_total_gb": round(mem_info.get("total_gb", 0), 2),
+ }
+ )
# CPU & Memory
memory = psutil.virtual_memory()
@@ -203,6 +213,7 @@ async def get_hardware_info():
# ============ Serve Frontend (Optional) ============
+
def setup_frontend(app: FastAPI, build_path: Path):
"""Mount frontend static files (optional)"""
if not build_path.exists():
@@ -211,15 +222,15 @@ def setup_frontend(app: FastAPI, build_path: Path):
# Mount assets
assets_dir = build_path / "assets"
if assets_dir.exists():
- app.mount("/assets", StaticFiles(directory=assets_dir), name="assets")
+ app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
@app.get("/")
async def serve_root():
content = (build_path / "index.html").read_bytes()
return Response(
- content=content,
- media_type="text/html",
- headers={"Cache-Control": "no-cache, no-store, must-revalidate"},
+ content = content,
+ media_type = "text/html",
+ headers = {"Cache-Control": "no-cache, no-store, must-revalidate"},
)
@app.get("/{full_path:path}")
@@ -230,8 +241,8 @@ def setup_frontend(app: FastAPI, build_path: Path):
file_path = (build_path / full_path).resolve()
# Block path traversal β ensure resolved path stays inside build_path
- if not str(file_path).startswith(str(build_path.resolve())):
- return Response(status_code=403)
+ if not file_path.is_relative_to(build_path.resolve()):
+ return Response(status_code = 403)
if file_path.is_file():
return FileResponse(file_path)
@@ -239,10 +250,9 @@ def setup_frontend(app: FastAPI, build_path: Path):
# Serve index.html as bytes β avoids Content-Length mismatch
content = (build_path / "index.html").read_bytes()
return Response(
- content=content,
- media_type="text/html",
- headers={"Cache-Control": "no-cache, no-store, must-revalidate"},
+ content = content,
+ media_type = "text/html",
+ headers = {"Cache-Control": "no-cache, no-store, must-revalidate"},
)
return True
-
diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py
index 13ab08c50b..4a53418439 100644
--- a/studio/backend/models/__init__.py
+++ b/studio/backend/models/__init__.py
@@ -4,6 +4,7 @@
"""
Pydantic models for API request/response schemas
"""
+
from .training import (
TrainingStartRequest,
TrainingJobResponse,
diff --git a/studio/backend/models/auth.py b/studio/backend/models/auth.py
index 47428a791d..c12d15617e 100644
--- a/studio/backend/models/auth.py
+++ b/studio/backend/models/auth.py
@@ -4,28 +4,38 @@
"""
Pydantic schemas for Authentication API
"""
+
from pydantic import BaseModel, Field
class AuthSetupRequest(BaseModel):
"""First-time setup: create the initial admin user + password."""
- setup_token: str = Field(..., description="One-time setup token printed to the server console")
- username: str = Field(..., description="Admin username")
- password: str = Field(..., min_length=8, description="Admin password (minimum 8 characters)")
+
+ setup_token: str = Field(
+ ..., description = "One-time setup token printed to the server console"
+ )
+ username: str = Field(..., description = "Admin username")
+ password: str = Field(
+ ..., min_length = 8, description = "Admin password (minimum 8 characters)"
+ )
class AuthLoginRequest(BaseModel):
"""Login payload: username/password to obtain a JWT."""
- username: str = Field(..., description="Username")
- password: str = Field(..., description="Password")
+
+ username: str = Field(..., description = "Username")
+ password: str = Field(..., description = "Password")
class RefreshTokenRequest(BaseModel):
"""Refresh token payload to obtain new access + refresh tokens."""
- refresh_token: str = Field(..., description="Refresh token from a previous login or refresh")
+
+ refresh_token: str = Field(
+ ..., description = "Refresh token from a previous login or refresh"
+ )
class AuthStatusResponse(BaseModel):
"""Indicate whether auth has been initialized."""
- initialized: bool = Field(..., description="True if auth setup has been completed")
+ initialized: bool = Field(..., description = "True if auth setup has been completed")
diff --git a/studio/backend/models/data_recipe.py b/studio/backend/models/data_recipe.py
index d7758cec36..6992572b00 100644
--- a/studio/backend/models/data_recipe.py
+++ b/studio/backend/models/data_recipe.py
@@ -13,13 +13,13 @@ from pydantic import BaseModel, Field
class RecipePayload(BaseModel):
- recipe: dict[str, Any] = Field(default_factory=dict)
+ recipe: dict[str, Any] = Field(default_factory = dict)
run: dict[str, Any] | None = None
ui: dict[str, Any] | None = None
class PreviewResponse(BaseModel):
- dataset: list[dict[str, Any]] = Field(default_factory=list)
+ dataset: list[dict[str, Any]] = Field(default_factory = list)
processor_artifacts: dict[str, Any] | None = None
analysis: dict[str, Any] | None = None
@@ -32,7 +32,7 @@ class ValidateError(BaseModel):
class ValidateResponse(BaseModel):
valid: bool
- errors: list[ValidateError] = Field(default_factory=list)
+ errors: list[ValidateError] = Field(default_factory = list)
raw_detail: str | None = None
@@ -41,42 +41,42 @@ class JobCreateResponse(BaseModel):
class SeedInspectRequest(BaseModel):
- dataset_name: str = Field(min_length=1)
+ dataset_name: str = Field(min_length = 1)
hf_token: str | None = None
subset: str | None = None
split: str | None = "train"
- preview_size: int = Field(default=10, ge=1, le=50)
+ preview_size: int = Field(default = 10, ge = 1, le = 50)
class SeedInspectUploadRequest(BaseModel):
- filename: str = Field(min_length=1)
- content_base64: str = Field(min_length=1)
- preview_size: int = Field(default=10, ge=1, le=50)
+ filename: str = Field(min_length = 1)
+ content_base64: str = Field(min_length = 1)
+ preview_size: int = Field(default = 10, ge = 1, le = 50)
seed_source_type: str | None = None
- unstructured_chunk_size: int | None = Field(default=None, ge=1, le=20000)
- unstructured_chunk_overlap: int | None = Field(default=None, ge=0, le=20000)
+ unstructured_chunk_size: int | None = Field(default = None, ge = 1, le = 20000)
+ unstructured_chunk_overlap: int | None = Field(default = None, ge = 0, le = 20000)
class SeedInspectResponse(BaseModel):
dataset_name: str
resolved_path: str
- columns: list[str] = Field(default_factory=list)
- preview_rows: list[dict[str, Any]] = Field(default_factory=list)
+ columns: list[str] = Field(default_factory = list)
+ preview_rows: list[dict[str, Any]] = Field(default_factory = list)
split: str | None = None
subset: str | None = None
class McpToolsListRequest(BaseModel):
- mcp_providers: list[dict[str, Any]] = Field(default_factory=list)
- timeout_sec: float | None = Field(default=None, gt=0)
+ mcp_providers: list[dict[str, Any]] = Field(default_factory = list)
+ timeout_sec: float | None = Field(default = None, gt = 0)
class McpToolsProviderResult(BaseModel):
name: str
- tools: list[str] = Field(default_factory=list)
+ tools: list[str] = Field(default_factory = list)
error: str | None = None
class McpToolsListResponse(BaseModel):
- providers: list[McpToolsProviderResult] = Field(default_factory=list)
- duplicate_tools: dict[str, list[str]] = Field(default_factory=dict)
+ providers: list[McpToolsProviderResult] = Field(default_factory = list)
+ duplicate_tools: dict[str, list[str]] = Field(default_factory = dict)
diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py
index 06d555eabd..f20d6f2d15 100644
--- a/studio/backend/models/datasets.py
+++ b/studio/backend/models/datasets.py
@@ -4,6 +4,7 @@
"""
Dataset-related Pydantic models for API requests and responses.
"""
+
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field, model_validator
@@ -11,13 +12,14 @@ from pydantic import BaseModel, Field, model_validator
class CheckFormatRequest(BaseModel):
"""Request for dataset format check"""
+
dataset_name: str # HuggingFace dataset name or local path
is_vlm: bool = False
hf_token: Optional[str] = None
subset: Optional[str] = None
train_split: Optional[str] = "train"
- @model_validator(mode="before")
+ @model_validator(mode = "before")
@classmethod
def _compat_split(cls, values: Any) -> Any:
"""Accept legacy 'split' field as alias for 'train_split'."""
@@ -28,6 +30,7 @@ class CheckFormatRequest(BaseModel):
class CheckFormatResponse(BaseModel):
"""Response for dataset format check"""
+
requires_manual_mapping: bool
detected_format: str
columns: List[str]
@@ -46,6 +49,7 @@ class CheckFormatResponse(BaseModel):
class AiAssistMappingRequest(BaseModel):
"""Request for LLM-assisted column classification (user-triggered)."""
+
columns: List[str]
samples: List[Dict[str, Any]] # Preview rows already loaded in the dialog
dataset_name: Optional[str] = None # For LLM context
@@ -56,6 +60,7 @@ class AiAssistMappingRequest(BaseModel):
class AiAssistMappingResponse(BaseModel):
"""Response from LLM-assisted column classification and conversion advice."""
+
success: bool
suggested_mapping: Optional[Dict[str, str]] = None
warning: Optional[str] = None
@@ -69,8 +74,9 @@ class AiAssistMappingResponse(BaseModel):
class UploadDatasetResponse(BaseModel):
"""Response with stored dataset path for training."""
- filename: str = Field(..., description="Original filename")
- stored_path: str = Field(..., description="Absolute path stored on backend")
+
+ filename: str = Field(..., description = "Original filename")
+ stored_path: str = Field(..., description = "Absolute path stored on backend")
class LocalDatasetItem(BaseModel):
@@ -90,4 +96,4 @@ class LocalDatasetItem(BaseModel):
class LocalDatasetsResponse(BaseModel):
- datasets: List[LocalDatasetItem] = Field(default_factory=list)
+ datasets: List[LocalDatasetItem] = Field(default_factory = list)
diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py
index 9625123721..a86596f199 100644
--- a/studio/backend/models/export.py
+++ b/studio/backend/models/export.py
@@ -4,6 +4,7 @@
"""
Pydantic schemas for Export API.
"""
+
from pydantic import BaseModel, Field
from typing import List, Optional, Literal, Dict, Any
@@ -11,20 +12,20 @@ from typing import List, Optional, Literal, Dict, Any
class LoadCheckpointRequest(BaseModel):
"""Request for loading a checkpoint into the export backend."""
- checkpoint_path: str = Field(..., description="Path to the checkpoint directory")
+ checkpoint_path: str = Field(..., description = "Path to the checkpoint directory")
max_seq_length: int = Field(
2048,
- ge=128,
- le=32768,
- description="Maximum sequence length for loading the model",
+ ge = 128,
+ le = 32768,
+ description = "Maximum sequence length for loading the model",
)
load_in_4bit: bool = Field(
True,
- description="Whether to load the model in 4-bit quantization",
+ description = "Whether to load the model in 4-bit quantization",
)
trust_remote_code: bool = Field(
False,
- description="Allow loading models with custom code. Only enable for checkpoints/base models you trust.",
+ description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.",
)
@@ -33,26 +34,26 @@ class ExportStatusResponse(BaseModel):
current_checkpoint: Optional[str] = Field(
None,
- description="Path to the currently loaded checkpoint, if any",
+ description = "Path to the currently loaded checkpoint, if any",
)
is_vision: bool = Field(
False,
- description="True if the loaded checkpoint is a vision model",
+ description = "True if the loaded checkpoint is a vision model",
)
is_peft: bool = Field(
False,
- description="True if the loaded checkpoint is a PEFT (LoRA) model",
+ description = "True if the loaded checkpoint is a PEFT (LoRA) model",
)
class ExportOperationResponse(BaseModel):
"""Generic response for export operations."""
- success: bool = Field(..., description="True if the operation succeeded")
- message: str = Field(..., description="Human-readable status or error message")
+ success: bool = Field(..., description = "True if the operation succeeded")
+ message: str = Field(..., description = "Human-readable status or error message")
details: Optional[Dict[str, Any]] = Field(
- default=None,
- description="Optional extra details about the operation",
+ default = None,
+ description = "Optional extra details about the operation",
)
@@ -61,27 +62,27 @@ class ExportCommonOptions(BaseModel):
save_directory: str = Field(
...,
- description="Local directory where the exported artifacts will be written",
+ description = "Local directory where the exported artifacts will be written",
)
push_to_hub: bool = Field(
False,
- description="If True, also push the exported model to the Hugging Face Hub",
+ description = "If True, also push the exported model to the Hugging Face Hub",
)
repo_id: Optional[str] = Field(
None,
- description="Hugging Face Hub repository ID (username/model-name)",
+ description = "Hugging Face Hub repository ID (username/model-name)",
)
hf_token: Optional[str] = Field(
None,
- description="Hugging Face access token used for Hub operations",
+ description = "Hugging Face access token used for Hub operations",
)
private: bool = Field(
False,
- description="If True, create a private repository on the Hub (where applicable)",
+ description = "If True, create a private repository on the Hub (where applicable)",
)
base_model_id: Optional[str] = Field(
None,
- description="HuggingFace model ID of the base model (for model card metadata)",
+ description = "HuggingFace model ID of the base model (for model card metadata)",
)
@@ -90,7 +91,7 @@ class ExportMergedModelRequest(ExportCommonOptions):
format_type: Literal["16-bit (FP16)", "4-bit (FP4)"] = Field(
"16-bit (FP16)",
- description="Export precision / format for the merged model",
+ description = "Export precision / format for the merged model",
)
@@ -98,7 +99,6 @@ class ExportBaseModelRequest(ExportCommonOptions):
"""Request for exporting a non-PEFT (base) model."""
# Uses fields from ExportCommonOptions only
- pass
class ExportGGUFRequest(BaseModel):
@@ -106,23 +106,23 @@ class ExportGGUFRequest(BaseModel):
save_directory: str = Field(
...,
- description="Directory where GGUF files will be saved",
+ description = "Directory where GGUF files will be saved",
)
quantization_method: str = Field(
"Q4_K_M",
- description='GGUF quantization method (e.g. "Q4_K_M")',
+ description = 'GGUF quantization method (e.g. "Q4_K_M")',
)
push_to_hub: bool = Field(
False,
- description="If True, also push GGUF artifacts to the Hugging Face Hub",
+ description = "If True, also push GGUF artifacts to the Hugging Face Hub",
)
repo_id: Optional[str] = Field(
None,
- description="Hugging Face Hub repository ID for GGUF upload",
+ description = "Hugging Face Hub repository ID for GGUF upload",
)
hf_token: Optional[str] = Field(
None,
- description="Hugging Face token for GGUF upload",
+ description = "Hugging Face token for GGUF upload",
)
@@ -130,6 +130,3 @@ class ExportLoRAAdapterRequest(ExportCommonOptions):
"""Request for exporting only the LoRA adapter (not merged)."""
# Uses fields from ExportCommonOptions only
- pass
-
-
diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index e9e581d1bd..46925fb3db 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -4,6 +4,7 @@
"""
Pydantic schemas for Inference API
"""
+
from __future__ import annotations
import time
@@ -15,21 +16,29 @@ from pydantic import BaseModel, Discriminator, Field, Tag
class LoadRequest(BaseModel):
"""Request to load a model for inference"""
- model_path: str = Field(..., description="Model identifier or local path")
- hf_token: Optional[str] = Field(None, description="HuggingFace token for gated models")
- max_seq_length: int = Field(4096, ge=128, le=32768, description="Maximum sequence length")
- load_in_4bit: bool = Field(True, description="Load model in 4-bit quantization")
- is_lora: bool = Field(False, description="Whether this is a LoRA adapter")
- gguf_variant: Optional[str] = Field(None, description="GGUF quantization variant (e.g. 'Q4_K_M')")
+
+ model_path: str = Field(..., description = "Model identifier or local path")
+ hf_token: Optional[str] = Field(
+ None, description = "HuggingFace token for gated models"
+ )
+ max_seq_length: int = Field(
+ 4096, ge = 128, le = 32768, description = "Maximum sequence length"
+ )
+ load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization")
+ is_lora: bool = Field(False, description = "Whether this is a LoRA adapter")
+ gguf_variant: Optional[str] = Field(
+ None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
+ )
trust_remote_code: bool = Field(
False,
- description="Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
+ description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
)
class UnloadRequest(BaseModel):
"""Request to unload a model"""
- model_path: str = Field(..., description="Model identifier to unload")
+
+ model_path: str = Field(..., description = "Model identifier to unload")
class ValidateModelRequest(BaseModel):
@@ -39,9 +48,14 @@ class ValidateModelRequest(BaseModel):
This does NOT actually load weights into GPU memory.
"""
- model_path: str = Field(..., description="Model identifier or local path")
- hf_token: Optional[str] = Field(None, description="HuggingFace token for gated models")
- gguf_variant: Optional[str] = Field(None, description="GGUF quantization variant (e.g. 'Q4_K_M')")
+
+ model_path: str = Field(..., description = "Model identifier or local path")
+ hf_token: Optional[str] = Field(
+ None, description = "HuggingFace token for gated models"
+ )
+ gguf_variant: Optional[str] = Field(
+ None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
+ )
class ValidateModelResponse(BaseModel):
@@ -51,58 +65,99 @@ class ValidateModelResponse(BaseModel):
valid == True means ModelConfig.from_identifier() succeeded and basic
introspection (GGUF / LoRA / vision flags) is available.
"""
- valid: bool = Field(..., description="Whether the model identifier looks valid")
- message: str = Field(..., description="Human-readable validation message")
- identifier: Optional[str] = Field(None, description="Resolved model identifier")
- display_name: Optional[str] = Field(None, description="Display name derived from identifier")
- is_gguf: bool = Field(False, description="Whether this is a GGUF model (llama.cpp)")
- is_lora: bool = Field(False, description="Whether this is a LoRA adapter")
- is_vision: bool = Field(False, description="Whether this is a vision-capable model")
+
+ valid: bool = Field(..., description = "Whether the model identifier looks valid")
+ message: str = Field(..., description = "Human-readable validation message")
+ identifier: Optional[str] = Field(None, description = "Resolved model identifier")
+ display_name: Optional[str] = Field(
+ None, description = "Display name derived from identifier"
+ )
+ is_gguf: bool = Field(False, description = "Whether this is a GGUF model (llama.cpp)")
+ is_lora: bool = Field(False, description = "Whether this is a LoRA adapter")
+ is_vision: bool = Field(False, description = "Whether this is a vision-capable model")
class GenerateRequest(BaseModel):
"""Request for text generation (legacy /generate/stream endpoint)"""
- messages: List[dict] = Field(..., description="Chat messages in OpenAI format")
- system_prompt: str = Field("You are a helpful AI assistant.", description="System prompt")
- temperature: float = Field(0.7, ge=0.0, le=2.0, description="Sampling temperature")
- top_p: float = Field(0.9, ge=0.0, le=1.0, description="Top-p sampling")
- top_k: int = Field(40, ge=-1, le=100, description="Top-k sampling")
- max_new_tokens: int = Field(2048, ge=1, le=4096, description="Maximum tokens to generate")
- repetition_penalty: float = Field(1.1, ge=1.0, le=2.0, description="Repetition penalty")
- image_base64: Optional[str] = Field(None, description="Base64 encoded image for vision models")
+
+ messages: List[dict] = Field(..., description = "Chat messages in OpenAI format")
+ system_prompt: str = Field(
+ "You are a helpful AI assistant.", description = "System prompt"
+ )
+ temperature: float = Field(0.7, ge = 0.0, le = 2.0, description = "Sampling temperature")
+ top_p: float = Field(0.9, ge = 0.0, le = 1.0, description = "Top-p sampling")
+ top_k: int = Field(40, ge = -1, le = 100, description = "Top-k sampling")
+ max_new_tokens: int = Field(
+ 2048, ge = 1, le = 4096, description = "Maximum tokens to generate"
+ )
+ repetition_penalty: float = Field(
+ 1.1, ge = 1.0, le = 2.0, description = "Repetition penalty"
+ )
+ image_base64: Optional[str] = Field(
+ None, description = "Base64 encoded image for vision models"
+ )
class LoadResponse(BaseModel):
"""Response after loading a model"""
- status: str = Field(..., description="Load status")
- model: str = Field(..., description="Model identifier")
- display_name: str = Field(..., description="Display name of the model")
- is_vision: bool = Field(False, description="Whether model is a vision model")
- is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
- is_gguf: bool = Field(False, description="Whether model is a GGUF model (llama.cpp)")
- is_audio: bool = Field(False, description="Whether model is a TTS audio model")
- audio_type: Optional[str] = Field(None, description="Audio codec type: snac, csm, bicodec, dac")
- has_audio_input: bool = Field(False, description="Whether model accepts audio input (ASR)")
- inference: dict = Field(..., description="Inference parameters (temperature, top_p, top_k, min_p)")
+
+ status: str = Field(..., description = "Load status")
+ model: str = Field(..., description = "Model identifier")
+ display_name: str = Field(..., description = "Display name of the model")
+ is_vision: bool = Field(False, description = "Whether model is a vision model")
+ is_lora: bool = Field(False, description = "Whether model is a LoRA adapter")
+ is_gguf: bool = Field(
+ False, description = "Whether model is a GGUF model (llama.cpp)"
+ )
+ is_audio: bool = Field(False, description = "Whether model is a TTS audio model")
+ audio_type: Optional[str] = Field(
+ None, description = "Audio codec type: snac, csm, bicodec, dac"
+ )
+ has_audio_input: bool = Field(
+ False, description = "Whether model accepts audio input (ASR)"
+ )
+ inference: dict = Field(
+ ..., description = "Inference parameters (temperature, top_p, top_k, min_p)"
+ )
class UnloadResponse(BaseModel):
"""Response after unloading a model"""
- status: str = Field(..., description="Unload status")
- model: str = Field(..., description="Model identifier that was unloaded")
+
+ status: str = Field(..., description = "Unload status")
+ model: str = Field(..., description = "Model identifier that was unloaded")
class InferenceStatusResponse(BaseModel):
"""Current inference backend status"""
- active_model: Optional[str] = Field(None, description="Currently active model identifier")
- is_vision: bool = Field(False, description="Whether the active model is a vision model")
- is_gguf: bool = Field(False, description="Whether the active model is a GGUF model (llama.cpp)")
- gguf_variant: Optional[str] = Field(None, description="GGUF quantization variant (e.g. Q4_K_M)")
- is_audio: bool = Field(False, description="Whether the active model is a TTS audio model")
- audio_type: Optional[str] = Field(None, description="Audio codec type: snac, csm, bicodec, dac")
- has_audio_input: bool = Field(False, description="Whether model accepts audio input (ASR)")
- loading: List[str] = Field(default_factory=list, description="Models currently being loaded")
- loaded: List[str] = Field(default_factory=list, description="Models currently loaded")
+
+ active_model: Optional[str] = Field(
+ None, description = "Currently active model identifier"
+ )
+ is_vision: bool = Field(
+ False, description = "Whether the active model is a vision model"
+ )
+ is_gguf: bool = Field(
+ False, description = "Whether the active model is a GGUF model (llama.cpp)"
+ )
+ gguf_variant: Optional[str] = Field(
+ None, description = "GGUF quantization variant (e.g. Q4_K_M)"
+ )
+ is_audio: bool = Field(
+ False, description = "Whether the active model is a TTS audio model"
+ )
+ audio_type: Optional[str] = Field(
+ None, description = "Audio codec type: snac, csm, bicodec, dac"
+ )
+ has_audio_input: bool = Field(
+ False, description = "Whether model accepts audio input (ASR)"
+ )
+ loading: List[str] = Field(
+ default_factory = list, description = "Models currently being loaded"
+ )
+ loaded: List[str] = Field(
+ default_factory = list, description = "Models currently loaded"
+ )
# =====================================================================
@@ -112,20 +167,24 @@ class InferenceStatusResponse(BaseModel):
# ββ Multimodal content parts (OpenAI vision format) ββββββββββββββ
+
class TextContentPart(BaseModel):
"""Text content part in a multimodal message."""
+
type: Literal["text"]
text: str
class ImageUrl(BaseModel):
"""Image URL object β supports data URIs and remote URLs."""
- url: str = Field(..., description="data:image/png;base64,... or https://...")
+
+ url: str = Field(..., description = "data:image/png;base64,... or https://...")
detail: Optional[Literal["auto", "low", "high"]] = "auto"
class ImageContentPart(BaseModel):
"""Image content part in a multimodal message."""
+
type: Literal["image_url"]
image_url: ImageUrl
@@ -148,6 +207,7 @@ ContentPart = Annotated[
# ββ Messages βββββββββββββββββββββββββββββββββββββββββββββββββββββ
+
class ChatMessage(BaseModel):
"""
A single message in the conversation.
@@ -155,8 +215,13 @@ class ChatMessage(BaseModel):
``content`` may be a plain string (text-only) or a list of
content parts for multimodal messages (OpenAI vision format).
"""
- role: Literal["system", "user", "assistant"] = Field(..., description="Message role")
- content: Union[str, list[ContentPart]] = Field(..., description="Message content (string or multimodal parts)")
+
+ role: Literal["system", "user", "assistant"] = Field(
+ ..., description = "Message role"
+ )
+ content: Union[str, list[ContentPart]] = Field(
+ ..., description = "Message content (string or multimodal parts)"
+ )
class ChatCompletionRequest(BaseModel):
@@ -165,22 +230,36 @@ class ChatCompletionRequest(BaseModel):
Extensions (non-OpenAI fields) are marked with 'x-unsloth'.
"""
- model: str = Field("default", description="Model identifier (informational; the active model is used)")
- messages: list[ChatMessage] = Field(..., description="Conversation messages")
- stream: bool = Field(True, description="Whether to stream the response via SSE")
- temperature: float = Field(0.7, ge=0.0, le=2.0)
- top_p: float = Field(0.9, ge=0.0, le=1.0)
- max_tokens: Optional[int] = Field(2048, ge=1, le=4096, description="Maximum tokens to generate")
+
+ model: str = Field(
+ "default",
+ description = "Model identifier (informational; the active model is used)",
+ )
+ messages: list[ChatMessage] = Field(..., description = "Conversation messages")
+ stream: bool = Field(True, description = "Whether to stream the response via SSE")
+ temperature: float = Field(0.7, ge = 0.0, le = 2.0)
+ top_p: float = Field(0.9, ge = 0.0, le = 1.0)
+ max_tokens: Optional[int] = Field(
+ 2048, ge = 1, le = 4096, description = "Maximum tokens to generate"
+ )
# ββ Unsloth extensions (ignored by standard OpenAI clients) ββ
- top_k: int = Field(40, ge=-1, le=100, description="[x-unsloth] Top-k sampling")
- min_p: float = Field(0.0, ge=0.0, le=1.0, description="[x-unsloth] Min-p sampling threshold")
- repetition_penalty: float = Field(1.1, ge=1.0, le=2.0, description="[x-unsloth] Repetition penalty")
- image_base64: Optional[str] = Field(None, description="[x-unsloth] Base64-encoded image for vision models")
- audio_base64: Optional[str] = Field(None, description="[x-unsloth] Base64-encoded WAV for audio-input models (ASR)")
+ top_k: int = Field(40, ge = -1, le = 100, description = "[x-unsloth] Top-k sampling")
+ min_p: float = Field(
+ 0.0, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold"
+ )
+ repetition_penalty: float = Field(
+ 1.1, ge = 1.0, le = 2.0, description = "[x-unsloth] Repetition penalty"
+ )
+ image_base64: Optional[str] = Field(
+ None, description = "[x-unsloth] Base64-encoded image for vision models"
+ )
+ audio_base64: Optional[str] = Field(
+ None, description = "[x-unsloth] Base64-encoded WAV for audio-input models (ASR)"
+ )
use_adapter: Optional[Union[bool, str]] = Field(
None,
- description=(
+ description = (
"[x-unsloth] Adapter control for compare mode. "
"null = no change (default), "
"false = disable adapters (base model), "
@@ -195,12 +274,14 @@ class ChatCompletionRequest(BaseModel):
class ChoiceDelta(BaseModel):
"""Delta content for a streaming chunk."""
+
role: Optional[str] = None
content: Optional[str] = None
class ChunkChoice(BaseModel):
"""A single choice in a streaming chunk."""
+
index: int = 0
delta: ChoiceDelta
finish_reason: Optional[Literal["stop", "length"]] = None
@@ -208,9 +289,10 @@ class ChunkChoice(BaseModel):
class ChatCompletionChunk(BaseModel):
"""A single SSE chunk in OpenAI streaming format."""
- id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
+
+ id: str = Field(default_factory = lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
- created: int = Field(default_factory=lambda: int(time.time()))
+ created: int = Field(default_factory = lambda: int(time.time()))
model: str = "default"
choices: list[ChunkChoice]
@@ -220,12 +302,14 @@ class ChatCompletionChunk(BaseModel):
class CompletionMessage(BaseModel):
"""The assistant's complete response message."""
+
role: Literal["assistant"] = "assistant"
content: str
class CompletionChoice(BaseModel):
"""A single choice in a non-streaming response."""
+
index: int = 0
message: CompletionMessage
finish_reason: Literal["stop", "length"] = "stop"
@@ -233,6 +317,7 @@ class CompletionChoice(BaseModel):
class CompletionUsage(BaseModel):
"""Token usage statistics (approximate)."""
+
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
@@ -240,9 +325,10 @@ class CompletionUsage(BaseModel):
class ChatCompletion(BaseModel):
"""Non-streaming chat completion response."""
- id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
+
+ id: str = Field(default_factory = lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
object: Literal["chat.completion"] = "chat.completion"
- created: int = Field(default_factory=lambda: int(time.time()))
+ created: int = Field(default_factory = lambda: int(time.time()))
model: str = "default"
choices: list[CompletionChoice]
- usage: CompletionUsage = Field(default_factory=CompletionUsage)
+ usage: CompletionUsage = Field(default_factory = CompletionUsage)
diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py
index 0f87d4543c..6634f69385 100644
--- a/studio/backend/models/models.py
+++ b/studio/backend/models/models.py
@@ -4,6 +4,7 @@
"""
Pydantic schemas for Model Management API
"""
+
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any, Literal
@@ -13,123 +14,169 @@ ModelType = Literal["text", "vision", "audio", "embeddings"]
class CheckpointInfo(BaseModel):
"""Information about a discovered checkpoint directory."""
- display_name: str = Field(..., description="User-friendly checkpoint name (folder name)")
- path: str = Field(..., description="Full path to the checkpoint directory")
- loss: Optional[float] = Field(None, description="Training loss at this checkpoint")
+ display_name: str = Field(
+ ..., description = "User-friendly checkpoint name (folder name)"
+ )
+ path: str = Field(..., description = "Full path to the checkpoint directory")
+ loss: Optional[float] = Field(None, description = "Training loss at this checkpoint")
class ModelCheckpoints(BaseModel):
"""A training run and its associated checkpoints."""
- name: str = Field(..., description="Training run folder name")
+ name: str = Field(..., description = "Training run folder name")
checkpoints: List[CheckpointInfo] = Field(
- default_factory=list,
- description="List of checkpoints for this training run (final + intermediate)",
+ default_factory = list,
+ description = "List of checkpoints for this training run (final + intermediate)",
)
base_model: Optional[str] = Field(
None,
- description="Base model name from adapter_config.json or config.json",
+ description = "Base model name from adapter_config.json or config.json",
)
peft_type: Optional[str] = Field(
None,
- description="PEFT type (e.g. LORA) if adapter training, None for full fine-tune",
+ description = "PEFT type (e.g. LORA) if adapter training, None for full fine-tune",
)
lora_rank: Optional[int] = Field(
None,
- description="LoRA rank (r) if applicable",
+ description = "LoRA rank (r) if applicable",
)
class CheckpointListResponse(BaseModel):
"""Response for listing available checkpoints in an outputs directory."""
- outputs_dir: str = Field(..., description="Directory that was scanned")
+ outputs_dir: str = Field(..., description = "Directory that was scanned")
models: List[ModelCheckpoints] = Field(
- default_factory=list,
- description="List of training runs with their checkpoints",
+ default_factory = list,
+ description = "List of training runs with their checkpoints",
)
class ModelDetails(BaseModel):
"""Detailed model configuration and metadata - can be used for both list and detail views"""
- id: str = Field(..., description="Model identifier")
- model_name: Optional[str] = Field(None, description="Model identifier (alias for id, for backward compatibility)")
- name: Optional[str] = Field(None, description="Display name for the model")
- config: Optional[Dict[str, Any]] = Field(None, description="Model configuration dictionary")
- is_vision: bool = Field(False, description="Whether model is a vision model")
- is_embedding: bool = Field(False, description="Whether model is an embedding/sentence-transformer model")
- is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
- is_gguf: bool = Field(False, description="Whether model is a GGUF model (llama.cpp format)")
- is_audio: bool = Field(False, description="Whether model is a TTS audio model")
- audio_type: Optional[str] = Field(None, description="Audio codec type: snac, csm, bicodec, dac")
- has_audio_input: bool = Field(False, description="Whether model accepts audio input (ASR)")
- model_type: Optional[ModelType] = Field(None, description="Collapsed model modality: text, vision, audio, or embeddings")
- base_model: Optional[str] = Field(None, description="Base model if this is a LoRA adapter")
+
+ id: str = Field(..., description = "Model identifier")
+ model_name: Optional[str] = Field(
+ None, description = "Model identifier (alias for id, for backward compatibility)"
+ )
+ name: Optional[str] = Field(None, description = "Display name for the model")
+ config: Optional[Dict[str, Any]] = Field(
+ None, description = "Model configuration dictionary"
+ )
+ is_vision: bool = Field(False, description = "Whether model is a vision model")
+ is_embedding: bool = Field(
+ False, description = "Whether model is an embedding/sentence-transformer model"
+ )
+ is_lora: bool = Field(False, description = "Whether model is a LoRA adapter")
+ is_gguf: bool = Field(
+ False, description = "Whether model is a GGUF model (llama.cpp format)"
+ )
+ is_audio: bool = Field(False, description = "Whether model is a TTS audio model")
+ audio_type: Optional[str] = Field(
+ None, description = "Audio codec type: snac, csm, bicodec, dac"
+ )
+ has_audio_input: bool = Field(
+ False, description = "Whether model accepts audio input (ASR)"
+ )
+ model_type: Optional[ModelType] = Field(
+ None, description = "Collapsed model modality: text, vision, audio, or embeddings"
+ )
+ base_model: Optional[str] = Field(
+ None, description = "Base model if this is a LoRA adapter"
+ )
class LoRAInfo(BaseModel):
"""LoRA adapter or exported model information"""
- display_name: str = Field(..., description="Display name for the LoRA")
- adapter_path: str = Field(..., description="Path to the LoRA adapter or exported model")
- base_model: Optional[str] = Field(None, description="Base model identifier")
- source: Optional[str] = Field(None, description="'training' or 'exported'")
- export_type: Optional[str] = Field(None, description="'lora', 'merged', or 'gguf' (for exports)")
+
+ display_name: str = Field(..., description = "Display name for the LoRA")
+ adapter_path: str = Field(
+ ..., description = "Path to the LoRA adapter or exported model"
+ )
+ base_model: Optional[str] = Field(None, description = "Base model identifier")
+ source: Optional[str] = Field(None, description = "'training' or 'exported'")
+ export_type: Optional[str] = Field(
+ None, description = "'lora', 'merged', or 'gguf' (for exports)"
+ )
class LoRAScanResponse(BaseModel):
"""Response schema for scanning trained LoRA adapters"""
- loras: List[LoRAInfo] = Field(default_factory=list, description="List of found LoRA adapters")
- outputs_dir: str = Field(..., description="Directory that was scanned")
+
+ loras: List[LoRAInfo] = Field(
+ default_factory = list, description = "List of found LoRA adapters"
+ )
+ outputs_dir: str = Field(..., description = "Directory that was scanned")
class ModelListResponse(BaseModel):
"""Response schema for listing models"""
- models: List[ModelDetails] = Field(default_factory=list, description="List of models")
- default_models: List[str] = Field(default_factory=list, description="List of default model IDs")
+
+ models: List[ModelDetails] = Field(
+ default_factory = list, description = "List of models"
+ )
+ default_models: List[str] = Field(
+ default_factory = list, description = "List of default model IDs"
+ )
class GgufVariantDetail(BaseModel):
"""A single GGUF quantization variant in a HuggingFace repo."""
- filename: str = Field(..., description="GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')")
- quant: str = Field(..., description="Quantization label (e.g., 'Q4_K_M')")
- size_bytes: int = Field(0, description="File size in bytes")
+
+ filename: str = Field(
+ ..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')"
+ )
+ quant: str = Field(..., description = "Quantization label (e.g., 'Q4_K_M')")
+ size_bytes: int = Field(0, description = "File size in bytes")
class GgufVariantsResponse(BaseModel):
"""Response for listing GGUF quantization variants in a HuggingFace repo."""
- repo_id: str = Field(..., description="HuggingFace repo ID")
- variants: List[GgufVariantDetail] = Field(default_factory=list, description="Available GGUF variants")
- has_vision: bool = Field(False, description="Whether the model has vision support (mmproj files)")
- default_variant: Optional[str] = Field(None, description="Recommended default quantization variant")
+
+ repo_id: str = Field(..., description = "HuggingFace repo ID")
+ variants: List[GgufVariantDetail] = Field(
+ default_factory = list, description = "Available GGUF variants"
+ )
+ has_vision: bool = Field(
+ False, description = "Whether the model has vision support (mmproj files)"
+ )
+ default_variant: Optional[str] = Field(
+ None, description = "Recommended default quantization variant"
+ )
class LocalModelInfo(BaseModel):
"""Discovered local model candidate."""
- id: str = Field(..., description="Identifier to use for loading/training")
- display_name: str = Field(..., description="Display label")
- path: str = Field(..., description="Local path where model data was discovered")
+
+ id: str = Field(..., description = "Identifier to use for loading/training")
+ display_name: str = Field(..., description = "Display label")
+ path: str = Field(..., description = "Local path where model data was discovered")
source: Literal["models_dir", "hf_cache"] = Field(
...,
- description="Discovery source",
+ description = "Discovery source",
)
model_id: Optional[str] = Field(
None,
- description="HF repo id for cached models, e.g. org/model",
+ description = "HF repo id for cached models, e.g. org/model",
)
updated_at: Optional[float] = Field(
None,
- description="Unix timestamp of latest observed update",
+ description = "Unix timestamp of latest observed update",
)
class LocalModelListResponse(BaseModel):
"""Response schema for listing local/cached models."""
- models_dir: str = Field(..., description="Directory scanned for custom local models")
+
+ models_dir: str = Field(
+ ..., description = "Directory scanned for custom local models"
+ )
hf_cache_dir: Optional[str] = Field(
None,
- description="HF cache root that was scanned",
+ description = "HF cache root that was scanned",
)
models: List[LocalModelInfo] = Field(
- default_factory=list,
- description="Discovered local/cached models",
+ default_factory = list,
+ description = "Discovered local/cached models",
)
diff --git a/studio/backend/models/responses.py b/studio/backend/models/responses.py
index 9e426d2158..3081f67422 100644
--- a/studio/backend/models/responses.py
+++ b/studio/backend/models/responses.py
@@ -5,45 +5,63 @@
Pydantic response schemas for endpoints that previously returned raw dicts.
These are small response models for training and model management routes.
"""
+
from pydantic import BaseModel, Field
from typing import Optional, List
# --- Training route response models ---
+
class TrainingStopResponse(BaseModel):
"""Response for stopping a training job"""
- status: str = Field(..., description="Current status: 'stopped' or 'idle'")
- message: str = Field(..., description="Human-readable status message")
+
+ status: str = Field(..., description = "Current status: 'stopped' or 'idle'")
+ message: str = Field(..., description = "Human-readable status message")
class TrainingMetricsResponse(BaseModel):
"""Response for training metrics history"""
- loss_history: List[float] = Field(default_factory=list, description="Loss values per step")
- lr_history: List[float] = Field(default_factory=list, description="Learning rate per step")
- step_history: List[int] = Field(default_factory=list, description="Step numbers")
- grad_norm_history: List[float] = Field(default_factory=list, description="Gradient norm values")
- grad_norm_step_history: List[int] = Field(default_factory=list, description="Step numbers for gradient norm values")
- current_loss: Optional[float] = Field(None, description="Most recent loss value")
- current_lr: Optional[float] = Field(None, description="Most recent learning rate")
- current_step: Optional[int] = Field(None, description="Most recent step number")
+
+ loss_history: List[float] = Field(
+ default_factory = list, description = "Loss values per step"
+ )
+ lr_history: List[float] = Field(
+ default_factory = list, description = "Learning rate per step"
+ )
+ step_history: List[int] = Field(default_factory = list, description = "Step numbers")
+ grad_norm_history: List[float] = Field(
+ default_factory = list, description = "Gradient norm values"
+ )
+ grad_norm_step_history: List[int] = Field(
+ default_factory = list, description = "Step numbers for gradient norm values"
+ )
+ current_loss: Optional[float] = Field(None, description = "Most recent loss value")
+ current_lr: Optional[float] = Field(None, description = "Most recent learning rate")
+ current_step: Optional[int] = Field(None, description = "Most recent step number")
# --- Model management route response models ---
+
class LoRABaseModelResponse(BaseModel):
"""Response for getting a LoRA's base model"""
- lora_path: str = Field(..., description="Path to the LoRA adapter")
- base_model: str = Field(..., description="Base model identifier")
+
+ lora_path: str = Field(..., description = "Path to the LoRA adapter")
+ base_model: str = Field(..., description = "Base model identifier")
class VisionCheckResponse(BaseModel):
"""Response for checking if a model is a vision model"""
- model_name: str = Field(..., description="Model identifier")
- is_vision: bool = Field(..., description="Whether the model is a vision model")
+
+ model_name: str = Field(..., description = "Model identifier")
+ is_vision: bool = Field(..., description = "Whether the model is a vision model")
class EmbeddingCheckResponse(BaseModel):
"""Response for checking if a model is an embedding model"""
- model_name: str = Field(..., description="Model identifier")
- is_embedding: bool = Field(..., description="Whether the model is an embedding/sentence-transformer model")
+
+ model_name: str = Field(..., description = "Model identifier")
+ is_embedding: bool = Field(
+ ..., description = "Whether the model is an embedding/sentence-transformer model"
+ )
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 8ea9eeec42..35bd9744a9 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -4,44 +4,63 @@
"""
Pydantic schemas for Training API
"""
+
from pydantic import BaseModel, Field, model_validator
from typing import Any, Optional, List, Dict, Literal
class TrainingStartRequest(BaseModel):
"""Request schema for starting training"""
+
# Model parameters
- model_name: str = Field(..., description="Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')")
- training_type: str = Field(..., description="Training type: 'LoRA/QLoRA' or 'Full Finetuning'")
- hf_token: Optional[str] = Field(None, description="HuggingFace token")
- load_in_4bit: bool = Field(True, description="Load model in 4-bit quantization")
- max_seq_length: int = Field(2048, description="Maximum sequence length")
+ model_name: str = Field(
+ ..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')"
+ )
+ training_type: str = Field(
+ ..., description = "Training type: 'LoRA/QLoRA' or 'Full Finetuning'"
+ )
+ hf_token: Optional[str] = Field(None, description = "HuggingFace token")
+ load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization")
+ max_seq_length: int = Field(2048, description = "Maximum sequence length")
trust_remote_code: bool = Field(
False,
- description="Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
+ description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
)
# Dataset parameters
- hf_dataset: Optional[str] = Field(None, description="HuggingFace dataset identifier")
- local_datasets: List[str] = Field(default_factory=list, description="List of local dataset paths")
- format_type: str = Field(..., description="Dataset format type")
+ hf_dataset: Optional[str] = Field(
+ None, description = "HuggingFace dataset identifier"
+ )
+ local_datasets: List[str] = Field(
+ default_factory = list, description = "List of local dataset paths"
+ )
+ format_type: str = Field(..., description = "Dataset format type")
subset: Optional[str] = None
- train_split: Optional[str] = Field("train", description="Training split name")
- eval_split: Optional[str] = Field(None, description="Eval split name. None = auto-detect")
- eval_steps: float = Field(0.00, description="Fraction of total steps between evals (0-1)")
- dataset_slice_start: Optional[int] = Field(None, description="Inclusive start row index for dataset slicing")
- dataset_slice_end: Optional[int] = Field(None, description="Inclusive end row index for dataset slicing")
+ train_split: Optional[str] = Field("train", description = "Training split name")
+ eval_split: Optional[str] = Field(
+ None, description = "Eval split name. None = auto-detect"
+ )
+ eval_steps: float = Field(
+ 0.00, description = "Fraction of total steps between evals (0-1)"
+ )
+ dataset_slice_start: Optional[int] = Field(
+ None, description = "Inclusive start row index for dataset slicing"
+ )
+ dataset_slice_end: Optional[int] = Field(
+ None, description = "Inclusive end row index for dataset slicing"
+ )
- @model_validator(mode="before")
+ @model_validator(mode = "before")
@classmethod
def _compat_split(cls, values: Any) -> Any:
"""Accept legacy 'split' field as alias for 'train_split'."""
if isinstance(values, dict) and "split" in values:
values.setdefault("train_split", values.pop("split"))
return values
+
custom_format_mapping: Optional[Dict[str, Any]] = Field(
None,
- description=(
+ description = (
"User-provided column-to-role mapping, e.g. {'image': 'image', 'caption': 'text'} "
"for VLM or {'instruction': 'user', 'output': 'assistant'} for LLM. "
"Enhanced format includes __system_prompt, __user_template, "
@@ -49,59 +68,77 @@ class TrainingStartRequest(BaseModel):
),
)
# Training parameters
- num_epochs: int = Field(1, description="Number of training epochs")
- learning_rate: str = Field("2e-4", description="Learning rate")
- batch_size: int = Field(1, description="Batch size")
- gradient_accumulation_steps: int = Field(1, description="Gradient accumulation steps")
- warmup_steps: Optional[int] = Field(None, description="Warmup steps")
- warmup_ratio: Optional[float] = Field(None, description="Warmup ratio")
- max_steps: Optional[int] = Field(None, description="Maximum training steps")
- save_steps: int = Field(100, description="Steps between checkpoints")
- weight_decay: float = Field(0.01, description="Weight decay")
- random_seed: int = Field(42, description="Random seed")
- packing: bool = Field(False, description="Enable sequence packing")
- optim: str = Field("adamw_8bit", description="Optimizer")
- lr_scheduler_type: str = Field("linear", description="Learning rate scheduler type")
+ num_epochs: int = Field(1, description = "Number of training epochs")
+ learning_rate: str = Field("2e-4", description = "Learning rate")
+ batch_size: int = Field(1, description = "Batch size")
+ gradient_accumulation_steps: int = Field(
+ 1, description = "Gradient accumulation steps"
+ )
+ warmup_steps: Optional[int] = Field(None, description = "Warmup steps")
+ warmup_ratio: Optional[float] = Field(None, description = "Warmup ratio")
+ max_steps: Optional[int] = Field(None, description = "Maximum training steps")
+ save_steps: int = Field(100, description = "Steps between checkpoints")
+ weight_decay: float = Field(0.01, description = "Weight decay")
+ random_seed: int = Field(42, description = "Random seed")
+ packing: bool = Field(False, description = "Enable sequence packing")
+ optim: str = Field("adamw_8bit", description = "Optimizer")
+ lr_scheduler_type: str = Field("linear", description = "Learning rate scheduler type")
# LoRA parameters
- use_lora: bool = Field(True, description="Use LoRA (derived from training_type)")
- lora_r: int = Field(16, description="LoRA rank")
- lora_alpha: int = Field(16, description="LoRA alpha")
- lora_dropout: float = Field(0.0, description="LoRA dropout")
- target_modules: List[str] = Field(default_factory=list, description="Target modules for LoRA")
- gradient_checkpointing: str = Field("", description="Gradient checkpointing setting")
- use_rslora: bool = Field(False, description="Use RSLoRA")
- use_loftq: bool = Field(False, description="Use LoftQ")
- train_on_completions: bool = Field(False, description="Train on completions only")
+ use_lora: bool = Field(True, description = "Use LoRA (derived from training_type)")
+ lora_r: int = Field(16, description = "LoRA rank")
+ lora_alpha: int = Field(16, description = "LoRA alpha")
+ lora_dropout: float = Field(0.0, description = "LoRA dropout")
+ target_modules: List[str] = Field(
+ default_factory = list, description = "Target modules for LoRA"
+ )
+ gradient_checkpointing: str = Field(
+ "", description = "Gradient checkpointing setting"
+ )
+ use_rslora: bool = Field(False, description = "Use RSLoRA")
+ use_loftq: bool = Field(False, description = "Use LoftQ")
+ train_on_completions: bool = Field(False, description = "Train on completions only")
# Vision-specific LoRA parameters
- finetune_vision_layers: bool = Field(False, description="Finetune vision layers")
- finetune_language_layers: bool = Field(False, description="Finetune language layers")
- finetune_attention_modules: bool = Field(False, description="Finetune attention modules")
- finetune_mlp_modules: bool = Field(False, description="Finetune MLP modules")
- is_dataset_image: bool = Field(False, description="Whether the dataset contains image data")
- is_dataset_audio: bool = Field(False, description="Whether the dataset contains audio data")
- is_embedding: bool = Field(False, description="Whether model is an embedding/sentence-transformer model")
+ finetune_vision_layers: bool = Field(False, description = "Finetune vision layers")
+ finetune_language_layers: bool = Field(
+ False, description = "Finetune language layers"
+ )
+ finetune_attention_modules: bool = Field(
+ False, description = "Finetune attention modules"
+ )
+ finetune_mlp_modules: bool = Field(False, description = "Finetune MLP modules")
+ is_dataset_image: bool = Field(
+ False, description = "Whether the dataset contains image data"
+ )
+ is_dataset_audio: bool = Field(
+ False, description = "Whether the dataset contains audio data"
+ )
+ is_embedding: bool = Field(
+ False, description = "Whether model is an embedding/sentence-transformer model"
+ )
# Logging parameters
- enable_wandb: bool = Field(False, description="Enable Weights & Biases logging")
- wandb_token: Optional[str] = Field(None, description="W&B token")
- wandb_project: Optional[str] = Field(None, description="W&B project name")
- enable_tensorboard: bool = Field(False, description="Enable TensorBoard logging")
- tensorboard_dir: Optional[str] = Field(None, description="TensorBoard directory")
+ enable_wandb: bool = Field(False, description = "Enable Weights & Biases logging")
+ wandb_token: Optional[str] = Field(None, description = "W&B token")
+ wandb_project: Optional[str] = Field(None, description = "W&B project name")
+ enable_tensorboard: bool = Field(False, description = "Enable TensorBoard logging")
+ tensorboard_dir: Optional[str] = Field(None, description = "TensorBoard directory")
class TrainingJobResponse(BaseModel):
"""Immediate response when training is initiated"""
- job_id: str = Field(..., description="Unique training job identifier")
- status: Literal["queued", "error"] = Field(..., description="Initial job status")
- message: str = Field(..., description="Human-readable status message")
- error: Optional[str] = Field(None, description="Error details if status is 'error'")
+
+ job_id: str = Field(..., description = "Unique training job identifier")
+ status: Literal["queued", "error"] = Field(..., description = "Initial job status")
+ message: str = Field(..., description = "Human-readable status message")
+ error: Optional[str] = Field(None, description = "Error details if status is 'error'")
class TrainingStatus(BaseModel):
"""Current training job status - works for streaming or polling"""
- job_id: str = Field(..., description="Training job identifier")
+
+ job_id: str = Field(..., description = "Training job identifier")
phase: Literal[
"idle",
"loading_model",
@@ -110,31 +147,49 @@ class TrainingStatus(BaseModel):
"training",
"completed",
"error",
- "stopped"
- ] = Field(..., description="Current phase of training pipeline")
- is_training_running: bool = Field(..., description="True if training loop is actively running")
- eval_enabled: bool = Field(False, description="True if evaluation dataset is configured for this training run")
- message: str = Field(..., description="Human-readable status message")
- error: Optional[str] = Field(None, description="Error details if phase is 'error'")
- details: Optional[dict] = Field(None, description="Phase-specific info, e.g. {'model_size': '8B'}")
+ "stopped",
+ ] = Field(..., description = "Current phase of training pipeline")
+ is_training_running: bool = Field(
+ ..., description = "True if training loop is actively running"
+ )
+ eval_enabled: bool = Field(
+ False,
+ description = "True if evaluation dataset is configured for this training run",
+ )
+ message: str = Field(..., description = "Human-readable status message")
+ error: Optional[str] = Field(None, description = "Error details if phase is 'error'")
+ details: Optional[dict] = Field(
+ None, description = "Phase-specific info, e.g. {'model_size': '8B'}"
+ )
metric_history: Optional[dict] = Field(
None,
- description="Full metric history arrays for chart recovery after SSE reconnection. "
- "Keys: 'steps', 'loss', 'lr', 'grad_norm', 'grad_norm_steps' β each a list of numeric values.",
+ description = "Full metric history arrays for chart recovery after SSE reconnection. "
+ "Keys: 'steps', 'loss', 'lr', 'grad_norm', 'grad_norm_steps' β each a list of numeric values.",
)
class TrainingProgress(BaseModel):
"""Training progress metrics - for streaming or polling"""
- job_id: str = Field(..., description="Training job identifier")
- step: int = Field(..., description="Current training step")
- total_steps: int = Field(..., description="Total training steps")
- loss: float = Field(..., description="Current loss value")
- learning_rate: float = Field(..., description="Current learning rate")
- progress_percent: float = Field(..., description="Progress percentage (0.0 to 100.0)")
- epoch: Optional[float] = Field(None, description="Current epoch")
- elapsed_seconds: Optional[float] = Field(None, description="Time elapsed since training started")
- eta_seconds: Optional[float] = Field(None, description="Estimated time remaining")
- grad_norm: Optional[float] = Field(None, description="L2 norm of gradients, computed before gradient clipping")
- num_tokens: Optional[int] = Field(None, description="Total number of tokens processed so far")
- eval_loss: Optional[float] = Field(None, description="Eval loss from the most recent evaluation step")
+
+ job_id: str = Field(..., description = "Training job identifier")
+ step: int = Field(..., description = "Current training step")
+ total_steps: int = Field(..., description = "Total training steps")
+ loss: float = Field(..., description = "Current loss value")
+ learning_rate: float = Field(..., description = "Current learning rate")
+ progress_percent: float = Field(
+ ..., description = "Progress percentage (0.0 to 100.0)"
+ )
+ epoch: Optional[float] = Field(None, description = "Current epoch")
+ elapsed_seconds: Optional[float] = Field(
+ None, description = "Time elapsed since training started"
+ )
+ eta_seconds: Optional[float] = Field(None, description = "Estimated time remaining")
+ grad_norm: Optional[float] = Field(
+ None, description = "L2 norm of gradients, computed before gradient clipping"
+ )
+ num_tokens: Optional[int] = Field(
+ None, description = "Total number of tokens processed so far"
+ )
+ eval_loss: Optional[float] = Field(
+ None, description = "Eval loss from the most recent evaluation step"
+ )
diff --git a/studio/backend/models/users.py b/studio/backend/models/users.py
index 7d3fbb4eba..8e982ed9f9 100644
--- a/studio/backend/models/users.py
+++ b/studio/backend/models/users.py
@@ -12,7 +12,6 @@ from pydantic import BaseModel, Field
class Token(BaseModel):
"""Authentication token model with access and refresh tokens."""
- access_token: str = Field(..., description="JWT access token (60 min expiry)")
- refresh_token: str = Field(..., description="Opaque refresh token (7 day expiry)")
- token_type: str = Field(..., description="Token type, always 'bearer'")
-
+ access_token: str = Field(..., description = "JWT access token (60 min expiry)")
+ refresh_token: str = Field(..., description = "Opaque refresh token (7 day expiry)")
+ token_type: str = Field(..., description = "Token type, always 'bearer'")
diff --git a/__init__.py b/studio/backend/plugins/__init__.py
similarity index 100%
rename from __init__.py
rename to studio/backend/plugins/__init__.py
diff --git a/studio/backend/plugins/data-designer-unstructured-seed/__init__.py b/studio/backend/plugins/data-designer-unstructured-seed/__init__.py
new file mode 100644
index 0000000000..32014236c6
--- /dev/null
+++ b/studio/backend/plugins/data-designer-unstructured-seed/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
diff --git a/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml b/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml
index 5b83e5468a..f826ffd992 100644
--- a/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml
+++ b/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
[build-system]
requires = ["setuptools>=68", "wheel"]
build-backend = "setuptools.build_meta"
diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py
index 9e81a967b7..80f51b2a24 100644
--- a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py
+++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py
@@ -36,9 +36,9 @@ def build_unstructured_preview_rows(
chunk_overlap: Any,
) -> list[dict[str, str]]:
parquet_path, rows = materialize_unstructured_seed_dataset(
- source_path=source_path,
- chunk_size=chunk_size,
- chunk_overlap=chunk_overlap,
+ source_path = source_path,
+ chunk_size = chunk_size,
+ chunk_overlap = chunk_overlap,
)
count = max(0, int(preview_size))
if rows:
@@ -47,12 +47,14 @@ def build_unstructured_preview_rows(
try:
import pandas as pd
except ImportError as exc: # pragma: no cover
- raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc
+ raise RuntimeError(
+ f"pandas is required for unstructured seed processing: {exc}"
+ ) from exc
dataframe = pd.read_parquet(parquet_path).head(count)
return [
{"chunk_text": str(value.get("chunk_text", "")).strip()}
- for value in dataframe.to_dict(orient="records")
+ for value in dataframe.to_dict(orient = "records")
if str(value.get("chunk_text", "")).strip()
]
@@ -69,9 +71,9 @@ def materialize_unstructured_seed_dataset(
size, overlap = resolve_chunking(chunk_size, chunk_overlap)
key = _compute_cache_key(
- source_path=resolved,
- chunk_size=size,
- chunk_overlap=overlap,
+ source_path = resolved,
+ chunk_size = size,
+ chunk_overlap = overlap,
)
parquet_path = _CACHE_DIR / f"{key}.parquet"
if parquet_path.exists():
@@ -79,9 +81,9 @@ def materialize_unstructured_seed_dataset(
text = load_unstructured_text_file(resolved)
chunks = split_text_into_chunks(
- text=text,
- chunk_size=size,
- chunk_overlap=overlap,
+ text = text,
+ chunk_size = size,
+ chunk_overlap = overlap,
)
if not chunks:
raise ValueError("No text found in unstructured seed source.")
@@ -91,10 +93,12 @@ def materialize_unstructured_seed_dataset(
try:
import pandas as pd
except ImportError as exc: # pragma: no cover
- raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc
+ raise RuntimeError(
+ f"pandas is required for unstructured seed processing: {exc}"
+ ) from exc
tmp_path = _CACHE_DIR / f"{key}.tmp.parquet"
- pd.DataFrame(rows).to_parquet(tmp_path, index=False)
+ pd.DataFrame(rows).to_parquet(tmp_path, index = False)
tmp_path.replace(parquet_path)
return parquet_path, rows
@@ -104,7 +108,7 @@ def load_unstructured_text_file(path: Path) -> str:
if ext not in {".txt", ".md"}:
raise ValueError(f"Unsupported unstructured seed file type: {ext}")
- raw = path.read_text(encoding="utf-8", errors="ignore")
+ raw = path.read_text(encoding = "utf-8", errors = "ignore")
return normalize_unstructured_text(raw)
diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py
index 1d6defa4df..e0a0392a69 100644
--- a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py
+++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py
@@ -15,11 +15,11 @@ from .chunking import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE, resolve_chunkin
class UnstructuredSeedSource(SeedSource):
seed_type: Literal["unstructured"] = "unstructured"
- path: str = Field(..., min_length=1)
+ path: str = Field(..., min_length = 1)
chunk_size: int = DEFAULT_CHUNK_SIZE
chunk_overlap: int = DEFAULT_CHUNK_OVERLAP
- @field_validator("path", mode="after")
+ @field_validator("path", mode = "after")
@classmethod
def _validate_path(cls, value: str) -> str:
path = Path(value).expanduser()
@@ -27,13 +27,13 @@ class UnstructuredSeedSource(SeedSource):
raise ValueError(f"Unstructured seed path is not a file: {path}")
return value
- @field_validator("chunk_size", mode="after")
+ @field_validator("chunk_size", mode = "after")
@classmethod
def _validate_chunk_size(cls, value: int) -> int:
size, _ = resolve_chunking(value, 0)
return size
- @field_validator("chunk_overlap", mode="after")
+ @field_validator("chunk_overlap", mode = "after")
@classmethod
def _validate_chunk_overlap(cls, value: int, info) -> int:
size = info.data.get("chunk_size", cls.model_fields["chunk_size"].default)
diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py
index acde1909c2..8a3deb9b92 100644
--- a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py
+++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py
@@ -18,8 +18,8 @@ class UnstructuredSeedReader(SeedReader[UnstructuredSeedSource]):
def get_dataset_uri(self) -> str:
path, _ = materialize_unstructured_seed_dataset(
- source_path=Path(self.source.path),
- chunk_size=self.source.chunk_size,
- chunk_overlap=self.source.chunk_overlap,
+ source_path = Path(self.source.path),
+ chunk_size = self.source.chunk_size,
+ chunk_overlap = self.source.chunk_overlap,
)
return str(path)
diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py
index de8c00d8b7..6f0d7ffd49 100644
--- a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py
+++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py
@@ -4,7 +4,7 @@
from data_designer.plugins.plugin import Plugin, PluginType
unstructured_seed_plugin = Plugin(
- impl_qualified_name="data_designer_unstructured_seed.impl.UnstructuredSeedReader",
- config_qualified_name="data_designer_unstructured_seed.config.UnstructuredSeedSource",
- plugin_type=PluginType.SEED_READER,
+ impl_qualified_name = "data_designer_unstructured_seed.impl.UnstructuredSeedReader",
+ config_qualified_name = "data_designer_unstructured_seed.config.UnstructuredSeedSource",
+ plugin_type = PluginType.SEED_READER,
)
diff --git a/studio/backend/requirements/__init__.py b/studio/backend/requirements/__init__.py
new file mode 100644
index 0000000000..32014236c6
--- /dev/null
+++ b/studio/backend/requirements/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
diff --git a/studio/backend/requirements/single-env/patch_metadata.py b/studio/backend/requirements/single-env/patch_metadata.py
index b5d4bfcaad..7bcaa56ac8 100644
--- a/studio/backend/requirements/single-env/patch_metadata.py
+++ b/studio/backend/requirements/single-env/patch_metadata.py
@@ -1,6 +1,7 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
"""Relax strict metadata pins so pip check matches known working single-env stack.
Why:
@@ -48,13 +49,13 @@ def metadata_path(dist_name: str) -> Path | None:
def patch_file(path: Path) -> bool:
- original = path.read_text(encoding="utf-8")
+ original = path.read_text(encoding = "utf-8")
updated = original
for pattern, repl in PATCHES:
updated = pattern.sub(repl, updated)
if updated == original:
return False
- path.write_text(updated, encoding="utf-8")
+ path.write_text(updated, encoding = "utf-8")
return True
diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py
index f7508cb6ff..83cc41e83b 100644
--- a/studio/backend/routes/auth.py
+++ b/studio/backend/routes/auth.py
@@ -4,6 +4,7 @@
"""
Authentication API routes
"""
+
from fastapi import APIRouter, HTTPException, status
import secrets
@@ -25,7 +26,7 @@ from auth.authentication import (
router = APIRouter()
-@router.get("/status", response_model=AuthStatusResponse)
+@router.get("/status", response_model = AuthStatusResponse)
async def auth_status() -> AuthStatusResponse:
"""
Check whether auth has already been initialized.
@@ -33,10 +34,10 @@ async def auth_status() -> AuthStatusResponse:
- initialized = False -> frontend should show "Set admin password" screen.
- initialized = True -> frontend should show normal login.
"""
- return AuthStatusResponse(initialized=storage.is_initialized())
+ return AuthStatusResponse(initialized = storage.is_initialized())
-@router.post("/setup", response_model=Token, status_code=status.HTTP_201_CREATED)
+@router.post("/setup", response_model = Token, status_code = status.HTTP_201_CREATED)
async def setup_auth(payload: AuthSetupRequest) -> Token:
"""
First-time setup: create the admin user and a JWT secret.
@@ -46,15 +47,15 @@ async def setup_auth(payload: AuthSetupRequest) -> Token:
"""
if storage.is_initialized():
raise HTTPException(
- status_code=status.HTTP_400_BAD_REQUEST,
- detail="Auth is already initialized.",
+ status_code = status.HTTP_400_BAD_REQUEST,
+ detail = "Auth is already initialized.",
)
# Validate the one-time setup token
if not storage.consume_setup_token(payload.setup_token):
raise HTTPException(
- status_code=status.HTTP_403_FORBIDDEN,
- detail="Invalid or expired setup token.",
+ status_code = status.HTTP_403_FORBIDDEN,
+ detail = "Invalid or expired setup token.",
)
# Generate a strong random JWT secret for this installation
@@ -63,34 +64,34 @@ async def setup_auth(payload: AuthSetupRequest) -> Token:
# Create user + generate tokens atomically β rollback if anything fails
try:
storage.create_initial_user(
- username=payload.username,
- password=payload.password,
- jwt_secret=jwt_secret,
+ username = payload.username,
+ password = payload.password,
+ jwt_secret = jwt_secret,
)
# Reload JWT secret from DB (so authentication.py picks it up)
reload_secret()
# Issue access + refresh tokens for the new user
- access_token = create_access_token(subject=payload.username)
- refresh_token = create_refresh_token(subject=payload.username)
+ access_token = create_access_token(subject = payload.username)
+ refresh_token = create_refresh_token(subject = payload.username)
except Exception as e:
# Rollback: remove the user row so setup can be retried
storage.delete_user(payload.username)
raise HTTPException(
- status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
- detail=f"Setup failed (rolled back): {str(e)}",
+ status_code = status.HTTP_500_INTERNAL_SERVER_ERROR,
+ detail = f"Setup failed (rolled back): {str(e)}",
)
return Token(
- access_token=access_token,
- refresh_token=refresh_token,
- token_type="bearer",
+ access_token = access_token,
+ refresh_token = refresh_token,
+ token_type = "bearer",
)
-@router.post("/login", response_model=Token)
+@router.post("/login", response_model = Token)
async def login(payload: AuthLoginRequest) -> Token:
"""
Login with username/password and receive access + refresh tokens.
@@ -98,27 +99,27 @@ async def login(payload: AuthLoginRequest) -> Token:
record = storage.get_user_and_secret(payload.username)
if record is None:
raise HTTPException(
- status_code=status.HTTP_401_UNAUTHORIZED,
- detail="Incorrect username or password",
+ status_code = status.HTTP_401_UNAUTHORIZED,
+ detail = "Incorrect username or password",
)
salt, pwd_hash, _jwt_secret = record
if not hashing.verify_password(payload.password, salt, pwd_hash):
raise HTTPException(
- status_code=status.HTTP_401_UNAUTHORIZED,
- detail="Incorrect username or password",
+ status_code = status.HTTP_401_UNAUTHORIZED,
+ detail = "Incorrect username or password",
)
- access_token = create_access_token(subject=payload.username)
- refresh_token = create_refresh_token(subject=payload.username)
+ access_token = create_access_token(subject = payload.username)
+ refresh_token = create_refresh_token(subject = payload.username)
return Token(
- access_token=access_token,
- refresh_token=refresh_token,
- token_type="bearer",
+ access_token = access_token,
+ refresh_token = refresh_token,
+ token_type = "bearer",
)
-@router.post("/refresh", response_model=Token)
+@router.post("/refresh", response_model = Token)
async def refresh(payload: RefreshTokenRequest) -> Token:
"""
Exchange a valid refresh token for a new access token.
@@ -128,13 +129,12 @@ async def refresh(payload: RefreshTokenRequest) -> Token:
new_access_token = refresh_access_token(payload.refresh_token)
if new_access_token is None:
raise HTTPException(
- status_code=status.HTTP_401_UNAUTHORIZED,
- detail="Invalid or expired refresh token",
+ status_code = status.HTTP_401_UNAUTHORIZED,
+ detail = "Invalid or expired refresh token",
)
return Token(
- access_token=new_access_token,
- refresh_token=payload.refresh_token,
- token_type="bearer",
+ access_token = new_access_token,
+ refresh_token = payload.refresh_token,
+ token_type = "bearer",
)
-
diff --git a/studio/backend/routes/data_recipe/__init__.py b/studio/backend/routes/data_recipe/__init__.py
index 6ec530751e..c596f189d4 100644
--- a/studio/backend/routes/data_recipe/__init__.py
+++ b/studio/backend/routes/data_recipe/__init__.py
@@ -21,7 +21,7 @@ from .mcp import router as mcp_router
from .seed import router as seed_router
from .validate import router as validate_router
-router = APIRouter(dependencies=[Depends(get_current_subject)])
+router = APIRouter(dependencies = [Depends(get_current_subject)])
router.include_router(seed_router)
router.include_router(validate_router)
router.include_router(jobs_router)
diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py
index 02b8befab6..4661615338 100644
--- a/studio/backend/routes/data_recipe/jobs.py
+++ b/studio/backend/routes/data_recipe/jobs.py
@@ -21,25 +21,30 @@ def _normalize_run_name(value: Any) -> str | None:
if value is None:
return None
if not isinstance(value, str):
- raise HTTPException(status_code=400, detail="invalid run_name: must be a string")
+ raise HTTPException(
+ status_code = 400, detail = "invalid run_name: must be a string"
+ )
trimmed = value.strip()
if not trimmed:
return None
return trimmed[:120]
-@router.post("/jobs", response_class=JSONResponse, response_model=JobCreateResponse)
+@router.post("/jobs", response_class = JSONResponse, response_model = JobCreateResponse)
def create_job(payload: RecipePayload):
recipe = payload.recipe
if not recipe.get("columns"):
- raise HTTPException(status_code=400, detail="Recipe must include columns.")
+ raise HTTPException(status_code = 400, detail = "Recipe must include columns.")
run: dict[str, Any] = payload.run or {}
run.pop("artifact_path", None)
run.pop("dataset_name", None)
execution_type = str(run.get("execution_type") or "full").strip().lower()
if execution_type not in {"preview", "full"}:
- raise HTTPException(status_code=400, detail="invalid execution_type: must be 'preview' or 'full'")
+ raise HTTPException(
+ status_code = 400,
+ detail = "invalid execution_type: must be 'preview' or 'full'",
+ )
run["execution_type"] = execution_type
run["run_name"] = _normalize_run_name(run.get("run_name"))
run_config_raw = run.get("run_config")
@@ -49,15 +54,17 @@ def create_job(payload: RecipePayload):
RunConfig.model_validate(run_config_raw)
except (ImportError, ValidationError, TypeError, ValueError) as exc:
- raise HTTPException(status_code=400, detail=f"invalid run_config: {exc}") from exc
+ raise HTTPException(
+ status_code = 400, detail = f"invalid run_config: {exc}"
+ ) from exc
mgr = get_job_manager()
try:
- job_id = mgr.start(recipe=recipe, run=run)
+ job_id = mgr.start(recipe = recipe, run = run)
except RuntimeError as exc:
- raise HTTPException(status_code=409, detail=str(exc)) from exc
+ raise HTTPException(status_code = 409, detail = str(exc)) from exc
except ValueError as exc:
- raise HTTPException(status_code=400, detail=str(exc)) from exc
+ raise HTTPException(status_code = 400, detail = str(exc)) from exc
return {"job_id": job_id}
@@ -67,7 +74,7 @@ def job_status(job_id: str):
mgr = get_job_manager()
state = mgr.get_status(job_id)
if state is None:
- raise HTTPException(status_code=404, detail="job not found")
+ raise HTTPException(status_code = 404, detail = "job not found")
return state
@@ -76,7 +83,7 @@ def current_job():
mgr = get_job_manager()
state = mgr.get_current_status()
if state is None:
- raise HTTPException(status_code=404, detail="no job")
+ raise HTTPException(status_code = 404, detail = "no job")
return state
@@ -85,7 +92,7 @@ def cancel_job(job_id: str):
mgr = get_job_manager()
ok = mgr.cancel(job_id)
if not ok:
- raise HTTPException(status_code=404, detail="job not found")
+ raise HTTPException(status_code = 404, detail = "job not found")
return mgr.get_status(job_id)
@@ -94,22 +101,22 @@ def job_analysis(job_id: str):
mgr = get_job_manager()
analysis = mgr.get_analysis(job_id)
if analysis is None:
- raise HTTPException(status_code=404, detail="analysis not ready")
+ raise HTTPException(status_code = 404, detail = "analysis not ready")
return analysis
@router.get("/jobs/{job_id}/dataset")
def job_dataset(
job_id: str,
- limit: int = Query(default=20, ge=1, le=500),
- offset: int = Query(default=0, ge=0),
+ limit: int = Query(default = 20, ge = 1, le = 500),
+ offset: int = Query(default = 0, ge = 0),
):
mgr = get_job_manager()
- result = mgr.get_dataset(job_id, limit=limit, offset=offset)
+ result = mgr.get_dataset(job_id, limit = limit, offset = offset)
if result is None:
- raise HTTPException(status_code=404, detail="dataset not ready")
+ raise HTTPException(status_code = 404, detail = "dataset not ready")
if "error" in result:
- raise HTTPException(status_code=422, detail=result["error"])
+ raise HTTPException(status_code = 422, detail = result["error"])
return {
"dataset": result["dataset"],
"total": result["total"],
@@ -136,9 +143,9 @@ async def job_events(request: Request, job_id: str):
except (TypeError, ValueError):
pass
- sub = mgr.subscribe(job_id, after_seq=after_seq)
+ sub = mgr.subscribe(job_id, after_seq = after_seq)
if sub is None:
- raise HTTPException(status_code=404, detail="job not found")
+ raise HTTPException(status_code = 404, detail = "job not found")
async def gen():
try:
@@ -148,11 +155,11 @@ async def job_events(request: Request, job_id: str):
while True:
if await request.is_disconnected():
break
- event = await sub.next_event(timeout_sec=1.0)
+ event = await sub.next_event(timeout_sec = 1.0)
if event is None:
continue
yield sub.format_sse(event)
finally:
mgr.unsubscribe(sub)
- return StreamingResponse(gen(), media_type="text/event-stream")
+ return StreamingResponse(gen(), media_type = "text/event-stream")
diff --git a/studio/backend/routes/data_recipe/mcp.py b/studio/backend/routes/data_recipe/mcp.py
index a106810a15..1f5c0f34e0 100644
--- a/studio/backend/routes/data_recipe/mcp.py
+++ b/studio/backend/routes/data_recipe/mcp.py
@@ -19,16 +19,16 @@ from models.data_recipe import (
router = APIRouter()
-@router.post("/mcp/tools", response_model=McpToolsListResponse)
+@router.post("/mcp/tools", response_model = McpToolsListResponse)
def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse:
try:
from data_designer.engine.mcp import io as mcp_io
except ImportError as exc:
return McpToolsListResponse(
- providers=[
+ providers = [
McpToolsProviderResult(
- name="",
- error=f"MCP dependencies unavailable: {exc}",
+ name = "",
+ error = f"MCP dependencies unavailable: {exc}",
)
]
)
@@ -42,29 +42,31 @@ def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse:
if len(built) != 1:
providers.append(
McpToolsProviderResult(
- name=provider_name,
- error="Unsupported MCP provider config.",
+ name = provider_name,
+ error = "Unsupported MCP provider config.",
)
)
continue
provider = built[0]
try:
- tools = mcp_io.list_tools(provider, timeout_sec=payload.timeout_sec)
- tool_names = sorted({tool.name for tool in tools if getattr(tool, "name", "")})
+ tools = mcp_io.list_tools(provider, timeout_sec = payload.timeout_sec)
+ tool_names = sorted(
+ {tool.name for tool in tools if getattr(tool, "name", "")}
+ )
for tool_name in tool_names:
tool_to_providers[tool_name].append(provider.name)
providers.append(
McpToolsProviderResult(
- name=provider.name,
- tools=tool_names,
+ name = provider.name,
+ tools = tool_names,
)
)
except Exception as exc:
providers.append(
McpToolsProviderResult(
- name=provider.name or provider_name,
- error=str(exc).strip() or "Failed to load tools.",
+ name = provider.name or provider_name,
+ error = str(exc).strip() or "Failed to load tools.",
)
)
@@ -75,6 +77,6 @@ def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse:
}
return McpToolsListResponse(
- providers=providers,
- duplicate_tools=duplicate_tools,
+ providers = providers,
+ duplicate_tools = duplicate_tools,
)
diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py
index cb8130e698..76e05be954 100644
--- a/studio/backend/routes/data_recipe/seed.py
+++ b/studio/backend/routes/data_recipe/seed.py
@@ -61,7 +61,7 @@ def _list_hf_data_files(*, dataset_name: str, token: str | None) -> list[str]:
return []
try:
api = HfApi()
- repo_files = api.list_repo_files(dataset_name, repo_type="dataset", token=token)
+ repo_files = api.list_repo_files(dataset_name, repo_type = "dataset", token = token)
return [file for file in repo_files if file.lower().endswith(DATA_EXTS)]
except (HfHubHTTPError, OSError, ValueError):
return []
@@ -86,10 +86,12 @@ def _select_best_file(data_files: list[str], split: str = DEFAULT_SPLIT) -> str
return (1, len(path))
return (2, len(path))
- return sorted(data_files, key=score)[0]
+ return sorted(data_files, key = score)[0]
-def _resolve_seed_hf_path(dataset_name: str, data_files: list[str], split: str = DEFAULT_SPLIT) -> str | None:
+def _resolve_seed_hf_path(
+ dataset_name: str, data_files: list[str], split: str = DEFAULT_SPLIT
+) -> str | None:
selected = _select_best_file(data_files, split)
if not selected:
return None
@@ -156,36 +158,42 @@ def _decode_base64_payload(content_base64: str) -> bytes:
if "," in raw and raw.lower().startswith("data:"):
raw = raw.split(",", 1)[1]
try:
- return base64.b64decode(raw, validate=True)
+ return base64.b64decode(raw, validate = True)
except binascii.Error as exc:
- raise HTTPException(status_code=400, detail="invalid base64 payload") from exc
+ raise HTTPException(status_code = 400, detail = "invalid base64 payload") from exc
-def _read_preview_rows_from_local_file(path: Path, preview_size: int) -> list[dict[str, Any]]:
+def _read_preview_rows_from_local_file(
+ path: Path, preview_size: int
+) -> list[dict[str, Any]]:
try:
import pandas as pd
except ImportError as exc:
- raise HTTPException(status_code=500, detail=f"seed inspect dependencies unavailable: {exc}") from exc
+ raise HTTPException(
+ status_code = 500, detail = f"seed inspect dependencies unavailable: {exc}"
+ ) from exc
ext = path.suffix.lower()
try:
if ext == ".csv":
- df = pd.read_csv(path, nrows=preview_size)
+ df = pd.read_csv(path, nrows = preview_size)
elif ext == ".jsonl":
- df = pd.read_json(path, lines=True).head(preview_size)
+ df = pd.read_json(path, lines = True).head(preview_size)
elif ext == ".json":
try:
df = pd.read_json(path).head(preview_size)
except ValueError:
- df = pd.read_json(path, lines=True).head(preview_size)
+ df = pd.read_json(path, lines = True).head(preview_size)
else:
- raise HTTPException(status_code=422, detail=f"unsupported file type: {ext}")
+ raise HTTPException(status_code = 422, detail = f"unsupported file type: {ext}")
except HTTPException:
raise
except (ValueError, OSError) as exc:
- raise HTTPException(status_code=422, detail=f"seed inspect failed: {exc}") from exc
+ raise HTTPException(
+ status_code = 422, detail = f"seed inspect failed: {exc}"
+ ) from exc
- rows = df.to_dict(orient="records")
+ rows = df.to_dict(orient = "records")
return _serialize_preview_rows(rows)
@@ -199,26 +207,33 @@ def _read_preview_rows_from_unstructured_file(
size, overlap = resolve_chunking(chunk_size, chunk_overlap)
try:
rows = build_unstructured_preview_rows(
- source_path=path,
- preview_size=preview_size,
- chunk_size=size,
- chunk_overlap=overlap,
+ source_path = path,
+ preview_size = preview_size,
+ chunk_size = size,
+ chunk_overlap = overlap,
)
except (FileNotFoundError, RuntimeError, ValueError, OSError) as exc:
- raise HTTPException(status_code=422, detail=f"seed inspect failed: {exc}") from exc
+ raise HTTPException(
+ status_code = 422, detail = f"seed inspect failed: {exc}"
+ ) from exc
return _serialize_preview_rows(rows)
-@router.post("/seed/inspect", response_model=SeedInspectResponse)
+@router.post("/seed/inspect", response_model = SeedInspectResponse)
def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
dataset_name = payload.dataset_name.strip()
if not dataset_name or dataset_name.count("/") < 1:
- raise HTTPException(status_code=400, detail="dataset_name must be a Hugging Face repo id like org/repo")
+ raise HTTPException(
+ status_code = 400,
+ detail = "dataset_name must be a Hugging Face repo id like org/repo",
+ )
try:
from datasets import load_dataset
except ImportError as exc:
- raise HTTPException(status_code=500, detail=f"seed inspect dependencies unavailable: {exc}") from exc
+ raise HTTPException(
+ status_code = 500, detail = f"seed inspect dependencies unavailable: {exc}"
+ ) from exc
split = _normalize_optional_text(payload.split) or DEFAULT_SPLIT
subset = _normalize_optional_text(payload.subset)
@@ -226,22 +241,22 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
preview_size = int(payload.preview_size)
preview_rows: list[dict[str, Any]] = []
- data_files = _list_hf_data_files(dataset_name=dataset_name, token=token)
+ data_files = _list_hf_data_files(dataset_name = dataset_name, token = token)
selected_file = _select_best_file(data_files, split)
if selected_file:
try:
single_file_kwargs = _build_stream_load_kwargs(
- dataset_name=dataset_name,
- split=split,
- subset=subset,
- token=token,
- data_file=selected_file,
+ dataset_name = dataset_name,
+ split = split,
+ subset = subset,
+ token = token,
+ data_file = selected_file,
)
preview_rows = _load_preview_rows(
- load_dataset_fn=load_dataset,
- load_kwargs=single_file_kwargs,
- preview_size=preview_size,
+ load_dataset_fn = load_dataset,
+ load_kwargs = single_file_kwargs,
+ preview_size = preview_size,
)
except (ValueError, OSError, RuntimeError):
preview_rows = []
@@ -249,21 +264,25 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
if not preview_rows:
try:
split_kwargs = _build_stream_load_kwargs(
- dataset_name=dataset_name,
- split=split,
- subset=subset,
- token=token,
+ dataset_name = dataset_name,
+ split = split,
+ subset = subset,
+ token = token,
)
preview_rows = _load_preview_rows(
- load_dataset_fn=load_dataset,
- load_kwargs=split_kwargs,
- preview_size=preview_size,
+ load_dataset_fn = load_dataset,
+ load_kwargs = split_kwargs,
+ preview_size = preview_size,
)
except (ValueError, OSError, RuntimeError) as exc:
- raise HTTPException(status_code=422, detail=f"seed inspect failed: {exc}") from exc
+ raise HTTPException(
+ status_code = 422, detail = f"seed inspect failed: {exc}"
+ ) from exc
if not preview_rows:
- raise HTTPException(status_code=422, detail="dataset appears empty or unreadable")
+ raise HTTPException(
+ status_code = 422, detail = "dataset appears empty or unreadable"
+ )
preview_rows = _serialize_preview_rows(preview_rows)
columns = _extract_columns(preview_rows)
@@ -272,19 +291,21 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
else:
resolved_path = _resolve_seed_hf_path(dataset_name, data_files, split)
if not resolved_path:
- raise HTTPException(status_code=422, detail="unable to resolve seed dataset path")
+ raise HTTPException(
+ status_code = 422, detail = "unable to resolve seed dataset path"
+ )
return SeedInspectResponse(
- dataset_name=dataset_name,
- resolved_path=resolved_path,
- columns=columns,
- preview_rows=preview_rows,
- split=split,
- subset=subset,
+ dataset_name = dataset_name,
+ resolved_path = resolved_path,
+ columns = columns,
+ preview_rows = preview_rows,
+ split = split,
+ subset = subset,
)
-@router.post("/seed/inspect-upload", response_model=SeedInspectResponse)
+@router.post("/seed/inspect-upload", response_model = SeedInspectResponse)
def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectResponse:
seed_source_type = _normalize_optional_text(payload.seed_source_type) or "local"
filename = _sanitize_filename(payload.filename)
@@ -292,18 +313,24 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons
if seed_source_type == "unstructured":
if ext not in UNSTRUCTURED_UPLOAD_EXTS:
allowed = ", ".join(sorted(UNSTRUCTURED_UPLOAD_EXTS))
- raise HTTPException(status_code=400, detail=f"unsupported file type: {ext}. allowed: {allowed}")
+ raise HTTPException(
+ status_code = 400,
+ detail = f"unsupported file type: {ext}. allowed: {allowed}",
+ )
else:
if ext not in LOCAL_UPLOAD_EXTS:
allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS))
- raise HTTPException(status_code=400, detail=f"unsupported file type: {ext}. allowed: {allowed}")
+ raise HTTPException(
+ status_code = 400,
+ detail = f"unsupported file type: {ext}. allowed: {allowed}",
+ )
file_bytes = _decode_base64_payload(payload.content_base64)
if not file_bytes:
- raise HTTPException(status_code=400, detail="empty upload payload")
+ raise HTTPException(status_code = 400, detail = "empty upload payload")
max_size_bytes = 50 * 1024 * 1024
if len(file_bytes) > max_size_bytes:
- raise HTTPException(status_code=413, detail="file too large (max 50MB)")
+ raise HTTPException(status_code = 413, detail = "file too large (max 50MB)")
ensure_dir(SEED_UPLOAD_DIR)
stored_name = f"{uuid4().hex}_{filename}"
@@ -312,10 +339,10 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons
if seed_source_type == "unstructured":
preview_rows = _read_preview_rows_from_unstructured_file(
- path=stored_path,
- preview_size=int(payload.preview_size),
- chunk_size=payload.unstructured_chunk_size,
- chunk_overlap=payload.unstructured_chunk_overlap,
+ path = stored_path,
+ preview_size = int(payload.preview_size),
+ chunk_size = payload.unstructured_chunk_size,
+ chunk_overlap = payload.unstructured_chunk_overlap,
)
else:
preview_rows = _read_preview_rows_from_local_file(
@@ -323,14 +350,16 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons
int(payload.preview_size),
)
if not preview_rows:
- raise HTTPException(status_code=422, detail="dataset appears empty or unreadable")
+ raise HTTPException(
+ status_code = 422, detail = "dataset appears empty or unreadable"
+ )
columns = _extract_columns(preview_rows)
return SeedInspectResponse(
- dataset_name=filename,
- resolved_path=str(stored_path),
- columns=columns,
- preview_rows=preview_rows,
- split=None,
- subset=None,
+ dataset_name = filename,
+ resolved_path = str(stored_path),
+ columns = columns,
+ preview_rows = preview_rows,
+ split = None,
+ subset = None,
)
diff --git a/studio/backend/routes/data_recipe/validate.py b/studio/backend/routes/data_recipe/validate.py
index f2858f97bd..a793a3b172 100644
--- a/studio/backend/routes/data_recipe/validate.py
+++ b/studio/backend/routes/data_recipe/validate.py
@@ -44,9 +44,9 @@ def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]:
_resolve_and_add_seed_columns(config, resource_provider.seed_reader)
_add_internal_row_id_column_if_needed(config)
violations = validate_data_designer_config(
- columns=config.columns,
- processor_configs=config.processors or [],
- allowed_references=_get_allowed_references(config),
+ columns = config.columns,
+ processor_configs = config.processors or [],
+ allowed_references = _get_allowed_references(config),
)
except (TypeError, ValueError, AttributeError):
return []
@@ -60,34 +60,34 @@ def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]:
message = str(violation.message).strip() or "Validation failed."
errors.append(
ValidateError(
- message=message,
- path=path,
- code=code,
+ message = message,
+ path = path,
+ code = code,
)
)
return errors
-@router.post("/validate", response_model=ValidateResponse)
+@router.post("/validate", response_model = ValidateResponse)
def validate(payload: RecipePayload) -> ValidateResponse:
recipe = payload.recipe
if not recipe.get("columns"):
return ValidateResponse(
- valid=False,
- errors=[ValidateError(message="Recipe must include columns.")],
+ valid = False,
+ errors = [ValidateError(message = "Recipe must include columns.")],
)
try:
validate_recipe(recipe)
except RuntimeError as exc:
- raise HTTPException(status_code=503, detail=str(exc)) from exc
+ raise HTTPException(status_code = 503, detail = str(exc)) from exc
except Exception as exc:
detail = str(exc).strip() or "Validation failed."
parsed_errors = _collect_validation_errors(recipe)
return ValidateResponse(
- valid=False,
- errors=parsed_errors or [ValidateError(message=detail)],
- raw_detail=detail,
+ valid = False,
+ errors = parsed_errors or [ValidateError(message = detail)],
+ raw_detail = detail,
)
- return ValidateResponse(valid=True)
+ return ValidateResponse(valid = True)
diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py
index 937de73512..415c30cd72 100644
--- a/studio/backend/routes/datasets.py
+++ b/studio/backend/routes/datasets.py
@@ -4,6 +4,7 @@
"""
Datasets API routes
"""
+
import base64
import io
import json
@@ -27,8 +28,6 @@ router = APIRouter()
logger = get_logger(__name__)
-
-
from models.datasets import (
AiAssistMappingRequest,
AiAssistMappingResponse,
@@ -53,9 +52,10 @@ def _serialize_preview_value(value):
try:
from PIL.Image import Image as PILImage
+
if isinstance(value, PILImage):
buffer = io.BytesIO()
- value.convert("RGB").save(buffer, format="JPEG", quality=85)
+ value.convert("RGB").save(buffer, format = "JPEG", quality = 85)
return {
"type": "image",
"mime": "image/jpeg",
@@ -88,10 +88,10 @@ def _serialize_preview_rows(rows):
# Tabular formats are preferred over archives for Tier 1 preview because
# archives (e.g. images.zip) may be loaded as ImageFolder datasets with
# synthetic columns (image/label) that don't match the real dataset schema.
-_TABULAR_EXTS = ('.parquet', '.json', '.jsonl', '.csv', '.tsv', '.arrow')
-_ARCHIVE_EXTS = ('.tar', '.tar.gz', '.tgz', '.gz', '.zst', '.zip', '.txt')
+_TABULAR_EXTS = (".parquet", ".json", ".jsonl", ".csv", ".tsv", ".arrow")
+_ARCHIVE_EXTS = (".tar", ".tar.gz", ".tgz", ".gz", ".zst", ".zip", ".txt")
DATA_EXTS = _TABULAR_EXTS + _ARCHIVE_EXTS
-LOCAL_FILE_EXTS = ('.json', '.jsonl', '.csv', '.parquet')
+LOCAL_FILE_EXTS = (".json", ".jsonl", ".csv", ".parquet")
LOCAL_UPLOAD_EXTS = {".csv", ".json", ".jsonl", ".parquet"}
LOCAL_DATASETS_ROOT = recipe_datasets_root()
DATASET_UPLOAD_DIR = dataset_uploads_root()
@@ -99,7 +99,7 @@ DATASET_UPLOAD_DIR = dataset_uploads_root()
def _safe_read_metadata(path: Path) -> dict | None:
try:
- payload = json.loads(path.read_text(encoding="utf-8"))
+ payload = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, ValueError, TypeError):
return None
if not isinstance(payload, dict):
@@ -200,30 +200,36 @@ def _build_local_dataset_items() -> list[LocalDatasetItem]:
items.append(
LocalDatasetItem(
- id=entry.name,
- label=entry.name,
- path=str(parquet_dir.resolve()),
- rows=rows,
- updated_at=updated_at,
- metadata=metadata_summary,
+ id = entry.name,
+ label = entry.name,
+ path = str(parquet_dir.resolve()),
+ rows = rows,
+ updated_at = updated_at,
+ metadata = metadata_summary,
)
)
- items.sort(key=lambda item: item.updated_at or 0, reverse=True)
+ items.sort(key = lambda item: item.updated_at or 0, reverse = True)
return items
-def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int):
+def _load_local_preview_slice(
+ *, dataset_path: Path, train_split: str, preview_size: int
+):
from datasets import load_dataset
if dataset_path.is_dir():
- parquet_dir = dataset_path / "parquet-files" if (dataset_path / "parquet-files").exists() else dataset_path
+ parquet_dir = (
+ dataset_path / "parquet-files"
+ if (dataset_path / "parquet-files").exists()
+ else dataset_path
+ )
parquet_files = sorted(parquet_dir.glob("*.parquet"))
if parquet_files:
dataset = load_dataset(
"parquet",
- data_files=[str(path) for path in parquet_files],
- split=train_split,
+ data_files = [str(path) for path in parquet_files],
+ split = train_split,
)
total_rows = len(dataset)
preview_slice = dataset.select(range(min(preview_size, total_rows)))
@@ -234,21 +240,22 @@ def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_s
candidate_files.extend(sorted(dataset_path.glob(f"*{ext}")))
if not candidate_files:
raise HTTPException(
- status_code=400,
- detail="Unsupported local dataset directory (expected parquet/json/jsonl/csv files)",
+ status_code = 400,
+ detail = "Unsupported local dataset directory (expected parquet/json/jsonl/csv files)",
)
dataset_path = candidate_files[0]
- if dataset_path.suffix in ['.json', '.jsonl']:
- dataset = load_dataset('json', data_files=str(dataset_path), split=train_split)
- elif dataset_path.suffix == '.csv':
- dataset = load_dataset('csv', data_files=str(dataset_path), split=train_split)
- elif dataset_path.suffix == '.parquet':
- dataset = load_dataset('parquet', data_files=str(dataset_path), split=train_split)
+ if dataset_path.suffix in [".json", ".jsonl"]:
+ dataset = load_dataset("json", data_files = str(dataset_path), split = train_split)
+ elif dataset_path.suffix == ".csv":
+ dataset = load_dataset("csv", data_files = str(dataset_path), split = train_split)
+ elif dataset_path.suffix == ".parquet":
+ dataset = load_dataset(
+ "parquet", data_files = str(dataset_path), split = train_split
+ )
else:
raise HTTPException(
- status_code=400,
- detail=f"Unsupported file format: {dataset_path.suffix}"
+ status_code = 400, detail = f"Unsupported file format: {dataset_path.suffix}"
)
total_rows = len(dataset)
@@ -263,7 +270,7 @@ def _sanitize_filename(filename: str) -> str:
return name
-@router.post("/upload", response_model=UploadDatasetResponse)
+@router.post("/upload", response_model = UploadDatasetResponse)
async def upload_dataset(
file: UploadFile,
current_subject: str = Depends(get_current_subject),
@@ -273,8 +280,8 @@ async def upload_dataset(
if ext not in LOCAL_UPLOAD_EXTS:
allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS))
raise HTTPException(
- status_code=400,
- detail=f"Unsupported file type: {ext}. Allowed: {allowed}",
+ status_code = 400,
+ detail = f"Unsupported file type: {ext}. Allowed: {allowed}",
)
max_size_bytes = 512 * 1024 * 1024
@@ -289,25 +296,27 @@ async def upload_dataset(
while chunk := await file.read(1024 * 1024):
size += len(chunk)
if size > max_size_bytes:
- stored_path.unlink(missing_ok=True)
- raise HTTPException(status_code=413, detail="File too large (max 512MB)")
+ stored_path.unlink(missing_ok = True)
+ raise HTTPException(
+ status_code = 413, detail = "File too large (max 512MB)"
+ )
f.write(chunk)
if size == 0:
- stored_path.unlink(missing_ok=True)
- raise HTTPException(status_code=400, detail="Empty upload payload")
+ stored_path.unlink(missing_ok = True)
+ raise HTTPException(status_code = 400, detail = "Empty upload payload")
- return UploadDatasetResponse(filename=filename, stored_path=str(stored_path))
+ return UploadDatasetResponse(filename = filename, stored_path = str(stored_path))
-@router.get("/local", response_model=LocalDatasetsResponse)
+@router.get("/local", response_model = LocalDatasetsResponse)
def list_local_datasets(
current_subject: str = Depends(get_current_subject),
) -> LocalDatasetsResponse:
- return LocalDatasetsResponse(datasets=_build_local_dataset_items())
+ return LocalDatasetsResponse(datasets = _build_local_dataset_items())
-@router.post("/check-format", response_model=CheckFormatResponse)
+@router.post("/check-format", response_model = CheckFormatResponse)
def check_format(
request: CheckFormatRequest,
current_subject: str = Depends(get_current_subject),
@@ -341,9 +350,9 @@ def check_format(
# ββ Local file ββββββββββββββββββββββββββββββββββββββββββ
train_split = request.train_split or "train"
preview_slice, total_rows = _load_local_preview_slice(
- dataset_path=dataset_path,
- train_split=train_split,
- preview_size=PREVIEW_SIZE,
+ dataset_path = dataset_path,
+ train_split = train_split,
+ preview_size = PREVIEW_SIZE,
)
else:
# ββ HuggingFace dataset βββββββββββββββββββββββββββββββββ
@@ -352,23 +361,32 @@ def check_format(
try:
from huggingface_hub import HfApi
+
api = HfApi()
repo_files = api.list_repo_files(
request.dataset_name,
- repo_type="dataset",
- token=request.hf_token or None,
+ repo_type = "dataset",
+ token = request.hf_token or None,
)
- data_files = [f for f in repo_files if any(f.endswith(ext) for ext in DATA_EXTS)]
+ data_files = [
+ f for f in repo_files if any(f.endswith(ext) for ext in DATA_EXTS)
+ ]
# Prefer tabular formats over archives (e.g. images.zip β ImageFolder
# with synthetic image/label columns that don't match the real schema).
- tabular_files = [f for f in data_files if any(f.endswith(ext) for ext in _TABULAR_EXTS)]
+ tabular_files = [
+ f
+ for f in data_files
+ if any(f.endswith(ext) for ext in _TABULAR_EXTS)
+ ]
candidates = tabular_files or data_files
# When a subset is specified, narrow to files whose name matches
# (e.g. subset="testmini" β prefer "testmini.parquet").
if request.subset and candidates:
- subset_matches = [f for f in candidates if request.subset in Path(f).stem]
+ subset_matches = [
+ f for f in candidates if request.subset in Path(f).stem
+ ]
if subset_matches:
candidates = subset_matches
@@ -394,7 +412,11 @@ def check_format(
if preview_slice is None:
# Tier 2: full streaming (resolves all files β slow for large repos)
logger.info("Tier 2: falling back to full streaming load_dataset")
- load_kwargs = {"path": request.dataset_name, "split": request.train_split, "streaming": True}
+ load_kwargs = {
+ "path": request.dataset_name,
+ "split": request.train_split,
+ "streaming": True,
+ }
if request.subset:
load_kwargs["name"] = request.subset
if request.hf_token:
@@ -405,17 +427,19 @@ def check_format(
rows = list(islice(streamed_ds, PREVIEW_SIZE))
if not rows:
raise HTTPException(
- status_code=400,
- detail="Dataset appears to be empty or could not be streamed"
+ status_code = 400,
+ detail = "Dataset appears to be empty or could not be streamed",
)
preview_slice = Dataset.from_list(rows)
total_rows = None
# Run lightweight format check on the preview slice
- result = check_dataset_format(preview_slice, is_vlm=request.is_vlm)
+ result = check_dataset_format(preview_slice, is_vlm = request.is_vlm)
- logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_image={result.get('is_image', False)}")
+ logger.info(
+ f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_image={result.get('is_image', False)}"
+ )
# Generate preview samples
preview_samples = None
@@ -428,13 +452,15 @@ def check_format(
try:
format_result = format_dataset(
preview_slice,
- format_type="auto",
- num_proc=1, # Only 10 preview rows β no need for multiprocessing
+ format_type = "auto",
+ num_proc = 1, # Only 10 preview rows β no need for multiprocessing
)
processed = format_result["dataset"]
preview_samples = _serialize_preview_rows(processed)
except Exception as e:
- logger.warning(f"Processed preview generation failed (non-fatal): {e}")
+ logger.warning(
+ f"Processed preview generation failed (non-fatal): {e}"
+ )
preview_samples = _serialize_preview_rows(preview_slice)
else:
preview_samples = _serialize_preview_rows(preview_slice)
@@ -445,7 +471,9 @@ def check_format(
if image_col and image_col in (result.get("columns") or []):
try:
sample_val = preview_slice[0][image_col]
- if isinstance(sample_val, str) and sample_val.startswith(("http://", "https://")):
+ if isinstance(sample_val, str) and sample_val.startswith(
+ ("http://", "https://")
+ ):
url_warning = (
"This dataset contains image URLs instead of embedded images. "
"Images will be downloaded during training, which may be slow for large datasets."
@@ -456,33 +484,32 @@ def check_format(
pass
return CheckFormatResponse(
- requires_manual_mapping=result["requires_manual_mapping"],
- detected_format=result["detected_format"],
- columns=result["columns"],
- is_image=result.get("is_image", False),
- is_audio=result.get("is_audio", False),
- multimodal_columns=result.get("multimodal_columns"),
- suggested_mapping=result.get("suggested_mapping"),
- detected_image_column=result.get("detected_image_column"),
- detected_audio_column=result.get("detected_audio_column"),
- detected_text_column=result.get("detected_text_column"),
- detected_speaker_column=result.get("detected_speaker_column"),
- preview_samples=preview_samples,
- total_rows=total_rows,
- warning=warning,
+ requires_manual_mapping = result["requires_manual_mapping"],
+ detected_format = result["detected_format"],
+ columns = result["columns"],
+ is_image = result.get("is_image", False),
+ is_audio = result.get("is_audio", False),
+ multimodal_columns = result.get("multimodal_columns"),
+ suggested_mapping = result.get("suggested_mapping"),
+ detected_image_column = result.get("detected_image_column"),
+ detected_audio_column = result.get("detected_audio_column"),
+ detected_text_column = result.get("detected_text_column"),
+ detected_speaker_column = result.get("detected_speaker_column"),
+ preview_samples = preview_samples,
+ total_rows = total_rows,
+ warning = warning,
)
except HTTPException:
raise
except Exception as e:
- logger.error(f"Error checking dataset format: {e}", exc_info=True)
+ logger.error(f"Error checking dataset format: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to check dataset format: {str(e)}"
+ status_code = 500, detail = f"Failed to check dataset format: {str(e)}"
)
-@router.post("/ai-assist-mapping", response_model=AiAssistMappingResponse)
+@router.post("/ai-assist-mapping", response_model = AiAssistMappingResponse)
def ai_assist_mapping(
request: AiAssistMappingRequest,
current_subject: str = Depends(get_current_subject),
@@ -507,35 +534,32 @@ def ai_assist_mapping(
]
result = llm_conversion_advisor(
- column_names=request.columns,
- samples=truncated,
- dataset_name=request.dataset_name,
- hf_token=request.hf_token,
- model_name=request.model_name,
- model_type=request.model_type,
+ column_names = request.columns,
+ samples = truncated,
+ dataset_name = request.dataset_name,
+ hf_token = request.hf_token,
+ model_name = request.model_name,
+ model_type = request.model_type,
)
if result and result.get("success"):
return AiAssistMappingResponse(
- success=True,
- suggested_mapping=result.get("suggested_mapping"),
- system_prompt=result.get("system_prompt"),
- user_template=result.get("user_template"),
- assistant_template=result.get("assistant_template"),
- label_mapping=result.get("label_mapping"),
- dataset_type=result.get("dataset_type"),
- is_conversational=result.get("is_conversational"),
- user_notification=result.get("user_notification"),
+ success = True,
+ suggested_mapping = result.get("suggested_mapping"),
+ system_prompt = result.get("system_prompt"),
+ user_template = result.get("user_template"),
+ assistant_template = result.get("assistant_template"),
+ label_mapping = result.get("label_mapping"),
+ dataset_type = result.get("dataset_type"),
+ is_conversational = result.get("is_conversational"),
+ user_notification = result.get("user_notification"),
)
return AiAssistMappingResponse(
- success=False,
- warning="AI could not determine column roles. Please assign them manually.",
+ success = False,
+ warning = "AI could not determine column roles. Please assign them manually.",
)
except Exception as e:
- logger.error(f"AI assist mapping failed: {e}", exc_info=True)
- raise HTTPException(
- status_code=500,
- detail=f"AI assist failed: {str(e)}"
- )
+ logger.error(f"AI assist mapping failed: {e}", exc_info = True)
+ raise HTTPException(status_code = 500, detail = f"AI assist failed: {str(e)}")
diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py
index 4daca71143..3e60eaaf20 100644
--- a/studio/backend/routes/export.py
+++ b/studio/backend/routes/export.py
@@ -43,11 +43,7 @@ router = APIRouter()
logger = get_logger(__name__)
-
-
-
-
-@router.post("/load-checkpoint", response_model=ExportOperationResponse)
+@router.post("/load-checkpoint", response_model = ExportOperationResponse)
async def load_checkpoint(
request: LoadCheckpointRequest,
current_subject: str = Depends(get_current_subject),
@@ -65,6 +61,7 @@ async def load_checkpoint(
# before loading the export checkpoint (they'd compete for VRAM).
try:
from core.inference import get_inference_backend
+
inf = get_inference_backend()
if inf.active_model_name:
logger.info(
@@ -79,6 +76,7 @@ async def load_checkpoint(
try:
from core.training import get_training_backend
+
trn = get_training_backend()
if trn.is_training_active():
logger.info("Stopping active training to free GPU memory for export")
@@ -88,35 +86,39 @@ async def load_checkpoint(
for _ in range(60): # up to 30s
if not trn.is_training_active():
break
- import time; time.sleep(0.5)
+ import time
+
+ time.sleep(0.5)
else:
- logger.warning("Training subprocess did not exit within 30s, proceeding anyway")
+ logger.warning(
+ "Training subprocess did not exit within 30s, proceeding anyway"
+ )
except Exception as e:
logger.warning("Could not stop training: %s", e)
backend = get_export_backend()
success, message = backend.load_checkpoint(
- checkpoint_path=request.checkpoint_path,
- max_seq_length=request.max_seq_length,
- load_in_4bit=request.load_in_4bit,
- trust_remote_code=request.trust_remote_code,
+ checkpoint_path = request.checkpoint_path,
+ max_seq_length = request.max_seq_length,
+ load_in_4bit = request.load_in_4bit,
+ trust_remote_code = request.trust_remote_code,
)
if not success:
- raise HTTPException(status_code=400, detail=message)
+ raise HTTPException(status_code = 400, detail = message)
- return ExportOperationResponse(success=True, message=message)
+ return ExportOperationResponse(success = True, message = message)
except HTTPException:
raise
except Exception as e:
- logger.error(f"Error loading checkpoint: {e}", exc_info=True)
+ logger.error(f"Error loading checkpoint: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to load checkpoint: {str(e)}",
+ status_code = 500,
+ detail = f"Failed to load checkpoint: {str(e)}",
)
-@router.post("/cleanup", response_model=ExportOperationResponse)
+@router.post("/cleanup", response_model = ExportOperationResponse)
async def cleanup_export_memory(
current_subject: str = Depends(get_current_subject),
):
@@ -131,25 +133,25 @@ async def cleanup_export_memory(
if not success:
raise HTTPException(
- status_code=500,
- detail="Memory cleanup failed. See server logs for details.",
+ status_code = 500,
+ detail = "Memory cleanup failed. See server logs for details.",
)
return ExportOperationResponse(
- success=True,
- message="Memory cleanup completed successfully",
+ success = True,
+ message = "Memory cleanup completed successfully",
)
except HTTPException:
raise
except Exception as e:
- logger.error(f"Error during export memory cleanup: {e}", exc_info=True)
+ logger.error(f"Error during export memory cleanup: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to cleanup export memory: {str(e)}",
+ status_code = 500,
+ detail = f"Failed to cleanup export memory: {str(e)}",
)
-@router.get("/status", response_model=ExportStatusResponse)
+@router.get("/status", response_model = ExportStatusResponse)
async def get_export_status(
current_subject: str = Depends(get_current_subject),
):
@@ -159,19 +161,19 @@ async def get_export_status(
try:
backend = get_export_backend()
return ExportStatusResponse(
- current_checkpoint=backend.current_checkpoint,
- is_vision=bool(getattr(backend, "is_vision", False)),
- is_peft=bool(getattr(backend, "is_peft", False)),
+ current_checkpoint = backend.current_checkpoint,
+ is_vision = bool(getattr(backend, "is_vision", False)),
+ is_peft = bool(getattr(backend, "is_peft", False)),
)
except Exception as e:
- logger.error(f"Error getting export status: {e}", exc_info=True)
+ logger.error(f"Error getting export status: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to get export status: {str(e)}",
+ status_code = 500,
+ detail = f"Failed to get export status: {str(e)}",
)
-@router.post("/export/merged", response_model=ExportOperationResponse)
+@router.post("/export/merged", response_model = ExportOperationResponse)
async def export_merged_model(
request: ExportMergedModelRequest,
current_subject: str = Depends(get_current_subject),
@@ -184,29 +186,29 @@ async def export_merged_model(
try:
backend = get_export_backend()
success, message = backend.export_merged_model(
- save_directory=request.save_directory,
- format_type=request.format_type,
- push_to_hub=request.push_to_hub,
- repo_id=request.repo_id,
- hf_token=request.hf_token,
- private=request.private,
+ save_directory = request.save_directory,
+ format_type = request.format_type,
+ push_to_hub = request.push_to_hub,
+ repo_id = request.repo_id,
+ hf_token = request.hf_token,
+ private = request.private,
)
if not success:
- raise HTTPException(status_code=400, detail=message)
+ raise HTTPException(status_code = 400, detail = message)
- return ExportOperationResponse(success=True, message=message)
+ return ExportOperationResponse(success = True, message = message)
except HTTPException:
raise
except Exception as e:
- logger.error(f"Error exporting merged model: {e}", exc_info=True)
+ logger.error(f"Error exporting merged model: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to export merged model: {str(e)}",
+ status_code = 500,
+ detail = f"Failed to export merged model: {str(e)}",
)
-@router.post("/export/base", response_model=ExportOperationResponse)
+@router.post("/export/base", response_model = ExportOperationResponse)
async def export_base_model(
request: ExportBaseModelRequest,
current_subject: str = Depends(get_current_subject),
@@ -219,29 +221,29 @@ async def export_base_model(
try:
backend = get_export_backend()
success, message = backend.export_base_model(
- save_directory=request.save_directory,
- push_to_hub=request.push_to_hub,
- repo_id=request.repo_id,
- hf_token=request.hf_token,
- private=request.private,
- base_model_id=request.base_model_id,
+ save_directory = request.save_directory,
+ push_to_hub = request.push_to_hub,
+ repo_id = request.repo_id,
+ hf_token = request.hf_token,
+ private = request.private,
+ base_model_id = request.base_model_id,
)
if not success:
- raise HTTPException(status_code=400, detail=message)
+ raise HTTPException(status_code = 400, detail = message)
- return ExportOperationResponse(success=True, message=message)
+ return ExportOperationResponse(success = True, message = message)
except HTTPException:
raise
except Exception as e:
- logger.error(f"Error exporting base model: {e}", exc_info=True)
+ logger.error(f"Error exporting base model: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to export base model: {str(e)}",
+ status_code = 500,
+ detail = f"Failed to export base model: {str(e)}",
)
-@router.post("/export/gguf", response_model=ExportOperationResponse)
+@router.post("/export/gguf", response_model = ExportOperationResponse)
async def export_gguf(
request: ExportGGUFRequest,
current_subject: str = Depends(get_current_subject),
@@ -254,28 +256,28 @@ async def export_gguf(
try:
backend = get_export_backend()
success, message = backend.export_gguf(
- save_directory=request.save_directory,
- quantization_method=request.quantization_method,
- push_to_hub=request.push_to_hub,
- repo_id=request.repo_id,
- hf_token=request.hf_token,
+ save_directory = request.save_directory,
+ quantization_method = request.quantization_method,
+ push_to_hub = request.push_to_hub,
+ repo_id = request.repo_id,
+ hf_token = request.hf_token,
)
if not success:
- raise HTTPException(status_code=400, detail=message)
+ raise HTTPException(status_code = 400, detail = message)
- return ExportOperationResponse(success=True, message=message)
+ return ExportOperationResponse(success = True, message = message)
except HTTPException:
raise
except Exception as e:
- logger.error(f"Error exporting GGUF model: {e}", exc_info=True)
+ logger.error(f"Error exporting GGUF model: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to export GGUF model: {str(e)}",
+ status_code = 500,
+ detail = f"Failed to export GGUF model: {str(e)}",
)
-@router.post("/export/lora", response_model=ExportOperationResponse)
+@router.post("/export/lora", response_model = ExportOperationResponse)
async def export_lora_adapter(
request: ExportLoRAAdapterRequest,
current_subject: str = Depends(get_current_subject),
@@ -288,24 +290,22 @@ async def export_lora_adapter(
try:
backend = get_export_backend()
success, message = backend.export_lora_adapter(
- save_directory=request.save_directory,
- push_to_hub=request.push_to_hub,
- repo_id=request.repo_id,
- hf_token=request.hf_token,
- private=request.private,
+ save_directory = request.save_directory,
+ push_to_hub = request.push_to_hub,
+ repo_id = request.repo_id,
+ hf_token = request.hf_token,
+ private = request.private,
)
if not success:
- raise HTTPException(status_code=400, detail=message)
+ raise HTTPException(status_code = 400, detail = message)
- return ExportOperationResponse(success=True, message=message)
+ return ExportOperationResponse(success = True, message = message)
except HTTPException:
raise
except Exception as e:
- logger.error(f"Error exporting LoRA adapter: {e}", exc_info=True)
+ logger.error(f"Error exporting LoRA adapter: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to export LoRA adapter: {str(e)}",
+ status_code = 500,
+ detail = f"Failed to export LoRA adapter: {str(e)}",
)
-
-
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 2f6626d7fb..7de7336946 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -4,6 +4,7 @@
"""
Inference API routes for model loading and text generation.
"""
+
import sys
import time
import uuid
@@ -18,7 +19,6 @@ import asyncio
import threading
-
# Add backend directory to path
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
@@ -69,15 +69,15 @@ router = APIRouter()
logger = get_logger(__name__)
-
# GGUF inference backend (llama-server)
_llama_cpp_backend = LlamaCppBackend()
+
def get_llama_cpp_backend() -> LlamaCppBackend:
return _llama_cpp_backend
-@router.post("/load", response_model=LoadResponse)
+@router.post("/load", response_model = LoadResponse)
async def load_model(
request: LoadRequest,
current_subject: str = Depends(get_current_subject),
@@ -98,15 +98,15 @@ async def load_model(
# Create config using clean factory method
# is_lora is auto-detected from adapter_config.json on disk/HF
config = ModelConfig.from_identifier(
- model_id=request.model_path,
- hf_token=request.hf_token,
- gguf_variant=request.gguf_variant,
+ model_id = request.model_path,
+ hf_token = request.hf_token,
+ gguf_variant = request.gguf_variant,
)
if not config:
raise HTTPException(
- status_code=400,
- detail=f"Invalid model identifier: {request.model_path}"
+ status_code = 400,
+ detail = f"Invalid model identifier: {request.model_path}",
)
# ββ GGUF path: load via llama-server ββββββββββββββββββββββ
@@ -116,34 +116,36 @@ async def load_model(
# Unload any active Unsloth model first to free VRAM
if unsloth_backend.active_model_name:
- logger.info(f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF")
+ logger.info(
+ f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF"
+ )
unsloth_backend.unload_model(unsloth_backend.active_model_name)
# Route to HF mode or local mode based on config
if config.gguf_hf_repo:
# HF mode: llama-server downloads via -hf "repo:quant"
success = llama_backend.load_model(
- hf_repo=config.gguf_hf_repo,
- hf_variant=config.gguf_variant,
- hf_token=request.hf_token,
- model_identifier=config.identifier,
- is_vision=config.is_vision,
- n_ctx=request.max_seq_length,
+ hf_repo = config.gguf_hf_repo,
+ hf_variant = config.gguf_variant,
+ hf_token = request.hf_token,
+ model_identifier = config.identifier,
+ is_vision = config.is_vision,
+ n_ctx = request.max_seq_length,
)
else:
# Local mode: llama-server loads via -m
success = llama_backend.load_model(
- gguf_path=config.gguf_file,
- mmproj_path=config.gguf_mmproj_file,
- model_identifier=config.identifier,
- is_vision=config.is_vision,
- n_ctx=request.max_seq_length,
+ gguf_path = config.gguf_file,
+ mmproj_path = config.gguf_mmproj_file,
+ model_identifier = config.identifier,
+ is_vision = config.is_vision,
+ n_ctx = request.max_seq_length,
)
if not success:
raise HTTPException(
- status_code=500,
- detail=f"Failed to load GGUF model: {config.display_name}"
+ status_code = 500,
+ detail = f"Failed to load GGUF model: {config.display_name}",
)
logger.info(f"Loaded GGUF model via llama-server: {config.identifier}")
@@ -151,13 +153,13 @@ async def load_model(
inference_config = load_inference_config(config.identifier)
return LoadResponse(
- status="loaded",
- model=config.identifier,
- display_name=config.display_name,
- is_vision=config.is_vision,
- is_lora=False,
- is_gguf=True,
- inference=inference_config,
+ status = "loaded",
+ model = config.identifier,
+ display_name = config.display_name,
+ is_vision = config.is_vision,
+ is_lora = False,
+ is_gguf = True,
+ inference = inference_config,
)
# ββ Standard path: load via Unsloth/transformers ββββββββββ
@@ -172,9 +174,12 @@ async def load_model(
# Shut down any export subprocess to free VRAM
try:
from core.export import get_export_backend
+
exp_backend = get_export_backend()
if exp_backend.current_checkpoint:
- logger.info("Shutting down export subprocess to free GPU memory for inference")
+ logger.info(
+ "Shutting down export subprocess to free GPU memory for inference"
+ )
exp_backend._shutdown_subprocess()
exp_backend.current_checkpoint = None
exp_backend.is_vision = False
@@ -189,6 +194,7 @@ async def load_model(
if config.is_lora and config.path:
import json
from pathlib import Path
+
adapter_cfg_path = Path(config.path) / "adapter_config.json"
if adapter_cfg_path.exists():
try:
@@ -208,10 +214,16 @@ async def load_model(
)
load_in_4bit = True
elif training_method:
- logger.info(f"Training method: {training_method}, load_in_4bit={load_in_4bit}")
+ logger.info(
+ f"Training method: {training_method}, load_in_4bit={load_in_4bit}"
+ )
else:
# No unsloth_training_method β fallback to base model name
- if config.base_model and "-bnb-4bit" not in config.base_model.lower() and load_in_4bit:
+ if (
+ config.base_model
+ and "-bnb-4bit" not in config.base_model.lower()
+ and load_in_4bit
+ ):
logger.info(
f"No unsloth_training_method in adapter_config.json. "
f"Base model '{config.base_model}' has no -bnb-4bit suffix β "
@@ -223,29 +235,30 @@ async def load_model(
# Load the model
success = backend.load_model(
- config=config,
- max_seq_length=request.max_seq_length,
- load_in_4bit=load_in_4bit,
- hf_token=request.hf_token,
- trust_remote_code=request.trust_remote_code,
+ config = config,
+ max_seq_length = request.max_seq_length,
+ load_in_4bit = load_in_4bit,
+ hf_token = request.hf_token,
+ trust_remote_code = request.trust_remote_code,
)
if not success:
# Check if YAML says this model needs trust_remote_code
if not request.trust_remote_code:
model_defaults = load_model_defaults(config.identifier)
- yaml_trust = model_defaults.get("inference", {}).get("trust_remote_code", False)
+ yaml_trust = model_defaults.get("inference", {}).get(
+ "trust_remote_code", False
+ )
if yaml_trust:
raise HTTPException(
- status_code=400,
- detail=(
+ status_code = 400,
+ detail = (
f"Model '{config.display_name}' requires trust_remote_code to be enabled. "
f"Please enable 'Trust remote code' in Chat Settings and try again."
),
)
raise HTTPException(
- status_code=500,
- detail=f"Failed to load model: {config.display_name}"
+ status_code = 500, detail = f"Failed to load model: {config.display_name}"
)
logger.info(f"Loaded model: {config.identifier}")
@@ -254,29 +267,26 @@ async def load_model(
inference_config = load_inference_config(config.identifier)
return LoadResponse(
- status="loaded",
- model=config.identifier,
- display_name=config.display_name,
- is_vision=config.is_vision,
- is_lora=config.is_lora,
- is_gguf=False,
- is_audio=config.is_audio,
- audio_type=config.audio_type,
- has_audio_input=config.has_audio_input,
- inference=inference_config,
+ status = "loaded",
+ model = config.identifier,
+ display_name = config.display_name,
+ is_vision = config.is_vision,
+ is_lora = config.is_lora,
+ is_gguf = False,
+ is_audio = config.is_audio,
+ audio_type = config.audio_type,
+ has_audio_input = config.has_audio_input,
+ inference = inference_config,
)
except HTTPException:
raise
except Exception as e:
- logger.error(f"Error loading model: {e}", exc_info=True)
- raise HTTPException(
- status_code=500,
- detail=f"Failed to load model: {str(e)}"
- )
+ logger.error(f"Error loading model: {e}", exc_info = True)
+ raise HTTPException(status_code = 500, detail = f"Failed to load model: {str(e)}")
-@router.post("/validate", response_model=ValidateModelResponse)
+@router.post("/validate", response_model = ValidateModelResponse)
async def validate_model(
request: ValidateModelRequest,
current_subject: str = Depends(get_current_subject),
@@ -289,38 +299,41 @@ async def validate_model(
"""
try:
config = ModelConfig.from_identifier(
- model_id=request.model_path,
- hf_token=request.hf_token,
- gguf_variant=request.gguf_variant,
+ model_id = request.model_path,
+ hf_token = request.hf_token,
+ gguf_variant = request.gguf_variant,
)
if not config:
raise HTTPException(
- status_code=400,
- detail=f"Invalid model identifier: {request.model_path}",
+ status_code = 400,
+ detail = f"Invalid model identifier: {request.model_path}",
)
return ValidateModelResponse(
- valid=True,
- message="Model identifier is valid.",
- identifier=config.identifier,
- display_name=getattr(config, "display_name", config.identifier),
- is_gguf=getattr(config, "is_gguf", False),
- is_lora=getattr(config, "is_lora", False),
- is_vision=getattr(config, "is_vision", False),
+ valid = True,
+ message = "Model identifier is valid.",
+ identifier = config.identifier,
+ display_name = getattr(config, "display_name", config.identifier),
+ is_gguf = getattr(config, "is_gguf", False),
+ is_lora = getattr(config, "is_lora", False),
+ is_vision = getattr(config, "is_vision", False),
)
except HTTPException:
raise
except Exception as e:
- logger.error(f"Error validating model identifier '{request.model_path}': {e}", exc_info=True)
+ logger.error(
+ f"Error validating model identifier '{request.model_path}': {e}",
+ exc_info = True,
+ )
raise HTTPException(
- status_code=400,
- detail=f"Invalid model: {str(e)}",
+ status_code = 400,
+ detail = f"Invalid model: {str(e)}",
)
-@router.post("/unload", response_model=UnloadResponse)
+@router.post("/unload", response_model = UnloadResponse)
async def unload_model(
request: UnloadRequest,
current_subject: str = Depends(get_current_subject),
@@ -332,23 +345,23 @@ async def unload_model(
try:
# Check if the GGUF backend has this model loaded
llama_backend = get_llama_cpp_backend()
- if llama_backend.is_loaded and llama_backend.model_identifier == request.model_path:
+ if (
+ llama_backend.is_loaded
+ and llama_backend.model_identifier == request.model_path
+ ):
llama_backend.unload_model()
logger.info(f"Unloaded GGUF model: {request.model_path}")
- return UnloadResponse(status="unloaded", model=request.model_path)
+ return UnloadResponse(status = "unloaded", model = request.model_path)
# Otherwise, unload from Unsloth backend
backend = get_inference_backend()
backend.unload_model(request.model_path)
logger.info(f"Unloaded model: {request.model_path}")
- return UnloadResponse(status="unloaded", model=request.model_path)
+ return UnloadResponse(status = "unloaded", model = request.model_path)
except Exception as e:
- logger.error(f"Error unloading model: {e}", exc_info=True)
- raise HTTPException(
- status_code=500,
- detail=f"Failed to unload model: {str(e)}"
- )
+ logger.error(f"Error unloading model: {e}", exc_info = True)
+ raise HTTPException(status_code = 500, detail = f"Failed to unload model: {str(e)}")
@router.post("/generate/stream")
@@ -358,17 +371,16 @@ async def generate_stream(
):
"""
Generate a chat response with Server-Sent Events (SSE) streaming.
-
+
For vision models, provide image_base64 with the base64-encoded image.
"""
backend = get_inference_backend()
-
+
if not backend.active_model_name:
raise HTTPException(
- status_code=400,
- detail="No model loaded. Call POST /inference/load first."
+ status_code = 400, detail = "No model loaded. Call POST /inference/load first."
)
-
+
# Decode image if provided (for vision models)
image = None
if request.image_base64:
@@ -376,58 +388,57 @@ async def generate_stream(
import base64
from PIL import Image
from io import BytesIO
-
+
# Check if current model supports vision
model_info = backend.models.get(backend.active_model_name, {})
if not model_info.get("is_vision"):
raise HTTPException(
- status_code=400,
- detail="Image provided but current model is text-only. Load a vision model."
+ status_code = 400,
+ detail = "Image provided but current model is text-only. Load a vision model.",
)
-
+
image_data = base64.b64decode(request.image_base64)
image = Image.open(BytesIO(image_data))
image = backend.resize_image(image)
-
+
except HTTPException:
raise
except Exception as e:
raise HTTPException(
- status_code=400,
- detail=f"Failed to decode image: {str(e)}"
+ status_code = 400, detail = f"Failed to decode image: {str(e)}"
)
-
+
async def stream():
try:
for chunk in backend.generate_chat_response(
- messages=request.messages,
- system_prompt=request.system_prompt,
- image=image,
- temperature=request.temperature,
- top_p=request.top_p,
- top_k=request.top_k,
- max_new_tokens=request.max_new_tokens,
- repetition_penalty=request.repetition_penalty,
+ messages = request.messages,
+ system_prompt = request.system_prompt,
+ image = image,
+ temperature = request.temperature,
+ top_p = request.top_p,
+ top_k = request.top_k,
+ max_new_tokens = request.max_new_tokens,
+ repetition_penalty = request.repetition_penalty,
):
yield f"data: {json.dumps({'content': chunk})}\n\n"
yield "data: [DONE]\n\n"
-
+
except Exception as e:
backend.reset_generation_state()
- logger.error(f"Error during generation: {e}", exc_info=True)
- yield f"data: {json.dumps({'error': str(e)})}\n\n"
-
+ logger.error(f"Error during generation: {e}", exc_info = True)
+ yield f"data: {json.dumps({'error': 'An internal error occurred'})}\n\n"
+
return StreamingResponse(
stream(),
- media_type="text/event-stream",
- headers={
+ media_type = "text/event-stream",
+ headers = {
"Cache-Control": "no-cache",
"Connection": "keep-alive",
- }
+ },
)
-@router.get("/status", response_model=InferenceStatusResponse)
+@router.get("/status", response_model = InferenceStatusResponse)
async def get_status(
current_subject: str = Depends(get_current_subject),
):
@@ -441,12 +452,12 @@ async def get_status(
# If a GGUF model is loaded via llama-server, report that
if llama_backend.is_loaded:
return InferenceStatusResponse(
- active_model=llama_backend.model_identifier,
- is_vision=llama_backend.is_vision,
- is_gguf=True,
- gguf_variant=llama_backend.hf_variant,
- loading=[],
- loaded=[llama_backend.model_identifier],
+ active_model = llama_backend.model_identifier,
+ is_vision = llama_backend.is_vision,
+ is_gguf = True,
+ gguf_variant = llama_backend.hf_variant,
+ loading = [],
+ loaded = [llama_backend.model_identifier],
)
# Otherwise, report Unsloth backend status
@@ -464,22 +475,19 @@ async def get_status(
has_audio_input = model_info.get("has_audio_input", False)
return InferenceStatusResponse(
- active_model=backend.active_model_name,
- is_vision=is_vision,
- is_gguf=False,
- is_audio=is_audio,
- audio_type=audio_type,
- has_audio_input=has_audio_input,
- loading=list(getattr(backend, 'loading_models', set())),
- loaded=list(backend.models.keys()),
+ active_model = backend.active_model_name,
+ is_vision = is_vision,
+ is_gguf = False,
+ is_audio = is_audio,
+ audio_type = audio_type,
+ has_audio_input = has_audio_input,
+ loading = list(getattr(backend, "loading_models", set())),
+ loaded = list(backend.models.keys()),
)
except Exception as e:
- logger.error(f"Error getting status: {e}", exc_info=True)
- raise HTTPException(
- status_code=500,
- detail=f"Failed to get status: {str(e)}"
- )
+ logger.error(f"Error getting status: {e}", exc_info = True)
+ raise HTTPException(status_code = 500, detail = f"Failed to get status: {str(e)}")
# =====================================================================
@@ -488,7 +496,11 @@ async def get_status(
@router.post("/audio/generate")
-async def generate_audio(payload: ChatCompletionRequest, request: Request, current_subject: str = Depends(get_current_subject)):
+async def generate_audio(
+ payload: ChatCompletionRequest,
+ request: Request,
+ current_subject: str = Depends(get_current_subject),
+):
"""
Generate audio (TTS) from the latest user message.
Returns a JSON response with base64-encoded WAV audio.
@@ -498,22 +510,24 @@ async def generate_audio(payload: ChatCompletionRequest, request: Request, curre
backend = get_inference_backend()
if not backend.active_model_name:
- raise HTTPException(status_code=400, detail="No model loaded.")
+ raise HTTPException(status_code = 400, detail = "No model loaded.")
model_info = backend.models.get(backend.active_model_name, {})
if not model_info.get("is_audio"):
- raise HTTPException(status_code=400, detail="Active model is not an audio model.")
+ raise HTTPException(
+ status_code = 400, detail = "Active model is not an audio model."
+ )
# Extract text from the last user message
_, chat_messages, _ = _extract_content_parts(payload.messages)
if not chat_messages:
- raise HTTPException(status_code=400, detail="No messages provided.")
+ raise HTTPException(status_code = 400, detail = "No messages provided.")
last_user_msg = next(
(m for m in reversed(chat_messages) if m["role"] == "user"), None
)
if not last_user_msg:
- raise HTTPException(status_code=400, detail="No user message found.")
+ raise HTTPException(status_code = 400, detail = "No user message found.")
text = last_user_msg["content"]
@@ -521,42 +535,46 @@ async def generate_audio(payload: ChatCompletionRequest, request: Request, curre
wav_bytes, sample_rate = await asyncio.get_event_loop().run_in_executor(
None,
lambda: backend.generate_audio_response(
- text=text,
- temperature=payload.temperature,
- top_p=payload.top_p,
- top_k=payload.top_k,
- min_p=payload.min_p,
- max_new_tokens=payload.max_tokens or 2048,
- repetition_penalty=payload.repetition_penalty,
- use_adapter=payload.use_adapter,
+ text = text,
+ temperature = payload.temperature,
+ top_p = payload.top_p,
+ top_k = payload.top_k,
+ min_p = payload.min_p,
+ max_new_tokens = payload.max_tokens or 2048,
+ repetition_penalty = payload.repetition_penalty,
+ use_adapter = payload.use_adapter,
),
)
audio_b64 = base64.b64encode(wav_bytes).decode("ascii")
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
- return JSONResponse(content={
- "id": completion_id,
- "object": "chat.completion.audio",
- "model": backend.active_model_name,
- "audio": {
- "data": audio_b64,
- "format": "wav",
- "sample_rate": sample_rate,
- },
- "choices": [{
- "index": 0,
- "message": {
- "role": "assistant",
- "content": f"[Generated audio from: \"{text[:100]}\"]",
+ return JSONResponse(
+ content = {
+ "id": completion_id,
+ "object": "chat.completion.audio",
+ "model": backend.active_model_name,
+ "audio": {
+ "data": audio_b64,
+ "format": "wav",
+ "sample_rate": sample_rate,
},
- "finish_reason": "stop",
- }],
- })
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": f'[Generated audio from: "{text[:100]}"]',
+ },
+ "finish_reason": "stop",
+ }
+ ],
+ }
+ )
except Exception as e:
- logger.error(f"Audio generation error: {e}", exc_info=True)
- raise HTTPException(status_code=500, detail=str(e))
+ logger.error(f"Audio generation error: {e}", exc_info = True)
+ raise HTTPException(status_code = 500, detail = str(e))
# =====================================================================
@@ -576,9 +594,9 @@ def _decode_audio_base64(b64: str) -> np.ndarray:
# torchaudio.load needs a file path or file-like object with format hint
# Write to a temp file so torchaudio can auto-detect the format
with tempfile.NamedTemporaryFile(
- suffix=".audio",
- delete=False,
- dir=str(ensure_dir(tmp_root())),
+ suffix = ".audio",
+ delete = False,
+ dir = str(ensure_dir(tmp_root())),
) as tmp:
tmp.write(raw)
tmp_path = tmp.name
@@ -589,11 +607,11 @@ def _decode_audio_base64(b64: str) -> np.ndarray:
# Convert to mono if stereo
if waveform.shape[0] > 1:
- waveform = waveform.mean(dim=0, keepdim=True)
+ waveform = waveform.mean(dim = 0, keepdim = True)
# Resample to 16kHz if needed
if sr != 16000:
- resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000)
+ resampler = torchaudio.transforms.Resample(orig_freq = sr, new_freq = 16000)
waveform = resampler(waveform)
return waveform.squeeze(0).numpy()
@@ -683,8 +701,8 @@ async def openai_chat_completions(
backend = get_inference_backend()
if not backend.active_model_name:
raise HTTPException(
- status_code=400,
- detail="No model loaded. Call POST /inference/load first.",
+ status_code = 400,
+ detail = "No model loaded. Call POST /inference/load first.",
)
model_name = backend.active_model_name or payload.model
@@ -697,8 +715,8 @@ async def openai_chat_completions(
# ββ Whisper without audio: return clear error ββ
if model_info.get("audio_type") == "whisper" and not payload.audio_base64:
raise HTTPException(
- status_code=400,
- detail="Whisper models require audio input. Please upload an audio file.",
+ status_code = 400,
+ detail = "Whisper models require audio input. Please upload an audio file.",
)
# ββ Audio INPUT path: decode WAV and route to audio input generation ββ
@@ -712,30 +730,38 @@ async def openai_chat_completions(
def audio_input_generate():
if model_info.get("audio_type") == "whisper":
return backend.generate_whisper_response(
- audio_array=audio_array,
- cancel_event=cancel_event,
+ audio_array = audio_array,
+ cancel_event = cancel_event,
)
return backend.generate_audio_input_response(
- messages=chat_messages,
- system_prompt=system_prompt,
- audio_array=audio_array,
- temperature=payload.temperature,
- top_p=payload.top_p,
- top_k=payload.top_k,
- min_p=payload.min_p,
- max_new_tokens=payload.max_tokens or 2048,
- repetition_penalty=payload.repetition_penalty,
- cancel_event=cancel_event,
+ messages = chat_messages,
+ system_prompt = system_prompt,
+ audio_array = audio_array,
+ temperature = payload.temperature,
+ top_p = payload.top_p,
+ top_k = payload.top_k,
+ min_p = payload.min_p,
+ max_new_tokens = payload.max_tokens or 2048,
+ repetition_penalty = payload.repetition_penalty,
+ cancel_event = cancel_event,
)
if payload.stream:
+
async def audio_input_stream():
try:
first_chunk = ChatCompletionChunk(
- id=completion_id, created=created, model=model_name,
- choices=[ChunkChoice(delta=ChoiceDelta(role="assistant"), finish_reason=None)],
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ ChunkChoice(
+ delta = ChoiceDelta(role = "assistant"),
+ finish_reason = None,
+ )
+ ],
)
- yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n"
+ yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n"
for chunk_text in audio_input_generate():
if await request.is_disconnected():
@@ -743,36 +769,60 @@ async def openai_chat_completions(
return
if chunk_text:
chunk = ChatCompletionChunk(
- id=completion_id, created=created, model=model_name,
- choices=[ChunkChoice(delta=ChoiceDelta(content=chunk_text), finish_reason=None)],
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ ChunkChoice(
+ delta = ChoiceDelta(content = chunk_text),
+ finish_reason = None,
+ )
+ ],
)
- yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n"
+ yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
final_chunk = ChatCompletionChunk(
- id=completion_id, created=created, model=model_name,
- choices=[ChunkChoice(delta=ChoiceDelta(), finish_reason="stop")],
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ ChunkChoice(delta = ChoiceDelta(), finish_reason = "stop")
+ ],
)
- yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
+ yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n"
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
cancel_event.set()
raise
except Exception as e:
- logger.error(f"Error during audio input streaming: {e}", exc_info=True)
- yield f"data: {json.dumps({'error': {'message': str(e), 'type': 'server_error'}})}\n\n"
+ logger.error(
+ f"Error during audio input streaming: {e}", exc_info = True
+ )
+ yield f"data: {json.dumps({'error': {'message': 'An internal error occurred', 'type': 'server_error'}})}\n\n"
return StreamingResponse(
audio_input_stream(),
- media_type="text/event-stream",
- headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
+ media_type = "text/event-stream",
+ headers = {
+ "Cache-Control": "no-cache",
+ "Connection": "keep-alive",
+ "X-Accel-Buffering": "no",
+ },
)
else:
full_text = "".join(audio_input_generate())
response = ChatCompletion(
- id=completion_id, created=created, model=model_name,
- choices=[CompletionChoice(message=CompletionMessage(content=full_text), finish_reason="stop")],
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ CompletionChoice(
+ message = CompletionMessage(content = full_text),
+ finish_reason = "stop",
+ )
+ ],
)
- return JSONResponse(content=response.model_dump())
+ return JSONResponse(content = response.model_dump())
# ββ Parse messages (handles multimodal content parts) βββββ
system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(
@@ -781,8 +831,8 @@ async def openai_chat_completions(
if not chat_messages:
raise HTTPException(
- status_code=400,
- detail="At least one non-system message is required.",
+ status_code = 400,
+ detail = "At least one non-system message is required.",
)
# ββ GGUF path: proxy to llama-server /v1/chat/completions ββ
@@ -791,8 +841,8 @@ async def openai_chat_completions(
image_b64 = extracted_image_b64 or payload.image_base64
if image_b64 and not llama_backend.is_vision:
raise HTTPException(
- status_code=400,
- detail="Image provided but current GGUF model does not support vision.",
+ status_code = 400,
+ detail = "Image provided but current GGUF model does not support vision.",
)
# Build message list with system prompt prepended
@@ -808,31 +858,34 @@ async def openai_chat_completions(
def gguf_generate():
return llama_backend.generate_chat_completion(
- messages=gguf_messages,
- image_b64=image_b64,
- temperature=payload.temperature,
- top_p=payload.top_p,
- top_k=payload.top_k,
- min_p=payload.min_p,
- max_tokens=payload.max_tokens or 2048,
- repetition_penalty=payload.repetition_penalty,
- cancel_event=cancel_event,
+ messages = gguf_messages,
+ image_b64 = image_b64,
+ temperature = payload.temperature,
+ top_p = payload.top_p,
+ top_k = payload.top_k,
+ min_p = payload.min_p,
+ max_tokens = payload.max_tokens or 2048,
+ repetition_penalty = payload.repetition_penalty,
+ cancel_event = cancel_event,
)
if payload.stream:
+
async def gguf_stream_chunks():
try:
# First chunk: role
first_chunk = ChatCompletionChunk(
- id=completion_id,
- created=created,
- model=model_name,
- choices=[ChunkChoice(
- delta=ChoiceDelta(role="assistant"),
- finish_reason=None,
- )],
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ ChunkChoice(
+ delta = ChoiceDelta(role = "assistant"),
+ finish_reason = None,
+ )
+ ],
)
- yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n"
+ yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n"
# Content chunks β llama backend yields cumulative text
prev_text = ""
@@ -840,48 +893,55 @@ async def openai_chat_completions(
if await request.is_disconnected():
cancel_event.set()
return
- new_text = cumulative[len(prev_text):]
+ new_text = cumulative[len(prev_text) :]
prev_text = cumulative
if not new_text:
continue
chunk = ChatCompletionChunk(
- id=completion_id,
- created=created,
- model=model_name,
- choices=[ChunkChoice(
- delta=ChoiceDelta(content=new_text),
- finish_reason=None,
- )],
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ ChunkChoice(
+ delta = ChoiceDelta(content = new_text),
+ finish_reason = None,
+ )
+ ],
)
- yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n"
+ yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
# Final chunk
final_chunk = ChatCompletionChunk(
- id=completion_id,
- created=created,
- model=model_name,
- choices=[ChunkChoice(
- delta=ChoiceDelta(),
- finish_reason="stop",
- )],
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ ChunkChoice(
+ delta = ChoiceDelta(),
+ finish_reason = "stop",
+ )
+ ],
)
- yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
+ yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n"
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
cancel_event.set()
raise
except Exception as e:
- logger.error(f"Error during GGUF streaming: {e}", exc_info=True)
+ logger.error(f"Error during GGUF streaming: {e}", exc_info = True)
error_chunk = {
- "error": {"message": str(e), "type": "server_error"},
+ "error": {
+ "message": "An internal error occurred",
+ "type": "server_error",
+ },
}
yield f"data: {json.dumps(error_chunk)}\n\n"
return StreamingResponse(
gguf_stream_chunks(),
- media_type="text/event-stream",
- headers={
+ media_type = "text/event-stream",
+ headers = {
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
@@ -894,19 +954,21 @@ async def openai_chat_completions(
full_text = token
response = ChatCompletion(
- id=completion_id,
- created=created,
- model=model_name,
- choices=[CompletionChoice(
- message=CompletionMessage(content=full_text),
- finish_reason="stop",
- )],
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ CompletionChoice(
+ message = CompletionMessage(content = full_text),
+ finish_reason = "stop",
+ )
+ ],
)
- return JSONResponse(content=response.model_dump())
+ return JSONResponse(content = response.model_dump())
except Exception as e:
- logger.error(f"Error during GGUF completion: {e}", exc_info=True)
- raise HTTPException(status_code=500, detail=str(e))
+ logger.error(f"Error during GGUF completion: {e}", exc_info = True)
+ raise HTTPException(status_code = 500, detail = str(e))
# ββ Standard Unsloth path βββββββββββββββββββββββββββββββββ
@@ -923,8 +985,8 @@ async def openai_chat_completions(
model_info = backend.models.get(backend.active_model_name, {})
if not model_info.get("is_vision"):
raise HTTPException(
- status_code=400,
- detail="Image provided but current model is text-only. Load a vision model.",
+ status_code = 400,
+ detail = "Image provided but current model is text-only. Load a vision model.",
)
image_data = base64.b64decode(image_b64)
@@ -934,52 +996,59 @@ async def openai_chat_completions(
except HTTPException:
raise
except Exception as e:
- raise HTTPException(status_code=400, detail=f"Failed to decode image: {e}")
+ raise HTTPException(status_code = 400, detail = f"Failed to decode image: {e}")
# Shared generation kwargs
gen_kwargs = dict(
- messages=chat_messages,
- system_prompt=system_prompt,
- image=image,
- temperature=payload.temperature,
- top_p=payload.top_p,
- top_k=payload.top_k,
- min_p=payload.min_p,
- max_new_tokens=payload.max_tokens or 2048,
- repetition_penalty=payload.repetition_penalty,
+ messages = chat_messages,
+ system_prompt = system_prompt,
+ image = image,
+ temperature = payload.temperature,
+ top_p = payload.top_p,
+ top_k = payload.top_k,
+ min_p = payload.min_p,
+ max_new_tokens = payload.max_tokens or 2048,
+ repetition_penalty = payload.repetition_penalty,
)
# Choose generation path (adapter-controlled or standard)
cancel_event = threading.Event()
if payload.use_adapter is not None:
+
def generate():
return backend.generate_with_adapter_control(
- use_adapter=payload.use_adapter,
- cancel_event=cancel_event,
+ use_adapter = payload.use_adapter,
+ cancel_event = cancel_event,
**gen_kwargs,
)
else:
+
def generate():
- return backend.generate_chat_response(cancel_event=cancel_event, **gen_kwargs)
+ return backend.generate_chat_response(
+ cancel_event = cancel_event, **gen_kwargs
+ )
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
created = int(time.time())
# ββ Streaming response ββββββββββββββββββββββββββββββββββββββββ
if payload.stream:
+
async def stream_chunks():
try:
first_chunk = ChatCompletionChunk(
- id=completion_id,
- created=created,
- model=model_name,
- choices=[ChunkChoice(
- delta=ChoiceDelta(role="assistant"),
- finish_reason=None,
- )],
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ ChunkChoice(
+ delta = ChoiceDelta(role = "assistant"),
+ finish_reason = None,
+ )
+ ],
)
- yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n"
+ yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n"
prev_text = ""
for cumulative in generate():
@@ -987,31 +1056,35 @@ async def openai_chat_completions(
cancel_event.set()
backend.reset_generation_state()
return
- new_text = cumulative[len(prev_text):]
+ new_text = cumulative[len(prev_text) :]
prev_text = cumulative
if not new_text:
continue
chunk = ChatCompletionChunk(
- id=completion_id,
- created=created,
- model=model_name,
- choices=[ChunkChoice(
- delta=ChoiceDelta(content=new_text),
- finish_reason=None,
- )],
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ ChunkChoice(
+ delta = ChoiceDelta(content = new_text),
+ finish_reason = None,
+ )
+ ],
)
- yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n"
+ yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
final_chunk = ChatCompletionChunk(
- id=completion_id,
- created=created,
- model=model_name,
- choices=[ChunkChoice(
- delta=ChoiceDelta(),
- finish_reason="stop",
- )],
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ ChunkChoice(
+ delta = ChoiceDelta(),
+ finish_reason = "stop",
+ )
+ ],
)
- yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
+ yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n"
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
@@ -1020,16 +1093,19 @@ async def openai_chat_completions(
raise
except Exception as e:
backend.reset_generation_state()
- logger.error(f"Error during OpenAI streaming: {e}", exc_info=True)
+ logger.error(f"Error during OpenAI streaming: {e}", exc_info = True)
error_chunk = {
- "error": {"message": str(e), "type": "server_error"},
+ "error": {
+ "message": "An internal error occurred",
+ "type": "server_error",
+ },
}
yield f"data: {json.dumps(error_chunk)}\n\n"
return StreamingResponse(
stream_chunks(),
- media_type="text/event-stream",
- headers={
+ media_type = "text/event-stream",
+ headers = {
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
@@ -1044,26 +1120,29 @@ async def openai_chat_completions(
full_text = token
response = ChatCompletion(
- id=completion_id,
- created=created,
- model=model_name,
- choices=[CompletionChoice(
- message=CompletionMessage(content=full_text),
- finish_reason="stop",
- )],
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ CompletionChoice(
+ message = CompletionMessage(content = full_text),
+ finish_reason = "stop",
+ )
+ ],
)
- return JSONResponse(content=response.model_dump())
+ return JSONResponse(content = response.model_dump())
except Exception as e:
backend.reset_generation_state()
- logger.error(f"Error during OpenAI completion: {e}", exc_info=True)
- raise HTTPException(status_code=500, detail=str(e))
+ logger.error(f"Error during OpenAI completion: {e}", exc_info = True)
+ raise HTTPException(status_code = 500, detail = str(e))
# =====================================================================
# OpenAI-Compatible Models Listing (/models β /v1/models)
# =====================================================================
+
@router.get("/models")
async def openai_list_models(
current_subject: str = Depends(get_current_subject),
@@ -1079,19 +1158,23 @@ async def openai_list_models(
# Check GGUF backend
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded:
- models.append({
- "id": llama_backend.model_identifier,
- "object": "model",
- "owned_by": "local",
- })
+ models.append(
+ {
+ "id": llama_backend.model_identifier,
+ "object": "model",
+ "owned_by": "local",
+ }
+ )
# Check Unsloth backend
backend = get_inference_backend()
if backend.active_model_name:
- models.append({
- "id": backend.active_model_name,
- "object": "model",
- "owned_by": "local",
- })
+ models.append(
+ {
+ "id": backend.active_model_name,
+ "object": "model",
+ "owned_by": "local",
+ }
+ )
return {"object": "list", "data": models}
diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index b6b36de346..99b42b491e 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -4,6 +4,8 @@
"""
Model Management API routes
"""
+
+import os
import sys
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Query
@@ -31,9 +33,18 @@ try:
list_gguf_variants,
ModelConfig,
)
- from utils.models.model_config import _pick_best_gguf, _extract_quant_label, is_audio_input_type
+ from utils.models.model_config import (
+ _pick_best_gguf,
+ _extract_quant_label,
+ is_audio_input_type,
+ )
from core.inference import get_inference_backend
- from utils.paths import outputs_root, exports_root, resolve_output_dir, resolve_export_dir
+ from utils.paths import (
+ outputs_root,
+ exports_root,
+ resolve_output_dir,
+ resolve_export_dir,
+ )
except ImportError:
# Fallback: try to import from parent directory
parent_backend = backend_path.parent / "backend"
@@ -50,9 +61,18 @@ except ImportError:
list_gguf_variants,
ModelConfig,
)
- from utils.models.model_config import _pick_best_gguf, _extract_quant_label, is_audio_input_type
+ from utils.models.model_config import (
+ _pick_best_gguf,
+ _extract_quant_label,
+ is_audio_input_type,
+ )
from core.inference import get_inference_backend
- from utils.paths import outputs_root, exports_root, resolve_output_dir, resolve_export_dir
+ from utils.paths import (
+ outputs_root,
+ exports_root,
+ resolve_output_dir,
+ resolve_export_dir,
+ )
from models import (
CheckpointInfo,
@@ -66,13 +86,19 @@ from models import (
ModelListResponse,
)
from models.models import GgufVariantDetail, GgufVariantsResponse, ModelType
-from models.responses import LoRABaseModelResponse, VisionCheckResponse, EmbeddingCheckResponse
+from models.responses import (
+ LoRABaseModelResponse,
+ VisionCheckResponse,
+ EmbeddingCheckResponse,
+)
router = APIRouter()
logger = get_logger(__name__)
-def derive_model_type(is_vision: bool, audio_type: Optional[str], is_embedding: bool = False) -> ModelType:
+def derive_model_type(
+ is_vision: bool, audio_type: Optional[str], is_embedding: bool = False
+) -> ModelType:
"""Collapse individual capability flags into a single model modality string."""
if is_embedding:
return "embeddings"
@@ -83,12 +109,11 @@ def derive_model_type(is_vision: bool, audio_type: Optional[str], is_embedding:
return "text"
-
-
def _resolve_hf_cache_dir() -> Path:
"""Resolve local HF cache root used by hub downloads."""
try:
from huggingface_hub.constants import HF_HUB_CACHE
+
return Path(HF_HUB_CACHE)
except Exception:
return Path.home() / ".cache" / "huggingface" / "hub"
@@ -117,11 +142,11 @@ def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]:
updated_at = None
found.append(
LocalModelInfo(
- id=str(child),
- display_name=child.name,
- path=str(child),
- source="models_dir",
- updated_at=updated_at,
+ id = str(child),
+ display_name = child.name,
+ path = str(child),
+ source = "models_dir",
+ updated_at = updated_at,
),
)
# Also scan for standalone .gguf files directly in the models directory
@@ -133,11 +158,11 @@ def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]:
updated_at = None
found.append(
LocalModelInfo(
- id=str(gguf_file),
- display_name=gguf_file.stem,
- path=str(gguf_file),
- source="models_dir",
- updated_at=updated_at,
+ id = str(gguf_file),
+ display_name = gguf_file.stem,
+ path = str(gguf_file),
+ source = "models_dir",
+ updated_at = updated_at,
),
)
@@ -153,7 +178,7 @@ def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]:
if not repo_dir.is_dir():
continue
- repo_name = repo_dir.name[len("models--"):]
+ repo_name = repo_dir.name[len("models--") :]
if not repo_name:
continue
model_id = repo_name.replace("--", "/")
@@ -165,28 +190,53 @@ def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]:
found.append(
LocalModelInfo(
- id=model_id,
- model_id=model_id,
- display_name=model_id.split("/")[-1],
- path=str(repo_dir),
- source="hf_cache",
- updated_at=updated_at,
+ id = model_id,
+ model_id = model_id,
+ display_name = model_id.split("/")[-1],
+ path = str(repo_dir),
+ source = "hf_cache",
+ updated_at = updated_at,
),
)
return found
-@router.get("/local", response_model=LocalModelListResponse)
+@router.get("/local", response_model = LocalModelListResponse)
async def list_local_models(
- models_dir: str = Query(default="./models", description="Directory to scan for local model folders"),
+ models_dir: str = Query(
+ default = "./models", description = "Directory to scan for local model folders"
+ ),
current_subject: str = Depends(get_current_subject),
):
"""
List local model candidates from custom models dir and HF cache.
"""
+ # Validate models_dir against an allowlist of trusted directories.
+ # Only the trusted Path objects are used for filesystem access -- the
+ # user-supplied string is only used for matching, never for path construction.
+ hf_cache_dir = _resolve_hf_cache_dir()
+ allowed_roots = [Path("./models").resolve(), hf_cache_dir]
+ try:
+ from utils.paths import studio_root, outputs_root
+
+ allowed_roots.extend([studio_root(), outputs_root()])
+ except Exception:
+ pass
+
+ requested = os.path.realpath(os.path.expanduser(models_dir))
+ models_root = None
+ for root in allowed_roots:
+ root_str = os.path.realpath(str(root))
+ if requested == root_str or requested.startswith(root_str + os.sep):
+ models_root = root # Use the trusted root, not the user-supplied path
+ break
+ if models_root is None:
+ raise HTTPException(
+ status_code = 403,
+ detail = "Directory not allowed",
+ )
+
try:
- models_root = Path(models_dir).expanduser().resolve()
- hf_cache_dir = _resolve_hf_cache_dir()
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
deduped: dict[str, LocalModelInfo] = {}
@@ -196,88 +246,80 @@ async def list_local_models(
models = sorted(
deduped.values(),
- key=lambda item: (item.updated_at or 0),
- reverse=True,
+ key = lambda item: (item.updated_at or 0),
+ reverse = True,
)
return LocalModelListResponse(
- models_dir=str(models_root),
- hf_cache_dir=str(hf_cache_dir),
- models=models,
+ models_dir = str(models_root),
+ hf_cache_dir = str(hf_cache_dir),
+ models = models,
)
except Exception as e:
- logger.error(f"Error listing local models: {e}", exc_info=True)
+ logger.error(f"Error listing local models: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to list local models: {str(e)}",
+ status_code = 500,
+ detail = f"Failed to list local models: {str(e)}",
)
-
-
@router.get("/list")
async def list_models(
current_subject: str = Depends(get_current_subject),
):
"""
List available models (default models and loaded models).
-
+
This endpoint returns the default models and any currently loaded models.
"""
try:
inference_backend = get_inference_backend()
-
+
# Get default models
default_models = inference_backend.default_models
-
+
# Get loaded models
loaded_models = []
for model_name, model_data in inference_backend.models.items():
_is_vision = model_data.get("is_vision", False)
_audio_type = model_data.get("audio_type")
model_info = ModelDetails(
- id=model_name,
- name=model_name.split("/")[-1] if "/" in model_name else model_name,
- is_vision=_is_vision,
- is_lora=model_data.get("is_lora", False),
- is_audio=model_data.get("is_audio", False),
- audio_type=_audio_type,
- has_audio_input=model_data.get("has_audio_input", False),
- model_type=derive_model_type(_is_vision, _audio_type),
+ id = model_name,
+ name = model_name.split("/")[-1] if "/" in model_name else model_name,
+ is_vision = _is_vision,
+ is_lora = model_data.get("is_lora", False),
+ is_audio = model_data.get("is_audio", False),
+ audio_type = _audio_type,
+ has_audio_input = model_data.get("has_audio_input", False),
+ model_type = derive_model_type(_is_vision, _audio_type),
)
loaded_models.append(model_info)
-
+
# Combine default and loaded models
all_models = []
seen_ids = set()
-
+
# Add default models
for model_id in default_models:
if model_id not in seen_ids:
model_info = ModelDetails(
- id=model_id,
- name=model_id.split("/")[-1] if "/" in model_id else model_id
+ id = model_id,
+ name = model_id.split("/")[-1] if "/" in model_id else model_id,
)
all_models.append(model_info)
seen_ids.add(model_id)
-
+
# Add loaded models
for model_info in loaded_models:
if model_info.id not in seen_ids:
all_models.append(model_info)
seen_ids.add(model_info.id)
-
- return ModelListResponse(
- models=all_models,
- default_models=default_models
- )
-
+
+ return ModelListResponse(models = all_models, default_models = default_models)
+
except Exception as e:
- logger.error(f"Error listing models: {e}", exc_info=True)
- raise HTTPException(
- status_code=500,
- detail=f"Failed to list models: {str(e)}"
- )
+ logger.error(f"Error listing models: {e}", exc_info = True)
+ raise HTTPException(status_code = 500, detail = f"Failed to list models: {str(e)}")
@router.get("/config/{model_name:path}")
@@ -293,18 +335,20 @@ async def get_model_config(
"""
try:
from utils.models.model_config import is_local_path
+
if not is_local_path(model_name):
model_name = model_name.lower()
-
+
logger.info(f"Getting model config for: {model_name}")
from utils.models.model_config import detect_audio_type
+
# Load model defaults from backend
config_dict = load_model_defaults(model_name)
# Detect model capabilities (pass HF token for gated models)
is_vision = is_vision_model(model_name)
- is_embedding = is_embedding_model(model_name, hf_token=hf_token)
- audio_type = detect_audio_type(model_name, hf_token=hf_token)
+ is_embedding = is_embedding_model(model_name, hf_token = hf_token)
+ audio_type = detect_audio_type(model_name, hf_token = hf_token)
# Check if it's a LoRA adapter
is_lora = False
@@ -316,33 +360,38 @@ async def get_model_config(
except Exception:
pass
- logger.info(f"Model config result for {model_name}: is_vision={is_vision}, is_embedding={is_embedding}, audio_type={audio_type}, is_lora={is_lora}")
- return ModelDetails(
- id=model_name,
- model_name=model_name,
- config=config_dict,
- is_vision=is_vision,
- is_embedding=is_embedding,
- is_lora=is_lora,
- is_audio=audio_type is not None,
- audio_type=audio_type,
- has_audio_input=is_audio_input_type(audio_type),
- model_type=derive_model_type(is_vision, audio_type, is_embedding),
- base_model=base_model,
+ logger.info(
+ f"Model config result for {model_name}: is_vision={is_vision}, is_embedding={is_embedding}, audio_type={audio_type}, is_lora={is_lora}"
)
-
+ return ModelDetails(
+ id = model_name,
+ model_name = model_name,
+ config = config_dict,
+ is_vision = is_vision,
+ is_embedding = is_embedding,
+ is_lora = is_lora,
+ is_audio = audio_type is not None,
+ audio_type = audio_type,
+ has_audio_input = is_audio_input_type(audio_type),
+ model_type = derive_model_type(is_vision, audio_type, is_embedding),
+ base_model = base_model,
+ )
+
except Exception as e:
- logger.error(f"Error getting model config: {e}", exc_info=True)
+ logger.error(f"Error getting model config: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to get model config: {str(e)}"
+ status_code = 500, detail = f"Failed to get model config: {str(e)}"
)
@router.get("/loras")
async def scan_loras(
- outputs_dir: str = Query(default=str(outputs_root()), description="Directory to scan for LoRA adapters"),
- exports_dir: str = Query(default=str(exports_root()), description="Directory to scan for exported models"),
+ outputs_dir: str = Query(
+ default = str(outputs_root()), description = "Directory to scan for LoRA adapters"
+ ),
+ exports_dir: str = Query(
+ default = str(exports_root()), description = "Directory to scan for exported models"
+ ),
current_subject: str = Depends(get_current_subject),
):
"""
@@ -357,102 +406,101 @@ async def scan_loras(
lora_list = []
# Scan training outputs
- trained_loras = scan_trained_loras(outputs_dir=resolved_outputs_dir)
+ trained_loras = scan_trained_loras(outputs_dir = resolved_outputs_dir)
for display_name, adapter_path in trained_loras:
base_model = get_base_model_from_lora(adapter_path)
- lora_list.append(LoRAInfo(
- display_name=display_name,
- adapter_path=adapter_path,
- base_model=base_model,
- source="training",
- ))
+ lora_list.append(
+ LoRAInfo(
+ display_name = display_name,
+ adapter_path = adapter_path,
+ base_model = base_model,
+ source = "training",
+ )
+ )
# Scan exported models (merged, LoRA, base β skips GGUF)
- exported = scan_exported_models(exports_dir=resolved_exports_dir)
+ exported = scan_exported_models(exports_dir = resolved_exports_dir)
for display_name, model_path, export_type, base_model in exported:
- lora_list.append(LoRAInfo(
- display_name=display_name,
- adapter_path=model_path,
- base_model=base_model,
- source="exported",
- export_type=export_type,
- ))
+ lora_list.append(
+ LoRAInfo(
+ display_name = display_name,
+ adapter_path = model_path,
+ base_model = base_model,
+ source = "exported",
+ export_type = export_type,
+ )
+ )
- return LoRAScanResponse(
- loras=lora_list,
- outputs_dir=resolved_outputs_dir
- )
+ return LoRAScanResponse(loras = lora_list, outputs_dir = resolved_outputs_dir)
except Exception as e:
- logger.error(f"Error scanning LoRAs: {e}", exc_info=True)
+ logger.error(f"Error scanning LoRAs: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to scan LoRA adapters: {str(e)}"
+ status_code = 500, detail = f"Failed to scan LoRA adapters: {str(e)}"
)
-@router.get("/loras/{lora_path:path}/base-model", response_model=LoRABaseModelResponse)
+@router.get("/loras/{lora_path:path}/base-model", response_model = LoRABaseModelResponse)
async def get_lora_base_model(
lora_path: str,
current_subject: str = Depends(get_current_subject),
):
"""
Get the base model for a LoRA adapter.
-
+
This endpoint wraps the backend get_base_model_from_lora function.
"""
try:
base_model = get_base_model_from_lora(lora_path)
-
+
if base_model is None:
raise HTTPException(
- status_code=404,
- detail=f"Could not determine base model for LoRA: {lora_path}"
+ status_code = 404,
+ detail = f"Could not determine base model for LoRA: {lora_path}",
)
-
+
return LoRABaseModelResponse(
- lora_path=lora_path,
- base_model=base_model,
+ lora_path = lora_path,
+ base_model = base_model,
)
-
+
except HTTPException:
raise
except Exception as e:
- logger.error(f"Error getting LoRA base model: {e}", exc_info=True)
+ logger.error(f"Error getting LoRA base model: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to get base model: {str(e)}"
+ status_code = 500, detail = f"Failed to get base model: {str(e)}"
)
-@router.get("/check-vision/{model_name:path}", response_model=VisionCheckResponse)
+@router.get("/check-vision/{model_name:path}", response_model = VisionCheckResponse)
async def check_vision_model(
model_name: str,
current_subject: str = Depends(get_current_subject),
):
"""
Check if a model is a vision model.
-
+
This endpoint wraps the backend is_vision_model function.
"""
try:
logger.info(f"Checking if vision model: {model_name}")
is_vision = is_vision_model(model_name)
-
+
logger.info(f"Vision check result for {model_name}: is_vision={is_vision}")
return VisionCheckResponse(
- model_name=model_name,
- is_vision=is_vision,
- )
-
- except Exception as e:
- logger.error(f"Error checking vision model: {e}", exc_info=True)
- raise HTTPException(
- status_code=500,
- detail=f"Failed to check vision model: {str(e)}"
+ model_name = model_name,
+ is_vision = is_vision,
)
-@router.get("/check-embedding/{model_name:path}", response_model=EmbeddingCheckResponse)
+ except Exception as e:
+ logger.error(f"Error checking vision model: {e}", exc_info = True)
+ raise HTTPException(
+ status_code = 500, detail = f"Failed to check vision model: {str(e)}"
+ )
+
+
+@router.get("/check-embedding/{model_name:path}", response_model = EmbeddingCheckResponse)
async def check_embedding_model(
model_name: str,
hf_token: Optional[str] = Query(None),
@@ -465,26 +513,31 @@ async def check_embedding_model(
"""
try:
logger.info(f"Checking if embedding model: {model_name}")
- is_embedding = is_embedding_model(model_name, hf_token=hf_token)
+ is_embedding = is_embedding_model(model_name, hf_token = hf_token)
- logger.info(f"Embedding check result for {model_name}: is_embedding={is_embedding}")
+ logger.info(
+ f"Embedding check result for {model_name}: is_embedding={is_embedding}"
+ )
return EmbeddingCheckResponse(
- model_name=model_name,
- is_embedding=is_embedding,
+ model_name = model_name,
+ is_embedding = is_embedding,
)
except Exception as e:
- logger.error(f"Error checking embedding model: {e}", exc_info=True)
+ logger.error(f"Error checking embedding model: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to check embedding model: {str(e)}"
+ status_code = 500, detail = f"Failed to check embedding model: {str(e)}"
)
-@router.get("/gguf-variants", response_model=GgufVariantsResponse)
+@router.get("/gguf-variants", response_model = GgufVariantsResponse)
async def get_gguf_variants(
- repo_id: str = Query(..., description="HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"),
- hf_token: Optional[str] = Query(None, description="HuggingFace token for private repos"),
+ repo_id: str = Query(
+ ..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"
+ ),
+ hf_token: Optional[str] = Query(
+ None, description = "HuggingFace token for private repos"
+ ),
current_subject: str = Depends(get_current_subject),
):
"""
@@ -495,7 +548,7 @@ async def get_gguf_variants(
default variant.
"""
try:
- variants, has_vision = list_gguf_variants(repo_id, hf_token=hf_token)
+ variants, has_vision = list_gguf_variants(repo_id, hf_token = hf_token)
# Determine default variant
filenames = [v.filename for v in variants]
@@ -503,32 +556,32 @@ async def get_gguf_variants(
default_variant = _extract_quant_label(best) if best else None
return GgufVariantsResponse(
- repo_id=repo_id,
- variants=[
+ repo_id = repo_id,
+ variants = [
GgufVariantDetail(
- filename=v.filename,
- quant=v.quant,
- size_bytes=v.size_bytes,
+ filename = v.filename,
+ quant = v.quant,
+ size_bytes = v.size_bytes,
)
for v in variants
],
- has_vision=has_vision,
- default_variant=default_variant,
+ has_vision = has_vision,
+ default_variant = default_variant,
)
except Exception as e:
- logger.error(f"Error listing GGUF variants for '{repo_id}': {e}", exc_info=True)
+ logger.error(f"Error listing GGUF variants for '{repo_id}': {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to list GGUF variants: {str(e)}",
+ status_code = 500,
+ detail = f"Failed to list GGUF variants: {str(e)}",
)
-@router.get("/checkpoints", response_model=CheckpointListResponse)
+@router.get("/checkpoints", response_model = CheckpointListResponse)
async def list_checkpoints(
outputs_dir: str = Query(
- default=str(outputs_root()),
- description="Directory to scan for checkpoints",
+ default = str(outputs_root()),
+ description = "Directory to scan for checkpoints",
),
current_subject: str = Depends(get_current_subject),
):
@@ -539,29 +592,29 @@ async def list_checkpoints(
"""
try:
resolved_outputs_dir = str(resolve_output_dir(outputs_dir))
- raw_models = scan_checkpoints(outputs_dir=resolved_outputs_dir)
+ raw_models = scan_checkpoints(outputs_dir = resolved_outputs_dir)
models = [
ModelCheckpoints(
- name=model_name,
- checkpoints=[
- CheckpointInfo(display_name=display_name, path=path, loss=loss)
+ name = model_name,
+ checkpoints = [
+ CheckpointInfo(display_name = display_name, path = path, loss = loss)
for display_name, path, loss in checkpoints
],
- base_model=metadata.get("base_model"),
- peft_type=metadata.get("peft_type"),
- lora_rank=metadata.get("lora_rank"),
+ base_model = metadata.get("base_model"),
+ peft_type = metadata.get("peft_type"),
+ lora_rank = metadata.get("lora_rank"),
)
for model_name, checkpoints, metadata in raw_models
]
return CheckpointListResponse(
- outputs_dir=resolved_outputs_dir,
- models=models,
+ outputs_dir = resolved_outputs_dir,
+ models = models,
)
except Exception as e:
- logger.error(f"Error listing checkpoints: {e}", exc_info=True)
+ logger.error(f"Error listing checkpoints: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to list checkpoints: {str(e)}",
+ status_code = 500,
+ detail = f"Failed to list checkpoints: {str(e)}",
)
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index c4010835ac..43b5212010 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -4,6 +4,7 @@
"""
Training API routes
"""
+
import sys
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Request
@@ -50,12 +51,11 @@ from pydantic import BaseModel as PydanticBaseModel
class TrainingStopRequest(PydanticBaseModel):
save: bool = True
+
router = APIRouter()
logger = get_logger(__name__)
-
-
@router.get("/hardware")
async def get_hardware_utilization(
current_subject: str = Depends(get_current_subject),
@@ -100,13 +100,13 @@ async def start_training(
if backend.is_training_active():
existing_job_id: Optional[str] = getattr(backend, "current_job_id", "")
return TrainingJobResponse(
- job_id=existing_job_id or job_id,
- status="error",
- message=(
+ job_id = existing_job_id or job_id,
+ status = "error",
+ message = (
"Training is already in progress. "
"Stop current training before starting a new one."
),
- error="Training already active",
+ error = "Training already active",
)
# Validate dataset paths if provided
@@ -128,8 +128,8 @@ async def start_training(
if missing_datasets:
missing_detail = "; ".join(missing_datasets[:3])
raise HTTPException(
- status_code=400,
- detail=f"Local dataset not found: {missing_detail}",
+ status_code = 400,
+ detail = f"Local dataset not found: {missing_detail}",
)
request.local_datasets = validated_datasets
@@ -167,7 +167,9 @@ async def start_training(
"lora_r": request.lora_r,
"lora_alpha": request.lora_alpha,
"lora_dropout": request.lora_dropout,
- "target_modules": request.target_modules if request.target_modules else None,
+ "target_modules": request.target_modules
+ if request.target_modules
+ else None,
"gradient_checkpointing": request.gradient_checkpointing.strip()
if request.gradient_checkpointing and request.gradient_checkpointing.strip()
else "unsloth",
@@ -194,15 +196,20 @@ async def start_training(
# net, consult the YAML directly so models that need it always get it.
if not training_kwargs["trust_remote_code"]:
model_defaults = load_model_defaults(request.model_name)
- yaml_trust = model_defaults.get("training", {}).get("trust_remote_code", False)
+ yaml_trust = model_defaults.get("training", {}).get(
+ "trust_remote_code", False
+ )
if yaml_trust:
- logger.info(f"YAML config sets trust_remote_code=True for {request.model_name}")
+ logger.info(
+ f"YAML config sets trust_remote_code=True for {request.model_name}"
+ )
training_kwargs["trust_remote_code"] = True
# Free GPU memory: shut down any running inference/export subprocesses
# before training starts (they'd compete for VRAM otherwise)
try:
from core.inference import get_inference_backend
+
inf_backend = get_inference_backend()
if inf_backend.active_model_name:
logger.info(
@@ -217,9 +224,12 @@ async def start_training(
try:
from core.export import get_export_backend
+
exp_backend = get_export_backend()
if exp_backend.current_checkpoint:
- logger.info("Shutting down export subprocess to free GPU memory for training")
+ logger.info(
+ "Shutting down export subprocess to free GPU memory for training"
+ )
exp_backend._shutdown_subprocess()
exp_backend.current_checkpoint = None
exp_backend.is_vision = False
@@ -233,28 +243,28 @@ async def start_training(
if not success:
progress_error = backend.trainer.training_progress.error
return TrainingJobResponse(
- job_id=job_id,
- status="error",
- message=progress_error or "Failed to start training subprocess",
- error=progress_error or "subprocess_start_failed",
+ job_id = job_id,
+ status = "error",
+ message = progress_error or "Failed to start training subprocess",
+ error = progress_error or "subprocess_start_failed",
)
return TrainingJobResponse(
- job_id=job_id,
- status="queued",
- message="Training job queued and starting in subprocess",
- error=None,
+ job_id = job_id,
+ status = "queued",
+ message = "Training job queued and starting in subprocess",
+ error = None,
)
except Exception as e:
- logger.error(f"Error starting training: {e}", exc_info=True)
+ logger.error(f"Error starting training: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to start training: {str(e)}",
+ status_code = 500,
+ detail = f"Failed to start training: {str(e)}",
)
-@router.post("/stop", response_model=TrainingStopResponse)
+@router.post("/stop", response_model = TrainingStopResponse)
async def stop_training(
body: TrainingStopRequest = TrainingStopRequest(),
current_subject: str = Depends(get_current_subject),
@@ -272,23 +282,21 @@ async def stop_training(
if not is_active:
return TrainingStopResponse(
- status="idle",
- message="No training job is currently running"
+ status = "idle", message = "No training job is currently running"
)
# Call backend stop method
- backend.stop_training(save=body.save)
+ backend.stop_training(save = body.save)
return TrainingStopResponse(
- status="stopped",
- message="Stop requested. Training will stop at the next safe step."
+ status = "stopped",
+ message = "Stop requested. Training will stop at the next safe step.",
)
except Exception as e:
- logger.error(f"Error stopping training: {e}", exc_info=True)
+ logger.error(f"Error stopping training: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to stop training: {str(e)}"
+ status_code = 500, detail = f"Failed to stop training: {str(e)}"
)
@@ -306,21 +314,30 @@ async def reset_training(
if is_active:
if backend._cancel_requested:
# Cancel (save=False) was requested β force-terminate so we can reset immediately
- logger.info("Force-terminating subprocess for immediate reset (cancel path)")
+ logger.info(
+ "Force-terminating subprocess for immediate reset (cancel path)"
+ )
backend.force_terminate()
else:
- logger.warning("Rejected reset while training active: is_active=%s", is_active)
+ logger.warning(
+ "Rejected reset while training active: is_active=%s", is_active
+ )
raise HTTPException(
- status_code=409,
- detail="Training is still running. Stop training and wait for it to finish before resetting.",
+ status_code = 409,
+ detail = "Training is still running. Stop training and wait for it to finish before resetting.",
)
logger.info("Reset training state: clearing runtime + metric history")
backend._should_stop = False # Clear stop flag so status returns to idle
backend.trainer._update_progress(
- is_training=False, is_completed=False, error=None,
- status_message="Ready to train", step=0, loss=0.0, epoch=0,
- total_steps=0,
+ is_training = False,
+ is_completed = False,
+ error = None,
+ status_message = "Ready to train",
+ step = 0,
+ loss = 0.0,
+ epoch = 0,
+ total_steps = 0,
)
backend.loss_history = []
backend.lr_history = []
@@ -331,10 +348,10 @@ async def reset_training(
except HTTPException:
raise
except Exception as e:
- logger.error(f"Error resetting training: {e}", exc_info=True)
+ logger.error(f"Error resetting training: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to reset training: {str(e)}",
+ status_code = 500,
+ detail = f"Failed to reset training: {str(e)}",
)
@@ -410,25 +427,24 @@ async def get_training_status(
}
return TrainingStatus(
- job_id=job_id,
- phase=phase,
- is_training_running=is_active,
- eval_enabled=backend.eval_enabled,
- message=status_message,
- error=error_message,
- details=details,
- metric_history=metric_history,
+ job_id = job_id,
+ phase = phase,
+ is_training_running = is_active,
+ eval_enabled = backend.eval_enabled,
+ message = status_message,
+ error = error_message,
+ details = details,
+ metric_history = metric_history,
)
-
+
except Exception as e:
- logger.error(f"Error getting training status: {e}", exc_info=True)
+ logger.error(f"Error getting training status: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to get training status: {str(e)}"
+ status_code = 500, detail = f"Failed to get training status: {str(e)}"
)
-@router.get("/metrics", response_model=TrainingMetricsResponse)
+@router.get("/metrics", response_model = TrainingMetricsResponse)
async def get_training_metrics(
current_subject: str = Depends(get_current_subject),
):
@@ -437,7 +453,7 @@ async def get_training_metrics(
"""
try:
backend = get_training_backend()
-
+
# Get metrics from backend
loss_history = backend.loss_history
lr_history = backend.lr_history
@@ -451,21 +467,20 @@ async def get_training_metrics(
current_step = step_history[-1] if step_history else None
return TrainingMetricsResponse(
- loss_history=loss_history,
- lr_history=lr_history,
- step_history=step_history,
- grad_norm_history=grad_norm_history,
- grad_norm_step_history=grad_norm_step_history,
- current_loss=current_loss,
- current_lr=current_lr,
- current_step=current_step,
+ loss_history = loss_history,
+ lr_history = lr_history,
+ step_history = step_history,
+ grad_norm_history = grad_norm_history,
+ grad_norm_step_history = grad_norm_step_history,
+ current_loss = current_loss,
+ current_lr = current_lr,
+ current_step = current_step,
)
-
+
except Exception as e:
- logger.error(f"Error getting training metrics: {e}", exc_info=True)
+ logger.error(f"Error getting training metrics: {e}", exc_info = True)
raise HTTPException(
- status_code=500,
- detail=f"Failed to get training metrics: {str(e)}"
+ status_code = 500, detail = f"Failed to get training metrics: {str(e)}"
)
@@ -518,29 +533,31 @@ async def stream_training_progress(
)
# Get actual values from progress object if available
- elapsed_seconds = getattr(progress, 'elapsed_seconds', None) if progress else None
- eta_seconds = getattr(progress, 'eta_seconds', None) if progress else None
+ elapsed_seconds = (
+ getattr(progress, "elapsed_seconds", None) if progress else None
+ )
+ eta_seconds = getattr(progress, "eta_seconds", None) if progress else None
grad_norm = grad_norm_override
if grad_norm is None and progress:
- grad_norm = getattr(progress, 'grad_norm', None)
- num_tokens = getattr(progress, 'num_tokens', None) if progress else None
+ grad_norm = getattr(progress, "grad_norm", None)
+ num_tokens = getattr(progress, "num_tokens", None) if progress else None
eval_loss = eval_loss_override
if eval_loss is None and progress:
- eval_loss = getattr(progress, 'eval_loss', None)
+ eval_loss = getattr(progress, "eval_loss", None)
return TrainingProgress(
- job_id=job_id,
- step=step,
- total_steps=total,
- loss=loss,
- learning_rate=learning_rate,
- progress_percent=progress_percent,
- epoch=epoch,
- elapsed_seconds=elapsed_seconds,
- eta_seconds=eta_seconds,
- grad_norm=grad_norm,
- num_tokens=num_tokens,
- eval_loss=eval_loss,
+ job_id = job_id,
+ step = step,
+ total_steps = total,
+ loss = loss,
+ learning_rate = learning_rate,
+ progress_percent = progress_percent,
+ epoch = epoch,
+ elapsed_seconds = elapsed_seconds,
+ eta_seconds = eta_seconds,
+ grad_norm = grad_norm,
+ num_tokens = num_tokens,
+ eval_loss = eval_loss,
)
def format_sse(
@@ -574,23 +591,37 @@ async def stream_training_progress(
}
for i, step_val in enumerate(backend.step_history):
if step_val > resume_from_step:
- loss_val = backend.loss_history[i] if i < len(backend.loss_history) else 0.0
- lr_val = backend.lr_history[i] if i < len(backend.lr_history) else 0.0
+ loss_val = (
+ backend.loss_history[i]
+ if i < len(backend.loss_history)
+ else 0.0
+ )
+ lr_val = (
+ backend.lr_history[i] if i < len(backend.lr_history) else 0.0
+ )
tp_replay = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
- total_replay = getattr(tp_replay, "total_steps", step_val) if tp_replay else step_val
- epoch_replay = getattr(tp_replay, "epoch", None) if tp_replay else None
+ total_replay = (
+ getattr(tp_replay, "total_steps", step_val)
+ if tp_replay
+ else step_val
+ )
+ epoch_replay = (
+ getattr(tp_replay, "epoch", None) if tp_replay else None
+ )
payload = build_progress(
step_val,
loss_val,
lr_val,
total_replay,
epoch_replay,
- progress=tp_replay,
- grad_norm_override=grad_norm_by_step.get(step_val),
+ progress = tp_replay,
+ grad_norm_override = grad_norm_by_step.get(step_val),
+ )
+ yield format_sse(
+ payload.model_dump_json(), event = "progress", event_id = step_val
)
- yield format_sse(payload.model_dump_json(), event="progress", event_id=step_val)
replayed += 1
if replayed:
logger.info(f"SSE reconnect: replayed {replayed} missed steps")
@@ -603,45 +634,62 @@ async def stream_training_progress(
initial_epoch = getattr(tp, "epoch", None) if tp else None
initial_progress = build_progress(
- step=0,
- loss=0.0,
- learning_rate=0.0,
- total_steps=initial_total_steps,
- epoch=initial_epoch,
- progress=tp,
+ step = 0,
+ loss = 0.0,
+ learning_rate = 0.0,
+ total_steps = initial_total_steps,
+ epoch = initial_epoch,
+ progress = tp,
+ )
+ yield format_sse(
+ initial_progress.model_dump_json(), event = "progress", event_id = 0
)
- yield format_sse(initial_progress.model_dump_json(), event="progress", event_id=0)
# If not active, send final state and exit
if not is_active:
if backend.step_history:
final_step = backend.step_history[-1]
- final_loss = backend.loss_history[-1] if backend.loss_history else 0.0
+ final_loss = (
+ backend.loss_history[-1] if backend.loss_history else 0.0
+ )
final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
final_total_steps = (
getattr(tp, "total_steps", final_step) if tp else final_step
)
final_epoch = getattr(tp, "epoch", None) if tp else None
- payload = build_progress(final_step, final_loss, final_lr, final_total_steps, final_epoch, progress=tp)
- yield format_sse(payload.model_dump_json(), event="complete", event_id=final_step)
+ payload = build_progress(
+ final_step,
+ final_loss,
+ final_lr,
+ final_total_steps,
+ final_epoch,
+ progress = tp,
+ )
+ yield format_sse(
+ payload.model_dump_json(), event = "complete", event_id = final_step
+ )
else:
yield format_sse(
- build_progress(-1, 0.0, 0.0, 0, progress=tp).model_dump_json(),
- event="complete",
- event_id=0,
+ build_progress(-1, 0.0, 0.0, 0, progress = tp).model_dump_json(),
+ event = "complete",
+ event_id = 0,
)
return
# ββ Live polling loop ββββββββββββββββββββββββββββββββββββ
last_step = resume_from_step if resume_from_step is not None else -1
no_update_count = 0
- max_no_updates = 1800 # Timeout after 30 minutes (large models need time for compilation)
+ max_no_updates = (
+ 1800 # Timeout after 30 minutes (large models need time for compilation)
+ )
while backend.is_training_active():
try:
if backend.step_history:
current_step = backend.step_history[-1]
- current_loss = backend.loss_history[-1] if backend.loss_history else 0.0
+ current_loss = (
+ backend.loss_history[-1] if backend.loss_history else 0.0
+ )
current_lr = backend.lr_history[-1] if backend.lr_history else 0.0
tp_inner = getattr(
getattr(backend, "trainer", None), "training_progress", None
@@ -651,7 +699,9 @@ async def stream_training_progress(
if tp_inner
else current_step
)
- current_epoch = getattr(tp_inner, "epoch", None) if tp_inner else None
+ current_epoch = (
+ getattr(tp_inner, "epoch", None) if tp_inner else None
+ )
# Only send if step changed
if current_step != last_step:
@@ -661,12 +711,12 @@ async def stream_training_progress(
current_lr,
current_total_steps,
current_epoch,
- progress=tp_inner,
+ progress = tp_inner,
)
yield format_sse(
progress_payload.model_dump_json(),
- event="progress",
- event_id=current_step,
+ event = "progress",
+ event_id = current_step,
)
last_step = current_step
no_update_count = 0
@@ -680,12 +730,12 @@ async def stream_training_progress(
current_lr,
current_total_steps,
current_epoch,
- progress=tp_inner,
+ progress = tp_inner,
)
yield format_sse(
heartbeat_payload.model_dump_json(),
- event="heartbeat",
- event_id=current_step,
+ event = "heartbeat",
+ event_id = current_step,
)
else:
# No steps yet, but training is active (model loading, etc.)
@@ -695,43 +745,53 @@ async def stream_training_progress(
# the frontend can show "Tokenizingβ¦" etc.
tp_prep = getattr(
getattr(backend, "trainer", None),
- "training_progress", None,
+ "training_progress",
+ None,
)
prep_total = (
- getattr(tp_prep, "total_steps", 0)
- if tp_prep else 0
+ getattr(tp_prep, "total_steps", 0) if tp_prep else 0
)
preparing_payload = build_progress(
- 0, 0.0, 0.0, prep_total, progress=tp_prep,
+ 0,
+ 0.0,
+ 0.0,
+ prep_total,
+ progress = tp_prep,
)
yield format_sse(
preparing_payload.model_dump_json(),
- event="heartbeat",
- event_id=0,
+ event = "heartbeat",
+ event_id = 0,
)
# Timeout check
if no_update_count > max_no_updates:
logger.warning("Progress stream timeout - no updates received")
- tp_timeout = getattr(getattr(backend, "trainer", None), "training_progress", None)
- timeout_payload = build_progress(last_step, 0.0, 0.0, 0, progress=tp_timeout)
+ tp_timeout = getattr(
+ getattr(backend, "trainer", None), "training_progress", None
+ )
+ timeout_payload = build_progress(
+ last_step, 0.0, 0.0, 0, progress = tp_timeout
+ )
yield format_sse(
timeout_payload.model_dump_json(),
- event="error",
- event_id=last_step if last_step >= 0 else 0,
+ event = "error",
+ event_id = last_step if last_step >= 0 else 0,
)
break
await asyncio.sleep(1) # Poll every second
except Exception as e:
- logger.error(f"Error in progress stream: {e}", exc_info=True)
- tp_error = getattr(getattr(backend, "trainer", None), "training_progress", None)
- error_payload = build_progress(0, 0.0, 0.0, 0, progress=tp_error)
+ logger.error(f"Error in progress stream: {e}", exc_info = True)
+ tp_error = getattr(
+ getattr(backend, "trainer", None), "training_progress", None
+ )
+ error_payload = build_progress(0, 0.0, 0.0, 0, progress = tp_error)
yield format_sse(
error_payload.model_dump_json(),
- event="error",
- event_id=last_step if last_step >= 0 else 0,
+ event = "error",
+ event_id = last_step if last_step >= 0 else 0,
)
break
@@ -739,9 +799,7 @@ async def stream_training_progress(
final_step = backend.step_history[-1] if backend.step_history else last_step
final_loss = backend.loss_history[-1] if backend.loss_history else 0.0
final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
- final_tp = getattr(
- getattr(backend, "trainer", None), "training_progress", None
- )
+ final_tp = getattr(getattr(backend, "trainer", None), "training_progress", None)
final_total_steps = (
getattr(final_tp, "total_steps", final_step) if final_tp else final_step
)
@@ -752,20 +810,20 @@ async def stream_training_progress(
final_lr,
final_total_steps,
final_epoch,
- progress=final_tp,
+ progress = final_tp,
)
yield format_sse(
final_payload.model_dump_json(),
- event="complete",
- event_id=final_step if final_step >= 0 else 0,
+ event = "complete",
+ event_id = final_step if final_step >= 0 else 0,
)
return StreamingResponse(
event_generator(),
- media_type="text/event-stream",
- headers={
+ media_type = "text/event-stream",
+ headers = {
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
- }
+ },
)
diff --git a/studio/backend/run.py b/studio/backend/run.py
index 32aae9831a..dde1335eae 100644
--- a/studio/backend/run.py
+++ b/studio/backend/run.py
@@ -5,11 +5,12 @@
Run script for Unsloth UI Backend.
Works independently and can be moved to any directory.
"""
+
import os
import sys
# Suppress annoying C-level dependency warnings globally (e.g. SwigPyPacked)
-os.environ["PYTHONWARNINGS"] = "ignore"
+os.environ["PYTHONWARNINGS"] = "ignore"
from pathlib import Path
@@ -19,6 +20,7 @@ if str(backend_dir) not in sys.path:
sys.path.insert(0, str(backend_dir))
from loggers import get_logger
+
logger = get_logger(__name__)
@@ -38,9 +40,9 @@ def _resolve_external_ip() -> str:
try:
req = urllib.request.Request(
"http://metadata.google.internal/computeMetadata/v1/instance/network-interfaces/0/access-configs/0/external-ip",
- headers={"Metadata-Flavor": "Google"},
+ headers = {"Metadata-Flavor": "Google"},
)
- with urllib.request.urlopen(req, timeout=1) as resp:
+ with urllib.request.urlopen(req, timeout = 1) as resp:
ip = resp.read().decode().strip()
if ip:
return ip
@@ -49,7 +51,7 @@ def _resolve_external_ip() -> str:
# 2. Try public IP service
try:
- with urllib.request.urlopen("https://ifconfig.me", timeout=3) as resp:
+ with urllib.request.urlopen("https://ifconfig.me", timeout = 3) as resp:
ip = resp.read().decode().strip()
if ip:
return ip
@@ -70,7 +72,7 @@ def _resolve_external_ip() -> str:
def run_server(
host: str = "0.0.0.0",
port: int = 8000,
- frontend_path: Path = "studio/frontend/dist",
+ frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist",
silent: bool = False,
):
"""
@@ -108,11 +110,13 @@ def run_server(
# Run server
def _run():
- config = uvicorn.Config(app, host=host, port=port, log_level="info", access_log=False)
+ config = uvicorn.Config(
+ app, host = host, port = port, log_level = "info", access_log = False
+ )
server = uvicorn.Server(config)
asyncio.run(server.serve())
- thread = Thread(target=_run, daemon=True)
+ thread = Thread(target = _run, daemon = True)
thread.start()
time.sleep(3)
@@ -135,24 +139,26 @@ def run_server(
if __name__ == "__main__":
import argparse
- parser = argparse.ArgumentParser(description="Run Unsloth UI Backend server")
- parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
- parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
+ parser = argparse.ArgumentParser(description = "Run Unsloth UI Backend server")
+ parser.add_argument("--host", default = "0.0.0.0", help = "Host to bind to")
+ parser.add_argument("--port", type = int, default = 8000, help = "Port to bind to")
parser.add_argument(
- "--frontend", type=str, default="studio/frontend/dist", help="Path to frontend build"
+ "--frontend",
+ type = str,
+ default = Path(__file__).resolve().parent.parent / "frontend" / "dist",
+ help = "Path to frontend build",
)
- parser.add_argument("--silent", action="store_true", help="Suppress output")
+ parser.add_argument("--silent", action = "store_true", help = "Suppress output")
args = parser.parse_args()
- frontend_path = Path(args.frontend) if args.frontend else None
- run_server(
- host=args.host, port=args.port, frontend_path=frontend_path, silent=args.silent
- )
+ kwargs = dict(host = args.host, port = args.port, silent = args.silent)
+ if args.frontend is not None:
+ kwargs["frontend_path"] = Path(args.frontend)
+ run_server(**kwargs)
# Keep running
import time
while True:
time.sleep(1)
-
diff --git a/studio/backend/state/__init__.py b/studio/backend/state/__init__.py
index e69de29bb2..32014236c6 100644
--- a/studio/backend/state/__init__.py
+++ b/studio/backend/state/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
diff --git a/studio/backend/tests/__init__.py b/studio/backend/tests/__init__.py
index e69de29bb2..32014236c6 100644
--- a/studio/backend/tests/__init__.py
+++ b/studio/backend/tests/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py
index 4b2bc38cb0..053e9b85d9 100644
--- a/studio/backend/tests/conftest.py
+++ b/studio/backend/tests/conftest.py
@@ -6,6 +6,7 @@ Shared pytest configuration for the backend test suite.
Ensures that the backend root is on sys.path so that
`import utils.utils` (and similar flat imports) resolve correctly.
"""
+
import sys
from pathlib import Path
diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py
index e896c9965c..3c33b33cb3 100644
--- a/studio/backend/tests/test_utils.py
+++ b/studio/backend/tests/test_utils.py
@@ -16,6 +16,7 @@ Run with:
cd studio/backend
python -m pytest tests/test_utils.py -v
"""
+
import platform
from unittest.mock import patch, MagicMock
@@ -24,18 +25,20 @@ import pytest
# --- Conditional framework imports ---
try:
import torch
+
HAS_TORCH = True
except ImportError:
HAS_TORCH = False
try:
import mlx.core as mx
+
HAS_MLX = True
except ImportError:
HAS_MLX = False
-needs_torch = pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not installed")
-needs_mlx = pytest.mark.skipif(not HAS_MLX, reason="MLX not installed")
+needs_torch = pytest.mark.skipif(not HAS_TORCH, reason = "PyTorch not installed")
+needs_mlx = pytest.mark.skipif(not HAS_MLX, reason = "MLX not installed")
from utils.hardware import (
get_device,
@@ -52,6 +55,7 @@ from utils.utils import format_error_message
# ========== Helpers ==========
+
def _actual_device() -> str:
"""Return the real device string for the current machine."""
if HAS_TORCH and torch.cuda.is_available():
@@ -69,6 +73,7 @@ def _reset_and_detect():
# ========== get_device() ==========
+
class TestGetDevice:
"""Tests for get_device() β should agree with the real hardware."""
@@ -89,28 +94,34 @@ class TestGetDevice:
@needs_torch
def test_returns_cuda_when_cuda_available(self):
- with patch("utils.hardware.hardware._has_torch", return_value=True), \
- patch("torch.cuda.is_available", return_value=True):
+ with (
+ patch("utils.hardware.hardware._has_torch", return_value = True),
+ patch("torch.cuda.is_available", return_value = True),
+ ):
assert _reset_and_detect() == DeviceType.CUDA
@needs_mlx
def test_returns_mlx_when_on_apple_silicon_with_mlx(self):
- with patch("utils.hardware.hardware._has_torch", return_value=False), \
- patch("utils.hardware.hardware.is_apple_silicon", return_value=True), \
- patch("utils.hardware.hardware._has_mlx", return_value=True):
+ with (
+ patch("utils.hardware.hardware._has_torch", return_value = False),
+ patch("utils.hardware.hardware.is_apple_silicon", return_value = True),
+ patch("utils.hardware.hardware._has_mlx", return_value = True),
+ ):
assert _reset_and_detect() == DeviceType.MLX
def test_returns_cpu_when_nothing_available(self):
- with patch("utils.hardware.hardware._has_torch", return_value=False), \
- patch("utils.hardware.hardware.is_apple_silicon", return_value=False), \
- patch("utils.hardware.hardware._has_mlx", return_value=False):
+ with (
+ patch("utils.hardware.hardware._has_torch", return_value = False),
+ patch("utils.hardware.hardware.is_apple_silicon", return_value = False),
+ patch("utils.hardware.hardware._has_mlx", return_value = False),
+ ):
assert _reset_and_detect() == DeviceType.CPU
# ========== is_apple_silicon() ==========
-class TestIsAppleSilicon:
+class TestIsAppleSilicon:
def test_returns_bool(self):
assert isinstance(is_apple_silicon(), bool)
@@ -136,6 +147,7 @@ class TestIsAppleSilicon:
# ========== clear_gpu_cache() ==========
+
class TestClearGpuCache:
"""clear_gpu_cache() must never raise, regardless of platform."""
@@ -144,9 +156,11 @@ class TestClearGpuCache:
@needs_torch
def test_calls_cuda_cache_when_cuda(self):
- with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \
- patch("torch.cuda.empty_cache") as mock_empty, \
- patch("torch.cuda.ipc_collect") as mock_ipc:
+ with (
+ patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
+ patch("torch.cuda.empty_cache") as mock_empty,
+ patch("torch.cuda.ipc_collect") as mock_ipc,
+ ):
clear_gpu_cache()
mock_empty.assert_called_once()
mock_ipc.assert_called_once()
@@ -154,18 +168,18 @@ class TestClearGpuCache:
@needs_mlx
def test_mlx_does_not_raise(self):
"""MLX cache clear is a no-op β should just succeed."""
- with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MLX):
+ with patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX):
clear_gpu_cache()
def test_noop_on_cpu(self):
- with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU):
+ with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU):
clear_gpu_cache()
# ========== get_gpu_memory_info() ==========
-class TestGetGpuMemoryInfo:
+class TestGetGpuMemoryInfo:
def test_returns_dict(self):
result = get_gpu_memory_info()
assert isinstance(result, dict)
@@ -183,8 +197,7 @@ class TestGetGpuMemoryInfo:
# --- When a GPU IS available ---
@pytest.mark.skipif(
- _actual_device() == "cpu",
- reason="No GPU available on this machine"
+ _actual_device() == "cpu", reason = "No GPU available on this machine"
)
def test_gpu_available_fields(self):
result = get_gpu_memory_info()
@@ -200,14 +213,16 @@ class TestGetGpuMemoryInfo:
@needs_torch
def test_cuda_path_returns_correct_fields(self):
mock_props = MagicMock()
- mock_props.total_memory = 16 * (1024 ** 3)
+ mock_props.total_memory = 16 * (1024**3)
mock_props.name = "NVIDIA Test GPU"
- with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \
- patch("torch.cuda.current_device", return_value=0), \
- patch("torch.cuda.get_device_properties", return_value=mock_props), \
- patch("torch.cuda.memory_allocated", return_value=4 * (1024 ** 3)), \
- patch("torch.cuda.memory_reserved", return_value=6 * (1024 ** 3)):
+ with (
+ patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
+ patch("torch.cuda.current_device", return_value = 0),
+ patch("torch.cuda.get_device_properties", return_value = mock_props),
+ patch("torch.cuda.memory_allocated", return_value = 4 * (1024**3)),
+ patch("torch.cuda.memory_reserved", return_value = 6 * (1024**3)),
+ ):
result = get_gpu_memory_info()
assert result["available"] is True
@@ -223,13 +238,15 @@ class TestGetGpuMemoryInfo:
@needs_mlx
def test_mlx_path_returns_correct_fields(self):
mock_psutil_mem = MagicMock()
- mock_psutil_mem.total = 32 * (1024 ** 3) # 32 GB unified
+ mock_psutil_mem.total = 32 * (1024**3) # 32 GB unified
mock_psutil = MagicMock()
mock_psutil.virtual_memory.return_value = mock_psutil_mem
- with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MLX), \
- patch.dict("sys.modules", {"psutil": mock_psutil}):
+ with (
+ patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX),
+ patch.dict("sys.modules", {"psutil": mock_psutil}),
+ ):
result = get_gpu_memory_info()
assert result["available"] is True
@@ -240,7 +257,7 @@ class TestGetGpuMemoryInfo:
# --- CPU-only path ---
def test_cpu_path_returns_unavailable(self):
- with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU):
+ with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU):
result = get_gpu_memory_info()
assert result["available"] is False
assert result["backend"] == "cpu"
@@ -249,8 +266,13 @@ class TestGetGpuMemoryInfo:
@needs_torch
def test_cuda_error_returns_unavailable(self):
- with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \
- patch("torch.cuda.current_device", side_effect=RuntimeError("CUDA init failed")):
+ with (
+ patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
+ patch(
+ "torch.cuda.current_device",
+ side_effect = RuntimeError("CUDA init failed"),
+ ),
+ ):
result = get_gpu_memory_info()
assert result["available"] is False
assert "error" in result
@@ -258,8 +280,8 @@ class TestGetGpuMemoryInfo:
# ========== log_gpu_memory() ==========
-class TestLogGpuMemory:
+class TestLogGpuMemory:
def test_does_not_raise(self):
log_gpu_memory("test")
@@ -275,8 +297,13 @@ class TestLogGpuMemory:
}
import structlog
from loggers import get_logger
- with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info), \
- caplog.at_level(logging.INFO, logger="utils.hardware.hardware"):
+
+ with (
+ patch(
+ "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
+ ),
+ caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"),
+ ):
log_gpu_memory("unit-test")
assert "unit-test" in caplog.text
@@ -287,8 +314,13 @@ class TestLogGpuMemory:
fake_info = {"available": False, "backend": "cpu"}
import structlog
from loggers import get_logger
- with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info), \
- caplog.at_level(logging.INFO, logger="utils.hardware.hardware"):
+
+ with (
+ patch(
+ "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
+ ),
+ caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"),
+ ):
log_gpu_memory("cpu-test")
assert "No GPU available" in caplog.text
@@ -296,8 +328,8 @@ class TestLogGpuMemory:
# ========== format_error_message() ==========
-class TestFormatErrorMessage:
+class TestFormatErrorMessage:
def test_not_found(self):
err = Exception("Repository not found for unsloth/test")
msg = format_error_message(err, "unsloth/test")
@@ -324,7 +356,7 @@ class TestFormatErrorMessage:
@needs_torch
def test_cuda_oom(self):
err = Exception("CUDA out of memory")
- with patch("utils.hardware.get_device", return_value=DeviceType.CUDA):
+ with patch("utils.hardware.get_device", return_value = DeviceType.CUDA):
msg = format_error_message(err, "big/model")
assert "GPU" in msg
assert "big/model" not in msg
@@ -335,7 +367,7 @@ class TestFormatErrorMessage:
@needs_mlx
def test_mlx_oom(self):
err = Exception("MLX backend out of memory")
- with patch("utils.hardware.get_device", return_value=DeviceType.MLX):
+ with patch("utils.hardware.get_device", return_value = DeviceType.MLX):
msg = format_error_message(err, "unsloth/huge-model")
assert "Apple Silicon" in msg
@@ -343,7 +375,7 @@ class TestFormatErrorMessage:
def test_cpu_oom(self):
err = Exception("not enough memory to allocate")
- with patch("utils.hardware.get_device", return_value=DeviceType.CPU):
+ with patch("utils.hardware.get_device", return_value = DeviceType.CPU):
msg = format_error_message(err, "any/model")
assert "system" in msg.lower()
diff --git a/studio/backend/utils/__init__.py b/studio/backend/utils/__init__.py
new file mode 100644
index 0000000000..32014236c6
--- /dev/null
+++ b/studio/backend/utils/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
diff --git a/studio/backend/utils/cache_cleanup.py b/studio/backend/utils/cache_cleanup.py
index 187f1c02cd..b3ffbcfc05 100644
--- a/studio/backend/utils/cache_cleanup.py
+++ b/studio/backend/utils/cache_cleanup.py
@@ -8,6 +8,7 @@ The unsloth_compiled_cache is created by unsloth_zoo/compiler.py during
FastModel.from_pretrained() and contains model-type-specific compiled Python
files. It should be cleared between model loads to avoid stale artefacts.
"""
+
import shutil
import structlog
from loggers import get_logger
@@ -16,8 +17,8 @@ from pathlib import Path
logger = get_logger(__name__)
# Possible locations where unsloth_compiled_cache may appear
-_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend
-_PROJECT_ROOT = _BACKEND_DIR.parent.parent # repo root
+_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend
+_PROJECT_ROOT = _BACKEND_DIR.parent.parent # repo root
_CACHE_DIRS = [
_BACKEND_DIR / "unsloth_compiled_cache",
@@ -31,4 +32,4 @@ def clear_unsloth_compiled_cache() -> None:
for cache_dir in _CACHE_DIRS:
if cache_dir.exists():
logger.info(f"Removing unsloth compiled cache: {cache_dir}")
- shutil.rmtree(cache_dir, ignore_errors=True)
+ shutil.rmtree(cache_dir, ignore_errors = True)
diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py
index ca2e7a8c98..23fc856772 100644
--- a/studio/backend/utils/datasets/chat_templates.py
+++ b/studio/backend/utils/datasets/chat_templates.py
@@ -61,7 +61,7 @@ def get_tokenizer_chat_template(tokenizer, model_name):
try:
tokenizer = get_chat_template(
tokenizer,
- chat_template=matched_template,
+ chat_template = matched_template,
)
except Exception as e:
logger.info(f"β οΈ Failed to apply Unsloth template '{matched_template}': {e}")
@@ -80,7 +80,7 @@ def get_tokenizer_chat_template(tokenizer, model_name):
try:
tokenizer = get_chat_template(
tokenizer,
- chat_template="chatml",
+ chat_template = "chatml",
)
except Exception as e:
logger.info(f"β οΈ Failed to apply default ChatML template: {e}")
@@ -119,14 +119,14 @@ def get_dataset_info_summary(dataset_info):
def apply_chat_template_to_dataset(
dataset_info,
tokenizer,
- model_name=None,
- custom_prompt_template=None,
- add_eos_token=False,
- remove_bos_prefix=False,
- custom_format_mapping=None,
- auto_detect_mapping=True,
- batch_size=1000,
- num_proc=None,
+ model_name = None,
+ custom_prompt_template = None,
+ add_eos_token = False,
+ remove_bos_prefix = False,
+ custom_format_mapping = None,
+ auto_detect_mapping = True,
+ batch_size = 1000,
+ num_proc = None,
):
"""
Applies chat template to dataset based on its format.
@@ -233,7 +233,7 @@ def apply_chat_template_to_dataset(
return result
try:
- dataset = dataset.map(_apply_custom_mapping, batched=True, batch_size=batch_size)
+ dataset = dataset.map(_apply_custom_mapping, batched = True, batch_size = batch_size)
# Update to use conversations format
final_format = "chatml_conversations"
chat_column = "conversations"
@@ -256,7 +256,7 @@ def apply_chat_template_to_dataset(
if not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template):
try:
from unsloth.chat_templates import get_chat_template
- tokenizer = get_chat_template(tokenizer, chat_template="alpaca")
+ tokenizer = get_chat_template(tokenizer, chat_template = "alpaca")
logger.info(f"π Set alpaca chat template on tokenizer for model saving")
except Exception as e:
logger.info(f"β οΈ Could not set alpaca template on tokenizer: {e}")
@@ -333,8 +333,8 @@ def apply_chat_template_to_dataset(
try:
text = tokenizer.apply_chat_template(
convo,
- tokenize=False,
- add_generation_prompt=False
+ tokenize = False,
+ add_generation_prompt = False
)
if remove_bos_prefix:
diff --git a/studio/backend/utils/datasets/data_collators.py b/studio/backend/utils/datasets/data_collators.py
index 3c0be0706c..04b24bbae3 100644
--- a/studio/backend/utils/datasets/data_collators.py
+++ b/studio/backend/utils/datasets/data_collators.py
@@ -12,11 +12,10 @@ import torch
from dataclasses import dataclass
from typing import Any, List, Optional, Union
from loggers import get_logger
+
logger = get_logger(__name__)
-
-
@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
"""
@@ -26,16 +25,23 @@ class DataCollatorSpeechSeq2SeqWithPadding:
masks padding in labels with -100, and strips leading BOS token.
Mirrors the collator from the Whisper.ipynb notebook.
"""
+
processor: Any
def __call__(self, features: List[dict]) -> dict:
- input_features = [{"input_features": feature["input_features"]} for feature in features]
- batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
+ input_features = [
+ {"input_features": feature["input_features"]} for feature in features
+ ]
+ batch = self.processor.feature_extractor.pad(
+ input_features, return_tensors = "pt"
+ )
label_features = [{"input_ids": feature["labels"]} for feature in features]
- labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
+ labels_batch = self.processor.tokenizer.pad(label_features, return_tensors = "pt")
- labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
+ labels = labels_batch["input_ids"].masked_fill(
+ labels_batch.attention_mask.ne(1), -100
+ )
if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
labels = labels[:, 1:]
@@ -54,6 +60,7 @@ class DeepSeekOCRDataCollator:
- Text tokenization
- Proper label masking for instruction fine-tuning
"""
+
processor: Any # Qwen2VLProcessor or similar
max_length: int = 2048
ignore_index: int = -100
@@ -86,7 +93,7 @@ class DeepSeekOCRDataCollator:
for item in content:
if isinstance(item, dict) and item.get("type") == "image":
img = item.get("image")
- if img is not None and hasattr(img, 'size'): # PIL Image
+ if img is not None and hasattr(img, "size"): # PIL Image
all_images.append(img)
# Process with the VL processor
@@ -94,19 +101,19 @@ class DeepSeekOCRDataCollator:
# Qwen2VL style processing
texts = [
self.processor.apply_chat_template(
- msgs, tokenize=False, add_generation_prompt=False
+ msgs, tokenize = False, add_generation_prompt = False
)
for msgs in all_messages
]
# Process with images
inputs = self.processor(
- text=texts,
- images=all_images if all_images else None,
- return_tensors="pt",
- padding=True,
- truncation=True,
- max_length=self.max_length,
+ text = texts,
+ images = all_images if all_images else None,
+ return_tensors = "pt",
+ padding = True,
+ truncation = True,
+ max_length = self.max_length,
)
# Create labels (mask input, keep output)
@@ -134,6 +141,7 @@ class VLMDataCollator:
- LLaVA
- Other VL models with compatible processors
"""
+
processor: Any
max_length: int = 2048
ignore_index: int = -100
@@ -163,26 +171,26 @@ class VLMDataCollator:
# Apply chat template
texts = [
self.processor.apply_chat_template(
- msgs, tokenize=False, add_generation_prompt=False
+ msgs, tokenize = False, add_generation_prompt = False
)
for msgs in all_messages
]
# Process inputs
inputs = self.processor(
- text=texts,
- images=all_images if all_images else None,
- return_tensors="pt",
- padding=True,
- truncation=True,
- max_length=self.max_length,
+ text = texts,
+ images = all_images if all_images else None,
+ return_tensors = "pt",
+ padding = True,
+ truncation = True,
+ max_length = self.max_length,
)
# Create labels
labels = inputs["input_ids"].clone()
# Mask padding
- if hasattr(self.processor, 'tokenizer'):
+ if hasattr(self.processor, "tokenizer"):
pad_token_id = self.processor.tokenizer.pad_token_id
else:
pad_token_id = self.processor.pad_token_id
diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py
index 9629c5aec4..9d15b86ca1 100644
--- a/studio/backend/utils/datasets/dataset_utils.py
+++ b/studio/backend/utils/datasets/dataset_utils.py
@@ -45,22 +45,21 @@ from .vlm_processing import generate_smart_vlm_instruction
from .data_collators import DeepSeekOCRDataCollator, VLMDataCollator
from .model_mappings import TEMPLATE_TO_MODEL_MAPPER
from loggers import get_logger
+
logger = get_logger(__name__)
-
-
def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
"""
Lightweight format check without processing - for frontend validation.
-
+
Use this to quickly determine if user needs to manually map columns
before calling the full format_and_template_dataset().
-
+
Args:
dataset: HuggingFace dataset
is_vlm: Whether this is a Vision-Language Model dataset
-
+
Returns:
dict: {
"requires_manual_mapping": bool - True if user must map columns,
@@ -71,8 +70,12 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
"detected_text_column": str or None - For VLM only,
}
"""
- columns = list(dataset.column_names) if hasattr(dataset, 'column_names') else list(next(iter(dataset)).keys())
-
+ columns = (
+ list(dataset.column_names)
+ if hasattr(dataset, "column_names")
+ else list(next(iter(dataset)).keys())
+ )
+
# Auto-detect multimodal data regardless of is_vlm flag
multimodal_info = detect_multimodal_dataset(dataset)
is_audio = multimodal_info.get("is_audio", False)
@@ -185,11 +188,17 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
**audio_fields,
}
+
# Normalise any format-specific role to canonical chatml (user/assistant/system)
_TO_CHATML = {
- "user": "user", "human": "user", "instruction": "user",
- "assistant": "assistant", "gpt": "assistant", "output": "assistant",
- "system": "system", "input": "system",
+ "user": "user",
+ "human": "user",
+ "instruction": "user",
+ "assistant": "assistant",
+ "gpt": "assistant",
+ "output": "assistant",
+ "system": "system",
+ "input": "system",
}
_CHATML_ROLE_ORDER = ("system", "user", "assistant")
_CHATML_TO_ALPACA = {"user": "instruction", "system": "input", "assistant": "output"}
@@ -232,11 +241,21 @@ def _apply_user_mapping(dataset, mapping: dict, batch_size: int = 1000):
for col in role_groups[chatml_role]:
if col in examples:
content = examples[col][i]
- convo.append({"role": chatml_role, "content": str(content) if content else ""})
+ convo.append(
+ {
+ "role": chatml_role,
+ "content": str(content) if content else "",
+ }
+ )
conversations.append(convo)
return {"conversations": conversations}
- return dataset.map(_convert, batched=True, batch_size=batch_size, remove_columns=dataset.column_names)
+ return dataset.map(
+ _convert,
+ batched = True,
+ batch_size = batch_size,
+ remove_columns = dataset.column_names,
+ )
def _extract_column_value(val, col: str, label_mapping: dict) -> str:
@@ -248,7 +267,7 @@ def _extract_column_value(val, col: str, label_mapping: dict) -> str:
inner = val["text"]
str_val = inner[0] if isinstance(inner, list) and inner else str(inner)
else:
- str_val = json.dumps(val, ensure_ascii=False)
+ str_val = json.dumps(val, ensure_ascii = False)
elif isinstance(val, list):
str_val = val[0] if len(val) == 1 else ", ".join(str(v) for v in val)
else:
@@ -286,6 +305,7 @@ def _apply_template_mapping(
role_groups[canonical].append(col)
import logging as _log
+
_log.getLogger(__name__).info(
f"Applying role mapping: sys={bool(system_prompt)}, "
f"user_cols={role_groups['user']}, asst_cols={role_groups['assistant']}, "
@@ -326,8 +346,10 @@ def _apply_template_mapping(
return {"conversations": conversations}
return dataset.map(
- _convert, batched=True, batch_size=batch_size,
- remove_columns=dataset.column_names,
+ _convert,
+ batched = True,
+ batch_size = batch_size,
+ remove_columns = dataset.column_names,
)
@@ -341,7 +363,11 @@ def _apply_user_mapping_alpaca(dataset, mapping: dict, batch_size: int = 1000):
Returns:
Dataset with instruction/input/output columns
"""
- col_for: dict[str, str | None] = {"instruction": None, "input": None, "output": None}
+ col_for: dict[str, str | None] = {
+ "instruction": None,
+ "input": None,
+ "output": None,
+ }
for col_name, role in mapping.items():
canonical = _TO_CHATML.get(role)
alpaca_field = _CHATML_TO_ALPACA.get(canonical) if canonical else None
@@ -352,25 +378,48 @@ def _apply_user_mapping_alpaca(dataset, mapping: dict, batch_size: int = 1000):
num = len(next(iter(examples.values())))
instructions, inputs, outputs = [], [], []
for i in range(num):
- for field, dest in (("instruction", instructions), ("input", inputs), ("output", outputs)):
+ for field, dest in (
+ ("instruction", instructions),
+ ("input", inputs),
+ ("output", outputs),
+ ):
col = col_for[field]
- val = str(examples[col][i]) if col and col in examples and examples[col][i] else ""
+ val = (
+ str(examples[col][i])
+ if col and col in examples and examples[col][i]
+ else ""
+ )
dest.append(val)
return {"instruction": instructions, "input": inputs, "output": outputs}
- return dataset.map(_convert, batched=True, batch_size=batch_size, remove_columns=dataset.column_names)
+ return dataset.map(
+ _convert,
+ batched = True,
+ batch_size = batch_size,
+ remove_columns = dataset.column_names,
+ )
def format_dataset(
dataset,
- format_type = "auto",
- tokenizer = None,
- aliases_for_system = ["system",],
- aliases_for_user = ["user", "human", "input",],
- aliases_for_assistant = ["gpt", "assistant", "output",],
- batch_size = 1000,
- num_proc = None,
- auto_detect_custom = True,
+ format_type = "auto",
+ tokenizer = None,
+ aliases_for_system = [
+ "system",
+ ],
+ aliases_for_user = [
+ "user",
+ "human",
+ "input",
+ ],
+ aliases_for_assistant = [
+ "gpt",
+ "assistant",
+ "output",
+ ],
+ batch_size = 1000,
+ num_proc = None,
+ auto_detect_custom = True,
custom_format_mapping = None,
):
"""
@@ -395,13 +444,17 @@ def format_dataset(
if custom_format_mapping:
try:
if format_type == "alpaca":
- mapped_dataset = _apply_user_mapping_alpaca(dataset, custom_format_mapping, batch_size)
+ mapped_dataset = _apply_user_mapping_alpaca(
+ dataset, custom_format_mapping, batch_size
+ )
final_format = "alpaca"
chat_column = None
else:
# auto / chatml / sharegpt / conversational β all produce chatml conversations
# (sharegpt is always standardized to role/content internally)
- mapped_dataset = _apply_user_mapping(dataset, custom_format_mapping, batch_size)
+ mapped_dataset = _apply_user_mapping(
+ dataset, custom_format_mapping, batch_size
+ )
final_format = "chatml_conversations"
chat_column = "conversations"
@@ -414,7 +467,9 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": [f"Applied user-provided column mapping ({format_type}): {custom_format_mapping}"]
+ "warnings": [
+ f"Applied user-provided column mapping ({format_type}): {custom_format_mapping}"
+ ],
}
except Exception as e:
return {
@@ -426,10 +481,9 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": [f"Failed to apply user mapping: {e}"]
+ "warnings": [f"Failed to apply user mapping: {e}"],
}
-
# Detect current format
detected = detect_dataset_format(dataset)
warnings = []
@@ -442,7 +496,6 @@ def format_dataset(
# AUTO MODE: Keep format but standardize if needed
if format_type == "auto":
-
# Alpaca - keep as is
if detected["format"] == "alpaca":
return {
@@ -454,16 +507,20 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": []
+ "warnings": [],
}
# ShareGPT - needs standardization
elif detected["format"] == "sharegpt":
try:
standardized = standardize_chat_format(
- dataset, tokenizer, aliases_for_system,
- aliases_for_user, aliases_for_assistant,
- batch_size, num_proc
+ dataset,
+ tokenizer,
+ aliases_for_system,
+ aliases_for_user,
+ aliases_for_assistant,
+ batch_size,
+ num_proc,
)
return {
"dataset": standardized,
@@ -474,7 +531,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": []
+ "warnings": [],
}
except Exception as e:
warnings.append(f"Failed to standardize ShareGPT format: {e}")
@@ -487,10 +544,14 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": warnings
+ "warnings": warnings,
}
- elif detected["format"] == "chatml" and detected["chat_column"] in ["conversations", "messages", "texts"]:
+ elif detected["format"] == "chatml" and detected["chat_column"] in [
+ "conversations",
+ "messages",
+ "texts",
+ ]:
return {
"dataset": dataset,
"detected_format": f"chatml_{detected['chat_column']}",
@@ -500,13 +561,14 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": warnings
+ "warnings": warnings,
}
-
# Unknown - try standardization, if fails pass as is
else:
- warnings.append(f"Unknown format detected. Keys found: {detected['sample_keys']}")
+ warnings.append(
+ f"Unknown format detected. Keys found: {detected['sample_keys']}"
+ )
# NEW: Try heuristic detection
if auto_detect_custom:
@@ -514,7 +576,6 @@ def format_dataset(
if custom_mapping:
warnings.append(f"Auto-detected column mapping: {custom_mapping}")
-
def _apply_auto_mapping(examples):
conversations = []
num_examples = len(examples[list(examples.keys())[0]])
@@ -523,25 +584,27 @@ def format_dataset(
all_columns = set(examples.keys())
mapped_columns = set(custom_mapping.keys())
preserved_columns = {
- col: examples[col]
- for col in all_columns - mapped_columns
+ col: examples[col] for col in all_columns - mapped_columns
}
for i in range(num_examples):
convo = []
- for target_role in ['system', 'user', 'assistant']:
+ for target_role in ["system", "user", "assistant"]:
for col_name, role in custom_mapping.items():
if role == target_role and col_name in examples:
content = examples[col_name][i]
if content and str(content).strip():
- convo.append({"role": role, "content": str(content)})
+ convo.append(
+ {"role": role, "content": str(content)}
+ )
conversations.append(convo)
return {"conversations": conversations, **preserved_columns}
-
try:
- dataset = dataset.map(_apply_auto_mapping, batched=True, batch_size=batch_size)
+ dataset = dataset.map(
+ _apply_auto_mapping, batched = True, batch_size = batch_size
+ )
return {
"dataset": dataset,
"detected_format": "unknown",
@@ -551,7 +614,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": warnings
+ "warnings": warnings,
}
except Exception as e:
warnings.append(f"Auto-detection failed: {e}")
@@ -560,9 +623,13 @@ def format_dataset(
if detected["chat_column"]:
try:
standardized = standardize_chat_format(
- dataset, tokenizer, aliases_for_system,
- aliases_for_user, aliases_for_assistant,
- batch_size, num_proc
+ dataset,
+ tokenizer,
+ aliases_for_system,
+ aliases_for_user,
+ aliases_for_assistant,
+ batch_size,
+ num_proc,
)
warnings.append("Successfully standardized unknown format")
return {
@@ -574,10 +641,12 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": warnings
+ "warnings": warnings,
}
except Exception as e:
- warnings.append(f"Could not standardize: {e}. Passing dataset as-is.")
+ warnings.append(
+ f"Could not standardize: {e}. Passing dataset as-is."
+ )
# Return as-is with warnings
return {
@@ -589,12 +658,11 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": warnings
+ "warnings": warnings,
}
# ALPACA MODE: Convert to Alpaca
elif format_type == "alpaca":
-
if detected["format"] == "alpaca":
return {
"dataset": dataset,
@@ -605,16 +673,20 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": []
+ "warnings": [],
}
elif detected["format"] in ["sharegpt", "chatml"]:
# First standardize if ShareGPT
if detected["format"] == "sharegpt":
dataset = standardize_chat_format(
- dataset, tokenizer, aliases_for_system,
- aliases_for_user, aliases_for_assistant,
- batch_size, num_proc
+ dataset,
+ tokenizer,
+ aliases_for_system,
+ aliases_for_user,
+ aliases_for_assistant,
+ batch_size,
+ num_proc,
)
# Then convert to Alpaca
@@ -628,7 +700,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": []
+ "warnings": [],
}
else:
@@ -642,12 +714,11 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": warnings
+ "warnings": warnings,
}
# CHATML MODE: Convert to ChatML
elif format_type in ["chatml", "conversational", "sharegpt"]:
-
if detected["format"] == "alpaca":
converted = convert_alpaca_to_chatml(dataset, batch_size, num_proc)
return {
@@ -659,14 +730,18 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": []
+ "warnings": [],
}
elif detected["format"] == "sharegpt":
standardized = standardize_chat_format(
- dataset, tokenizer, aliases_for_system,
- aliases_for_user, aliases_for_assistant,
- batch_size, num_proc
+ dataset,
+ tokenizer,
+ aliases_for_system,
+ aliases_for_user,
+ aliases_for_assistant,
+ batch_size,
+ num_proc,
)
return {
"dataset": standardized,
@@ -677,7 +752,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": []
+ "warnings": [],
}
elif detected["format"] == "chatml":
@@ -690,7 +765,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": []
+ "warnings": [],
}
else:
@@ -698,9 +773,13 @@ def format_dataset(
if detected["chat_column"]:
try:
standardized = standardize_chat_format(
- dataset, tokenizer, aliases_for_system,
- aliases_for_user, aliases_for_assistant,
- batch_size, num_proc
+ dataset,
+ tokenizer,
+ aliases_for_system,
+ aliases_for_user,
+ aliases_for_assistant,
+ batch_size,
+ num_proc,
)
return {
"dataset": standardized,
@@ -711,7 +790,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": warnings
+ "warnings": warnings,
}
except Exception as e:
warnings.append(f"Standardization failed: {e}")
@@ -725,7 +804,7 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
- "warnings": warnings
+ "warnings": warnings,
}
else:
@@ -737,25 +816,34 @@ def format_and_template_dataset(
model_name,
tokenizer,
is_vlm = False,
- format_type="auto",
+ format_type = "auto",
# VLM-specific parameters
- vlm_instruction=None, # Now optional - will auto-generate
- vlm_text_column=None,
- vlm_image_column=None,
- dataset_name=None,
-
- custom_prompt_template=None,
- add_eos_token=False,
- remove_bos_prefix=False,
- custom_format_mapping=None,
- auto_detect_custom=True,
- auto_detect_mapping=True,
- aliases_for_system=["system",],
- aliases_for_user=["user", "human", "input",],
- aliases_for_assistant=["gpt", "assistant", "output",],
- batch_size=1000,
- num_proc=None,
- progress_callback=None,
+ vlm_instruction = None, # Now optional - will auto-generate
+ vlm_text_column = None,
+ vlm_image_column = None,
+ dataset_name = None,
+ custom_prompt_template = None,
+ add_eos_token = False,
+ remove_bos_prefix = False,
+ custom_format_mapping = None,
+ auto_detect_custom = True,
+ auto_detect_mapping = True,
+ aliases_for_system = [
+ "system",
+ ],
+ aliases_for_user = [
+ "user",
+ "human",
+ "input",
+ ],
+ aliases_for_assistant = [
+ "gpt",
+ "assistant",
+ "output",
+ ],
+ batch_size = 1000,
+ num_proc = None,
+ progress_callback = None,
):
"""
Convenience function that combines format_dataset and apply_chat_template_to_dataset.
@@ -786,25 +874,27 @@ def format_and_template_dataset(
# Expect mapping like: {"image_col": "image", "caption_col": "text"}
user_vlm_image_column = None
user_vlm_text_column = None
-
+
for col, role in custom_format_mapping.items():
if role == "image":
user_vlm_image_column = col
elif role in ["text", "user", "caption", "assistant"]:
user_vlm_text_column = col
-
+
if user_vlm_image_column and user_vlm_text_column:
try:
dataset = convert_to_vlm_format(
dataset,
- instruction=vlm_instruction,
- text_column=user_vlm_text_column,
- image_column=user_vlm_image_column,
- dataset_name=dataset_name,
- progress_callback=progress_callback,
+ instruction = vlm_instruction,
+ text_column = user_vlm_text_column,
+ image_column = user_vlm_image_column,
+ dataset_name = dataset_name,
+ progress_callback = progress_callback,
)
- warnings.append(f"Applied user VLM mapping: image='{user_vlm_image_column}', text='{user_vlm_text_column}'")
-
+ warnings.append(
+ f"Applied user VLM mapping: image='{user_vlm_image_column}', text='{user_vlm_text_column}'"
+ )
+
return {
"dataset": dataset,
"detected_format": "user_mapped",
@@ -826,7 +916,9 @@ def format_and_template_dataset(
f"text='{user_vlm_text_column}') failed: {e} β "
f"falling back to auto-detection"
)
- logger.info(f"β οΈ User VLM mapping failed, falling back to auto-detection...")
+ logger.info(
+ f"β οΈ User VLM mapping failed, falling back to auto-detection..."
+ )
custom_format_mapping = None # clear so auto-detection runs below
else:
errors.append(
@@ -850,10 +942,13 @@ def format_and_template_dataset(
if vlm_structure["format"] == "vlm_messages_llava":
try:
dataset = convert_llava_to_vlm_format(dataset)
- warnings.append("Converted from Llava format (image indices) to standard VLM format")
+ warnings.append(
+ "Converted from Llava format (image indices) to standard VLM format"
+ )
except Exception as e:
errors.append(f"Failed to convert Llava format: {e}")
import traceback
+
traceback.print_exc()
return {
@@ -872,15 +967,18 @@ def format_and_template_dataset(
try:
dataset = convert_sharegpt_with_images_to_vlm_format(
dataset,
- image_column=vlm_structure["image_column"],
- messages_column=vlm_structure["messages_column"],
- dataset_name=dataset_name,
- progress_callback=progress_callback,
+ image_column = vlm_structure["image_column"],
+ messages_column = vlm_structure["messages_column"],
+ dataset_name = dataset_name,
+ progress_callback = progress_callback,
+ )
+ warnings.append(
+ "Converted from ShareGPT+image format to standard VLM format"
)
- warnings.append("Converted from ShareGPT+image format to standard VLM format")
except Exception as e:
errors.append(f"Failed to convert ShareGPT+image format: {e}")
import traceback
+
traceback.print_exc()
return {
@@ -910,14 +1008,18 @@ def format_and_template_dataset(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
+
friendly = llm_generate_dataset_warning(
- issues, dataset_name=dataset_name, modality="vision",
- column_names=columns,
+ issues,
+ dataset_name = dataset_name,
+ modality = "vision",
+ column_names = columns,
)
except Exception:
pass
errors.append(
- friendly or f"Could not auto-detect image/text columns. Found: {vlm_structure}. "
+ friendly
+ or f"Could not auto-detect image/text columns. Found: {vlm_structure}. "
)
return {
"dataset": dataset,
@@ -933,21 +1035,26 @@ def format_and_template_dataset(
try:
dataset = convert_to_vlm_format(
dataset,
- instruction=vlm_instruction,
- text_column=vlm_text_column,
- image_column=vlm_image_column,
- dataset_name=dataset_name,
- progress_callback=progress_callback,
+ instruction = vlm_instruction,
+ text_column = vlm_text_column,
+ image_column = vlm_image_column,
+ dataset_name = dataset_name,
+ progress_callback = progress_callback,
)
if vlm_instruction:
- warnings.append(f"Using user-provided instruction: '{vlm_instruction}'")
+ warnings.append(
+ f"Using user-provided instruction: '{vlm_instruction}'"
+ )
else:
- warnings.append("Auto-generated instruction based on dataset analysis")
+ warnings.append(
+ "Auto-generated instruction based on dataset analysis"
+ )
except Exception as e:
errors.append(f"Failed to convert to VLM format: {e}")
import traceback
+
traceback.print_exc()
return {
@@ -987,41 +1094,45 @@ def format_and_template_dataset(
# Step 1: Format the dataset
dataset_info = format_dataset(
dataset,
- format_type=format_type,
- tokenizer=tokenizer,
- auto_detect_custom=auto_detect_custom,
- custom_format_mapping=custom_format_mapping,
- aliases_for_system=aliases_for_system,
- aliases_for_user=aliases_for_user,
- aliases_for_assistant=aliases_for_assistant,
- batch_size=batch_size,
- num_proc=num_proc,
+ format_type = format_type,
+ tokenizer = tokenizer,
+ auto_detect_custom = auto_detect_custom,
+ custom_format_mapping = custom_format_mapping,
+ aliases_for_system = aliases_for_system,
+ aliases_for_user = aliases_for_user,
+ aliases_for_assistant = aliases_for_assistant,
+ batch_size = batch_size,
+ num_proc = num_proc,
)
# Step 2: Apply chat template
# Gemma emits a leading that must be stripped for text-only chatml/sharegpt.
- is_alpaca = format_type == "alpaca" or (format_type == "auto" and dataset_info["detected_format"] == "alpaca")
+ is_alpaca = format_type == "alpaca" or (
+ format_type == "auto" and dataset_info["detected_format"] == "alpaca"
+ )
is_gemma = "gemma" in model_name.lower()
if is_gemma and not dataset_info["is_image"] and not is_alpaca:
remove_bos_prefix = True
template_result = apply_chat_template_to_dataset(
- dataset_info=dataset_info,
- tokenizer=tokenizer,
- model_name=model_name,
- custom_prompt_template=custom_prompt_template,
- add_eos_token=add_eos_token,
- remove_bos_prefix=remove_bos_prefix,
- custom_format_mapping=custom_format_mapping,
- auto_detect_mapping=auto_detect_mapping,
- batch_size=batch_size,
- num_proc=num_proc,
+ dataset_info = dataset_info,
+ tokenizer = tokenizer,
+ model_name = model_name,
+ custom_prompt_template = custom_prompt_template,
+ add_eos_token = add_eos_token,
+ remove_bos_prefix = remove_bos_prefix,
+ custom_format_mapping = custom_format_mapping,
+ auto_detect_mapping = auto_detect_mapping,
+ batch_size = batch_size,
+ num_proc = num_proc,
)
# Step 3: Generate summary
summary = get_dataset_info_summary(dataset_info)
# Combine results
- all_warnings = dataset_info.get("warnings", []) + template_result.get("warnings", [])
+ all_warnings = dataset_info.get("warnings", []) + template_result.get(
+ "warnings", []
+ )
all_errors = template_result.get("errors", [])
# If format_dataset returned "unknown" but apply_chat_template rescued
diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py
index be707e6127..264789c41c 100644
--- a/studio/backend/utils/datasets/format_conversion.py
+++ b/studio/backend/utils/datasets/format_conversion.py
@@ -12,19 +12,28 @@ import os
from datasets import IterableDataset
from loggers import get_logger
+
logger = get_logger(__name__)
-
-
def standardize_chat_format(
dataset,
- tokenizer=None,
- aliases_for_system=["system",],
- aliases_for_user=["user", "human", "input",],
- aliases_for_assistant=["gpt", "assistant", "output",],
- batch_size=1000,
- num_proc=None,
+ tokenizer = None,
+ aliases_for_system = [
+ "system",
+ ],
+ aliases_for_user = [
+ "user",
+ "human",
+ "input",
+ ],
+ aliases_for_assistant = [
+ "gpt",
+ "assistant",
+ "output",
+ ],
+ batch_size = 1000,
+ num_proc = None,
):
"""
Our own standardization function that handles BOTH messages and conversations.
@@ -67,22 +76,25 @@ def standardize_chat_format(
return dataset # Unexpected structure
keys = list(uniques.keys())
- length_first = len(set(uniques[keys[0]]))
+ length_first = len(set(uniques[keys[0]]))
length_second = len(set(uniques[keys[1]]))
# Determine which is role and which is content
if length_first < length_second:
- role_key = keys[0]
+ role_key = keys[0]
content_key = keys[1]
else:
- role_key = keys[1]
+ role_key = keys[1]
content_key = keys[0]
# Mapping for aliases
aliases_mapping = {}
- for x in aliases_for_system: aliases_mapping[x] = "system"
- for x in aliases_for_user: aliases_mapping[x] = "user"
- for x in aliases_for_assistant: aliases_mapping[x] = "assistant"
+ for x in aliases_for_system:
+ aliases_mapping[x] = "system"
+ for x in aliases_for_user:
+ aliases_mapping[x] = "user"
+ for x in aliases_for_assistant:
+ aliases_mapping[x] = "assistant"
def _standardize_dataset(examples):
convos = examples[chat_column]
@@ -109,10 +121,9 @@ def standardize_chat_format(
return {chat_column: all_convos}
-
dataset_map_kwargs = {
- 'batched': True,
- 'batch_size': batch_size,
+ "batched": True,
+ "batch_size": batch_size,
}
if not isinstance(dataset, IterableDataset):
@@ -123,13 +134,13 @@ def standardize_chat_format(
else:
num_proc = safe_num_proc(num_proc)
- dataset_map_kwargs['num_proc'] = num_proc
- dataset_map_kwargs['desc'] = "Standardizing chat format"
+ dataset_map_kwargs["num_proc"] = num_proc
+ dataset_map_kwargs["desc"] = "Standardizing chat format"
return dataset.map(_standardize_dataset, **dataset_map_kwargs)
-def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
+def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None):
"""
Converts ChatML format (messages OR conversations) to Alpaca format.
Handles both standardized and ShareGPT formats.
@@ -142,10 +153,16 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
def _convert(examples):
# Auto-detect which column name is used
- chatml_data = examples.get("messages") or examples.get("conversations") or examples.get("texts")
+ chatml_data = (
+ examples.get("messages")
+ or examples.get("conversations")
+ or examples.get("texts")
+ )
if chatml_data is None:
- raise ValueError("No 'messages' or 'conversations' or 'texts' column found.")
+ raise ValueError(
+ "No 'messages' or 'conversations' or 'texts' column found."
+ )
instructions = []
outputs = []
@@ -172,15 +189,11 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
inputs.append("") # Alpaca typically has empty input
outputs.append(output)
- return {
- "instruction": instructions,
- "input": inputs,
- "output": outputs
- }
+ return {"instruction": instructions, "input": inputs, "output": outputs}
dataset_map_kwargs = {
- 'batched': True,
- 'batch_size': batch_size,
+ "batched": True,
+ "batch_size": batch_size,
}
if not isinstance(dataset, IterableDataset):
@@ -191,13 +204,13 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
else:
num_proc = safe_num_proc(num_proc)
- dataset_map_kwargs['num_proc'] = num_proc
- dataset_map_kwargs['desc'] = "Converting ChatML to Alpaca format"
+ dataset_map_kwargs["num_proc"] = num_proc
+ dataset_map_kwargs["desc"] = "Converting ChatML to Alpaca format"
return dataset.map(_convert, **dataset_map_kwargs)
-def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None):
+def convert_alpaca_to_chatml(dataset, batch_size = 1000, num_proc = None):
"""
Converts Alpaca format to ChatML format.
@@ -222,15 +235,15 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None):
# Build conversation in standard ChatML format
convo = [
{"role": "user", "content": user_content},
- {"role": "assistant", "content": output}
+ {"role": "assistant", "content": output},
]
conversations.append(convo)
return {"conversations": conversations}
dataset_map_kwargs = {
- 'batched': True,
- 'batch_size': batch_size,
+ "batched": True,
+ "batch_size": batch_size,
}
if not isinstance(dataset, IterableDataset):
@@ -241,8 +254,8 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None):
else:
num_proc = safe_num_proc(num_proc)
- dataset_map_kwargs['num_proc'] = num_proc
- dataset_map_kwargs['desc'] = "Converting Alpaca to ChatML format"
+ dataset_map_kwargs["num_proc"] = num_proc
+ dataset_map_kwargs["desc"] = "Converting Alpaca to ChatML format"
return dataset.map(_convert, **dataset_map_kwargs)
@@ -262,11 +275,11 @@ def _format_eta(seconds):
def convert_to_vlm_format(
dataset,
- instruction=None,
- text_column="text",
- image_column="image",
- dataset_name=None,
- progress_callback=None,
+ instruction = None,
+ text_column = "text",
+ image_column = "image",
+ dataset_name = None,
+ progress_callback = None,
):
"""
Converts simple {image, text} format to VLM messages format.
@@ -290,27 +303,31 @@ def convert_to_vlm_format(
def _notify(msg):
"""Send status update to the training overlay if callback is available."""
if progress_callback:
- progress_callback(status_message=msg)
+ progress_callback(status_message = msg)
# Generate smart instruction if not provided
if instruction is None:
instruction_info = generate_smart_vlm_instruction(
dataset,
- text_column=text_column,
- image_column=image_column,
- dataset_name=dataset_name,
+ text_column = text_column,
+ image_column = image_column,
+ dataset_name = dataset_name,
)
instruction = instruction_info["instruction"]
instruction_column = instruction_info.get("instruction_column")
uses_dynamic = instruction_info["uses_dynamic_instruction"]
- logger.info(f"π Auto-detected instruction type: {instruction_info['instruction_type']}")
+ logger.info(
+ f"π Auto-detected instruction type: {instruction_info['instruction_type']}"
+ )
logger.info(f"π Confidence: {instruction_info['confidence']:.2f}")
if not uses_dynamic:
logger.info(f"π Using instruction: '{instruction}'")
else:
- logger.info(f"π Using dynamic instructions from column: '{instruction_column}'")
+ logger.info(
+ f"π Using dynamic instructions from column: '{instruction_column}'"
+ )
else:
instruction_column = None
uses_dynamic = False
@@ -324,13 +341,17 @@ def convert_to_vlm_format(
if image_data.startswith(("http://", "https://")):
import fsspec
from io import BytesIO
- with fsspec.open(image_data, "rb", expand=True) as f:
+
+ with fsspec.open(image_data, "rb", expand = True) as f:
image_data = Image.open(BytesIO(f.read())).convert("RGB")
elif _image_lookup is not None and image_data in _image_lookup:
# Bare filename β resolve via HF repo lookup
from huggingface_hub import hf_hub_download
+
local_path = hf_hub_download(
- dataset_name, _image_lookup[image_data], repo_type="dataset",
+ dataset_name,
+ _image_lookup[image_data],
+ repo_type = "dataset",
)
image_data = Image.open(local_path).convert("RGB")
else:
@@ -340,6 +361,7 @@ def convert_to_vlm_format(
text_data = sample[text_column]
if isinstance(text_data, list) and len(text_data) > 0:
import random
+
text_data = random.choice(text_data)
# Get instruction (static or dynamic)
@@ -354,15 +376,10 @@ def convert_to_vlm_format(
"role": "user",
"content": [
{"type": "text", "text": current_instruction},
- {"type": "image", "image": image_data} # PIL object
- ]
+ {"type": "image", "image": image_data}, # PIL object
+ ],
},
- {
- "role": "assistant",
- "content": [
- {"type": "text", "text": text_data}
- ]
- }
+ {"role": "assistant", "content": [{"type": "text", "text": text_data}]},
]
# Return dict with messages
@@ -370,13 +387,15 @@ def convert_to_vlm_format(
total = len(dataset)
first_image = next(iter(dataset))[image_column]
- has_urls = isinstance(first_image, str) and first_image.startswith(("http://", "https://"))
+ has_urls = isinstance(first_image, str) and first_image.startswith(
+ ("http://", "https://")
+ )
# ββ Bare-filename detection: images stored as filenames (e.g. "img_001.png")
# that don't exist locally. Build a basenameβrepo_path lookup so we can
# resolve them via hf_hub_download during conversion.
_image_lookup = None
- _IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp', '.tiff')
+ _IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff")
if (
not has_urls
and isinstance(first_image, str)
@@ -385,18 +404,25 @@ def convert_to_vlm_format(
):
try:
from huggingface_hub import HfApi
+
_notify("Resolving image filenames from HF repo...")
- logger.info(f"π Image column contains bare filenames (e.g. '{first_image}') β building repo lookup...")
- repo_files = HfApi().list_repo_files(dataset_name, repo_type="dataset")
+ logger.info(
+ f"π Image column contains bare filenames (e.g. '{first_image}') β building repo lookup..."
+ )
+ repo_files = HfApi().list_repo_files(dataset_name, repo_type = "dataset")
_image_lookup = {
os.path.basename(f): f
for f in repo_files
if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS)
}
if first_image in _image_lookup:
- logger.info(f"β
Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' β '{_image_lookup[first_image]}')")
+ logger.info(
+ f"β
Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' β '{_image_lookup[first_image]}')"
+ )
else:
- logger.info(f"β οΈ Built lookup with {len(_image_lookup)} images but '{first_image}' not found β falling back to local open")
+ logger.info(
+ f"β οΈ Built lookup with {len(_image_lookup)} images but '{first_image}' not found β falling back to local open"
+ )
_image_lookup = None
except Exception as e:
logger.info(f"β οΈ Failed to build HF repo image lookup: {e}")
@@ -413,15 +439,19 @@ def convert_to_vlm_format(
num_workers = safe_num_proc()
_notify(f"Probing {PROBE_SIZE} image URLs with {num_workers} workers...")
- logger.info(f"π Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers...")
+ logger.info(
+ f"π Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers..."
+ )
probe_samples = [dataset[i] for i in range(PROBE_SIZE)]
probe_ok = 0
probe_fail = 0
probe_start = time.time()
- with ThreadPoolExecutor(max_workers=num_workers) as executor:
- futures = {executor.submit(_convert_single_sample, s): s for s in probe_samples}
+ with ThreadPoolExecutor(max_workers = num_workers) as executor:
+ futures = {
+ executor.submit(_convert_single_sample, s): s for s in probe_samples
+ }
for future in as_completed(futures):
try:
future.result()
@@ -443,9 +473,12 @@ def convert_to_vlm_format(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
+
friendly = llm_generate_dataset_warning(
- issues, dataset_name=dataset_name, modality="vision",
- column_names=[image_column, text_column],
+ issues,
+ dataset_name = dataset_name,
+ modality = "vision",
+ column_names = [image_column, text_column],
)
except Exception:
pass
@@ -471,7 +504,9 @@ def convert_to_vlm_format(
if probe_fail > 0:
info_msg += f" | {fail_rate:.0%} broken URLs will be skipped"
- logger.info(f"β
Probe passed: {probe_ok}/{probe_total} ok, {probe_fail} failed ({fail_rate:.0%}), {throughput:.1f} img/s")
+ logger.info(
+ f"β
Probe passed: {probe_ok}/{probe_total} ok, {probe_fail} failed ({fail_rate:.0%}), {throughput:.1f} img/s"
+ )
logger.info(f"β±οΈ Estimated time for {total:,} samples: ~{eta_str}")
_notify(info_msg)
@@ -496,8 +531,11 @@ def convert_to_vlm_format(
batch_end = min(batch_start + batch_size, total)
batch_samples = [dataset[i] for i in range(batch_start, batch_end)]
- with ThreadPoolExecutor(max_workers=num_workers) as executor:
- futures = {executor.submit(_convert_single_sample, s): i for i, s in enumerate(batch_samples)}
+ with ThreadPoolExecutor(max_workers = num_workers) as executor:
+ futures = {
+ executor.submit(_convert_single_sample, s): i
+ for i, s in enumerate(batch_samples)
+ }
batch_results = [None] * len(batch_samples)
for future in as_completed(futures):
idx = futures[future]
@@ -506,9 +544,13 @@ def convert_to_vlm_format(
except Exception as e:
failed_count += 1
if failed_count == 1:
- print(f"β οΈ First VLM conversion failure: {type(e).__name__}: {e}")
+ print(
+ f"β οΈ First VLM conversion failure: {type(e).__name__}: {e}"
+ )
if failed_count == 1:
- logger.info(f"β οΈ First VLM conversion failure: {type(e).__name__}: {e}")
+ logger.info(
+ f"β οΈ First VLM conversion failure: {type(e).__name__}: {e}"
+ )
converted_list.extend(r for r in batch_results if r is not None)
@@ -519,11 +561,13 @@ def convert_to_vlm_format(
remaining_time = (total - done) / rate if rate > 0 else 0
eta_str = _format_eta(remaining_time)
progress_msg = f"Downloading images: {done:,}/{total:,} ({done*100//total}%) | ~{eta_str} remaining | {failed_count} skipped"
- logger.info(f" [{done}/{total}] {rate:.1f} img/s, {failed_count} failed, ETA {eta_str}")
+ logger.info(
+ f" [{done}/{total}] {rate:.1f} img/s, {failed_count} failed, ETA {eta_str}"
+ )
_notify(progress_msg)
else:
# Sequential conversion for local/embedded images (fast, no I/O bottleneck)
- pbar = tqdm(dataset, total=total, desc="Converting VLM samples", unit="sample")
+ pbar = tqdm(dataset, total = total, desc = "Converting VLM samples", unit = "sample")
for sample in pbar:
try:
converted_list.append(_convert_single_sample(sample))
@@ -534,13 +578,17 @@ def convert_to_vlm_format(
print(f"β οΈ First VLM conversion failure: {type(e).__name__}: {e}")
if failed_count == 1:
# Log the first failure to aid debugging
- logger.info(f"β οΈ First VLM conversion failure: {type(e).__name__}: {e}")
- pbar.set_postfix(ok=len(converted_list), failed=failed_count, refresh=False)
+ logger.info(
+ f"β οΈ First VLM conversion failure: {type(e).__name__}: {e}"
+ )
+ pbar.set_postfix(ok = len(converted_list), failed = failed_count, refresh = False)
pbar.close()
if failed_count > 0:
fail_rate = failed_count / total
- logger.info(f"β οΈ Skipped {failed_count}/{total} ({fail_rate:.0%}) samples with broken/unreachable images")
+ logger.info(
+ f"β οΈ Skipped {failed_count}/{total} ({fail_rate:.0%}) samples with broken/unreachable images"
+ )
# For datasets that skipped the probe (small URL datasets), check fail rate now
if has_urls and fail_rate >= MAX_FAIL_RATE:
issues = [
@@ -550,9 +598,12 @@ def convert_to_vlm_format(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
+
friendly = llm_generate_dataset_warning(
- issues, dataset_name=dataset_name, modality="vision",
- column_names=[image_column, text_column],
+ issues,
+ dataset_name = dataset_name,
+ modality = "vision",
+ column_names = [image_column, text_column],
)
except Exception:
pass
@@ -573,14 +624,18 @@ def convert_to_vlm_format(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
+
friendly = llm_generate_dataset_warning(
- issues, dataset_name=dataset_name, modality="vision",
- column_names=[image_column, text_column],
+ issues,
+ dataset_name = dataset_name,
+ modality = "vision",
+ column_names = [image_column, text_column],
)
except Exception:
pass
raise ValueError(
- friendly or (
+ friendly
+ or (
f"All {total} samples failed during VLM conversion β no usable images found. "
"This dataset may contain only image URLs that are no longer accessible."
)
@@ -595,10 +650,10 @@ def convert_to_vlm_format(
def convert_sharegpt_with_images_to_vlm_format(
dataset,
- image_column="image",
- messages_column="conversations",
- dataset_name=None,
- progress_callback=None,
+ image_column = "image",
+ messages_column = "conversations",
+ dataset_name = None,
+ progress_callback = None,
):
"""
Converts ShareGPT/ChatML datasets that have a separate image column and
@@ -619,16 +674,18 @@ def convert_sharegpt_with_images_to_vlm_format(
from PIL import Image
from tqdm import tqdm
- _IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp', '.tiff')
+ _IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff")
_ROLE_MAP = {
- "human": "user", "user": "user",
- "gpt": "assistant", "assistant": "assistant",
+ "human": "user",
+ "user": "user",
+ "gpt": "assistant",
+ "assistant": "assistant",
"system": "system",
}
def _notify(msg):
if progress_callback:
- progress_callback(status_message=msg)
+ progress_callback(status_message = msg)
# ββ Resolve image loading strategy (same 3-tier as convert_to_vlm_format) ββ
total = len(dataset)
@@ -643,9 +700,12 @@ def convert_sharegpt_with_images_to_vlm_format(
):
try:
from huggingface_hub import HfApi
+
_notify("Resolving image filenames from HF repo...")
- logger.info(f"π Image column contains bare filenames (e.g. '{first_image}') β building repo lookup...")
- repo_files = HfApi().list_repo_files(dataset_name, repo_type="dataset")
+ logger.info(
+ f"π Image column contains bare filenames (e.g. '{first_image}') β building repo lookup..."
+ )
+ repo_files = HfApi().list_repo_files(dataset_name, repo_type = "dataset")
_image_lookup = {
os.path.basename(f): f
for f in repo_files
@@ -656,9 +716,13 @@ def convert_sharegpt_with_images_to_vlm_format(
if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS):
_image_lookup[f] = f
if first_image in _image_lookup:
- logger.info(f"β
Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' β '{_image_lookup[first_image]}')")
+ logger.info(
+ f"β
Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' β '{_image_lookup[first_image]}')"
+ )
else:
- logger.info(f"β οΈ Built lookup with {len(_image_lookup)} images but '{first_image}' not found β falling back to local open")
+ logger.info(
+ f"β οΈ Built lookup with {len(_image_lookup)} images but '{first_image}' not found β falling back to local open"
+ )
_image_lookup = None
except Exception as e:
logger.info(f"β οΈ Failed to build HF repo image lookup: {e}")
@@ -666,25 +730,32 @@ def convert_sharegpt_with_images_to_vlm_format(
def _resolve_image(image_data):
"""Resolve image data to a PIL Image object."""
- if hasattr(image_data, 'size') and hasattr(image_data, 'mode'):
+ if hasattr(image_data, "size") and hasattr(image_data, "mode"):
return image_data # Already PIL
if isinstance(image_data, str):
if image_data.startswith(("http://", "https://")):
import fsspec
from io import BytesIO
- with fsspec.open(image_data, "rb", expand=True) as f:
+
+ with fsspec.open(image_data, "rb", expand = True) as f:
return Image.open(BytesIO(f.read())).convert("RGB")
elif _image_lookup is not None and image_data in _image_lookup:
from huggingface_hub import hf_hub_download
+
local_path = hf_hub_download(
- dataset_name, _image_lookup[image_data], repo_type="dataset",
+ dataset_name,
+ _image_lookup[image_data],
+ repo_type = "dataset",
)
return Image.open(local_path).convert("RGB")
else:
return Image.open(image_data).convert("RGB")
- if isinstance(image_data, dict) and ("bytes" in image_data or "path" in image_data):
+ if isinstance(image_data, dict) and (
+ "bytes" in image_data or "path" in image_data
+ ):
if image_data.get("bytes"):
from io import BytesIO
+
return Image.open(BytesIO(image_data["bytes"])).convert("RGB")
if image_data.get("path"):
return Image.open(image_data["path"]).convert("RGB")
@@ -726,7 +797,7 @@ def convert_sharegpt_with_images_to_vlm_format(
converted_list = []
failed_count = 0
- pbar = tqdm(dataset, total=total, desc="Converting ShareGPT+image", unit="sample")
+ pbar = tqdm(dataset, total = total, desc = "Converting ShareGPT+image", unit = "sample")
for sample in pbar:
try:
converted_list.append(_convert_single_sample(sample))
@@ -734,11 +805,13 @@ def convert_sharegpt_with_images_to_vlm_format(
failed_count += 1
if failed_count == 1:
logger.info(f"β οΈ First conversion failure: {type(e).__name__}: {e}")
- pbar.set_postfix(ok=len(converted_list), failed=failed_count, refresh=False)
+ pbar.set_postfix(ok = len(converted_list), failed = failed_count, refresh = False)
pbar.close()
if failed_count > 0:
- logger.info(f"β οΈ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples")
+ logger.info(
+ f"β οΈ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples"
+ )
if len(converted_list) == 0:
raise ValueError(
@@ -764,7 +837,9 @@ def convert_llava_to_vlm_format(dataset):
"""
from PIL import Image
- logger.info(f"π Converting {len(dataset)} samples from Llava format to standard VLM format...")
+ logger.info(
+ f"π Converting {len(dataset)} samples from Llava format to standard VLM format..."
+ )
def _convert_single_sample(sample):
"""Convert a single llava sample to standard VLM format."""
@@ -787,10 +862,12 @@ def convert_llava_to_vlm_format(dataset):
if isinstance(pil_image, str):
pil_image = Image.open(pil_image).convert("RGB")
- new_content.append({
- "type": "image",
- "image": pil_image # Actual PIL object
- })
+ new_content.append(
+ {
+ "type": "image",
+ "image": pil_image, # Actual PIL object
+ }
+ )
else:
# No index, try to use first image
if len(images) > 0:
@@ -798,22 +875,13 @@ def convert_llava_to_vlm_format(dataset):
if isinstance(pil_image, str):
pil_image = Image.open(pil_image).convert("RGB")
- new_content.append({
- "type": "image",
- "image": pil_image
- })
+ new_content.append({"type": "image", "image": pil_image})
elif item["type"] == "text":
# Keep text as-is (only type + text)
- new_content.append({
- "type": "text",
- "text": item.get("text", "")
- })
+ new_content.append({"type": "text", "text": item.get("text", "")})
- new_messages.append({
- "role": msg["role"],
- "content": new_content
- })
+ new_messages.append({"role": msg["role"], "content": new_content})
return {"messages": new_messages}
diff --git a/studio/backend/utils/datasets/format_detection.py b/studio/backend/utils/datasets/format_detection.py
index e03e85e594..7b70ff3a76 100644
--- a/studio/backend/utils/datasets/format_detection.py
+++ b/studio/backend/utils/datasets/format_detection.py
@@ -13,7 +13,10 @@ import re
def _keyword_in_column(keyword: str, col_name: str) -> bool:
"""Word-boundary keyword match to avoid false positives like 'pic' in 'topic'."""
- return re.search(r'\b' + re.escape(keyword) + r'\b', col_name, re.IGNORECASE) is not None
+ return (
+ re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE)
+ is not None
+ )
def detect_dataset_format(dataset):
@@ -37,7 +40,7 @@ def detect_dataset_format(dataset):
"format": "alpaca",
"chat_column": None,
"needs_standardization": False,
- "sample_keys": []
+ "sample_keys": [],
}
# Check for chat-based formats (messages or conversations)
@@ -65,7 +68,7 @@ def detect_dataset_format(dataset):
"format": "sharegpt",
"chat_column": chat_column,
"needs_standardization": True,
- "sample_keys": list(msg_keys)
+ "sample_keys": list(msg_keys),
}
# ChatML uses "role" and "content"
@@ -74,7 +77,7 @@ def detect_dataset_format(dataset):
"format": "chatml",
"chat_column": chat_column,
"needs_standardization": False,
- "sample_keys": list(msg_keys)
+ "sample_keys": list(msg_keys),
}
# Unknown structure but has chat column
@@ -83,7 +86,7 @@ def detect_dataset_format(dataset):
"format": "unknown",
"chat_column": chat_column,
"needs_standardization": None,
- "sample_keys": list(msg_keys)
+ "sample_keys": list(msg_keys),
}
except Exception as e:
return {
@@ -91,7 +94,7 @@ def detect_dataset_format(dataset):
"chat_column": chat_column,
"needs_standardization": None,
"sample_keys": [],
- "error": str(e)
+ "error": str(e),
}
# No recognized format
@@ -99,7 +102,7 @@ def detect_dataset_format(dataset):
"format": "unknown",
"chat_column": None,
"needs_standardization": None,
- "sample_keys": []
+ "sample_keys": [],
}
@@ -120,49 +123,86 @@ def detect_custom_format_heuristic(dataset):
# Keywords
assistant_words = [
- 'output', 'answer', 'response', 'assistant', 'completion',
- 'expected', 'recommendation', 'reply', 'result', 'target',
- 'solution', 'explanation', 'solve'
+ "output",
+ "answer",
+ "response",
+ "assistant",
+ "completion",
+ "expected",
+ "recommendation",
+ "reply",
+ "result",
+ "target",
+ "solution",
+ "explanation",
+ "solve",
]
# Split into high/low priority
user_words_high_priority = [
- 'input', 'question', 'query', 'prompt', 'instruction',
- 'request', 'snippet', 'user', 'text',
- 'problem', 'exercise'
+ "input",
+ "question",
+ "query",
+ "prompt",
+ "instruction",
+ "request",
+ "snippet",
+ "user",
+ "text",
+ "problem",
+ "exercise",
]
- user_words_low_priority = ['task'] # Ambiguous - can be user OR system
+ user_words_low_priority = ["task"] # Ambiguous - can be user OR system
user_words = user_words_high_priority + user_words_low_priority
system_words = [
- 'system', 'context', 'description', 'persona', 'role',
- 'template', 'task' # Also in system
+ "system",
+ "context",
+ "description",
+ "persona",
+ "role",
+ "template",
+ "task", # Also in system
]
# Metadata columns to ignore
metadata_exact_match = {
- 'id', 'idx', 'index', 'key', 'timestamp', 'date',
- 'metadata', 'source', 'kind', 'type', 'category',
- 'score', 'label', 'tag', 'inference_mode'
+ "id",
+ "idx",
+ "index",
+ "key",
+ "timestamp",
+ "date",
+ "metadata",
+ "source",
+ "kind",
+ "type",
+ "category",
+ "score",
+ "label",
+ "tag",
+ "inference_mode",
}
metadata_prefix_patterns = [
- 'problem_type', 'problem_source',
- 'generation_model', 'pass_rate',
+ "problem_type",
+ "problem_source",
+ "generation_model",
+ "pass_rate",
]
priority_patterns = {
- 'generated': 100,
- 'gen_': 90,
- 'model_': 80,
- 'predicted': 70,
- 'completion': 60,
+ "generated": 100,
+ "gen_": 90,
+ "model_": 80,
+ "predicted": 70,
+ "completion": 60,
}
def has_keyword(col_name, keywords):
"""Check if any keyword appears in column name."""
col_lower = col_name.lower()
- col_normalized = col_lower.replace('_', '').replace('-', '').replace(' ', '')
+ col_normalized = col_lower.replace("_", "").replace("-", "").replace(" ", "")
for keyword in keywords:
if keyword in col_lower or keyword in col_normalized:
@@ -180,13 +220,16 @@ def detect_custom_format_heuristic(dataset):
return True
for pattern in metadata_prefix_patterns:
- if col_lower.startswith(pattern.split('_')[0] + '_') and col_lower != pattern:
- if '_' in col_lower:
- prefix = col_lower.split('_')[0]
- if prefix in ['generation', 'pass', 'inference']:
+ if (
+ col_lower.startswith(pattern.split("_")[0] + "_")
+ and col_lower != pattern
+ ):
+ if "_" in col_lower:
+ prefix = col_lower.split("_")[0]
+ if prefix in ["generation", "pass", "inference"]:
return True
- if len(col_lower) <= 2 and not col_lower in ['qa', 'q', 'a']:
+ if len(col_lower) <= 2 and not col_lower in ["qa", "q", "a"]:
return True
return False
@@ -221,16 +264,18 @@ def detect_custom_format_heuristic(dataset):
score += 10
# Penalize ambiguous keywords when scoring for user
- if role_type == 'user':
+ if role_type == "user":
col_lower = col_name.lower()
# If column is ONLY "task" (or task_xxx), give it lower priority for user role
- if 'task' in col_lower and not any(kw in col_lower for kw in user_words_high_priority):
+ if "task" in col_lower and not any(
+ kw in col_lower for kw in user_words_high_priority
+ ):
score -= 15 # Significant penalty so other user columns win
priority_bonus = get_priority_score(col_name)
score += priority_bonus
- if role_type in ['assistant', 'user']:
+ if role_type in ["assistant", "user"]:
avg_length = get_content_length(col_name)
if num_candidates > 1:
@@ -256,20 +301,24 @@ def detect_custom_format_heuristic(dataset):
content_columns = [col for col in all_columns if not is_metadata(col)]
# Count candidates first
- assistant_potential = [col for col in content_columns if has_keyword(col, assistant_words)]
+ assistant_potential = [
+ col for col in content_columns if has_keyword(col, assistant_words)
+ ]
user_potential = [col for col in content_columns if has_keyword(col, user_words)]
# STEP 1: Find best ASSISTANT column
assistant_candidates = []
for col in assistant_potential:
- score = score_column(col, assistant_words, 'assistant', len(assistant_potential))
+ score = score_column(
+ col, assistant_words, "assistant", len(assistant_potential)
+ )
if score > 0:
assistant_candidates.append((col, score))
if assistant_candidates:
- assistant_candidates.sort(key=lambda x: x[1], reverse=True)
+ assistant_candidates.sort(key = lambda x: x[1], reverse = True)
assistant_col = assistant_candidates[0][0]
- mapping[assistant_col] = 'assistant'
+ mapping[assistant_col] = "assistant"
else:
assistant_col = None
@@ -278,14 +327,14 @@ def detect_custom_format_heuristic(dataset):
for col in user_potential:
if col == assistant_col:
continue
- score = score_column(col, user_words, 'user', len(user_potential))
+ score = score_column(col, user_words, "user", len(user_potential))
if score > 0:
user_candidates.append((col, score))
if user_candidates:
- user_candidates.sort(key=lambda x: x[1], reverse=True)
+ user_candidates.sort(key = lambda x: x[1], reverse = True)
user_col = user_candidates[0][0]
- mapping[user_col] = 'user'
+ mapping[user_col] = "user"
else:
user_col = None
@@ -296,7 +345,7 @@ def detect_custom_format_heuristic(dataset):
for col in remaining_columns:
if has_keyword(col, system_words):
# Found a system match in remaining columns
- mapping[col] = 'system'
+ mapping[col] = "system"
system_col = col
break
@@ -309,22 +358,22 @@ def detect_custom_format_heuristic(dataset):
# If no strong keyword match, decide based on what's missing
if not has_keyword(remaining_col, user_words + assistant_words):
- mapping[remaining_col] = 'system'
+ mapping[remaining_col] = "system"
elif user_col is None:
# No user column yet, assign this as user
- mapping[remaining_col] = 'user'
+ mapping[remaining_col] = "user"
else:
# Already have user + assistant, treat as system context
- mapping[remaining_col] = 'system'
+ mapping[remaining_col] = "system"
# VALIDATION: Ensure we have at least user + assistant
- has_user = any(role == 'user' for role in mapping.values())
- has_assistant = any(role == 'assistant' for role in mapping.values())
+ has_user = any(role == "user" for role in mapping.values())
+ has_assistant = any(role == "assistant" for role in mapping.values())
if not has_user and len(remaining_columns) > 0:
for col in remaining_columns:
if col not in mapping:
- mapping[col] = 'user'
+ mapping[col] = "user"
has_user = True
break
@@ -358,14 +407,27 @@ def detect_multimodal_dataset(dataset):
# Keywords that indicate image data
image_keywords = [
- 'image', 'img', 'pixel',
- 'jpg', 'jpeg', 'png', 'webp', 'bmp', 'gif', 'tiff', 'svg',
- 'photo', 'pic', 'picture', 'visual',
- 'file_name', 'filename',
+ "image",
+ "img",
+ "pixel",
+ "jpg",
+ "jpeg",
+ "png",
+ "webp",
+ "bmp",
+ "gif",
+ "tiff",
+ "svg",
+ "photo",
+ "pic",
+ "picture",
+ "visual",
+ "file_name",
+ "filename",
]
# Keywords that indicate audio data
- audio_keywords = ['audio', 'speech', 'wav', 'waveform', 'sound']
+ audio_keywords = ["audio", "speech", "wav", "waveform", "sound"]
multimodal_columns = []
audio_columns = []
@@ -419,7 +481,7 @@ def detect_multimodal_dataset(dataset):
# Detect text column for audio datasets
detected_text_col = None
if audio_columns:
- text_keywords = ['text', 'sentence', 'transcript', 'transcription', 'label']
+ text_keywords = ["text", "sentence", "transcript", "transcription", "label"]
for col_name in column_names:
if col_name.lower() in text_keywords:
detected_text_col = col_name
@@ -430,7 +492,7 @@ def detect_multimodal_dataset(dataset):
# Detect speaker_id column for TTS datasets (CSM, Orpheus, Spark)
detected_speaker_col = None
if audio_columns:
- speaker_keywords = ['source', 'speaker', 'speaker_id']
+ speaker_keywords = ["source", "speaker", "speaker_id"]
for col_name in column_names:
if col_name.lower() in speaker_keywords:
detected_speaker_col = col_name
@@ -456,6 +518,7 @@ def _is_image_value(value) -> bool:
# PIL Image instance
try:
from PIL.Image import Image as PILImage
+
if isinstance(value, PILImage):
return True
except ImportError:
@@ -470,7 +533,9 @@ def _is_image_value(value) -> bool:
if "bytes" in value and "path" in value:
# Check path extension to exclude audio files
path = value.get("path") or ""
- if isinstance(path, str) and any(path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS):
+ if isinstance(path, str) and any(
+ path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS
+ ):
return False
return True
@@ -479,11 +544,13 @@ def _is_image_value(value) -> bool:
return _has_image_header(value)
# String that looks like an image file path or URL
- _IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp', '.tiff', '.svg')
+ _IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".svg")
if isinstance(value, str) and len(value) < 1000:
lower = value.strip().lower()
# Image URL (http://... ending in image extension)
- if lower.startswith(("http://", "https://")) and any(lower.split("?")[0].endswith(ext) for ext in _IMAGE_EXTS):
+ if lower.startswith(("http://", "https://")) and any(
+ lower.split("?")[0].endswith(ext) for ext in _IMAGE_EXTS
+ ):
return True
# Image file path (relative or absolute path ending in image extension)
if any(lower.endswith(ext) for ext in _IMAGE_EXTS):
@@ -493,7 +560,15 @@ def _is_image_value(value) -> bool:
_AUDIO_EXTENSIONS = (
- ".wav", ".mp3", ".flac", ".ogg", ".opus", ".m4a", ".aac", ".wma", ".webm",
+ ".wav",
+ ".mp3",
+ ".flac",
+ ".ogg",
+ ".opus",
+ ".m4a",
+ ".aac",
+ ".wma",
+ ".webm",
)
@@ -509,7 +584,9 @@ def _is_audio_value(value) -> bool:
# Undecoded/streaming β {"bytes": b"...", "path": "some.wav"}
if "bytes" in value or "path" in value:
path = value.get("path") or ""
- if isinstance(path, str) and any(path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS):
+ if isinstance(path, str) and any(
+ path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS
+ ):
return True
return False
@@ -520,19 +597,19 @@ def _has_image_header(data: bytes) -> bool:
if len(data) < 4:
return False
# JPEG
- if data[:2] == b'\xff\xd8':
+ if data[:2] == b"\xff\xd8":
return True
# PNG
- if data[:4] == b'\x89PNG':
+ if data[:4] == b"\x89PNG":
return True
# GIF
- if data[:3] == b'GIF':
+ if data[:3] == b"GIF":
return True
# WebP
- if data[:4] == b'RIFF' and len(data) >= 12 and data[8:12] == b'WEBP':
+ if data[:4] == b"RIFF" and len(data) >= 12 and data[8:12] == b"WEBP":
return True
# BMP
- if data[:2] == b'BM':
+ if data[:2] == b"BM":
return True
return False
@@ -568,10 +645,13 @@ def detect_vlm_dataset_structure(dataset):
if isinstance(content, list) and len(content) > 0:
if isinstance(content[0], dict) and "type" in content[0]:
-
# Check for llava format
- has_index = any('index' in item for item in content if isinstance(item, dict))
- has_images_column = 'images' in column_names
+ has_index = any(
+ "index" in item
+ for item in content
+ if isinstance(item, dict)
+ )
+ has_images_column = "images" in column_names
if has_index and has_images_column:
return {
@@ -583,7 +663,11 @@ def detect_vlm_dataset_structure(dataset):
}
# Standard VLM format
- has_image = any('image' in item for item in content if isinstance(item, dict))
+ has_image = any(
+ "image" in item
+ for item in content
+ if isinstance(item, dict)
+ )
if has_image:
return {
"format": "vlm_messages",
@@ -637,26 +721,65 @@ def detect_vlm_dataset_structure(dataset):
# Define metadata patterns to EXCLUDE
metadata_patterns = {
- 'suffixes': ['_id', '_url', '_name', '_filename', '_uri', '_link', '_key', '_index'],
- 'prefixes': ['id_', 'url_', 'name_', 'filename_', 'uri_', 'link_', 'key_', 'index_'],
+ "suffixes": [
+ "_id",
+ "_url",
+ "_name",
+ "_filename",
+ "_uri",
+ "_link",
+ "_key",
+ "_index",
+ ],
+ "prefixes": [
+ "id_",
+ "url_",
+ "name_",
+ "filename_",
+ "uri_",
+ "link_",
+ "key_",
+ "index_",
+ ],
}
# Image-related keywords
- image_keywords = ['image', 'img', 'photo', 'picture', 'pic', 'visual', 'scan', 'file_name', 'filename']
+ image_keywords = [
+ "image",
+ "img",
+ "photo",
+ "picture",
+ "pic",
+ "visual",
+ "scan",
+ "file_name",
+ "filename",
+ ]
# Text-related keywords
- text_keywords = ['text', 'caption', 'captions', 'description', 'answer', 'output', 'response', 'label']
+ text_keywords = [
+ "text",
+ "caption",
+ "captions",
+ "description",
+ "answer",
+ "output",
+ "response",
+ "label",
+ ]
def is_metadata_column(col_name):
"""Check if column name looks like metadata."""
col_lower = col_name.lower()
# Check suffixes
- if any(col_lower.endswith(suffix) for suffix in metadata_patterns['suffixes']):
+ if any(col_lower.endswith(suffix) for suffix in metadata_patterns["suffixes"]):
return True
# Check prefixes
- if any(col_lower.startswith(prefix) for prefix in metadata_patterns['prefixes']):
+ if any(
+ col_lower.startswith(prefix) for prefix in metadata_patterns["prefixes"]
+ ):
return True
return False
@@ -664,11 +787,13 @@ def detect_vlm_dataset_structure(dataset):
def _score_image_candidate(col, sample_value):
"""Score a candidate image column by how resolvable its value is."""
# PIL Image object (highest priority - already loaded)
- if hasattr(sample_value, 'size') and hasattr(sample_value, 'mode'):
+ if hasattr(sample_value, "size") and hasattr(sample_value, "mode"):
return 100
# Dict with image data (bytes/path from HF Image feature)
- if isinstance(sample_value, dict) and ('bytes' in sample_value or 'path' in sample_value):
+ if isinstance(sample_value, dict) and (
+ "bytes" in sample_value or "path" in sample_value
+ ):
return 75
if isinstance(sample_value, str):
@@ -693,13 +818,16 @@ def detect_vlm_dataset_structure(dataset):
# Local file β check it exists
if not sample_value.startswith(("http://", "https://")):
- return os.path.exists(sample_value) # bare filenames return False here, that's OK
+ return os.path.exists(
+ sample_value
+ ) # bare filenames return False here, that's OK
# URL β quick HEAD request with short timeout
try:
import urllib.request
- req = urllib.request.Request(sample_value, method="HEAD")
- resp = urllib.request.urlopen(req, timeout=3)
+
+ req = urllib.request.Request(sample_value, method = "HEAD")
+ resp = urllib.request.urlopen(req, timeout = 3)
return resp.status < 400
except Exception:
return False
@@ -732,7 +860,7 @@ def detect_vlm_dataset_structure(dataset):
if not candidates:
return None
- candidates.sort(key=lambda x: x[1], reverse=True)
+ candidates.sort(key = lambda x: x[1], reverse = True)
# Single candidate or top candidate is PIL/dict β no probing needed
if len(candidates) == 1 or candidates[0][1] >= 75:
@@ -766,14 +894,18 @@ def detect_vlm_dataset_structure(dataset):
# Longer text = higher priority (likely content, not just a label)
priority = min(len(sample_value), 1000) # Cap at 1000
candidates.append((col, priority))
- elif isinstance(sample_value, list) and len(sample_value) > 0 and isinstance(sample_value[0], str):
+ elif (
+ isinstance(sample_value, list)
+ and len(sample_value) > 0
+ and isinstance(sample_value[0], str)
+ ):
# List of strings (e.g. captions list) β lower priority than plain strings
priority = min(len(sample_value[0]), 1000) // 2
candidates.append((col, priority))
# Return highest priority candidate
if candidates:
- candidates.sort(key=lambda x: x[1], reverse=True)
+ candidates.sort(key = lambda x: x[1], reverse = True)
return candidates[0][0]
return None
diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py
index 556caba2f5..6e8c25cba5 100644
--- a/studio/backend/utils/datasets/llm_assist.py
+++ b/studio/backend/utils/datasets/llm_assist.py
@@ -42,33 +42,45 @@ def precache_helper_gguf():
return
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
- variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
+ variant = os.environ.get(
+ "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
+ )
try:
from huggingface_hub import HfApi, hf_hub_download
+ from huggingface_hub.utils import disable_progress_bars, enable_progress_bars
+
+ disable_progress_bars()
+ logging.getLogger("huggingface_hub").setLevel(logging.WARNING)
# Find the GGUF file matching the variant
api = HfApi()
- files = api.list_repo_files(repo, repo_type="model")
+ files = api.list_repo_files(repo, repo_type = "model")
gguf_files = [f for f in files if f.endswith(".gguf")]
# Find all GGUF files matching the variant (may be split into shards)
variant_lower = variant.lower().replace("-", "_")
matching = sorted(
- f for f in gguf_files
- if variant_lower in f.lower().replace("-", "_")
+ f for f in gguf_files if variant_lower in f.lower().replace("-", "_")
)
if matching:
- logger.info(f"Pre-caching helper GGUF: {repo}/{matching[0]}"
- + (f" (+{len(matching) - 1} shards)" if len(matching) > 1 else ""))
+ logger.info(
+ f"Pre-caching helper GGUF: {repo}/{matching[0]}"
+ + (f" (+{len(matching) - 1} shards)" if len(matching) > 1 else "")
+ )
for target in matching:
- hf_hub_download(repo_id=repo, filename=target)
+ hf_hub_download(repo_id = repo, filename = target)
logger.info(f"Helper GGUF cached: {len(matching)} file(s)")
else:
logger.warning(f"No GGUF matching variant '{variant}' in {repo}")
except Exception as e:
logger.warning(f"Failed to pre-cache helper GGUF: {e}")
+ finally:
+ try:
+ enable_progress_bars()
+ except Exception as e:
+ pass
def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
@@ -81,7 +93,9 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
return None
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
- variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
+ variant = os.environ.get(
+ "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
+ )
backend = None
try:
@@ -89,15 +103,14 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
backend = LlamaCppBackend()
logger.info(f"Loading helper model: {repo} ({variant})")
- print(f"π€ Loading helper model: {repo} ({variant})...")
ok = backend.load_model(
- hf_repo=repo,
- hf_variant=variant,
- model_identifier=f"helper:{repo}:{variant}",
- is_vision=False,
- n_ctx=2048,
- n_gpu_layers=-1,
+ hf_repo = repo,
+ hf_variant = variant,
+ model_identifier = f"helper:{repo}:{variant}",
+ is_vision = False,
+ n_ctx = 2048,
+ n_gpu_layers = -1,
)
if not ok:
logger.warning("Helper model failed to start")
@@ -106,12 +119,12 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
messages = [{"role": "user", "content": prompt}]
cumulative = ""
for text in backend.generate_chat_completion(
- messages=messages,
- temperature=0.1,
- top_p=0.9,
- top_k=20,
- max_tokens=max_tokens,
- repetition_penalty=1.0,
+ messages = messages,
+ temperature = 0.1,
+ top_p = 0.9,
+ top_k = 20,
+ max_tokens = max_tokens,
+ repetition_penalty = 1.0,
):
cumulative = text # cumulative β last value is full text
@@ -127,7 +140,7 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
if backend is not None:
try:
backend.unload_model()
- print("π€ Helper model unloaded")
+ logger.info("Helper model unloaded")
except Exception:
pass
@@ -176,7 +189,7 @@ def llm_generate_vlm_instruction(
"Respond with ONLY the instruction sentence, nothing else."
)
- result = _run_with_helper(prompt, max_tokens=100)
+ result = _run_with_helper(prompt, max_tokens = 100)
if not result:
return None
@@ -187,7 +200,7 @@ def llm_generate_vlm_instruction(
logger.warning(f"Helper model returned unusable instruction: {instruction!r}")
return None
- print(f"π€ LLM-generated instruction: {instruction}")
+ logger.info(f"LLM-generated instruction: {instruction}")
return {
"instruction": instruction,
"confidence": 0.85,
@@ -231,7 +244,7 @@ def llm_classify_columns(
'Example: {"question": "user", "answer": "assistant", "id": "metadata"}'
)
- result = _run_with_helper(prompt, max_tokens=200)
+ result = _run_with_helper(prompt, max_tokens = 200)
if not result:
return None
@@ -248,6 +261,7 @@ def llm_classify_columns(
except json.JSONDecodeError:
# Try to find JSON object in the response
import re
+
match = re.search(r"\{[^}]+\}", text)
if match:
try:
@@ -266,7 +280,11 @@ def llm_classify_columns(
valid_roles = {"user", "assistant", "system", "metadata"}
cleaned = {}
for col, role in mapping.items():
- if col in column_names and isinstance(role, str) and role.lower() in valid_roles:
+ if (
+ col in column_names
+ and isinstance(role, str)
+ and role.lower() in valid_roles
+ ):
cleaned[col] = role.lower()
if not cleaned:
@@ -278,7 +296,7 @@ def llm_classify_columns(
logger.warning(f"Helper model mapping missing user/assistant: {cleaned}")
return None
- print(f"π€ LLM-classified columns: {cleaned}")
+ logger.info(f"LLM-classified columns: {cleaned}")
return cleaned
@@ -319,7 +337,7 @@ def llm_generate_dataset_warning(
"Keep it under 3 sentences. Be specific about the dataset."
)
- result = _run_with_helper(prompt, max_tokens=200)
+ result = _run_with_helper(prompt, max_tokens = 200)
if not result:
return None
@@ -328,7 +346,7 @@ def llm_generate_dataset_warning(
if len(warning) < 10 or len(warning) > 500:
return None
- print(f"π€ LLM-generated warning: {warning}")
+ logger.info(f"LLM-generated warning: {warning}")
return warning
@@ -369,18 +387,16 @@ def _parse_json_response(text: str) -> Optional[dict]:
return None
-def _generate_with_backend(
- backend, messages: list[dict], max_tokens: int = 512
-) -> str:
+def _generate_with_backend(backend, messages: list[dict], max_tokens: int = 512) -> str:
"""Run one chat completion on an already-loaded backend. Returns raw text."""
cumulative = ""
for text in backend.generate_chat_completion(
- messages=messages,
- temperature=0.1,
- top_p=0.9,
- top_k=20,
- max_tokens=max_tokens,
- repetition_penalty=1.0,
+ messages = messages,
+ temperature = 0.1,
+ top_p = 0.9,
+ top_k = 20,
+ max_tokens = max_tokens,
+ repetition_penalty = 1.0,
):
cumulative = text
return cumulative.strip()
@@ -398,7 +414,7 @@ def fetch_hf_dataset_card(
try:
from huggingface_hub import DatasetCard
- card = DatasetCard.load(dataset_name, token=hf_token)
+ card = DatasetCard.load(dataset_name, token = hf_token)
readme = card.text or ""
# Truncate at sentence boundary
@@ -413,14 +429,21 @@ def fetch_hf_dataset_card(
metadata = {}
if card.data:
for key in (
- "task_categories", "task_ids", "language",
- "size_categories", "tags", "license", "pretty_name",
+ "task_categories",
+ "task_ids",
+ "language",
+ "size_categories",
+ "tags",
+ "license",
+ "pretty_name",
):
val = getattr(card.data, key, None)
if val is not None:
metadata[key] = val
- logger.info(f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields")
+ logger.info(
+ f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields"
+ )
return readme, metadata
except Exception as e:
@@ -447,30 +470,31 @@ def _run_multi_pass_advisor(
return None
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
- variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
+ variant = os.environ.get(
+ "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
+ )
backend = None
try:
from core.inference.llama_cpp import LlamaCppBackend
backend = LlamaCppBackend()
- print(f"π€ Loading advisor model: {repo} ({variant})...")
+ logger.info(f"Loading advisor model: {repo} ({variant})")
t0 = time.monotonic()
ok = backend.load_model(
- hf_repo=repo,
- hf_variant=variant,
- model_identifier=f"advisor:{repo}:{variant}",
- is_vision=False,
- n_ctx=2048,
- n_gpu_layers=-1,
+ hf_repo = repo,
+ hf_variant = variant,
+ model_identifier = f"advisor:{repo}:{variant}",
+ is_vision = False,
+ n_ctx = 2048,
+ n_gpu_layers = -1,
)
if not ok:
logger.warning("Advisor model failed to start")
return None
- print(f"π€ Advisor model loaded in {time.monotonic() - t0:.1f}s")
-
+ logger.info(f"Advisor model loaded in {time.monotonic() - t0:.1f}s")
# ββ Format samples ββ
samples_text = ""
for i, row in enumerate(samples[:5], 1):
@@ -478,8 +502,9 @@ def _run_multi_pass_advisor(
samples_text += f"Row {i}:\n" + "\n".join(parts) + "\n"
metadata_str = (
- json.dumps(dataset_metadata, indent=2, default=str)[:500]
- if dataset_metadata else "N/A"
+ json.dumps(dataset_metadata, indent = 2, default = str)[:500]
+ if dataset_metadata
+ else "N/A"
)
card_excerpt = (dataset_card or "")[:1200] or "N/A"
@@ -489,13 +514,14 @@ def _run_multi_pass_advisor(
if model_name:
try:
from utils.models.model_config import load_model_config
- config = load_model_config(model_name, use_auth=True, token=hf_token)
+
+ config = load_model_config(model_name, use_auth = True, token = hf_token)
archs = getattr(config, "architectures", [])
if archs and "Gemma3nForConditionalGeneration" in archs:
is_gemma_3n = True
except Exception:
is_gemma_3n = "gemma-3n" in model_name.lower()
-
+
if model_type == "audio" and not is_gemma_3n:
target_hints = (
"\n\nHINT: The user is training an AUDIO model. The dataset MUST contain "
@@ -514,7 +540,7 @@ def _run_multi_pass_advisor(
)
# ββ Pass 1: Classify ββ
- print("π€ Pass 1: Classifying dataset...", flush=True)
+ logger.info("Pass 1: Classifying dataset...")
t1 = time.monotonic()
messages1 = [
{
@@ -559,9 +585,9 @@ def _run_multi_pass_advisor(
Respond with ONLY the JSON object. No markdown, no explanation."""),
},
]
- raw1 = _generate_with_backend(backend, messages1, max_tokens=256)
+ raw1 = _generate_with_backend(backend, messages1, max_tokens = 256)
pass1 = _parse_json_response(raw1)
- print(f"π€ Pass 1 done ({time.monotonic() - t1:.1f}s): {pass1}", flush=True)
+ logger.info(f"Pass 1 done ({time.monotonic() - t1:.1f}s): {pass1}")
if not pass1:
logger.warning(f"Advisor Pass 1 failed to produce JSON: {raw1[:200]}")
@@ -580,7 +606,9 @@ def _run_multi_pass_advisor(
}
# ββ Pass 2: Map columns to roles ββ
- print("π€ Pass 2: Mapping columns to roles...", flush=True)
+ logger.info("Pass 2: Mapping columns to roles...")
+
+
t2 = time.monotonic()
messages2 = [
{
@@ -592,13 +620,13 @@ def _run_multi_pass_advisor(
'- "user" = This column contains INPUT that the model will receive as a prompt.\n'
'- "assistant" = This column contains OUTPUT that the model should learn to generate.\n\n'
"CRITICAL RULES:\n"
- "1. There MUST be at least one column assigned to \"user\" AND at least one "
- "column assigned to \"assistant\". Never assign all columns to the same role.\n"
+ '1. There MUST be at least one column assigned to "user" AND at least one '
+ 'column assigned to "assistant". Never assign all columns to the same role.\n'
"2. The column that contains the TARGET or OUTPUT or ANSWER or LABEL must "
- "ALWAYS be assigned to \"assistant\". This is the thing the model should learn "
+ 'ALWAYS be assigned to "assistant". This is the thing the model should learn '
"to produce.\n"
"3. The columns that contain the SOURCE or INPUT or CONTEXT or QUESTION must "
- "be assigned to \"user\". This is what the model receives.\n"
+ 'be assigned to "user". This is what the model receives.\n'
'4. Metadata columns like "id", "index", "source", "url", "date" should be '
'set to "skip".\n\n'
"You must respond with ONLY a valid JSON object."
@@ -611,7 +639,7 @@ def _run_multi_pass_advisor(
Here is a dataset that has been classified:
CLASSIFICATION:
- {json.dumps(pass1, indent=2)}
+ {json.dumps(pass1, indent = 2)}
COLUMNS AVAILABLE: {columns}
@@ -659,9 +687,9 @@ def _run_multi_pass_advisor(
Respond with ONLY the JSON object."""),
},
]
- raw2 = _generate_with_backend(backend, messages2, max_tokens=512)
+ raw2 = _generate_with_backend(backend, messages2, max_tokens = 512)
pass2 = _parse_json_response(raw2)
- print(f"π€ Pass 2 done ({time.monotonic() - t2:.1f}s): {pass2}", flush=True)
+ logger.info(f"Pass 2 done ({time.monotonic() - t2:.1f}s): {pass2}")
if not pass2:
logger.warning(f"Advisor Pass 2 failed to produce JSON: {raw2[:200]}")
@@ -674,10 +702,7 @@ def _run_multi_pass_advisor(
# Validate: must have at least one user AND one assistant
roles_present = set(column_roles.values())
if "user" not in roles_present or "assistant" not in roles_present:
- print(
- f"π€ Pass 2 sanity fail: missing user or assistant role: {column_roles}",
- flush=True,
- )
+ logger.warning(f"Pass 2 sanity fail: missing user or assistant role: {column_roles}")
return None # triggers fallback to simple classification
# ββ Pass 3: System prompt (non-conversational datasets only) ββ
@@ -686,7 +711,7 @@ def _run_multi_pass_advisor(
is_conv = pass1.get("is_conversational", False)
if not is_conv:
- print("π€ Pass 3: Generating system prompt...", flush=True)
+ logger.info("Pass 3: Generating system prompt...")
t3 = time.monotonic()
# Format label mapping info for the prompt
@@ -726,8 +751,8 @@ def _run_multi_pass_advisor(
Write ONLY the system prompt text. No quotes, no labels, no explanation around it."""),
},
]
- raw3 = _generate_with_backend(backend, messages3, max_tokens=256)
- print(f"π€ Pass 3 done ({time.monotonic() - t3:.1f}s): {raw3[:200] if raw3 else None}", flush=True)
+ raw3 = _generate_with_backend(backend, messages3, max_tokens = 256)
+ logger.info(f"Pass 3 done ({time.monotonic() - t3:.1f}s): {raw3[:200] if raw3 else None}")
if raw3:
# Pass 3 returns raw text, not JSON β clean it up
@@ -746,15 +771,14 @@ def _run_multi_pass_advisor(
note_parts = [f"This is a {dtype} dataset (not conversational)."]
if desc:
note_parts.append(desc)
- note_parts.append("Columns have been mapped to conversation roles. You can adjust the mapping if needed.")
+ note_parts.append(
+ "Columns have been mapped to conversation roles. You can adjust the mapping if needed."
+ )
user_notification = " ".join(note_parts)
total_time = time.monotonic() - t0
- print(
- f"π€ Advisor complete ({total_time:.1f}s): type={dtype}, "
- f"mapping={suggested_mapping}, sys_prompt={bool(sys_prompt)}, label_map={bool(label_map)}",
- flush=True,
- )
+ logger.info(f"Advisor complete ({total_time:.1f}s): type={dtype}, mapping={suggested_mapping}, sys_prompt={bool(sys_prompt)}, label_map={bool(label_map)}")
+
return {
"success": True,
@@ -774,7 +798,7 @@ def _run_multi_pass_advisor(
if backend is not None:
try:
backend.unload_model()
- print("π€ Advisor model unloaded")
+ logger.info("Advisor model unloaded")
except Exception:
pass
@@ -805,18 +829,18 @@ def llm_conversion_advisor(
# Try multi-pass advisor
result = _run_multi_pass_advisor(
- columns=column_names,
- samples=samples,
- dataset_name=dataset_name,
- dataset_card=dataset_card,
- dataset_metadata=dataset_metadata,
- model_name=model_name,
- model_type=model_type,
- hf_token=hf_token,
+ columns = column_names,
+ samples = samples,
+ dataset_name = dataset_name,
+ dataset_card = dataset_card,
+ dataset_metadata = dataset_metadata,
+ model_name = model_name,
+ model_type = model_type,
+ hf_token = hf_token,
)
if result and result.get("success"):
- print(f"π€ Conversion advisor succeeded: type={result.get('dataset_type')}")
+ logger.info(f"Conversion advisor succeeded: type={result.get('dataset_type')}")
return result
# Fallback: simple column classification
@@ -826,7 +850,8 @@ def llm_conversion_advisor(
return {
"success": True,
"suggested_mapping": {
- col: role for col, role in simple_mapping.items()
+ col: role
+ for col, role in simple_mapping.items()
if role in ("user", "assistant", "system")
},
"dataset_type": None,
diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py
index 31c2da28dc..96b0c50876 100644
--- a/studio/backend/utils/datasets/model_mappings.py
+++ b/studio/backend/utils/datasets/model_mappings.py
@@ -8,7 +8,6 @@ This module contains the mapping dictionaries that associate model names
with their corresponding chat templates and response markers.
"""
-
TEMPLATE_TO_MODEL_MAPPER = {
"phi-3.5": (
"unsloth/Phi-3.5-mini-instruct-bnb-4bit",
@@ -407,14 +406,11 @@ MODEL_TO_TEMPLATE_MAPPER = {}
for key, values in TEMPLATE_TO_MODEL_MAPPER.items():
for value in values:
MODEL_TO_TEMPLATE_MAPPER[value] = key
- pass
# Get lowercased
lowered_key = key.lower()
for value in values:
MODEL_TO_TEMPLATE_MAPPER[value.lower()] = lowered_key
- pass
-pass
TEMPLATE_TO_RESPONSES_MAPPER = {
@@ -531,4 +527,3 @@ TEMPLATE_TO_RESPONSES_MAPPER = {
"response": "<|assistant|>",
},
}
-
diff --git a/studio/backend/utils/datasets/vlm_processing.py b/studio/backend/utils/datasets/vlm_processing.py
index 63aeec514e..7b63152ede 100644
--- a/studio/backend/utils/datasets/vlm_processing.py
+++ b/studio/backend/utils/datasets/vlm_processing.py
@@ -14,9 +14,9 @@ from itertools import islice
def generate_smart_vlm_instruction(
dataset,
- text_column="text",
- image_column="image",
- dataset_name=None,
+ text_column = "text",
+ image_column = "image",
+ dataset_name = None,
):
"""
Generate smart, context-aware instruction for VLM datasets using heuristics.
@@ -66,11 +66,12 @@ def generate_smart_vlm_instruction(
# OCR / Transcription
"ocr": {
"keywords": ["ocr", "transcribe", "transcript"],
- "content_hints": [r"[A-Za-z\u0600-\u06FF]{10,}"], # Long text passages (Latin/Arabic)
+ "content_hints": [
+ r"[A-Za-z\u0600-\u06FF]{10,}"
+ ], # Long text passages (Latin/Arabic)
"instruction": "Transcribe all the text shown in this image.",
"confidence": 0.9,
},
-
# LaTeX / Math
"latex": {
"keywords": ["latex", "math", "formula", "equation"],
@@ -78,7 +79,6 @@ def generate_smart_vlm_instruction(
"instruction": "Convert this image to LaTeX notation.",
"confidence": 0.95,
},
-
# Caption / Description
"caption": {
"keywords": ["caption", "description", "describe"],
@@ -86,15 +86,21 @@ def generate_smart_vlm_instruction(
"instruction": "Provide a detailed description of this image.",
"confidence": 0.85,
},
-
# Medical / Radiology
"medical": {
- "keywords": ["medical", "radiology", "xray", "ct", "mri", "scan", "diagnosis"],
+ "keywords": [
+ "medical",
+ "radiology",
+ "xray",
+ "ct",
+ "mri",
+ "scan",
+ "diagnosis",
+ ],
"content_hints": [r"\b(lesion|radiograph|patient|diagnosis|findings)\b"],
"instruction": "Analyze this medical image and describe the key findings.",
"confidence": 0.9,
},
-
# Code / Programming
"code": {
"keywords": ["code", "program", "function", "algorithm"],
@@ -102,7 +108,6 @@ def generate_smart_vlm_instruction(
"instruction": "Explain what this code visualization shows.",
"confidence": 0.85,
},
-
# Chart / Graph
"chart": {
"keywords": ["chart", "graph", "plot", "visualization", "diagram"],
@@ -110,7 +115,6 @@ def generate_smart_vlm_instruction(
"instruction": "Describe this chart or graph, including key data points and trends.",
"confidence": 0.85,
},
-
# Document / Text Recognition
"document": {
"keywords": ["document", "page", "paragraph", "article"],
@@ -132,7 +136,9 @@ def generate_smart_vlm_instruction(
score += 0.5
# Check dataset name if provided
- if dataset_name and any(keyword in dataset_name.lower() for keyword in task_info["keywords"]):
+ if dataset_name and any(
+ keyword in dataset_name.lower() for keyword in task_info["keywords"]
+ ):
score += 0.3
# Check content patterns
@@ -186,7 +192,7 @@ def generate_smart_vlm_instruction(
row = {}
for col in s:
val = s[col]
- if hasattr(val, 'size') and hasattr(val, 'mode'): # PIL Image
+ if hasattr(val, "size") and hasattr(val, "mode"): # PIL Image
row[col] = ""
elif isinstance(val, list):
row[col] = str(val)[:300]
@@ -195,15 +201,15 @@ def generate_smart_vlm_instruction(
sample_rows.append(row)
llm_result = llm_generate_vlm_instruction(
- column_names=list(column_names),
- samples=sample_rows,
- dataset_name=dataset_name,
+ column_names = list(column_names),
+ samples = sample_rows,
+ dataset_name = dataset_name,
)
if llm_result and llm_result.get("instruction"):
print(
f"\n[DEBUG] LLM-assisted VLM instruction generated: "
f"'{llm_result['instruction']}' (confidence={llm_result.get('confidence', 'N/A')})\n",
- flush=True,
+ flush = True,
)
return {
"instruction": llm_result["instruction"],
@@ -214,6 +220,7 @@ def generate_smart_vlm_instruction(
}
except Exception as e:
import logging
+
logging.getLogger(__name__).debug(f"LLM-assisted instruction skipped: {e}")
# ===== LEVEL 5: Generic Fallback =====
diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py
index 600084f5f2..c1f8b62010 100644
--- a/studio/backend/utils/hardware/__init__.py
+++ b/studio/backend/utils/hardware/__init__.py
@@ -4,6 +4,7 @@
"""
Hardware detection and GPU utilities
"""
+
from .hardware import (
DeviceType,
DEVICE,
@@ -21,17 +22,17 @@ from .hardware import (
)
__all__ = [
- 'DeviceType',
- 'DEVICE',
- 'detect_hardware',
- 'get_device',
- 'is_apple_silicon',
- 'clear_gpu_cache',
- 'get_gpu_memory_info',
- 'log_gpu_memory',
- 'get_gpu_summary',
- 'get_package_versions',
- 'get_gpu_utilization',
- 'get_physical_gpu_count',
- 'safe_num_proc',
+ "DeviceType",
+ "DEVICE",
+ "detect_hardware",
+ "get_device",
+ "is_apple_silicon",
+ "clear_gpu_cache",
+ "get_gpu_memory_info",
+ "log_gpu_memory",
+ "get_gpu_summary",
+ "get_package_versions",
+ "get_gpu_utilization",
+ "get_physical_gpu_count",
+ "safe_num_proc",
]
diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index be8e2eb9c0..c743dcd897 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -15,6 +15,7 @@ Usage:
import torch
...
"""
+
import platform
import structlog
from loggers import get_logger
@@ -26,11 +27,13 @@ logger = get_logger(__name__)
# ========== Device Enum ==========
+
class DeviceType(str, Enum):
"""Supported compute backends. Inherits from str so it serializes cleanly in JSON."""
+
CUDA = "cuda"
- MLX = "mlx"
- CPU = "cpu"
+ MLX = "mlx"
+ CPU = "cpu"
# ========== Global State (set once by detect_hardware) ==========
@@ -40,6 +43,7 @@ DEVICE: Optional[DeviceType] = None
# ========== Detection ==========
+
def is_apple_silicon() -> bool:
"""Check if running on Apple Silicon hardware (pure platform check, no ML imports)."""
return platform.system() == "Darwin" and platform.machine() == "arm64"
@@ -49,6 +53,7 @@ def _has_torch() -> bool:
"""Check if PyTorch is importable."""
try:
import torch
+
return True
except ImportError:
return False
@@ -58,6 +63,7 @@ def _has_mlx() -> bool:
"""Check if MLX is importable."""
try:
import mlx.core
+
return True
except ImportError:
return False
@@ -80,6 +86,7 @@ def detect_hardware() -> DeviceType:
# --- CUDA: try PyTorch ---
if _has_torch():
import torch
+
if torch.cuda.is_available():
DEVICE = DeviceType.CUDA
device_name = torch.cuda.get_device_properties(0).name
@@ -101,6 +108,7 @@ def detect_hardware() -> DeviceType:
# ========== Convenience helpers ==========
+
def get_device() -> DeviceType:
"""
Return the detected device. Auto-detects if detect_hardware() hasn't been called yet.
@@ -118,12 +126,14 @@ def clear_gpu_cache():
Safe to call on any platform β no-ops gracefully.
"""
import gc
+
gc.collect()
device = get_device()
if device == DeviceType.CUDA:
import torch
+
torch.cuda.synchronize()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
@@ -144,6 +154,7 @@ def get_gpu_memory_info() -> Dict[str, Any]:
if device == DeviceType.CUDA:
try:
import torch
+
idx = torch.cuda.current_device()
props = torch.cuda.get_device_properties(idx)
@@ -215,6 +226,7 @@ def log_gpu_memory(context: str):
# ========== GPU Summary & Package Versions ==========
+
def get_gpu_summary() -> Dict[str, Any]:
"""
Return a compact summary of the primary GPU.
@@ -256,6 +268,7 @@ def get_package_versions() -> Dict[str, Optional[str]]:
# CUDA toolkit version bundled with torch
try:
import torch
+
versions["cuda"] = getattr(torch.version, "cuda", None)
except Exception:
versions["cuda"] = None
@@ -265,6 +278,7 @@ def get_package_versions() -> Dict[str, Optional[str]]:
# ========== Live GPU Utilization (nvidia-smi) ==========
+
def get_gpu_utilization() -> Dict[str, Any]:
"""
Return a live snapshot of GPU utilization via ``nvidia-smi``.
@@ -312,9 +326,9 @@ def get_gpu_utilization() -> Dict[str, Any]:
"memory.used,memory.total,power.draw,power.limit",
"--format=csv,noheader,nounits",
],
- capture_output=True,
- text=True,
- timeout=5,
+ capture_output = True,
+ text = True,
+ timeout = 5,
)
if result.returncode == 0 and result.stdout.strip():
@@ -360,7 +374,9 @@ def get_gpu_utilization() -> Dict[str, Any]:
power_limit = smi_data.get("power_limit")
vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None
- vram_total_gb = round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
+ vram_total_gb = (
+ round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
+ )
vram_pct = (
round((vram_used_mb / vram_total_mb) * 100, 1)
if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0
@@ -395,6 +411,7 @@ def get_gpu_utilization() -> Dict[str, Any]:
_physical_gpu_count: Optional[int] = None
+
def get_physical_gpu_count() -> int:
"""
Return the number of physical NVIDIA GPUs on the machine.
@@ -409,9 +426,12 @@ def get_physical_gpu_count() -> int:
try:
import subprocess
+
result = subprocess.run(
["nvidia-smi", "-L"],
- capture_output=True, text=True, timeout=5,
+ capture_output = True,
+ text = True,
+ timeout = 5,
)
if result.returncode == 0 and result.stdout.strip():
_physical_gpu_count = len(result.stdout.strip().splitlines())
diff --git a/studio/backend/utils/inference/__init__.py b/studio/backend/utils/inference/__init__.py
index 15645ef89a..9dc9c08767 100644
--- a/studio/backend/utils/inference/__init__.py
+++ b/studio/backend/utils/inference/__init__.py
@@ -4,7 +4,7 @@
"""
Inference utility functions
"""
+
from utils.inference.inference_config import load_inference_config
__all__ = ["load_inference_config"]
-
diff --git a/studio/backend/utils/inference/inference_config.py b/studio/backend/utils/inference/inference_config.py
index 3e4727e4e3..5b92efce6e 100644
--- a/studio/backend/utils/inference/inference_config.py
+++ b/studio/backend/utils/inference/inference_config.py
@@ -7,6 +7,7 @@ Inference configuration loading utilities.
This module provides functions to load inference parameters (temperature, top_p, top_k, min_p)
from model YAML configuration files, with fallback to default.yaml.
"""
+
from pathlib import Path
from typing import Dict, Any
import yaml
@@ -21,15 +22,15 @@ logger = get_logger(__name__)
def load_inference_config(model_identifier: str) -> Dict[str, Any]:
"""
Load inference configuration parameters for a model.
-
+
This function loads inference parameters (temperature, top_p, top_k, min_p) from the
model's YAML configuration file using the same mapping logic as the /config endpoint.
If a parameter is missing from the model's config, it falls back to the value in
default.yaml.
-
+
Args:
model_identifier: Model identifier (e.g., "unsloth/llama-3-8b-bnb-4bit")
-
+
Returns:
Dictionary containing inference parameters:
{
@@ -41,30 +42,33 @@ def load_inference_config(model_identifier: str) -> Dict[str, Any]:
"""
# Load model defaults to get inference parameters
model_defaults = load_model_defaults(model_identifier)
-
+
# Load default.yaml for fallback values
script_dir = Path(__file__).parent.parent.parent
defaults_dir = script_dir / "assets" / "configs" / "model_defaults"
default_config_path = defaults_dir / "default.yaml"
-
+
default_inference = {}
if default_config_path.exists():
try:
- with open(default_config_path, 'r', encoding='utf-8') as f:
+ with open(default_config_path, "r", encoding = "utf-8") as f:
default_config = yaml.safe_load(f) or {}
default_inference = default_config.get("inference", {})
except Exception as e:
logger.warning(f"Failed to load default.yaml: {e}")
-
+
# Extract inference parameters from model config, fallback to defaults
model_inference = model_defaults.get("inference", {})
inference_config = {
- "temperature": model_inference.get("temperature", default_inference.get("temperature", 0.7)),
+ "temperature": model_inference.get(
+ "temperature", default_inference.get("temperature", 0.7)
+ ),
"top_p": model_inference.get("top_p", default_inference.get("top_p", 0.95)),
"top_k": model_inference.get("top_k", default_inference.get("top_k", -1)),
"min_p": model_inference.get("min_p", default_inference.get("min_p", 0.01)),
- "trust_remote_code": model_inference.get("trust_remote_code", default_inference.get("trust_remote_code", False)),
+ "trust_remote_code": model_inference.get(
+ "trust_remote_code", default_inference.get("trust_remote_code", False)
+ ),
}
-
- return inference_config
+ return inference_config
diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py
index 46ec1362ff..82236d8013 100644
--- a/studio/backend/utils/models/__init__.py
+++ b/studio/backend/utils/models/__init__.py
@@ -4,6 +4,7 @@
"""
Model and LoRA configuration handling
"""
+
from .model_config import (
ModelConfig,
GgufVariantInfo,
@@ -24,20 +25,20 @@ from .model_config import (
from .checkpoints import scan_checkpoints
__all__ = [
- 'ModelConfig',
- 'GgufVariantInfo',
- 'is_vision_model',
- 'is_embedding_model',
- 'detect_audio_type',
- 'is_audio_input_type',
- 'VALID_AUDIO_TYPES',
- 'scan_trained_loras',
- 'scan_exported_models',
- 'load_model_defaults',
- 'get_base_model_from_lora',
- 'load_model_config',
- 'list_gguf_variants',
- 'MODEL_NAME_MAPPING',
- 'UI_STATUS_INDICATORS',
- 'scan_checkpoints',
+ "ModelConfig",
+ "GgufVariantInfo",
+ "is_vision_model",
+ "is_embedding_model",
+ "detect_audio_type",
+ "is_audio_input_type",
+ "VALID_AUDIO_TYPES",
+ "scan_trained_loras",
+ "scan_exported_models",
+ "load_model_defaults",
+ "get_base_model_from_lora",
+ "load_model_config",
+ "list_gguf_variants",
+ "MODEL_NAME_MAPPING",
+ "UI_STATUS_INDICATORS",
+ "scan_checkpoints",
]
diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py
index f33d170816..a7cb80f338 100644
--- a/studio/backend/utils/models/checkpoints.py
+++ b/studio/backend/utils/models/checkpoints.py
@@ -4,6 +4,7 @@
"""
Checkpoint scanning utilities for discovering training runs and their checkpoints.
"""
+
import json
import structlog
from loggers import get_logger
@@ -86,7 +87,9 @@ def scan_checkpoints(
name_part = parts[0]
idx = name_part.find("_")
if idx > 0:
- metadata["base_model"] = name_part[:idx] + "/" + name_part[idx + 1:]
+ metadata["base_model"] = (
+ name_part[:idx] + "/" + name_part[idx + 1 :]
+ )
else:
metadata["base_model"] = name_part
@@ -109,13 +112,19 @@ def scan_checkpoints(
# Assign the last checkpoint's loss to the main adapter entry
if len(checkpoints) > 1:
last_checkpoint_loss = checkpoints[-1][2]
- checkpoints[0] = (checkpoints[0][0], checkpoints[0][1], last_checkpoint_loss)
+ checkpoints[0] = (
+ checkpoints[0][0],
+ checkpoints[0][1],
+ last_checkpoint_loss,
+ )
models.append((item.name, checkpoints, metadata))
- logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)")
+ logger.debug(
+ f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)"
+ )
# Sort by modification time (newest first)
- models.sort(key=lambda x: Path(x[1][0][1]).stat().st_mtime, reverse=True)
+ models.sort(key = lambda x: Path(x[1][0][1]).stat().st_mtime, reverse = True)
logger.info(f"Found {len(models)} training runs in {outputs_dir}")
return models
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index 21aa1daae1..960e263cb2 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -4,6 +4,7 @@
"""
Model and LoRA configuration handling
"""
+
from transformers import AutoConfig
from dataclasses import dataclass
from typing import Optional, Dict, Any
@@ -77,7 +78,6 @@ MODEL_NAME_MAPPING = {
"unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml": [
"unsloth/ERNIE-4.5-VL-28B-A3B-PT",
],
-
"tiiuae_Falcon-H1-0.5B-Instruct.yaml": [
"tiiuae/Falcon-H1-0.5B-Instruct",
"unsloth/Falcon-H1-0.5B-Instruct",
@@ -132,7 +132,6 @@ MODEL_NAME_MAPPING = {
"unsloth/gpt-oss-20b-unsloth-bnb-4bit",
"unsloth/gpt-oss-20b-BF16",
],
-
"unsloth_gpt-oss-120b.yaml": [
"openai/gpt-oss-120b",
"unsloth/gpt-oss-120b-unsloth-bnb-4bit",
@@ -169,7 +168,6 @@ MODEL_NAME_MAPPING = {
"unsloth/Meta-Llama-3.1-405B-bnb-4bit",
"meta-llama/Meta-Llama-3.1-405B",
],
-
"unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml": [
"unsloth/Meta-Llama-3.1-8B-Instruct-unsloth-bnb-4bit",
"unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
@@ -229,10 +227,9 @@ MODEL_NAME_MAPPING = {
"unsloth/Mistral-Nemo-Base-2407-bnb-4bit",
"unsloth/Mistral-Nemo-Base-2407",
"mistralai/Mistral-Nemo-Base-2407",
- "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit",
- "unsloth/Mistral-Nemo-Instruct-2407",
+ "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit",
+ "unsloth/Mistral-Nemo-Instruct-2407",
"mistralai/Mistral-Nemo-Instruct-2407",
-
],
"unsloth_Mistral-Small-Instruct-2409.yaml": [
"unsloth/Mistral-Small-Instruct-2409-bnb-4bit",
@@ -384,7 +381,10 @@ for canonical_file, model_names in MODEL_NAME_MAPPING.items():
for model_name in model_names:
_REVERSE_MODEL_MAPPING[model_name.lower()] = canonical_file
-def load_model_config(model_name: str, use_auth: bool = False, token: Optional[str] = None):
+
+def load_model_config(
+ model_name: str, use_auth: bool = False, token: Optional[str] = None
+):
"""
Load model config with optional authentication control.
"""
@@ -392,32 +392,30 @@ def load_model_config(model_name: str, use_auth: bool = False, token: Optional[s
if token:
# Explicit token provided - use it
return AutoConfig.from_pretrained(
- model_name,
- trust_remote_code=True,
- token=token
+ model_name, trust_remote_code = True, token = token
)
if not use_auth:
# Load without any authentication (for public model checks)
with without_hf_auth():
return AutoConfig.from_pretrained(
- model_name,
- trust_remote_code=True,
- token=None
+ model_name, trust_remote_code = True, token = None
)
# Use default authentication (cached tokens)
- return AutoConfig.from_pretrained(
- model_name,
- trust_remote_code=True
- )
+ return AutoConfig.from_pretrained(model_name, trust_remote_code = True)
# VLM architecture suffixes and known VLM model_type values.
_VLM_ARCH_SUFFIXES = ("ForConditionalGeneration", "ForVisionText2Text")
_VLM_MODEL_TYPES = {
- 'phi3_v', 'llava', 'llava_next', 'llava_onevision',
- 'internvl_chat', 'cogvlm2', 'minicpmv',
+ "phi3_v",
+ "llava",
+ "llava_next",
+ "llava_onevision",
+ "internvl_chat",
+ "cogvlm2",
+ "minicpmv",
}
# Pre-computed .venv_t5 path and backend dir for subprocess version switching.
@@ -426,7 +424,7 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent)
# Inline script executed in a subprocess with transformers 5.x activated.
# Receives model_name and token via argv, prints JSON result to stdout.
-_VISION_CHECK_SCRIPT = r'''
+_VISION_CHECK_SCRIPT = r"""
import sys, os, json
os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -472,10 +470,12 @@ try:
except Exception as exc:
logger.info(json.dumps({"error": str(exc)}))
sys.exit(1)
-'''
+"""
-def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None) -> bool:
+def _is_vision_model_subprocess(
+ model_name: str, hf_token: Optional[str] = None
+) -> bool:
"""Run is_vision_model check in a subprocess with transformers 5.x.
Same pattern as training/inference workers: spawn a clean subprocess
@@ -486,16 +486,26 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
try:
result = subprocess.run(
- [sys.executable, "-c", _VISION_CHECK_SCRIPT,
- _VENV_T5_DIR, _BACKEND_DIR, model_name, token_arg],
- capture_output=True, text=True, timeout=60,
+ [
+ sys.executable,
+ "-c",
+ _VISION_CHECK_SCRIPT,
+ _VENV_T5_DIR,
+ _BACKEND_DIR,
+ model_name,
+ token_arg,
+ ],
+ capture_output = True,
+ text = True,
+ timeout = 60,
)
if result.returncode != 0:
stderr = result.stderr.strip()
logger.warning(
"Vision check subprocess failed for '%s': %s",
- model_name, stderr or result.stdout.strip(),
+ model_name,
+ stderr or result.stdout.strip(),
)
return False
@@ -503,7 +513,8 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
if "error" in data:
logger.warning(
"Vision check subprocess error for '%s': %s",
- model_name, data["error"],
+ model_name,
+ data["error"],
)
return False
@@ -511,7 +522,10 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
logger.info(
"Vision check (subprocess, transformers 5.x) for '%s': "
"model_type=%s, architectures=%s, is_vision=%s",
- model_name, data.get("model_type"), data.get("architectures"), is_vlm,
+ model_name,
+ data.get("model_type"),
+ data.get("architectures"),
+ is_vlm,
)
return is_vlm
@@ -540,52 +554,54 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
# because AutoConfig in the main process (transformers 4.57.x) doesn't
# recognize their architectures.
from utils.transformers_version import needs_transformers_5
+
if needs_transformers_5(model_name):
logger.info(
"Model '%s' needs transformers 5.x β checking vision via subprocess",
model_name,
)
- return _is_vision_model_subprocess(model_name, hf_token=hf_token)
+ return _is_vision_model_subprocess(model_name, hf_token = hf_token)
try:
- config = load_model_config(model_name, use_auth=True, token=hf_token)
+ config = load_model_config(model_name, use_auth = True, token = hf_token)
# Exclude audio-only models that share ForConditionalGeneration suffix
# (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration)
- _audio_only_model_types = {'csm', 'whisper'}
- model_type = getattr(config, 'model_type', None)
+ _audio_only_model_types = {"csm", "whisper"}
+ model_type = getattr(config, "model_type", None)
if model_type in _audio_only_model_types:
return False
# Check 1: Architecture class name patterns
- if hasattr(config, 'architectures'):
- is_vlm = any(
- x.endswith(_VLM_ARCH_SUFFIXES)
- for x in config.architectures
- )
+ if hasattr(config, "architectures"):
+ is_vlm = any(x.endswith(_VLM_ARCH_SUFFIXES) for x in config.architectures)
if is_vlm:
- logger.info(f"Model {model_name} detected as VLM: architecture {config.architectures}")
+ logger.info(
+ f"Model {model_name} detected as VLM: architecture {config.architectures}"
+ )
return True
# Check 2: Has vision_config (most VLMs: LLaVA, Gemma-3, Qwen2-VL, etc.)
- if hasattr(config, 'vision_config'):
+ if hasattr(config, "vision_config"):
logger.info(f"Model {model_name} detected as VLM: has vision_config")
return True
# Check 3: Has img_processor (Phi-3.5 Vision uses this instead of vision_config)
- if hasattr(config, 'img_processor'):
+ if hasattr(config, "img_processor"):
logger.info(f"Model {model_name} detected as VLM: has img_processor")
return True
# Check 4: Has image_token_index (common in VLMs for image placeholder tokens)
- if hasattr(config, 'image_token_index'):
+ if hasattr(config, "image_token_index"):
logger.info(f"Model {model_name} detected as VLM: has image_token_index")
return True
# Check 5: Known VLM model_type values that may not match above checks
- if hasattr(config, 'model_type'):
+ if hasattr(config, "model_type"):
if config.model_type in _VLM_MODEL_TYPES:
- logger.info(f"Model {model_name} detected as VLM: model_type={config.model_type}")
+ logger.info(
+ f"Model {model_name} detected as VLM: model_type={config.model_type}"
+ )
return True
return False
@@ -595,19 +611,20 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
return False
-VALID_AUDIO_TYPES = ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm')
+VALID_AUDIO_TYPES = ("snac", "csm", "bicodec", "dac", "whisper", "audio_vlm")
# Cache detection results per session to avoid repeated API calls
_audio_detection_cache: Dict[str, Optional[str]] = {}
# Tokenizer token patterns β audio_type (all 6 types detected from tokenizer_config.json)
_AUDIO_TOKEN_PATTERNS = {
- 'csm': lambda tokens: '<|AUDIO|>' in tokens and '<|audio_eos|>' in tokens,
- 'whisper': lambda tokens: '<|startoftranscript|>' in tokens,
- 'audio_vlm': lambda tokens: '' in tokens,
- 'bicodec': lambda tokens: any(t.startswith('<|bicodec_') for t in tokens),
- 'dac': lambda tokens: '<|audio_start|>' in tokens and '<|audio_end|>' in tokens,
- 'snac': lambda tokens: sum(1 for t in tokens if t.startswith(' 10000,
+ "csm": lambda tokens: "<|AUDIO|>" in tokens and "<|audio_eos|>" in tokens,
+ "whisper": lambda tokens: "<|startoftranscript|>" in tokens,
+ "audio_vlm": lambda tokens: "" in tokens,
+ "bicodec": lambda tokens: any(t.startswith("<|bicodec_") for t in tokens),
+ "dac": lambda tokens: "<|audio_start|>" in tokens and "<|audio_end|>" in tokens,
+ "snac": lambda tokens: sum(1 for t in tokens if t.startswith(" 10000,
}
@@ -631,17 +648,20 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option
return result
-def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]:
+def _detect_audio_from_tokenizer(
+ model_name: str, hf_token: Optional[str] = None
+) -> Optional[str]:
"""Detect audio type from tokenizer special tokens (for LLM-based audio models).
First checks local HF cache, then fetches tokenizer_config.json from HuggingFace.
Checks added_tokens_decoder for distinctive patterns.
"""
+
def _check_token_patterns(tok_config: dict) -> Optional[str]:
- added = tok_config.get('added_tokens_decoder', {})
+ added = tok_config.get("added_tokens_decoder", {})
if not added:
return None
- token_contents = [v.get('content', '') for v in added.values()]
+ token_contents = [v.get("content", "") for v in added.values()]
for audio_type, check_fn in _AUDIO_TOKEN_PATTERNS.items():
if check_fn(token_contents):
return audio_type
@@ -650,6 +670,7 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
# 1) Check local HF cache first (works for gated/offline models)
try:
from huggingface_hub.constants import HF_HUB_CACHE
+
cache_dir = Path(HF_HUB_CACHE)
repo_dir_name = f"models--{model_name.replace('/', '--')}"
repo_dir = cache_dir / repo_dir_name
@@ -657,7 +678,10 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
snapshots_dir = repo_dir / "snapshots"
if snapshots_dir.exists():
for snapshot in snapshots_dir.iterdir():
- for tok_path in ['tokenizer_config.json', 'LLM/tokenizer_config.json']:
+ for tok_path in [
+ "tokenizer_config.json",
+ "LLM/tokenizer_config.json",
+ ]:
tok_file = snapshot / tok_path
if tok_file.exists():
tok_config = json.loads(tok_file.read_text())
@@ -672,16 +696,16 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
import requests
import os
- paths_to_try = ['tokenizer_config.json', 'LLM/tokenizer_config.json']
+ paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"]
# Use provided token, or fall back to env
- token = hf_token or os.environ.get('HF_TOKEN')
+ token = hf_token or os.environ.get("HF_TOKEN")
headers = {}
if token:
- headers['Authorization'] = f'Bearer {token}'
+ headers["Authorization"] = f"Bearer {token}"
for tok_path in paths_to_try:
url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}"
- resp = requests.get(url, headers=headers, timeout=15)
+ resp = requests.get(url, headers = headers, timeout = 15)
if not resp.ok:
continue
@@ -692,7 +716,9 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
return None
except Exception as e:
- logger.debug(f"Could not detect audio type from tokenizer for {model_name}: {e}")
+ logger.debug(
+ f"Could not detect audio type from tokenizer for {model_name}: {e}"
+ )
return None
@@ -701,7 +727,7 @@ def is_audio_input_type(audio_type: Optional[str]) -> bool:
Whisper (ASR) and audio_vlm (Gemma3n) accept audio input.
"""
- return audio_type in ('whisper', 'audio_vlm')
+ return audio_type in ("whisper", "audio_vlm")
def _is_mmproj(filename: str) -> bool:
@@ -756,7 +782,8 @@ def detect_gguf_model(path: str) -> Optional[str]:
if p.is_dir():
gguf_files = sorted(
(f for f in p.glob("*.gguf") if not _is_mmproj(f.name)),
- key=lambda f: f.stat().st_size, reverse=True,
+ key = lambda f: f.stat().st_size,
+ reverse = True,
)
if gguf_files:
return str(gguf_files[0].resolve())
@@ -767,9 +794,18 @@ def detect_gguf_model(path: str) -> Optional[str]:
# Preferred GGUF quantization levels, in descending priority.
# Q4_K_M is a good default: small, fast, acceptable quality.
_GGUF_QUANT_PREFERENCE = [
- "Q4_K_M", "Q4_K_S", "Q5_K_M", "Q5_K_S",
- "Q6_K", "Q8_0", "Q3_K_M", "Q3_K_L", "Q2_K",
- "F16", "BF16", "F32",
+ "Q4_K_M",
+ "Q4_K_S",
+ "Q5_K_M",
+ "Q5_K_S",
+ "Q6_K",
+ "Q8_0",
+ "Q3_K_M",
+ "Q3_K_L",
+ "Q2_K",
+ "F16",
+ "BF16",
+ "F32",
]
@@ -797,9 +833,10 @@ def _pick_best_gguf(filenames: list[str]) -> Optional[str]:
@dataclass
class GgufVariantInfo:
"""A single GGUF quantization variant from a HuggingFace repo."""
- filename: str # e.g., "gemma-3-4b-it-Q4_K_M.gguf"
- quant: str # e.g., "Q4_K_M" (extracted from filename)
- size_bytes: int # file size
+
+ filename: str # e.g., "gemma-3-4b-it-Q4_K_M.gguf"
+ quant: str # e.g., "Q4_K_M" (extracted from filename)
+ size_bytes: int # file size
def _extract_quant_label(filename: str) -> str:
@@ -815,21 +852,23 @@ def _extract_quant_label(filename: str) -> str:
"MXFP4_MOE/model-MXFP4_MOE-0001.gguf"β "MXFP4_MOE"
"""
import re
+
# Use only the basename (rfilename may include directory)
basename = filename.rsplit("/", 1)[-1]
# Strip .gguf and any shard suffix (-00001-of-00010)
- stem = re.sub(r'-\d{3,}-of-\d{3,}', '', basename.rsplit(".", 1)[0])
+ stem = re.sub(r"-\d{3,}-of-\d{3,}", "", basename.rsplit(".", 1)[0])
# Match known quantization patterns
match = re.search(
- r'(UD-)?' # Optional UD- prefix (Ultra Discrete)
- r'(MXFP[0-9]+(?:_[A-Z0-9]+)*' # MXFP variants: MXFP4, MXFP4_MOE
- r'|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?' # IQ variants: IQ4_XS, IQ4_NL, IQ1_S
- r'|TQ[0-9]+_[0-9]+' # Ternary quant: TQ1_0, TQ2_0
- r'|Q[0-9]+_K_[A-Z]+' # K-quant: Q4_K_M, Q3_K_S
- r'|Q[0-9]+_[0-9]+' # Standard: Q8_0, Q5_1
- r'|Q[0-9]+_K' # Short K-quant: Q6_K
- r'|BF16|F16|F32)', # Full precision
- stem, re.IGNORECASE,
+ r"(UD-)?" # Optional UD- prefix (Ultra Discrete)
+ r"(MXFP[0-9]+(?:_[A-Z0-9]+)*" # MXFP variants: MXFP4, MXFP4_MOE
+ r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?" # IQ variants: IQ4_XS, IQ4_NL, IQ1_S
+ r"|TQ[0-9]+_[0-9]+" # Ternary quant: TQ1_0, TQ2_0
+ r"|Q[0-9]+_K_[A-Z]+" # K-quant: Q4_K_M, Q3_K_S
+ r"|Q[0-9]+_[0-9]+" # Standard: Q8_0, Q5_1
+ r"|Q[0-9]+_K" # Short K-quant: Q6_K
+ r"|BF16|F16|F32)", # Full precision
+ stem,
+ re.IGNORECASE,
)
if match:
prefix = match.group(1) or ""
@@ -853,11 +892,11 @@ def list_gguf_variants(
"""
from huggingface_hub import model_info as hf_model_info
- info = hf_model_info(repo_id, token=hf_token, files_metadata=True)
+ info = hf_model_info(repo_id, token = hf_token, files_metadata = True)
variants: list[GgufVariantInfo] = []
has_vision = False
- quant_totals: dict[str, int] = {} # quant -> total bytes
+ quant_totals: dict[str, int] = {} # quant -> total bytes
quant_first_file: dict[str, str] = {} # quant -> first filename (for display)
for sibling in info.siblings:
@@ -877,11 +916,13 @@ def list_gguf_variants(
quant_first_file[quant] = fname
for quant, total_size in quant_totals.items():
- variants.append(GgufVariantInfo(
- filename=quant_first_file[quant],
- quant=quant,
- size_bytes=total_size,
- ))
+ variants.append(
+ GgufVariantInfo(
+ filename = quant_first_file[quant],
+ quant = quant,
+ size_bytes = total_size,
+ )
+ )
return variants, has_vision
@@ -898,7 +939,7 @@ def detect_gguf_model_remote(
try:
from huggingface_hub import model_info as hf_model_info
- info = hf_model_info(repo_id, token=hf_token)
+ info = hf_model_info(repo_id, token = hf_token)
repo_files = [s.rfilename for s in info.siblings]
return _pick_best_gguf(repo_files)
except Exception as e:
@@ -919,9 +960,9 @@ def download_gguf_file(
from huggingface_hub import hf_hub_download
local_path = hf_hub_download(
- repo_id=repo_id,
- filename=filename,
- token=hf_token,
+ repo_id = repo_id,
+ filename = filename,
+ token = hf_token,
)
return local_path
@@ -964,7 +1005,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
try:
from huggingface_hub import model_info as hf_model_info
- info = hf_model_info(model_name, token=hf_token)
+ info = hf_model_info(model_name, token = hf_token)
tags = set(info.tags or [])
pipeline_tag = info.pipeline_tag or ""
@@ -1024,16 +1065,21 @@ def scan_trained_loras(outputs_dir: str = str(outputs_root())) -> List[Tuple[str
logger.debug(f"Found trained LoRA: {display_name}")
# Sort by modification time (newest first)
- trained_loras.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True)
+ trained_loras.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True)
- logger.info(f"Found {len(trained_loras)} trained LoRA adapters in {outputs_dir}")
+ logger.info(
+ f"Found {len(trained_loras)} trained LoRA adapters in {outputs_dir}"
+ )
return trained_loras
except Exception as e:
logger.error(f"Error scanning outputs folder: {e}")
return []
-def scan_exported_models(exports_dir: str = str(exports_root())) -> List[Tuple[str, str, str, Optional[str]]]:
+
+def scan_exported_models(
+ exports_dir: str = str(exports_root()),
+) -> List[Tuple[str, str, str, Optional[str]]]:
"""
Scan exports folder for exported models (merged, LoRA, GGUF).
@@ -1082,9 +1128,8 @@ def scan_exported_models(exports_dir: str = str(exports_root())) -> List[Tuple[s
adapter_config = checkpoint_dir / "adapter_config.json"
config_file = checkpoint_dir / "config.json"
- has_weights = (
- any(checkpoint_dir.glob("*.safetensors"))
- or any(checkpoint_dir.glob("*.bin"))
+ has_weights = any(checkpoint_dir.glob("*.safetensors")) or any(
+ checkpoint_dir.glob("*.bin")
)
has_gguf = any(checkpoint_dir.glob("*.gguf"))
@@ -1134,7 +1179,9 @@ def scan_exported_models(exports_dir: str = str(exports_root())) -> List[Tuple[s
# Fallback: read base model from the original training run's
# adapter_config.json in ./outputs/{run_name}/
if not base_model:
- outputs_adapter_cfg = resolve_output_dir(run_dir.name) / "adapter_config.json"
+ outputs_adapter_cfg = (
+ resolve_output_dir(run_dir.name) / "adapter_config.json"
+ )
try:
if outputs_adapter_cfg.exists():
cfg = json.loads(outputs_adapter_cfg.read_text())
@@ -1147,7 +1194,7 @@ def scan_exported_models(exports_dir: str = str(exports_root())) -> List[Tuple[s
results.append((display_name, model_path, export_type, base_model))
logger.debug(f"Found exported model: {display_name} ({export_type})")
- results.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True)
+ results.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True)
logger.info(f"Found {len(results)} exported models in {exports_dir}")
return results
@@ -1177,11 +1224,13 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
# Try adapter_config.json first
adapter_config_path = lora_path_obj / "adapter_config.json"
if adapter_config_path.exists():
- with open(adapter_config_path, 'r') as f:
+ with open(adapter_config_path, "r") as f:
config = json.load(f)
base_model = config.get("base_model_name_or_path")
if base_model:
- logger.info(f"Detected base model from adapter_config.json: {base_model}")
+ logger.info(
+ f"Detected base model from adapter_config.json: {base_model}"
+ )
return base_model
# Fallback: try training_args.bin (requires torch)
@@ -1189,10 +1238,13 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
if training_args_path.exists():
try:
import torch
+
training_args = torch.load(training_args_path)
- if hasattr(training_args, 'model_name_or_path'):
+ if hasattr(training_args, "model_name_or_path"):
base_model = training_args.model_name_or_path
- logger.info(f"Detected base model from training_args.bin: {base_model}")
+ logger.info(
+ f"Detected base model from training_args.bin: {base_model}"
+ )
return base_model
except Exception as e:
logger.warning(f"Could not load training_args.bin: {e}")
@@ -1216,22 +1268,23 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
except Exception as e:
logger.error(f"Error reading base model from LoRA config: {e}")
return None
-pass
+
# Status indicators that appear in UI dropdowns
UI_STATUS_INDICATORS = [" (Ready)", " (Loading...)", " (Active)", "β "]
+
def load_model_defaults(model_name: str) -> Dict[str, Any]:
"""
Load default training parameters for a model from YAML file.
-
+
Args:
model_name: Model identifier (e.g., "unsloth/Meta-Llama-3.1-8B-bnb-4bit")
-
+
Returns:
Dictionary with default parameters from YAML file, or empty dict if not found
-
- The function looks for a YAML file in configs/model_defaults/ (including subfolders)
+
+ The function looks for a YAML file in configs/model_defaults/ (including subfolders)
based on the model name or its aliases from MODEL_NAME_MAPPING.
If no specific file exists, it falls back to default.yaml.
"""
@@ -1239,22 +1292,26 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
# Get the script directory to locate configs
script_dir = Path(__file__).parent.parent.parent
defaults_dir = script_dir / "assets" / "configs" / "model_defaults"
-
+
# First, check if model is in the mapping
if model_name.lower() in _REVERSE_MODEL_MAPPING:
canonical_file = _REVERSE_MODEL_MAPPING[model_name.lower()]
# Search in subfolders and root
for config_path in defaults_dir.rglob(canonical_file):
if config_path.is_file():
- with open(config_path, 'r', encoding='utf-8') as f:
+ with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
- logger.info(f"Loaded model defaults from {config_path} (via mapping)")
+ logger.info(
+ f"Loaded model defaults from {config_path} (via mapping)"
+ )
return config
-
+
# If model_name is a local path (e.g. /home/.../Spark-TTS-0.5B/LLM from
# adapter_config.json), try matching the last 1-2 path components against
# the registry (e.g. "Spark-TTS-0.5B/LLM").
- if model_name not in _REVERSE_MODEL_MAPPING and (model_name.startswith("/") or model_name.startswith(".")):
+ if model_name not in _REVERSE_MODEL_MAPPING and (
+ model_name.startswith("/") or model_name.startswith(".")
+ ):
parts = Path(model_name).parts
for depth in [2, 1]:
if len(parts) >= depth:
@@ -1263,9 +1320,11 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
canonical_file = _REVERSE_MODEL_MAPPING[suffix]
for config_path in defaults_dir.rglob(canonical_file):
if config_path.is_file():
- with open(config_path, 'r', encoding='utf-8') as f:
+ with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
- logger.info(f"Loaded model defaults from {config_path} (via path suffix '{suffix}')")
+ logger.info(
+ f"Loaded model defaults from {config_path} (via path suffix '{suffix}')"
+ )
return config
# Try exact model name match (for backward compatibility)
@@ -1273,48 +1332,58 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
# Search in subfolders and root
for config_path in defaults_dir.rglob(model_filename):
if config_path.is_file():
- with open(config_path, 'r', encoding='utf-8') as f:
+ with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded model defaults from {config_path}")
return config
-
+
# Fall back to default.yaml
default_config_path = defaults_dir / "default.yaml"
if default_config_path.exists():
- with open(default_config_path, 'r', encoding='utf-8') as f:
+ with open(default_config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded default model defaults from {default_config_path}")
return config
-
+
logger.warning(f"No default config found for model {model_name}")
return {}
-
+
except Exception as e:
logger.error(f"Error loading model defaults for {model_name}: {e}")
return {}
+
@dataclass
class ModelConfig:
"""Configuration for a model to load"""
- identifier: str # Clean model identifier (org/name or path)
- display_name: str # Original UI display name
- path: str # Normalized filesystem path
- is_local: bool # Is this a local file vs HF model?
- is_cached: bool # Is this already in HF cache?
- is_vision: bool # Is this a vision model?
- is_lora: bool # Is this a lora adapter?
- is_gguf: bool = False # Is this a GGUF model?
- is_audio: bool = False # Is this a TTS audio model?
- audio_type: Optional[str] = None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac'
+
+ identifier: str # Clean model identifier (org/name or path)
+ display_name: str # Original UI display name
+ path: str # Normalized filesystem path
+ is_local: bool # Is this a local file vs HF model?
+ is_cached: bool # Is this already in HF cache?
+ is_vision: bool # Is this a vision model?
+ is_lora: bool # Is this a lora adapter?
+ is_gguf: bool = False # Is this a GGUF model?
+ is_audio: bool = False # Is this a TTS audio model?
+ audio_type: Optional[str] = (
+ None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac'
+ )
has_audio_input: bool = False # Accepts audio input (ASR/speech understanding)
gguf_file: Optional[str] = None # Full path to the .gguf file (local mode)
- gguf_mmproj_file: Optional[str] = None # Full path to the mmproj .gguf file (vision projection)
- gguf_hf_repo: Optional[str] = None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
+ gguf_mmproj_file: Optional[str] = (
+ None # Full path to the mmproj .gguf file (vision projection)
+ )
+ gguf_hf_repo: Optional[str] = (
+ None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
+ )
gguf_variant: Optional[str] = None # Quantization variant (e.g. "Q4_K_M")
base_model: Optional[str] = None # Base model (for LoRAs)
@classmethod
- def from_lora_path(cls, lora_path: str, hf_token: Optional[str] = None) -> Optional['ModelConfig']:
+ def from_lora_path(
+ cls, lora_path: str, hf_token: Optional[str] = None
+ ) -> Optional["ModelConfig"]:
"""
Create ModelConfig from a local LoRA adapter path.
@@ -1341,26 +1410,26 @@ class ModelConfig:
return None
# Check if base model is vision
- is_vision = is_vision_model(base_model, hf_token=hf_token)
+ is_vision = is_vision_model(base_model, hf_token = hf_token)
# Check if base model is audio
- audio_type = detect_audio_type(base_model, hf_token=hf_token)
+ audio_type = detect_audio_type(base_model, hf_token = hf_token)
display_name = lora_path_obj.name
identifier = lora_path # Use path as identifier for local LoRAs
return cls(
- identifier=identifier,
- display_name=display_name,
- path=lora_path,
- is_local=True,
- is_cached=True, # Local LoRAs are always "cached"
- is_vision=is_vision,
- is_lora=True,
- is_audio=audio_type is not None and audio_type != 'audio_vlm',
- audio_type=audio_type,
- has_audio_input=is_audio_input_type(audio_type),
- base_model=base_model,
+ identifier = identifier,
+ display_name = display_name,
+ path = lora_path,
+ is_local = True,
+ is_cached = True, # Local LoRAs are always "cached"
+ is_vision = is_vision,
+ is_lora = True,
+ is_audio = audio_type is not None and audio_type != "audio_vlm",
+ audio_type = audio_type,
+ has_audio_input = is_audio_input_type(audio_type),
+ base_model = base_model,
)
except Exception as e:
@@ -1374,7 +1443,7 @@ class ModelConfig:
hf_token: Optional[str] = None,
is_lora: bool = False,
gguf_variant: Optional[str] = None,
- ) -> Optional['ModelConfig']:
+ ) -> Optional["ModelConfig"]:
"""
Create ModelConfig from a clean model identifier.
@@ -1432,7 +1501,7 @@ class ModelConfig:
try:
meta = json.loads(meta_path.read_text())
base = meta.get("base_model")
- if base and is_vision_model(base, hf_token=hf_token):
+ if base and is_vision_model(base, hf_token = hf_token):
base_is_vision = True
logger.info(f"GGUF base model '{base}' is a vision model")
except Exception as e:
@@ -1444,27 +1513,30 @@ class ModelConfig:
gguf_is_vision = True
logger.info(f"Detected mmproj for vision: {mmproj_file}")
elif base_is_vision:
- logger.warning(f"Base model is vision but no mmproj file found in {gguf_dir}")
+ logger.warning(
+ f"Base model is vision but no mmproj file found in {gguf_dir}"
+ )
return cls(
- identifier=identifier,
- display_name=display_name,
- path=path,
- is_local=True,
- is_cached=True,
- is_vision=gguf_is_vision,
- is_lora=False,
- is_gguf=True,
- gguf_file=gguf_file,
- gguf_mmproj_file=mmproj_file,
+ identifier = identifier,
+ display_name = display_name,
+ path = path,
+ is_local = True,
+ is_cached = True,
+ is_vision = gguf_is_vision,
+ is_lora = False,
+ is_gguf = True,
+ gguf_file = gguf_file,
+ gguf_mmproj_file = mmproj_file,
)
else:
# Check if the HF repo contains GGUF files
- gguf_filename = detect_gguf_model_remote(identifier, hf_token=hf_token)
+ gguf_filename = detect_gguf_model_remote(identifier, hf_token = hf_token)
if gguf_filename:
# Preflight: verify llama-server binary exists BEFORE user waits
# for a multi-GB download that llama-server handles natively
from core.inference.llama_cpp import LlamaCppBackend
+
if not LlamaCppBackend._find_llama_server_binary():
raise RuntimeError(
"llama-server binary not found β cannot load GGUF models. "
@@ -1472,7 +1544,7 @@ class ModelConfig:
)
# Use list_gguf_variants() to detect vision & resolve variant
- variants, has_vision = list_gguf_variants(identifier, hf_token=hf_token)
+ variants, has_vision = list_gguf_variants(identifier, hf_token = hf_token)
variant = gguf_variant
if not variant:
# Auto-select best quantization
@@ -1489,17 +1561,17 @@ class ModelConfig:
f"variant={variant}, vision={has_vision}"
)
return cls(
- identifier=identifier,
- display_name=display_name,
- path=identifier,
- is_local=False,
- is_cached=False,
- is_vision=has_vision,
- is_lora=False,
- is_gguf=True,
- gguf_file=None,
- gguf_hf_repo=identifier,
- gguf_variant=variant,
+ identifier = identifier,
+ display_name = display_name,
+ path = identifier,
+ is_local = False,
+ is_cached = False,
+ is_vision = has_vision,
+ is_lora = False,
+ is_gguf = True,
+ gguf_file = None,
+ gguf_hf_repo = identifier,
+ gguf_variant = variant,
)
# Auto-detect LoRA for local paths (check adapter_config.json on disk)
@@ -1507,20 +1579,25 @@ class ModelConfig:
detected_base = get_base_model_from_lora(path)
if detected_base:
is_lora = True
- logger.info(f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})")
-
+ logger.info(
+ f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})"
+ )
+
# Auto-detect LoRA for remote HF models (check repo file listing)
if not is_lora and not is_local:
try:
from huggingface_hub import model_info as hf_model_info
- info = hf_model_info(identifier, token=hf_token)
+
+ info = hf_model_info(identifier, token = hf_token)
repo_files = [s.rfilename for s in info.siblings]
if "adapter_config.json" in repo_files:
is_lora = True
logger.info(f"Auto-detected remote LoRA adapter: '{identifier}'")
except Exception as e:
- logger.debug(f"Could not check remote LoRA status for '{identifier}': {e}")
-
+ logger.debug(
+ f"Could not check remote LoRA status for '{identifier}': {e}"
+ )
+
# Handle LoRA adapters
base_model = None
if is_lora:
@@ -1531,15 +1608,20 @@ class ModelConfig:
# Remote LoRA: download adapter_config.json from HF
try:
from huggingface_hub import hf_hub_download
- config_path = hf_hub_download(identifier, "adapter_config.json", token=hf_token)
- with open(config_path, 'r') as f:
+
+ config_path = hf_hub_download(
+ identifier, "adapter_config.json", token = hf_token
+ )
+ with open(config_path, "r") as f:
adapter_config = json.load(f)
base_model = adapter_config.get("base_model_name_or_path")
if base_model:
logger.info(f"Resolved remote LoRA base model: '{base_model}'")
except Exception as e:
- logger.warning(f"Could not download adapter_config.json for '{identifier}': {e}")
-
+ logger.warning(
+ f"Could not download adapter_config.json for '{identifier}': {e}"
+ )
+
if not base_model:
logger.warning(f"Could not determine base model for LoRA '{path}'")
return None
@@ -1547,34 +1629,35 @@ class ModelConfig:
else:
check_model = identifier
- vision = is_vision_model(check_model, hf_token=hf_token)
- audio_type_val = detect_audio_type(check_model, hf_token=hf_token)
+ vision = is_vision_model(check_model, hf_token = hf_token)
+ audio_type_val = detect_audio_type(check_model, hf_token = hf_token)
has_audio_in = is_audio_input_type(audio_type_val)
display_name = Path(path).name if is_local else identifier.split("/")[-1]
return cls(
- identifier=identifier,
- display_name=display_name,
- path=path,
- is_local=is_local,
- is_cached=is_model_cached(identifier) if not is_local else True,
- is_vision=vision,
- is_lora=is_lora,
- is_audio=audio_type_val is not None and audio_type_val != 'audio_vlm',
- audio_type=audio_type_val,
- has_audio_input=has_audio_in,
- base_model=base_model,
+ identifier = identifier,
+ display_name = display_name,
+ path = path,
+ is_local = is_local,
+ is_cached = is_model_cached(identifier) if not is_local else True,
+ is_vision = vision,
+ is_lora = is_lora,
+ is_audio = audio_type_val is not None and audio_type_val != "audio_vlm",
+ audio_type = audio_type_val,
+ has_audio_input = has_audio_in,
+ base_model = base_model,
)
-
@classmethod
- def from_ui_selection(cls,
- dropdown_value: Optional[str],
- search_value: Optional[str],
- local_models: list = None,
- hf_token: Optional[str] = None,
- is_lora: bool = False) -> Optional['ModelConfig']:
+ def from_ui_selection(
+ cls,
+ dropdown_value: Optional[str],
+ search_value: Optional[str],
+ local_models: list = None,
+ hf_token: Optional[str] = None,
+ is_lora: bool = False,
+ ) -> Optional["ModelConfig"]:
"""
Create a universal ModelConfig from UI dropdown/search selections.
Handles base models and LoRA adapters.
@@ -1592,7 +1675,9 @@ class ModelConfig:
# Use the correct 'local_models' parameter to resolve display names
if " (Active)" in selected or " (Ready)" in selected:
- clean_display_name = selected.replace(" (Active)", "").replace(" (Ready)", "")
+ clean_display_name = selected.replace(" (Active)", "").replace(
+ " (Ready)", ""
+ )
if local_models:
for local_display, local_path in local_models:
if local_display == clean_display_name:
@@ -1621,25 +1706,28 @@ class ModelConfig:
# For a LoRA, we MUST find its base model.
base_model = get_base_model_from_lora(path)
if not base_model:
- logger.warning(f"Could not determine base model for LoRA '{path}'. Cannot create config.")
- return None # Cannot proceed without a base model
+ logger.warning(
+ f"Could not determine base model for LoRA '{path}'. Cannot create config."
+ )
+ return None # Cannot proceed without a base model
# A LoRA's vision capability is determined by its base model.
- is_vision = is_vision_model(base_model, hf_token=hf_token)
+ is_vision = is_vision_model(base_model, hf_token = hf_token)
else:
# For a base model, just check its own vision status.
- is_vision = is_vision_model(identifier, hf_token=hf_token)
+ is_vision = is_vision_model(identifier, hf_token = hf_token)
from utils.paths import is_model_cached
+
is_cached = is_model_cached(identifier) if not is_local else True
return cls(
- identifier=identifier,
- display_name=display_name,
- path=path,
- is_local=is_local,
- is_cached=is_cached,
- is_vision=is_vision,
- is_lora=is_lora,
- base_model=base_model, # This will be None for base models, and populated for LoRAs
+ identifier = identifier,
+ display_name = display_name,
+ path = path,
+ is_local = is_local,
+ is_cached = is_cached,
+ is_vision = is_vision,
+ is_lora = is_lora,
+ base_model = base_model, # This will be None for base models, and populated for LoRAs
)
diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py
index 3850b57638..507fb1106b 100644
--- a/studio/backend/utils/paths/__init__.py
+++ b/studio/backend/utils/paths/__init__.py
@@ -4,6 +4,7 @@
"""
Path utilities for model and dataset handling
"""
+
from .path_utils import normalize_path, is_local_path, is_model_cached, get_cache_path
from .storage_roots import (
studio_root,
@@ -21,6 +22,7 @@ from .storage_roots import (
oxc_validator_tmp_root,
tensorboard_root,
ensure_dir,
+ ensure_studio_directories,
resolve_under_root,
resolve_output_dir,
resolve_export_dir,
@@ -29,28 +31,29 @@ from .storage_roots import (
)
__all__ = [
- 'normalize_path',
- 'is_local_path',
- 'is_model_cached',
- 'get_cache_path',
- 'studio_root',
- 'assets_root',
- 'datasets_root',
- 'dataset_uploads_root',
- 'recipe_datasets_root',
- 'outputs_root',
- 'exports_root',
- 'auth_root',
- 'auth_db_path',
- 'tmp_root',
- 'seed_uploads_root',
- 'unstructured_seed_cache_root',
- 'oxc_validator_tmp_root',
- 'tensorboard_root',
- 'ensure_dir',
- 'resolve_under_root',
- 'resolve_output_dir',
- 'resolve_export_dir',
- 'resolve_tensorboard_dir',
- 'resolve_dataset_path',
+ "normalize_path",
+ "is_local_path",
+ "is_model_cached",
+ "get_cache_path",
+ "studio_root",
+ "assets_root",
+ "datasets_root",
+ "dataset_uploads_root",
+ "recipe_datasets_root",
+ "outputs_root",
+ "exports_root",
+ "auth_root",
+ "auth_db_path",
+ "tmp_root",
+ "seed_uploads_root",
+ "unstructured_seed_cache_root",
+ "oxc_validator_tmp_root",
+ "tensorboard_root",
+ "ensure_dir",
+ "ensure_studio_directories",
+ "resolve_under_root",
+ "resolve_output_dir",
+ "resolve_export_dir",
+ "resolve_tensorboard_dir",
+ "resolve_dataset_path",
]
diff --git a/studio/backend/utils/paths/path_utils.py b/studio/backend/utils/paths/path_utils.py
index f93262c2cb..1d6a952399 100644
--- a/studio/backend/utils/paths/path_utils.py
+++ b/studio/backend/utils/paths/path_utils.py
@@ -4,6 +4,7 @@
"""
Path utilities for model and dataset handling
"""
+
import os
from pathlib import Path
from typing import Optional
@@ -25,14 +26,14 @@ def normalize_path(path: str) -> str:
return path
# Handle Windows drive letters (C:\\ or c:\\)
- if len(path) >= 3 and path[1] == ':' and path[2] in ('\\', '/'):
+ if len(path) >= 3 and path[1] == ":" and path[2] in ("\\", "/"):
drive = path[0].lower()
- rest = path[3:].replace('\\', '/')
- return f'/mnt/{drive}/{rest}'
+ rest = path[3:].replace("\\", "/")
+ return f"/mnt/{drive}/{rest}"
# Already Unix-style or relative
- return path.replace('\\', '/')
-pass
+ return path.replace("\\", "/")
+
def is_local_path(path: str) -> bool:
"""
@@ -53,26 +54,26 @@ def is_local_path(path: str) -> bool:
pass
# Obvious HF patterns
- if path.count('/') == 1 and not path.startswith(('/', '.', '~')):
+ if path.count("/") == 1 and not path.startswith(("/", ".", "~")):
return False # Looks like org/model format
# Filesystem indicators
return (
- path.startswith(('/', '.', '~')) or # Unix absolute/relative
- ':' in path or # Windows drive or URL
- '\\' in path or # Windows separator
- os.path.isabs(path) # System-absolute
+ path.startswith(("/", ".", "~")) # Unix absolute/relative
+ or ":" in path # Windows drive or URL
+ or "\\" in path # Windows separator
+ or os.path.isabs(path) # System-absolute
)
-pass
+
def get_cache_path(model_name: str) -> Optional[Path]:
"""Get HuggingFace cache path for a model if it exists."""
- cache_dir = Path.home() / '.cache' / 'huggingface' / 'hub'
+ cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
model_cache_name = model_name.replace("/", "--")
- model_cache_path = cache_dir / f'models--{model_cache_name}'
+ model_cache_path = cache_dir / f"models--{model_cache_name}"
return model_cache_path if model_cache_path.exists() else None
-pass
+
def is_model_cached(model_name: str) -> bool:
"""Check if model is downloaded in HuggingFace cache."""
@@ -81,9 +82,8 @@ def is_model_cached(model_name: str) -> bool:
return False
# Check for actual model files
- for suffix in ['.safetensors', '.bin', '.json']:
- if list(cache_path.rglob(f'*{suffix}')):
+ for suffix in [".safetensors", ".bin", ".json"]:
+ if list(cache_path.rglob(f"*{suffix}")):
return True
return False
-pass
diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py
index a005d604a8..d9adce2105 100644
--- a/studio/backend/utils/paths/storage_roots.py
+++ b/studio/backend/utils/paths/storage_roots.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
from __future__ import annotations
from pathlib import Path
@@ -61,11 +64,29 @@ def tensorboard_root() -> Path:
def ensure_dir(path: Path) -> Path:
- path.mkdir(parents=True, exist_ok=True)
+ path.mkdir(parents = True, exist_ok = True)
return path
-def _clean_relative_path(path_value: str, *, strip_prefixes: tuple[str, ...] = ()) -> Path:
+def ensure_studio_directories() -> None:
+ """Create all standard studio directories on startup."""
+ for dir_fn in (
+ studio_root,
+ assets_root,
+ datasets_root,
+ dataset_uploads_root,
+ recipe_datasets_root,
+ outputs_root,
+ exports_root,
+ auth_root,
+ tensorboard_root,
+ ):
+ ensure_dir(dir_fn())
+
+
+def _clean_relative_path(
+ path_value: str, *, strip_prefixes: tuple[str, ...] = ()
+) -> Path:
path = Path(path_value).expanduser()
parts = [part for part in path.parts if part not in ("", ".")]
while parts and parts[0] in strip_prefixes:
@@ -86,31 +107,31 @@ def resolve_under_root(
if path.is_absolute():
return path
- cleaned = _clean_relative_path(str(path), strip_prefixes=strip_prefixes)
+ cleaned = _clean_relative_path(str(path), strip_prefixes = strip_prefixes)
return root / cleaned
def resolve_output_dir(path_value: str | None = None) -> Path:
return resolve_under_root(
path_value,
- root=outputs_root(),
- strip_prefixes=("outputs",),
+ root = outputs_root(),
+ strip_prefixes = ("outputs",),
)
def resolve_export_dir(path_value: str | None = None) -> Path:
return resolve_under_root(
path_value,
- root=exports_root(),
- strip_prefixes=("exports",),
+ root = exports_root(),
+ strip_prefixes = ("exports",),
)
def resolve_tensorboard_dir(path_value: str | None = None) -> Path:
return resolve_under_root(
path_value,
- root=tensorboard_root(),
- strip_prefixes=("runs", "tensorboard"),
+ root = tensorboard_root(),
+ strip_prefixes = ("runs", "tensorboard"),
)
diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py
index 4cc7217ccb..58fbc35e33 100644
--- a/studio/backend/utils/transformers_version.py
+++ b/studio/backend/utils/transformers_version.py
@@ -33,7 +33,6 @@ from pathlib import Path
logger = get_logger(__name__)
-
# ---------------------------------------------------------------------------
# Detection
# ---------------------------------------------------------------------------
@@ -41,12 +40,12 @@ logger = get_logger(__name__)
# Lowercase substrings β if ANY appears anywhere in the lowered model name,
# we need transformers 5.x.
TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = (
- "ministral-3-", # Ministral-3-{3,8,14}B-{Instruct,Reasoning,Base}-2512
- "glm-4.7-flash", # GLM-4.7-Flash
- "qwen3-30b-a3b", # Qwen3-30B-A3B-Instruct-2507 and variants
- "qwen3.5", # Qwen3.5 family (35B-A3B, etc.)
- "qwen3-next", # Qwen3-Next and variants
- "tiny_qwen3_moe", # imdatta0/tiny_qwen3_moe_2.8B_0.7B
+ "ministral-3-", # Ministral-3-{3,8,14}B-{Instruct,Reasoning,Base}-2512
+ "glm-4.7-flash", # GLM-4.7-Flash
+ "qwen3-30b-a3b", # Qwen3-30B-A3B-Instruct-2507 and variants
+ "qwen3.5", # Qwen3.5 family (35B-A3B, etc.)
+ "qwen3-next", # Qwen3-Next and variants
+ "tiny_qwen3_moe", # imdatta0/tiny_qwen3_moe_2.8B_0.7B
)
# Versions
@@ -77,7 +76,8 @@ def _resolve_base_model(model_name: str) -> str:
if base:
logger.info(
"Resolved LoRA adapter '%s' β base model '%s'",
- model_name, base,
+ model_name,
+ base,
)
return base
except Exception as exc:
@@ -87,18 +87,21 @@ def _resolve_base_model(model_name: str) -> str:
if local_path.is_dir():
try:
from utils.models import get_base_model_from_lora
+
base = get_base_model_from_lora(model_name)
if base:
logger.info(
"Resolved LoRA adapter '%s' β base model '%s' "
"(via get_base_model_from_lora)",
- model_name, base,
+ model_name,
+ base,
)
return base
except Exception as exc:
logger.debug(
"get_base_model_from_lora failed for '%s': %s",
- model_name, exc,
+ model_name,
+ exc,
)
return model_name
@@ -115,6 +118,7 @@ def needs_transformers_5(model_name: str) -> bool:
# Version switching (in-process β used only by export)
# ---------------------------------------------------------------------------
+
def _get_in_memory_version() -> str | None:
"""Return the transformers version currently loaded in this process."""
tf = sys.modules.get("transformers")
@@ -153,7 +157,8 @@ def _purge_modules() -> int:
"""
importlib.invalidate_caches()
to_remove = [
- k for k in list(sys.modules.keys())
+ k
+ for k in list(sys.modules.keys())
if any(k == p or k.startswith(p + ".") for p in _PURGE_PREFIXES)
]
for key in to_remove:
@@ -167,15 +172,21 @@ def _ensure_venv_t5_exists() -> bool:
return True
logger.warning(".venv_t5 not found at %s β installing at runtime", _VENV_T5_DIR)
- os.makedirs(_VENV_T5_DIR, exist_ok=True)
+ os.makedirs(_VENV_T5_DIR, exist_ok = True)
for pkg in (f"transformers=={TRANSFORMERS_5_VERSION}", "huggingface_hub==1.3.0"):
cmd = [
- sys.executable, "-m", "pip", "install",
- "--target", _VENV_T5_DIR,
+ sys.executable,
+ "-m",
+ "pip",
+ "install",
+ "--target",
+ _VENV_T5_DIR,
"--no-deps",
pkg,
]
- result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+ result = subprocess.run(
+ cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True
+ )
if result.returncode != 0:
logger.error("pip install failed:\n%s", result.stdout)
return False
@@ -186,7 +197,9 @@ def _ensure_venv_t5_exists() -> bool:
def _activate_5x() -> None:
"""Prepend .venv_t5/ to sys.path, purge stale modules, reimport."""
if not _ensure_venv_t5_exists():
- raise RuntimeError(f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}")
+ raise RuntimeError(
+ f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}"
+ )
if _VENV_T5_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_DIR)
@@ -196,6 +209,7 @@ def _activate_5x() -> None:
logger.info("Purged %d cached modules", count)
import transformers
+
logger.info("Loaded transformers %s", transformers.__version__)
@@ -209,6 +223,7 @@ def _deactivate_5x() -> None:
logger.info("Purged %d cached modules", count)
import transformers
+
logger.info("Reverted to transformers %s", transformers.__version__)
@@ -236,7 +251,10 @@ def ensure_transformers_version(model_name: str) -> None:
logger.info(
"Version check for '%s' (resolved: '%s'): need=%s, in_memory=%s",
- model_name, resolved, target_version, in_memory,
+ model_name,
+ resolved,
+ target_version,
+ in_memory,
)
# --- Already correct? ---------------------------------------------------
@@ -245,7 +263,8 @@ def ensure_transformers_version(model_name: str) -> None:
if in_memory_major == target_major:
logger.info(
"transformers %s already loaded β correct for '%s'",
- in_memory, model_name,
+ in_memory,
+ model_name,
)
return
@@ -254,7 +273,9 @@ def ensure_transformers_version(model_name: str) -> None:
logger.info("Activating transformers %s via .venv_t5β¦", TRANSFORMERS_5_VERSION)
_activate_5x()
else:
- logger.info("Reverting to default transformers %sβ¦", TRANSFORMERS_DEFAULT_VERSION)
+ logger.info(
+ "Reverting to default transformers %sβ¦", TRANSFORMERS_DEFAULT_VERSION
+ )
_deactivate_5x()
final = _get_in_memory_version()
diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py
index 0cf3abfebf..4e61a5b969 100644
--- a/studio/backend/utils/utils.py
+++ b/studio/backend/utils/utils.py
@@ -4,6 +4,7 @@
"""
Shared backend utilities
"""
+
import os
import structlog
from loggers import get_logger
@@ -28,26 +29,26 @@ def without_hf_auth():
"""
# Save environment variables
saved_env = {}
- env_vars = ['HF_TOKEN', 'HUGGINGFACE_HUB_TOKEN', 'HF_HOME']
+ env_vars = ["HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_HOME"]
for var in env_vars:
if var in os.environ:
saved_env[var] = os.environ[var]
del os.environ[var]
# Save disable flag
- saved_disable = os.environ.get('HF_HUB_DISABLE_IMPLICIT_TOKEN')
- os.environ['HF_HUB_DISABLE_IMPLICIT_TOKEN'] = '1'
+ saved_disable = os.environ.get("HF_HUB_DISABLE_IMPLICIT_TOKEN")
+ os.environ["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "1"
# Move token files temporarily
token_files = []
token_locations = [
- Path.home() / '.cache' / 'huggingface' / 'token',
- Path.home() / '.huggingface' / 'token'
+ Path.home() / ".cache" / "huggingface" / "token",
+ Path.home() / ".huggingface" / "token",
]
for token_loc in token_locations:
if token_loc.exists():
- temp = tempfile.NamedTemporaryFile(delete=False)
+ temp = tempfile.NamedTemporaryFile(delete = False)
temp.close()
shutil.move(str(token_loc), temp.name)
token_files.append((token_loc, temp.name))
@@ -58,7 +59,7 @@ def without_hf_auth():
# Restore tokens
for original, temp in token_files:
try:
- original.parent.mkdir(parents=True, exist_ok=True)
+ original.parent.mkdir(parents = True, exist_ok = True)
shutil.move(temp, str(original))
except Exception as e:
logger.error(f"Failed to restore token {original}: {e}")
@@ -68,10 +69,10 @@ def without_hf_auth():
os.environ[var] = value
if saved_disable is not None:
- os.environ['HF_HUB_DISABLE_IMPLICIT_TOKEN'] = saved_disable
+ os.environ["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = saved_disable
else:
- os.environ.pop('HF_HUB_DISABLE_IMPLICIT_TOKEN', None)
-pass
+ os.environ.pop("HF_HUB_DISABLE_IMPLICIT_TOKEN", None)
+
def format_error_message(error: Exception, model_name: str) -> str:
"""
@@ -85,7 +86,7 @@ def format_error_message(error: Exception, model_name: str) -> str:
User-friendly error string
"""
error_str = str(error).lower()
- model_short = model_name.split('/')[-1] if '/' in model_name else model_name
+ model_short = model_name.split("/")[-1] if "/" in model_name else model_name
if "repository not found" in error_str or "404" in error_str:
return f"Model '{model_short}' not found. Check the model name."
@@ -99,12 +100,19 @@ def format_error_message(error: Exception, model_name: str) -> str:
if "invalid user token" in error_str:
return "Invalid HF token. Please check your token and try again."
- if "memory" in error_str or "cuda" in error_str or "mlx" in error_str or "out of memory" in error_str:
+ if (
+ "memory" in error_str
+ or "cuda" in error_str
+ or "mlx" in error_str
+ or "out of memory" in error_str
+ ):
from utils.hardware import get_device
+
device = get_device()
- device_label = {"cuda": "GPU", "mlx": "Apple Silicon GPU", "cpu": "system"}.get(device.value, "GPU")
+ device_label = {"cuda": "GPU", "mlx": "Apple Silicon GPU", "cpu": "system"}.get(
+ device.value, "GPU"
+ )
return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory."
# Generic fallback
return str(error)
-pass
diff --git a/studio/frontend/.gitignore b/studio/frontend/.gitignore
index 3483430dcf..bf7ac45ef1 100644
--- a/studio/frontend/.gitignore
+++ b/studio/frontend/.gitignore
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
# Logs
logs
*.log
diff --git a/studio/frontend/eslint.config.js b/studio/frontend/eslint.config.js
index 88dec236a8..e76e018697 100644
--- a/studio/frontend/eslint.config.js
+++ b/studio/frontend/eslint.config.js
@@ -1,3 +1,6 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
import js from "@eslint/js";
import reactHooks from "eslint-plugin-react-hooks";
import reactRefresh from "eslint-plugin-react-refresh";
diff --git a/studio/frontend/index.html b/studio/frontend/index.html
index c55ae77f62..4f81ffd4ff 100644
--- a/studio/frontend/index.html
+++ b/studio/frontend/index.html
@@ -1,4 +1,7 @@
+
+
+
diff --git a/studio/frontend/public/huggingface.svg b/studio/frontend/public/huggingface.svg
index 45459331bb..c48757ff87 100644
--- a/studio/frontend/public/huggingface.svg
+++ b/studio/frontend/public/huggingface.svg
@@ -1 +1,4 @@
+
+
+
\ No newline at end of file
diff --git a/studio/frontend/public/vite.svg b/studio/frontend/public/vite.svg
index e7b8dfb1b2..8f4c4a2a10 100644
--- a/studio/frontend/public/vite.svg
+++ b/studio/frontend/public/vite.svg
@@ -1 +1,4 @@
+
+
+
\ No newline at end of file
diff --git a/studio/frontend/src/assets/react.svg b/studio/frontend/src/assets/react.svg
index 6c87de9bb3..2b67796fea 100644
--- a/studio/frontend/src/assets/react.svg
+++ b/studio/frontend/src/assets/react.svg
@@ -1 +1,4 @@
+
+
+
\ No newline at end of file
diff --git a/studio/frontend/src/components/ui/collapsible.tsx b/studio/frontend/src/components/ui/collapsible.tsx
index 3cc4953cf3..3566eb9859 100644
--- a/studio/frontend/src/components/ui/collapsible.tsx
+++ b/studio/frontend/src/components/ui/collapsible.tsx
@@ -1,40 +1,40 @@
-// SPDX-License-Identifier: AGPL-3.0-only
-// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-
-import { cn } from "@/lib/utils";
-import { Collapsible as CollapsiblePrimitive } from "radix-ui";
-
-function Collapsible({
- ...props
-}: React.ComponentProps) {
- return ;
-}
-
-function CollapsibleTrigger({
- ...props
-}: React.ComponentProps) {
- return (
-
- );
-}
-
-function CollapsibleContent({
- className,
- ...props
-}: React.ComponentProps) {
- return (
-
- );
-}
-
-export { Collapsible, CollapsibleTrigger, CollapsibleContent };
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { cn } from "@/lib/utils";
+import { Collapsible as CollapsiblePrimitive } from "radix-ui";
+
+function Collapsible({
+ ...props
+}: React.ComponentProps) {
+ return ;
+}
+
+function CollapsibleTrigger({
+ ...props
+}: React.ComponentProps) {
+ return (
+
+ );
+}
+
+function CollapsibleContent({
+ className,
+ ...props
+}: React.ComponentProps) {
+ return (
+
+ );
+}
+
+export { Collapsible, CollapsibleTrigger, CollapsibleContent };
diff --git a/studio/frontend/src/components/ui/tooltip.tsx b/studio/frontend/src/components/ui/tooltip.tsx
index d5f1497041..f58d9f5309 100644
--- a/studio/frontend/src/components/ui/tooltip.tsx
+++ b/studio/frontend/src/components/ui/tooltip.tsx
@@ -1,62 +1,62 @@
-// SPDX-License-Identifier: AGPL-3.0-only
-// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-
-import { Tooltip as TooltipPrimitive } from "radix-ui";
-import type * as React from "react";
-
-import { cn } from "@/lib/utils";
-
-function TooltipProvider({
- delayDuration = 400,
- ...props
-}: React.ComponentProps) {
- return (
-
- );
-}
-
-function Tooltip({
- ...props
-}: React.ComponentProps) {
- return (
-
-
-
- );
-}
-
-function TooltipTrigger({
- ...props
-}: React.ComponentProps) {
- return ;
-}
-
-function TooltipContent({
- className,
- sideOffset = 0,
- children,
- ...props
-}: React.ComponentProps) {
- return (
-
-
- {children}
-
-
-
- );
-}
-
-export { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger };
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { Tooltip as TooltipPrimitive } from "radix-ui";
+import type * as React from "react";
+
+import { cn } from "@/lib/utils";
+
+function TooltipProvider({
+ delayDuration = 400,
+ ...props
+}: React.ComponentProps) {
+ return (
+
+ );
+}
+
+function Tooltip({
+ ...props
+}: React.ComponentProps) {
+ return (
+
+
+
+ );
+}
+
+function TooltipTrigger({
+ ...props
+}: React.ComponentProps) {
+ return ;
+}
+
+function TooltipContent({
+ className,
+ sideOffset = 0,
+ children,
+ ...props
+}: React.ComponentProps) {
+ return (
+
+
+ {children}
+
+
+
+ );
+}
+
+export { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger };
diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css
index 419883ebf6..397853a7b7 100644
--- a/studio/frontend/src/index.css
+++ b/studio/frontend/src/index.css
@@ -1,5 +1,5 @@
-/* SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0 */
-/* Copyright Β© 2025 Unsloth AI */
+/* SPDX-License-Identifier: AGPL-3.0-only */
+/* Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 */
@import "tailwindcss";
@import "tw-animate-css";
diff --git a/studio/frontend/vite.config.ts b/studio/frontend/vite.config.ts
index c8c9028b2e..6e9a03df95 100644
--- a/studio/frontend/vite.config.ts
+++ b/studio/frontend/vite.config.ts
@@ -1,3 +1,6 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
import path from "node:path";
import tailwindcss from "@tailwindcss/vite";
import react from "@vitejs/plugin-react";
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index a4fed429a2..9ce16cd7bb 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -1,4 +1,5 @@
#!/usr/bin/env python3
+
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
@@ -31,6 +32,7 @@ LOCAL_DD_UNSTRUCTURED_PLUGIN = (
# ββ Color support ββββββββββββββββββββββββββββββββββββββββββββββββββββββ
+
def _enable_colors() -> bool:
"""Try to enable ANSI color support. Returns True if available."""
if not hasattr(sys.stdout, "fileno"):
@@ -43,6 +45,7 @@ def _enable_colors() -> bool:
if IS_WINDOWS:
try:
import ctypes
+
kernel32 = ctypes.windll.kernel32
# Enable ENABLE_VIRTUAL_TERMINAL_PROCESSING (0x0004) on stdout
handle = kernel32.GetStdHandle(-11) # STD_OUTPUT_HANDLE
@@ -54,14 +57,18 @@ def _enable_colors() -> bool:
return False
return True # Unix terminals support ANSI by default
+
_HAS_COLOR = _enable_colors()
+
def _green(msg: str) -> str:
return f"\033[92m{msg}\033[0m" if _HAS_COLOR else msg
+
def _cyan(msg: str) -> str:
return f"\033[96m{msg}\033[0m" if _HAS_COLOR else msg
+
def _red(msg: str) -> str:
return f"\033[91m{msg}\033[0m" if _HAS_COLOR else msg
@@ -71,29 +78,33 @@ def run(label: str, cmd: list[str], *, quiet: bool = True) -> None:
print(_cyan(f" {label}..."))
result = subprocess.run(
cmd,
- stdout=subprocess.PIPE if quiet else None,
- stderr=subprocess.STDOUT if quiet else None,
+ stdout = subprocess.PIPE if quiet else None,
+ stderr = subprocess.STDOUT if quiet else None,
)
if result.returncode != 0:
print(_red(f"β {label} failed (exit code {result.returncode}):"))
if result.stdout:
- print(result.stdout.decode(errors="replace"))
+ print(result.stdout.decode(errors = "replace"))
sys.exit(result.returncode)
# Packages to skip on Windows (require special build steps)
-WINDOWS_SKIP_PACKAGES = {"open_spiel"}
+WINDOWS_SKIP_PACKAGES = {"open_spiel", "triton_kernels"}
def _filter_requirements(req: Path, skip: set[str]) -> Path:
"""Return a temp copy of a requirements file with certain packages removed."""
- lines = req.read_text(encoding="utf-8").splitlines(keepends=True)
+ lines = req.read_text(encoding = "utf-8").splitlines(keepends = True)
filtered = [
- line for line in lines
+ line
+ for line in lines
if not any(line.strip().lower().startswith(pkg) for pkg in skip)
]
tmp = tempfile.NamedTemporaryFile(
- mode="w", suffix=".txt", delete=False, encoding="utf-8",
+ mode = "w",
+ suffix = ".txt",
+ delete = False,
+ encoding = "utf-8",
)
tmp.writelines(filtered)
tmp.close()
@@ -121,8 +132,7 @@ def pip_install(
finally:
# Clean up temp file if we created one
if actual_req is not None and actual_req != req:
- actual_req.unlink(missing_ok=True)
-
+ actual_req.unlink(missing_ok = True)
def download_file(url: str, dest: Path) -> None:
@@ -134,7 +144,8 @@ def patch_package_file(package_name: str, relative_path: str, url: str) -> None:
"""Download a file from url and overwrite a file inside an installed package."""
result = subprocess.run(
[sys.executable, "-m", "pip", "show", package_name],
- capture_output=True, text=True,
+ capture_output = True,
+ text = True,
)
if result.returncode != 0:
print(_red(f" β οΈ Could not find package {package_name}, skipping patch"))
@@ -157,6 +168,7 @@ def patch_package_file(package_name: str, relative_path: str, url: str) -> None:
# ββ Main install sequence βββββββββββββββββββββββββββββββββββββββββββββ
+
def install_python_stack() -> int:
print(_cyan("ββ Installing Python stack ββ"))
@@ -165,80 +177,85 @@ def install_python_stack() -> int:
# 2. Core packages: unsloth-zoo + unsloth
pip_install(
- "Installing unsloth-zoo + unsloth",
+ "Installing base packages",
"--no-cache-dir",
- req=REQ_ROOT / "base.txt",
+ req = REQ_ROOT / "base.txt",
)
# 3. Extra dependencies
pip_install(
"Installing additional unsloth dependencies",
"--no-cache-dir",
- req=REQ_ROOT / "extras.txt",
+ req = REQ_ROOT / "extras.txt",
)
# 3b. Extra dependencies (no-deps) β audio model support etc.
pip_install(
"Installing extras (no-deps)",
- "--no-deps", "--no-cache-dir",
- req=REQ_ROOT / "extras-no-deps.txt",
+ "--no-deps",
+ "--no-cache-dir",
+ req = REQ_ROOT / "extras-no-deps.txt",
)
# 4. Overrides (torchao, transformers) β force-reinstall
pip_install(
- "Installing torchao + transformers overrides",
- "--force-reinstall", "--no-cache-dir",
- req=REQ_ROOT / "overrides.txt",
+ "Installing dependency overrides",
+ "--force-reinstall",
+ "--no-cache-dir",
+ req = REQ_ROOT / "overrides.txt",
)
# 5. Triton kernels (no-deps, from source)
- pip_install(
- "Installing triton kernels",
- "--no-deps", "--no-cache-dir",
- req=REQ_ROOT / "triton-kernels.txt",
- constrain=False,
- )
+ if not IS_WINDOWS:
+ pip_install(
+ "Installing triton kernels",
+ "--no-deps",
+ "--no-cache-dir",
+ req = REQ_ROOT / "triton-kernels.txt",
+ constrain = False,
+ )
- # 6. Patch: override llama_cpp.py with fix from unsloth-zoo feature/llama-cpp-windows-support branch
- patch_package_file(
- "unsloth-zoo",
- os.path.join("unsloth_zoo", "llama_cpp.py"),
- "https://raw.githubusercontent.com/unslothai/unsloth-zoo/refs/heads/main/unsloth_zoo/llama_cpp.py",
- )
+ # # 6. Patch: override llama_cpp.py with fix from unsloth-zoo feature/llama-cpp-windows-support branch
+ # patch_package_file(
+ # "unsloth-zoo",
+ # os.path.join("unsloth_zoo", "llama_cpp.py"),
+ # "https://raw.githubusercontent.com/unslothai/unsloth-zoo/refs/heads/main/unsloth_zoo/llama_cpp.py",
+ # )
- # 7a. Patch: override vision.py with fix from unsloth PR #4091
- patch_package_file(
- "unsloth",
- os.path.join("unsloth", "models", "vision.py"),
- "https://raw.githubusercontent.com/unslothai/unsloth/80e0108a684c882965a02a8ed851e3473c1145ab/unsloth/models/vision.py",
- )
+ # # 7a. Patch: override vision.py with fix from unsloth PR #4091
+ # patch_package_file(
+ # "unsloth",
+ # os.path.join("unsloth", "models", "vision.py"),
+ # "https://raw.githubusercontent.com/unslothai/unsloth/80e0108a684c882965a02a8ed851e3473c1145ab/unsloth/models/vision.py",
+ # )
- # 7b. Patch : override save.py with fix from feature/llama-cpp-windows-support
- patch_package_file(
- "unsloth",
- os.path.join("unsloth", "save.py"),
- "https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/main/unsloth/save.py",
- )
+ # # 7b. Patch : override save.py with fix from feature/llama-cpp-windows-support
+ # patch_package_file(
+ # "unsloth",
+ # os.path.join("unsloth", "save.py"),
+ # "https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/main/unsloth/save.py",
+ # )
# 8. Studio dependencies
pip_install(
"Installing studio dependencies",
"--no-cache-dir",
- req=REQ_ROOT / "studio.txt",
+ req = REQ_ROOT / "studio.txt",
)
# 9. Data-designer dependencies
pip_install(
- "Installing data-designer dependencies",
+ "Installing data-designer base dependencies",
"--no-cache-dir",
- req=SINGLE_ENV / "data-designer-deps.txt",
+ req = SINGLE_ENV / "data-designer-deps.txt",
)
# 10. Data-designer packages (no-deps to avoid conflicts)
pip_install(
"Installing data-designer",
- "--no-cache-dir", "--no-deps",
- req=SINGLE_ENV / "data-designer.txt",
+ "--no-cache-dir",
+ "--no-deps",
+ req = SINGLE_ENV / "data-designer.txt",
)
# 11. Local Data Designer seed plugin
@@ -251,9 +268,11 @@ def install_python_stack() -> int:
return 1
pip_install(
"Installing local data-designer unstructured plugin",
- "--no-cache-dir", "--no-deps",
- "-e", str(LOCAL_DD_UNSTRUCTURED_PLUGIN),
- constrain=False,
+ "--no-cache-dir",
+ "--no-deps",
+ "-e",
+ str(LOCAL_DD_UNSTRUCTURED_PLUGIN),
+ constrain = False,
)
# 12. Patch metadata for single-env compatibility
@@ -265,7 +284,8 @@ def install_python_stack() -> int:
# 13. Final check (silent; third-party conflicts are expected)
subprocess.run(
[sys.executable, "-m", "pip", "check"],
- stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
+ stdout = subprocess.DEVNULL,
+ stderr = subprocess.DEVNULL,
)
print(_green("β
Python dependencies installed"))
diff --git a/studio/setup.bat b/studio/setup.bat
index ef16abd263..72d1acb141 100644
--- a/studio/setup.bat
+++ b/studio/setup.bat
@@ -1,2 +1,5 @@
@echo off
+REM SPDX-License-Identifier: AGPL-3.0-only
+REM Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
powershell -ExecutionPolicy Bypass -File "%~dp0setup.ps1" %*
diff --git a/studio/setup.ps1 b/studio/setup.ps1
index 8b662e1998..bf0ba754f0 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -1,6 +1,6 @@
#Requires -Version 5.1
-# SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0
-# Copyright Β© 2025 Unsloth AI
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
<#
.SYNOPSIS
Full environment setup for Unsloth Studio on Windows (bundled version).
@@ -653,7 +653,7 @@ Write-Host "[OK] Using $PythonCmd ($(& $PythonCmd --version 2>&1))" -ForegroundC
# Always create a .venv for isolation -- even for pip installs.
# Created in the repo root (parent of studio/).
-$VenvDir = Join-Path (Split-Path -Parent $PSScriptRoot) ".venv"
+$VenvDir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv"
if (-not (Test-Path $VenvDir)) {
Write-Host " Creating virtual environment at $VenvDir..." -ForegroundColor Cyan
& $PythonCmd -m venv $VenvDir
@@ -726,7 +726,7 @@ $ErrorActionPreference = $prevEAP
# The training subprocess just prepends .venv_t5/ to sys.path β instant switch.
Write-Host ""
Write-Host " Pre-installing transformers 5.x for newer model support..." -ForegroundColor Cyan
-$VenvT5Dir = Join-Path (Split-Path -Parent $PSScriptRoot) ".venv_t5"
+$VenvT5Dir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv_t5"
if (Test-Path $VenvT5Dir) { Remove-Item -Recurse -Force $VenvT5Dir }
New-Item -ItemType Directory -Path $VenvT5Dir -Force | Out-Null
$prevEAP_t5 = $ErrorActionPreference
@@ -956,55 +956,6 @@ if (Test-Path $LlamaServerBin) {
}
}
-# ============================================
-# Add shell aliases (PowerShell profile + cmd batch files)
-# ============================================
-Write-Host ""
-$RepoDir = Split-Path -Parent $PSScriptRoot
-$VenvPython = Join-Path $RepoDir ".venv\Scripts\python.exe"
-$CliScript = Join-Path $RepoDir "cli.py"
-$FrontendDist = Join-Path $PSScriptRoot "frontend\dist"
-$AliasAdded = $false
-
-# --- PowerShell profile: add functions ---
-$ProfileDir = Split-Path $PROFILE -Parent
-if (-not (Test-Path $ProfileDir)) { New-Item -ItemType Directory -Path $ProfileDir -Force | Out-Null }
-if (-not (Test-Path $PROFILE)) { New-Item -ItemType File -Path $PROFILE -Force | Out-Null }
-
-if (-not (Select-String -Path $PROFILE -Pattern "unsloth-studio" -Quiet -ErrorAction SilentlyContinue)) {
- $block = @"
-
-# Unsloth Studio launcher
-function unsloth-studio { & "$VenvPython" "$CliScript" studio -f "$FrontendDist" @args }
-function unsloth-ui { & "$VenvPython" "$CliScript" studio -f "$FrontendDist" @args }
-"@
- Add-Content -Path $PROFILE -Value $block
- Write-Host "[OK] Aliases 'unsloth-studio' and 'unsloth-ui' added to $PROFILE" -ForegroundColor Green
- $AliasAdded = $true
-} else {
- Write-Host "[OK] Aliases 'unsloth-studio' and 'unsloth-ui' already exist in $PROFILE" -ForegroundColor Green
-}
-
-# --- cmd.exe: create batch files and ensure they're on PATH ---
-$BatDir = Join-Path $RepoDir ".venv\Scripts"
-foreach ($name in @("unsloth-studio", "unsloth-ui")) {
- $batPath = Join-Path $BatDir "$name.bat"
- if (-not (Test-Path $batPath)) {
- Set-Content -Path $batPath -Value "@echo off`r`n`"$VenvPython`" `"$CliScript`" studio -f `"$FrontendDist`" %*"
- }
-}
-# Persist .venv\Scripts to User PATH so commands work in new cmd.exe terminals without activation
-$userPath = [Environment]::GetEnvironmentVariable('Path', 'User')
-if (-not $userPath -or $userPath -notlike "*$BatDir*") {
- if ($userPath) {
- [Environment]::SetEnvironmentVariable('Path', "$BatDir;$userPath", 'User')
- } else {
- [Environment]::SetEnvironmentVariable('Path', "$BatDir", 'User')
- }
- Write-Host " Persisted $BatDir to User PATH" -ForegroundColor Gray
-}
-Write-Host "[OK] Batch launchers created (works from any new cmd.exe or PowerShell)" -ForegroundColor Green
-
# ============================================
# Done
# ============================================
@@ -1012,8 +963,7 @@ Write-Host ""
Write-Host "+===============================================+" -ForegroundColor Green
Write-Host "| Setup Complete! |" -ForegroundColor Green
Write-Host "| |" -ForegroundColor Green
-Write-Host "| IMPORTANT: Open a NEW terminal, then run: |" -ForegroundColor Yellow
-Write-Host "| |" -ForegroundColor Green
+Write-Host "| Launch with: |" -ForegroundColor Green
Write-Host "| unsloth studio -H 0.0.0.0 -p 8000 |" -ForegroundColor Green
Write-Host "| |" -ForegroundColor Green
Write-Host "+===============================================+" -ForegroundColor Green
diff --git a/studio/setup.sh b/studio/setup.sh
index 45fedc8882..d42f772397 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -1,6 +1,6 @@
#!/usr/bin/env bash
-# SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0
-# Copyright Β© 2025 Unsloth AI
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
set -euo pipefail
@@ -322,57 +322,6 @@ rm -rf "$LLAMA_CPP_DIR"
fi
}
-# ββ 9. Add shell alias (skip in Colab) ββ
-# Note: venv activation does NOT persist across terminal sessions.
-# This alias hardcodes the venv python path so users don't need to activate.
-if [ "$IS_COLAB" = false ]; then
-echo ""
-REPO_DIR="$REPO_ROOT"
-STUDIO_VENV_PY="$HOME/.unsloth/studio/.venv/bin/python"
-
-# Detect the user's default shell and pick the right rc file
-USER_SHELL="$(basename "${SHELL:-/bin/bash}")"
-case "$USER_SHELL" in
- zsh)
- SHELL_RC="$HOME/.zshrc"
- ALIAS_BLOCK="alias unsloth-studio='${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio'
-alias unsloth-ui='${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio'"
- ;;
- fish)
- SHELL_RC="$HOME/.config/fish/config.fish"
- ALIAS_BLOCK="alias unsloth-studio '${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio'
-alias unsloth-ui '${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio'"
- ;;
- ksh)
- SHELL_RC="$HOME/.kshrc"
- ALIAS_BLOCK="alias unsloth-studio='${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio'
-alias unsloth-ui='${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio'"
- ;;
- *)
- SHELL_RC="$HOME/.bashrc"
- ALIAS_BLOCK="alias unsloth-studio='${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio'
-alias unsloth-ui='${STUDIO_VENV_PY} ${REPO_DIR}/cli.py studio'"
- ;;
-esac
-
-echo " Detected shell: $USER_SHELL β $SHELL_RC"
-
-ALIAS_ADDED=false
-if ! grep -qF "unsloth-studio" "$SHELL_RC" 2>/dev/null; then
- mkdir -p "$(dirname "$SHELL_RC")" # needed for fish's nested config path
- cat >> "$SHELL_RC" < list[dict | str]:
- """Extract `data:` payloads from raw SSE text. Returns dicts or raw strings."""
- results = []
- for line in raw.split("\n"):
- if line.startswith("data: "):
- payload = line[len("data: "):]
- if payload == "[DONE]":
- results.append("[DONE]")
- else:
- try:
- results.append(json.loads(payload))
- except json.JSONDecodeError:
- results.append(payload)
- return results
-
-
-@pytest.fixture()
-def client():
- yield TestClient(app)
-
-
-# =====================================================================
-# Streaming tests
-# =====================================================================
-
-
-class TestStreamingChunkFormat:
- """Each SSE chunk must match the OpenAI chat.completion.chunk schema."""
-
- def test_chunks_have_required_fields(self, client: TestClient):
- mock_backend = _make_mock_backend(tokens=["Hi"])
- with patch("routes.inference.get_inference_backend", return_value=mock_backend):
- resp = client.post(
- "/api/inference/chat/completions",
- json={
- "messages": [{"role": "user", "content": "Hello"}],
- "stream": True,
- },
- )
-
- assert resp.status_code == 200
- chunks = _parse_sse_data(resp.text)
-
- # Filter to actual chunk dicts (not [DONE])
- json_chunks = [c for c in chunks if isinstance(c, dict) and "choices" in c]
- assert len(json_chunks) >= 2 # role chunk + content chunk(s) + final
-
- for chunk in json_chunks:
- assert "id" in chunk
- assert chunk["object"] == "chat.completion.chunk"
- assert "created" in chunk
- assert "model" in chunk
- assert len(chunk["choices"]) == 1
- assert "delta" in chunk["choices"][0]
-
- def test_first_chunk_has_role(self, client: TestClient):
- mock_backend = _make_mock_backend(tokens=["Hi"])
- with patch("routes.inference.get_inference_backend", return_value=mock_backend):
- resp = client.post(
- "/api/inference/chat/completions",
- json={"messages": [{"role": "user", "content": "Hello"}]},
- )
-
- chunks = [c for c in _parse_sse_data(resp.text) if isinstance(c, dict) and "choices" in c]
- first = chunks[0]
- assert first["choices"][0]["delta"].get("role") == "assistant"
-
- def test_last_chunk_has_stop_finish_reason(self, client: TestClient):
- mock_backend = _make_mock_backend(tokens=["Done"])
- with patch("routes.inference.get_inference_backend", return_value=mock_backend):
- resp = client.post(
- "/api/inference/chat/completions",
- json={"messages": [{"role": "user", "content": "Hello"}]},
- )
-
- chunks = [c for c in _parse_sse_data(resp.text) if isinstance(c, dict) and "choices" in c]
- last = chunks[-1]
- assert last["choices"][0]["finish_reason"] == "stop"
- # Delta should be empty on the final chunk
- assert last["choices"][0]["delta"].get("content") is None
-
- def test_stream_ends_with_done(self, client: TestClient):
- mock_backend = _make_mock_backend(tokens=["x"])
- with patch("routes.inference.get_inference_backend", return_value=mock_backend):
- resp = client.post(
- "/api/inference/chat/completions",
- json={"messages": [{"role": "user", "content": "Hello"}]},
- )
-
- all_data = _parse_sse_data(resp.text)
- assert all_data[-1] == "[DONE]"
-
- def test_consistent_id_across_chunks(self, client: TestClient):
- mock_backend = _make_mock_backend(tokens=["a", "b", "c"])
- with patch("routes.inference.get_inference_backend", return_value=mock_backend):
- resp = client.post(
- "/api/inference/chat/completions",
- json={"messages": [{"role": "user", "content": "Hello"}]},
- )
-
- chunks = [c for c in _parse_sse_data(resp.text) if isinstance(c, dict) and "choices" in c]
- ids = set(c["id"] for c in chunks)
- assert len(ids) == 1, "All chunks should share the same completion ID"
-
-
-class TestStreamingHeaders:
- """Verify response headers for SSE proxy compatibility."""
-
- def test_headers(self, client: TestClient):
- mock_backend = _make_mock_backend(tokens=["x"])
- with patch("routes.inference.get_inference_backend", return_value=mock_backend):
- resp = client.post(
- "/api/inference/chat/completions",
- json={"messages": [{"role": "user", "content": "Hello"}]},
- )
-
- assert resp.headers["content-type"].startswith("text/event-stream")
- assert resp.headers.get("cache-control") == "no-cache"
- assert resp.headers.get("x-accel-buffering") == "no"
-
-
-# =====================================================================
-# Non-streaming tests
-# =====================================================================
-
-
-class TestNonStreaming:
- """When stream=false, return a single ChatCompletion JSON object."""
-
- def test_returns_json_object(self, client: TestClient):
- mock_backend = _make_mock_backend(tokens=["Full response text"])
- with patch("routes.inference.get_inference_backend", return_value=mock_backend):
- resp = client.post(
- "/api/inference/chat/completions",
- json={
- "messages": [{"role": "user", "content": "Hello"}],
- "stream": False,
- },
- )
-
- assert resp.status_code == 200
- body = resp.json()
- assert body["object"] == "chat.completion"
- assert body["choices"][0]["message"]["role"] == "assistant"
- assert body["choices"][0]["message"]["content"] == "Full response text"
- assert body["choices"][0]["finish_reason"] == "stop"
-
- def test_non_streaming_has_model(self, client: TestClient):
- mock_backend = _make_mock_backend(tokens=["x"], active_model="my-model")
- with patch("routes.inference.get_inference_backend", return_value=mock_backend):
- resp = client.post(
- "/api/inference/chat/completions",
- json={
- "messages": [{"role": "user", "content": "Hi"}],
- "stream": False,
- },
- )
-
- body = resp.json()
- assert body["model"] == "my-model"
-
-
-# =====================================================================
-# System prompt extraction
-# =====================================================================
-
-
-class TestSystemPromptExtraction:
- """System messages should be extracted and passed as system_prompt."""
-
- def test_system_message_extracted(self, client: TestClient):
- mock_backend = _make_mock_backend(tokens=["ok"])
- with patch("routes.inference.get_inference_backend", return_value=mock_backend):
- client.post(
- "/api/inference/chat/completions",
- json={
- "messages": [
- {"role": "system", "content": "You are a pirate."},
- {"role": "user", "content": "Hello"},
- ],
- "stream": False,
- },
- )
-
- # Check that generate_chat_response was called with the correct system_prompt
- call_kwargs = mock_backend.generate_chat_response.call_args[1]
- assert call_kwargs["system_prompt"] == "You are a pirate."
- # System message should NOT be in the chat_messages list
- assert all(m["role"] != "system" for m in call_kwargs["messages"])
-
- def test_default_system_prompt_when_none(self, client: TestClient):
- mock_backend = _make_mock_backend(tokens=["ok"])
- with patch("routes.inference.get_inference_backend", return_value=mock_backend):
- client.post(
- "/api/inference/chat/completions",
- json={
- "messages": [{"role": "user", "content": "Hello"}],
- "stream": False,
- },
- )
-
- call_kwargs = mock_backend.generate_chat_response.call_args[1]
- assert call_kwargs["system_prompt"] == "You are a helpful AI assistant."
-
-
-# =====================================================================
-# Error handling
-# =====================================================================
-
-
-class TestErrorHandling:
- """Validate error responses for bad requests."""
-
- def test_no_model_loaded(self, client: TestClient):
- mock_backend = _make_mock_backend()
- mock_backend.active_model_name = None
- with patch("routes.inference.get_inference_backend", return_value=mock_backend):
- resp = client.post(
- "/api/inference/chat/completions",
- json={"messages": [{"role": "user", "content": "Hi"}]},
- )
-
- assert resp.status_code == 400
- assert "No model loaded" in resp.json()["detail"]
-
- def test_only_system_messages_rejected(self, client: TestClient):
- mock_backend = _make_mock_backend()
- with patch("routes.inference.get_inference_backend", return_value=mock_backend):
- resp = client.post(
- "/api/inference/chat/completions",
- json={
- "messages": [{"role": "system", "content": "You are a bot."}],
- },
- )
-
- assert resp.status_code == 400
- assert "non-system message" in resp.json()["detail"]
diff --git a/studio/tests/test_remote_lora_detection.py b/studio/tests/test_remote_lora_detection.py
deleted file mode 100644
index d14ad25d6d..0000000000
--- a/studio/tests/test_remote_lora_detection.py
+++ /dev/null
@@ -1,43 +0,0 @@
-# SPDX-License-Identifier: AGPL-3.0-only
-# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-
-"""
-Test remote LoRA adapter detection via HuggingFace Hub API.
-
-Verifies that we can detect whether a remote HF model is a LoRA adapter
-by checking for adapter_config.json in the repo file listing.
-"""
-import pytest
-from huggingface_hub import model_info
-
-
-def is_remote_lora_adapter(model_id: str, hf_token: str = None) -> bool:
- """
- Check if a remote HuggingFace model is a LoRA adapter
- by looking for adapter_config.json in its repo files.
- """
- try:
- info = model_info(model_id, token=hf_token)
- filenames = [s.rfilename for s in info.siblings]
- return "adapter_config.json" in filenames
- except Exception:
- return False
-
-
-class TestRemoteLoRADetection:
- """Test remote LoRA adapter detection via HF Hub API."""
-
- def test_lora_adapter_detected(self):
- """edbeeching/llama-se-rl-adapter is a known LoRA adapter on HF."""
- result = is_remote_lora_adapter("edbeeching/llama-se-rl-adapter")
- assert result is True, "Expected edbeeching/llama-se-rl-adapter to be detected as a LoRA adapter"
-
- def test_base_model_not_detected_as_lora(self):
- """google/gemma-3-4b-it is a full base model, not a LoRA adapter."""
- result = is_remote_lora_adapter("google/gemma-3-4b-it")
- assert result is False, "Expected google/gemma-3-4b-it to NOT be detected as a LoRA adapter"
-
- def test_nonexistent_model_returns_false(self):
- """A nonexistent model should return False, not raise."""
- result = is_remote_lora_adapter("this-org-does-not-exist/fake-model-12345")
- assert result is False, "Expected nonexistent model to return False"
diff --git a/studio/tests/test_sse_progress.py b/studio/tests/test_sse_progress.py
deleted file mode 100644
index b7ee9c9b8f..0000000000
--- a/studio/tests/test_sse_progress.py
+++ /dev/null
@@ -1,323 +0,0 @@
-# SPDX-License-Identifier: AGPL-3.0-only
-# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-
-"""
-Tests for the SSE training progress endpoint and status fallback.
-
-Validates:
- - SSE spec compliance: `retry:`, `id:`, `event:` fields
- - Named event types: progress, heartbeat, complete, error
- - Last-Event-ID reconnection and history replay
- - /status metric_history fallback (Option B)
-
-All tests mock the training backend and bypass auth.
-"""
-import sys
-from pathlib import Path
-from typing import Optional
-from unittest.mock import MagicMock, patch, PropertyMock
-import re
-
-import pytest
-
-# ββ Path setup ββββββββββββββββββββββββββββββββββββββββββββββββββββ
-# Add backend root so bare `from routesβ¦`, `from modelsβ¦` etc. resolve.
-_backend_root = Path(__file__).resolve().parent.parent / "backend"
-if str(_backend_root) not in sys.path:
- sys.path.insert(0, str(_backend_root))
-
-from fastapi.testclient import TestClient
-from main import app
-from auth.authentication import get_current_subject
-
-
-# ββ Fixtures ββββββββββββββββββββββββββββββββββββββββββββββββββββββ
-
-def _bypass_auth():
- """Dependency override that skips real JWT validation."""
- return "test-user"
-
-
-def _make_mock_backend(
- *,
- is_active: bool = False,
- step_history: list | None = None,
- loss_history: list | None = None,
- lr_history: list | None = None,
- total_steps: int = 100,
- epoch: int | None = 1,
- job_id: str = "job_test_001",
-):
- """Build a lightweight mock that quacks like TrainingBackend."""
- backend = MagicMock()
- backend.current_job_id = job_id
- backend.step_history = step_history or []
- backend.loss_history = loss_history or []
- backend.lr_history = lr_history or []
- backend.is_training_active.return_value = is_active
- backend._training_thread = None
-
- # trainer.training_progress / get_training_progress()
- tp = MagicMock()
- tp.total_steps = total_steps
- tp.epoch = epoch
- tp.step = step_history[-1] if step_history else 0
- tp.loss = loss_history[-1] if loss_history else 0.0
- tp.learning_rate = lr_history[-1] if lr_history else 0.0
- tp.status_message = "Training..."
- tp.error = None
- tp.is_completed = not is_active and bool(step_history)
-
- backend.trainer = MagicMock()
- backend.trainer.training_progress = tp
- backend.trainer.get_training_progress.return_value = tp
-
- return backend
-
-
-@pytest.fixture()
-def client():
- """TestClient with auth bypassed."""
- app.dependency_overrides[get_current_subject] = _bypass_auth
- yield TestClient(app)
- app.dependency_overrides.clear()
-
-
-# ββ SSE Parsing Helpers βββββββββββββββββββββββββββββββββββββββββββ
-
-def parse_sse_events(raw: str) -> list[dict]:
- """
- Parse raw SSE text into a list of event dicts.
-
- Each dict has optional keys: 'id', 'event', 'data', 'retry'.
- """
- events: list[dict] = []
- current: dict = {}
-
- for line in raw.split("\n"):
- if line.startswith("retry:"):
- # retry is a standalone directive, not part of a normal event
- events.append({"retry": line.split(":", 1)[1].strip()})
- continue
- if line.startswith("id:"):
- current["id"] = line.split(":", 1)[1].strip()
- elif line.startswith("event:"):
- current["event"] = line.split(":", 1)[1].strip()
- elif line.startswith("data:"):
- current["data"] = line.split(":", 1)[1].strip()
- elif line == "" and current:
- events.append(current)
- current = {}
-
- if current:
- events.append(current)
- return events
-
-
-# =====================================================================
-# Option A β /api/train/progress (SSE)
-# =====================================================================
-
-
-class TestSSERetryDirective:
- """The first thing the stream emits must be `retry: 3000`."""
-
- def test_retry_is_first_event(self, client: TestClient):
- mock_backend = _make_mock_backend(is_active=False)
- with patch("routes.training.get_training_backend", return_value=mock_backend):
- resp = client.get("/api/train/progress")
-
- assert resp.status_code == 200
- assert resp.headers["content-type"].startswith("text/event-stream")
-
- events = parse_sse_events(resp.text)
- assert len(events) >= 1
- assert events[0] == {"retry": "3000"}
-
-
-class TestSSEEventFields:
- """Every non-retry event must include `id:`, `event:`, and `data:` fields."""
-
- def test_events_have_id_and_event_type(self, client: TestClient):
- mock_backend = _make_mock_backend(
- is_active=False,
- step_history=[1, 2, 3],
- loss_history=[2.0, 1.5, 1.0],
- lr_history=[1e-4, 1e-4, 1e-4],
- total_steps=3,
- )
- with patch("routes.training.get_training_backend", return_value=mock_backend):
- resp = client.get("/api/train/progress")
-
- events = parse_sse_events(resp.text)
- data_events = [e for e in events if "data" in e]
-
- assert len(data_events) >= 1
- for evt in data_events:
- assert "id" in evt, f"Missing `id:` field in event: {evt}"
- assert "event" in evt, f"Missing `event:` field in event: {evt}"
- assert "data" in evt
-
-
-class TestSSENamedEventTypes:
- """Events use the correct named types: progress, complete, heartbeat, error."""
-
- def test_idle_sends_progress_then_complete(self, client: TestClient):
- mock_backend = _make_mock_backend(
- is_active=False,
- step_history=[10],
- loss_history=[1.5],
- lr_history=[1e-4],
- total_steps=10,
- )
- with patch("routes.training.get_training_backend", return_value=mock_backend):
- resp = client.get("/api/train/progress")
-
- events = parse_sse_events(resp.text)
- data_events = [e for e in events if "event" in e and e.get("event") != "retry"]
-
- event_types = [e["event"] for e in data_events]
- assert "progress" in event_types
- assert "complete" in event_types
-
- def test_no_history_sends_complete(self, client: TestClient):
- mock_backend = _make_mock_backend(is_active=False)
- with patch("routes.training.get_training_backend", return_value=mock_backend):
- resp = client.get("/api/train/progress")
-
- events = parse_sse_events(resp.text)
- data_events = [e for e in events if "event" in e]
- assert any(e["event"] == "complete" for e in data_events)
-
-
-class TestSSELastEventIDResume:
- """When `Last-Event-ID` header is sent, the server replays missed steps."""
-
- def test_replays_steps_after_last_event_id(self, client: TestClient):
- mock_backend = _make_mock_backend(
- is_active=False,
- step_history=[1, 2, 3, 4, 5],
- loss_history=[2.5, 2.0, 1.5, 1.2, 1.0],
- lr_history=[1e-4, 1e-4, 1e-4, 1e-4, 1e-4],
- total_steps=5,
- )
- with patch("routes.training.get_training_backend", return_value=mock_backend):
- resp = client.get(
- "/api/train/progress",
- headers={"Last-Event-ID": "2"},
- )
-
- events = parse_sse_events(resp.text)
- # Filter to progress events (replayed ones)
- progress_events = [e for e in events if e.get("event") == "progress"]
-
- # Steps 3, 4, 5 should have been replayed
- replayed_ids = [int(e["id"]) for e in progress_events]
- assert 3 in replayed_ids
- assert 4 in replayed_ids
- assert 5 in replayed_ids
- # Steps 1, 2 should NOT be replayed
- assert 1 not in replayed_ids
- assert 2 not in replayed_ids
-
- def test_no_replay_without_header(self, client: TestClient):
- """Without Last-Event-ID, should start fresh (initial progress event)."""
- mock_backend = _make_mock_backend(
- is_active=False,
- step_history=[1, 2, 3],
- loss_history=[2.0, 1.5, 1.0],
- lr_history=[1e-4, 1e-4, 1e-4],
- total_steps=3,
- )
- with patch("routes.training.get_training_backend", return_value=mock_backend):
- resp = client.get("/api/train/progress")
-
- events = parse_sse_events(resp.text)
- progress_events = [e for e in events if e.get("event") == "progress"]
-
- # Should have initial step=0 progress event
- assert any(e.get("id") == "0" for e in progress_events)
-
- def test_invalid_last_event_id_treated_as_fresh(self, client: TestClient):
- """Non-integer Last-Event-ID should be ignored gracefully."""
- mock_backend = _make_mock_backend(is_active=False)
- with patch("routes.training.get_training_backend", return_value=mock_backend):
- resp = client.get(
- "/api/train/progress",
- headers={"Last-Event-ID": "not-a-number"},
- )
-
- assert resp.status_code == 200
- events = parse_sse_events(resp.text)
- # Should still work β treated as a fresh connection
- assert any(e.get("event") == "progress" or e.get("event") == "complete" for e in events)
-
-
-class TestSSEResponseHeaders:
- """Verify SSE response headers for proxy compatibility."""
-
- def test_headers(self, client: TestClient):
- mock_backend = _make_mock_backend(is_active=False)
- with patch("routes.training.get_training_backend", return_value=mock_backend):
- resp = client.get("/api/train/progress")
-
- assert resp.headers["content-type"].startswith("text/event-stream")
- assert resp.headers.get("cache-control") == "no-cache"
- assert resp.headers.get("x-accel-buffering") == "no"
-
-
-# =====================================================================
-# Option B β /api/train/status (metric_history fallback)
-# =====================================================================
-
-
-class TestStatusMetricHistory:
- """The /status endpoint returns metric_history for chart recovery."""
-
- def test_metric_history_populated_when_history_exists(self, client: TestClient):
- mock_backend = _make_mock_backend(
- is_active=True,
- step_history=[1, 2, 3, 4, 5],
- loss_history=[2.5, 2.0, 1.5, 1.2, 1.0],
- lr_history=[1e-4, 1e-4, 1e-4, 1e-4, 1e-4],
- total_steps=10,
- )
- with patch("routes.training.get_training_backend", return_value=mock_backend):
- resp = client.get("/api/train/status")
-
- assert resp.status_code == 200
- body = resp.json()
-
- assert "metric_history" in body
- mh = body["metric_history"]
- assert mh is not None
- assert mh["steps"] == [1, 2, 3, 4, 5]
- assert mh["loss"] == [2.5, 2.0, 1.5, 1.2, 1.0]
- assert mh["lr"] == [1e-4, 1e-4, 1e-4, 1e-4, 1e-4]
-
- def test_metric_history_null_when_no_history(self, client: TestClient):
- mock_backend = _make_mock_backend(is_active=False)
- with patch("routes.training.get_training_backend", return_value=mock_backend):
- resp = client.get("/api/train/status")
-
- assert resp.status_code == 200
- body = resp.json()
- assert body["metric_history"] is None
-
- def test_status_still_returns_phase_and_details(self, client: TestClient):
- """Ensure adding metric_history didn't break existing fields."""
- mock_backend = _make_mock_backend(
- is_active=True,
- step_history=[5],
- loss_history=[1.5],
- lr_history=[1e-4],
- total_steps=100,
- )
- with patch("routes.training.get_training_backend", return_value=mock_backend):
- resp = client.get("/api/train/status")
-
- body = resp.json()
- assert body["phase"] == "training"
- assert body["is_training_running"] is True
- assert body["job_id"] == "job_test_001"
- assert "details" in body
-