diff --git a/cli/commands/studio.py b/cli/commands/studio.py index 4287c756a8..2489676844 100644 --- a/cli/commands/studio.py +++ b/cli/commands/studio.py @@ -1,19 +1,34 @@ +import time +from pathlib import Path +from typing import Optional + import typer def studio( port: int = typer.Option(8000, "--port", "-p", help="Port to run the UI server on."), host: str = typer.Option("0.0.0.0", "--host", "-H", help="Host address to bind to."), - share: bool = typer.Option(True, "--share", "-s", help="Create a public Gradio share link."), + frontend: Optional[Path] = typer.Option(None, "--frontend", "-f", help="Path to frontend build directory."), + silent: bool = typer.Option(False, "--silent", "-q", help="Suppress startup messages."), ): - """Launch the Unsloth web UI for training, inference, and export.""" - from app import demo, script_dir + """Launch the Unsloth web UI backend server.""" + from studio.backend.run import run_server - typer.echo(f"Starting Unsloth UI on http://{host}:{port}") + if not silent: + from studio.backend.run import _resolve_external_ip + display_host = _resolve_external_ip() if host == "0.0.0.0" else host + typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}") - demo.launch( - share=share, - server_port=port, - server_name=host, - favicon_path=f"{script_dir}/assets/favicon-32x32.png", + run_server( + host=host, + port=port, + frontend_path=frontend, + silent=silent, ) + + # Keep running until interrupted + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + typer.echo("\nShutting down...") diff --git a/setup.sh b/setup.sh index c493d357ee..2db0c7f5a5 100755 --- a/setup.sh +++ b/setup.sh @@ -212,38 +212,42 @@ USER_SHELL="$(basename "${SHELL:-/bin/bash}")" case "$USER_SHELL" in zsh) SHELL_RC="$HOME/.zshrc" - ALIAS_BLOCK="alias unsloth-ui='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py ui -f ${REPO_DIR}/studio/frontend/dist'" + ALIAS_BLOCK="alias unsloth-studio='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist' +alias unsloth-ui='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist'" ;; fish) SHELL_RC="$HOME/.config/fish/config.fish" # fish uses 'abbr' or 'function'; a simple alias works via 'alias' in config.fish - ALIAS_BLOCK="alias unsloth-ui '${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py ui -f ${REPO_DIR}/studio/frontend/dist'" + ALIAS_BLOCK="alias unsloth-studio '${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist' +alias unsloth-ui '${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist'" ;; ksh) SHELL_RC="$HOME/.kshrc" - ALIAS_BLOCK="alias unsloth-ui='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py ui -f ${REPO_DIR}/studio/frontend/dist'" + ALIAS_BLOCK="alias unsloth-studio='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist' +alias unsloth-ui='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist'" ;; *) # Default to bash for bash and any other POSIX-compatible shell SHELL_RC="$HOME/.bashrc" - ALIAS_BLOCK="alias unsloth-ui='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py ui -f ${REPO_DIR}/studio/frontend/dist'" + ALIAS_BLOCK="alias unsloth-studio='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist' +alias unsloth-ui='${REPO_DIR}/.venv/bin/python ${REPO_DIR}/cli.py studio -f ${REPO_DIR}/studio/frontend/dist'" ;; esac echo " Detected shell: $USER_SHELL → $SHELL_RC" ALIAS_ADDED=false -if ! grep -qF "unsloth-ui" "$SHELL_RC" 2>/dev/null; then +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" < Generator[str, None, None]: """ Generate response for text or vision models. - Acquires the generation lock. For adapter-controlled generation, - use generate_with_adapter_control() instead. + The generation lock is acquired by the background generation thread. """ - with self._generation_lock: - yield from self._generate_chat_response_inner( - messages=messages, - system_prompt=system_prompt, - image=image, - temperature=temperature, - top_p=top_p, - top_k=top_k, - min_p=min_p, - max_new_tokens=max_new_tokens, - repetition_penalty=repetition_penalty, - cancel_event=cancel_event, - ) + 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, + ) def _generate_chat_response_inner(self, messages: list, @@ -553,10 +558,14 @@ class InferenceBackend: min_p: float = 0.0, max_new_tokens: int = 256, repetition_penalty: float = 1.1, - cancel_event=None) -> Generator[str, None, None]: + cancel_event=None, + _adapter_state=None) -> Generator[str, None, None]: """ - Inner generation logic (no lock). Called by both generate_chat_response + Inner generation logic. Called by both generate_chat_response and generate_with_adapter_control. + + _adapter_state is passed to generate_stream/vision so the background + thread can toggle adapters under the generation lock. """ if not self.active_model_name: yield "Error: No active model" @@ -616,6 +625,7 @@ class InferenceBackend: 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, ) def _generate_vision_response(self, messages, system_prompt, image, @@ -677,6 +687,7 @@ class InferenceBackend: 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, @@ -686,16 +697,17 @@ class InferenceBackend: err: dict[str, str] = {} def generate_fn(): - try: - model.generate(**generation_kwargs) - except Exception as e: - err["msg"] = str(e) - logger.error(f"Vision generation error in thread: {e}") - finally: + with self._generation_lock: try: - streamer.end() - except Exception: - pass + model.generate(**generation_kwargs) + except Exception as e: + err["msg"] = str(e) + logger.error(f"Vision generation error in thread: {e}") + finally: + try: + streamer.end() + except Exception: + pass thread = threading.Thread(target=generate_fn) thread.start() @@ -741,8 +753,13 @@ class InferenceBackend: min_p: float = 0.0, max_new_tokens: int = 256, repetition_penalty: float = 1.1, - cancel_event=None) -> Generator[str, None, None]: - """Generate streaming text response (text models only).""" + 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 + before model.generate(), all under _generation_lock. + """ if not self.active_model_name: yield "Error: No active model" return @@ -773,7 +790,7 @@ class InferenceBackend: top_k=top_k, min_p=min_p, repetition_penalty=repetition_penalty, - do_sample=True, + 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, ) @@ -795,16 +812,19 @@ class InferenceBackend: ) def generate_fn(): - try: - model.generate(**generation_kwargs) - except Exception as e: - err["msg"] = str(e) - logger.error(f"Generation error: {e}") - finally: + with self._generation_lock: try: - streamer.end() - except Exception: - pass + if _adapter_state is not None: + self._apply_adapter_state(_adapter_state) + model.generate(**generation_kwargs) + except Exception as e: + err["msg"] = str(e) + logger.error(f"Generation error: {e}") + finally: + try: + streamer.end() + except Exception: + pass err: dict[str, str] = {} thread = threading.Thread(target=generate_fn) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index e93ba097d1..c36fa235ed 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -867,6 +867,7 @@ class UnslothTrainer: self.trainer, instruction_part=instruction_part, response_part=response_part, + num_proc=config_args.get("dataset_num_proc", max(1, os.cpu_count() // 4)), ) print("Train on responses only configured successfully\n") except Exception as e: diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 8b94ddefdc..39119f1123 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -70,33 +70,48 @@ def _serialize_preview_rows(rows): # --- Endpoints --- +# Recognized data-file extensions for the single-file fallback approach. +DATA_EXTS = ( + '.parquet', + '.json', '.jsonl', + '.csv', '.tsv', + '.txt', + '.arrow', + '.tar', '.tar.gz', '.tgz', + '.gz', '.zst', + '.zip', +) + + @router.post("/check-format", response_model=CheckFormatResponse) -async def check_format(request: CheckFormatRequest): +def check_format(request: CheckFormatRequest): """ Check if a dataset requires manual column mapping. - - This is a lightweight check that streams only the first N rows, - runs format detection, and (if processable) returns processed - preview samples. The full dataset is re-processed at training time. - - For HuggingFace datasets we use streaming mode so we never download - the entire dataset — only the rows we actually need are fetched. + + Strategy for HuggingFace datasets: + 1. list_repo_files → pick the first data file → load_dataset(data_files=[…]) + Avoids resolving thousands of files; typically ~2-4 s. + 2. Full streaming load_dataset as a last-resort fallback. + + Local files are loaded directly. + + Using a plain `def` (not async) so FastAPI runs this in a thread-pool, + preventing any blocking IO from freezing the event loop. """ try: from itertools import islice from datasets import Dataset, load_dataset from utils.datasets import format_dataset - + PREVIEW_SIZE = 10 - + logger.info(f"Checking format for dataset: {request.dataset_name}") - - # Load dataset + dataset_path = Path(request.dataset_name) total_rows = None - + if dataset_path.exists(): - # Local dataset — direct load is fine (files are local) + # ── Local file ────────────────────────────────────────── if dataset_path.suffix in ['.json', '.jsonl']: dataset = load_dataset('json', data_files=str(dataset_path), split=request.train_split) elif dataset_path.suffix == '.csv': @@ -111,54 +126,83 @@ async def check_format(request: CheckFormatRequest): total_rows = len(dataset) preview_slice = dataset.select(range(min(PREVIEW_SIZE, total_rows))) else: - # HuggingFace dataset — use STREAMING to avoid downloading everything - load_kwargs = {"path": request.dataset_name, "split": request.train_split, "streaming": True} - if request.subset: - load_kwargs["name"] = request.subset - if request.hf_token: - load_kwargs["token"] = request.hf_token - - streamed_ds = load_dataset(**load_kwargs) - - # Take only the first PREVIEW_SIZE rows from the stream - 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" + # ── HuggingFace dataset ───────────────────────────────── + # Tier 1: list_repo_files → load only the first data file + preview_slice = None + + 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, ) - - # Convert list-of-dicts into a proper Dataset for downstream compat - preview_slice = Dataset.from_list(rows) - # total_rows unknown in streaming mode + data_files = [f for f in repo_files if any(f.endswith(ext) for ext in DATA_EXTS)] + + if data_files: + first_file = data_files[0] + logger.info(f"Tier 1: loading single file {first_file}") + load_kwargs = { + "path": request.dataset_name, + "data_files": [first_file], + "split": "train", + "streaming": True, + } + if request.hf_token: + load_kwargs["token"] = request.hf_token + + streamed_ds = load_dataset(**load_kwargs) + rows = list(islice(streamed_ds, PREVIEW_SIZE)) + if rows: + preview_slice = Dataset.from_list(rows) + except Exception as e: + logger.warning(f"Tier 1 (single-file) failed: {e}") + + 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} + if request.subset: + load_kwargs["name"] = request.subset + if request.hf_token: + load_kwargs["token"] = request.hf_token + + streamed_ds = load_dataset(**load_kwargs) + + 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" + ) + + 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) - - logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}") - + + logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_multimodal={result.get('is_multimodal', False)}") + # Generate preview samples preview_samples = None if not result["requires_manual_mapping"]: - # Format detected — return processed preview try: format_result = format_dataset( preview_slice, format_type="auto", custom_format_mapping=result.get("suggested_mapping"), + 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}") - # Fall back to raw samples so frontend still has something preview_samples = _serialize_preview_rows(preview_slice) else: - # Format detection failed — return raw samples so user can - # see actual data and map columns in the frontend preview_samples = _serialize_preview_rows(preview_slice) - + return CheckFormatResponse( requires_manual_mapping=result["requires_manual_mapping"], detected_format=result["detected_format"], @@ -171,7 +215,7 @@ async def check_format(request: CheckFormatRequest): preview_samples=preview_samples, total_rows=total_rows, ) - + except HTTPException: raise except Exception as e: diff --git a/studio/backend/utils/datasets/format_detection.py b/studio/backend/utils/datasets/format_detection.py index ae1b6a388d..9283ea5d55 100644 --- a/studio/backend/utils/datasets/format_detection.py +++ b/studio/backend/utils/datasets/format_detection.py @@ -328,6 +328,11 @@ def detect_multimodal_dataset(dataset): """ Detects if dataset contains multimodal data (images/vision). + Two-pass approach: + 1. Column-name heuristic (fast): checks for keywords like 'image', 'img', 'pixel'. + 2. Value-type inspection (reliable): checks if actual values are PIL Images, + bytes with image headers, or HF Image-feature dicts. + Returns: dict: { "is_multimodal": bool, @@ -339,11 +344,16 @@ def detect_multimodal_dataset(dataset): column_names = list(sample.keys()) # Keywords that indicate multimodal/image data - multimodal_keywords = ['image', 'img', 'pixel'] + multimodal_keywords = [ + 'image', 'img', 'pixel', + 'jpg', 'jpeg', 'png', 'webp', 'bmp', 'gif', 'tiff', 'svg', + 'photo', 'pic', 'picture', 'visual', + ] multimodal_columns = [] modality_types = set() + # ── Pass 1: column-name heuristic ─────────────────────── for col_name in column_names: col_lower = col_name.lower() @@ -353,6 +363,17 @@ def detect_multimodal_dataset(dataset): modality_types.add(keyword) break # Don't check other keywords for this column + # ── Pass 2: inspect actual values ─────────────────────── + # Catches columns with non-obvious names (e.g. "jpg", "photo", "pic") + already_detected = set(multimodal_columns) + for col_name in column_names: + if col_name in already_detected: + continue + value = sample[col_name] + if _is_image_value(value): + multimodal_columns.append(col_name) + modality_types.add("image") + return { "is_multimodal": len(multimodal_columns) > 0, "multimodal_columns": multimodal_columns, @@ -360,6 +381,54 @@ def detect_multimodal_dataset(dataset): } +def _is_image_value(value) -> bool: + """Check if a single sample value looks like image data.""" + if value is None: + return False + + # PIL Image instance + try: + from PIL.Image import Image as PILImage + if isinstance(value, PILImage): + return True + except ImportError: + pass + + # HF datasets Image feature stores decoded images as PIL or dicts with + # {"bytes": b"...", "path": "..."} when not yet decoded + if isinstance(value, dict): + if "bytes" in value and "path" in value: + return True + + # Raw bytes with a known image magic header + if isinstance(value, (bytes, bytearray)): + return _has_image_header(value) + + return False + + +def _has_image_header(data: bytes) -> bool: + """Quick magic-byte check for common image formats.""" + if len(data) < 4: + return False + # JPEG + if data[:2] == b'\xff\xd8': + return True + # PNG + if data[:4] == b'\x89PNG': + return True + # GIF + if data[:3] == b'GIF': + return True + # WebP + if data[:4] == b'RIFF' and len(data) >= 12 and data[8:12] == b'WEBP': + return True + # BMP + if data[:2] == b'BM': + return True + return False + + def detect_vlm_dataset_structure(dataset): """ Detects if VLM dataset is: diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index c109c92ead..aaaca0eeb3 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -8,6 +8,7 @@ import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { Button } from "@/components/ui/button"; +import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { cn } from "@/lib/utils"; import { ActionBarMorePrimitive, @@ -38,7 +39,7 @@ import { RefreshCwIcon, SquareIcon, } from "lucide-react"; -import { type FC, useRef } from "react"; +import { type FC, useRef, useState } from "react"; export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ hideComposer, @@ -275,6 +276,32 @@ const AssistantMessage: FC = () => { ); }; +const COPY_RESET_MS = 2000; + +const CopyButton: FC = () => { + const aui = useAui(); + const [copied, setCopied] = useState(false); + const resetTimeoutRef = useRef | null>(null); + + const handleCopy = () => { + const text = aui.message().getCopyText(); + if (copyToClipboard(text)) { + setCopied(true); + if (resetTimeoutRef.current) clearTimeout(resetTimeoutRef.current); + resetTimeoutRef.current = setTimeout(() => { + setCopied(false); + resetTimeoutRef.current = null; + }, COPY_RESET_MS); + } + }; + + return ( + + {copied ? : } + + ); +}; + const AssistantActionBar: FC = () => { return ( { autohideFloat="single-branch" className="aui-assistant-action-bar-root col-start-3 row-start-2 -ml-1 flex gap-1 text-muted-foreground data-floating:absolute data-floating:rounded-md data-floating:border data-floating:bg-background data-floating:p-1 data-floating:shadow-sm" > - - - message.isCopied}> - - - !message.isCopied}> - - - - + @@ -352,16 +370,7 @@ const UserActionBar: FC = () => { autohide="not-last" className="aui-user-action-bar-root flex items-center" > - - - message.isCopied}> - - - !message.isCopied}> - - - - + diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index 8840b9363c..da60328d40 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -73,10 +73,26 @@ export const TARGET_MODULES = [ "down_proj", ]; +export const OPTIMIZER_OPTIONS: ReadonlyArray<{ value: string; label: string }> = [ + { value: "adamw_8bit", label: "AdamW 8-bit" }, + { value: "paged_adamw_8bit", label: "Paged AdamW 8-bit" }, + { value: "adamw_bnb_8bit", label: "AdamW BNB 8-bit" }, + { value: "paged_adamw_32bit", label: "Paged AdamW 32-bit" }, + { value: "adamw_torch", label: "AdamW (PyTorch)" }, + { value: "adamw_torch_fused", label: "AdamW (PyTorch Fused)" }, +]; + +export const LR_SCHEDULER_OPTIONS: ReadonlyArray<{ value: string; label: string }> = [ + { value: "linear", label: "Linear" }, + { value: "cosine", label: "Cosine" }, +]; + export const DEFAULT_HYPERPARAMS = { epochs: 3, contextLength: 2048, learningRate: 2e-4, + optimizerType: "adamw_8bit", + lrSchedulerType: "linear", loraRank: 16, loraAlpha: 32, loraDropout: 0.05, diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 8997570271..8e9e7df0e0 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -159,7 +159,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (abortSignal.aborted) return; warmupToastShown = true; toast.promise(firstTokenPromise, { - loading: "Warming up model", + loading: "Generating", success: "Generating", error: (err) => err instanceof Error && err.message ? err.message : "Generation failed", diff --git a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx index 0275e5a6da..484b10ddb4 100644 --- a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx @@ -35,6 +35,7 @@ import { import { useDebouncedValue, useHfDatasetSearch, + useHfTokenValidation, useInfiniteScroll, } from "@/hooks"; import { cn, formatCompact } from "@/lib/utils"; @@ -103,10 +104,14 @@ export function DatasetStep() { isLoading, isLoadingMore, fetchMore, + error: hfSearchError, } = useHfDatasetSearch(debouncedQuery, { accessToken: hfToken || undefined, }); + const { error: tokenValidationError, isChecking: isCheckingToken } = + useHfTokenValidation(hfToken); + const resultIds = useMemo(() => hfResults.map((r) => r.id), [hfResults]); const comboboxAnchorRef = useRef(null); @@ -179,6 +184,23 @@ export function DatasetStep() { onChange={(e) => setHfToken(e.target.value)} /> + {(tokenValidationError ?? hfSearchError) && ( +

+ {tokenValidationError ?? hfSearchError} + {" — "} + + Get or update token + +

+ )} + {isCheckingToken && ( +

Checking token…

+ )} diff --git a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx index 3f63775ff9..ce506b6085 100644 --- a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx @@ -34,6 +34,7 @@ import { MODEL_TYPE_TO_HF_TASK } from "@/config/training"; import { useDebouncedValue, useHfModelSearch, + useHfTokenValidation, useInfiniteScroll, } from "@/hooks"; import { formatCompact } from "@/lib/utils"; @@ -80,11 +81,15 @@ export function ModelSelectionStep() { isLoading, isLoadingMore, fetchMore, + error: hfSearchError, } = useHfModelSearch(debouncedQuery, { task, accessToken: hfToken || undefined, }); + const { error: tokenValidationError, isChecking: isCheckingToken } = + useHfTokenValidation(hfToken); + const resultIds = useMemo(() => hfResults.map((r) => r.id), [hfResults]); const comboboxAnchorRef = useRef(null); @@ -126,6 +131,23 @@ export function ModelSelectionStep() { onChange={(e) => setHfToken(e.target.value)} /> + {(tokenValidationError ?? hfSearchError) && ( +

+ {tokenValidationError ?? hfSearchError} + {" — "} + + Get or update token + +

+ )} + {isCheckingToken && ( +

Checking token…

+ )} diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx index 648adfbc89..e135f301af 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx @@ -63,6 +63,7 @@ export function DatasetPreviewDialog({ const mappingOk = !!manualMapping.input && !!manualMapping.output; const leftLabel = isVlm ? "Image" : "Input"; const rightLabel = isVlm ? "Text" : "Output"; + const isHfDataset = !!datasetName && datasetName.includes("/"); useEffect(() => { if (!manualMapping.input || !manualMapping.output) return; @@ -266,8 +267,13 @@ export function DatasetPreviewDialog({

- Loading preview... + {isHfDataset ? "Fetching dataset preview from Hugging Face..." : "Loading preview..."}

+ {isHfDataset && ( +

+ This may take a moment for large datasets +

+ )} )} diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index 97bd1304f0..d142ae5f54 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -25,6 +25,7 @@ import { import { useDebouncedValue, useHfDatasetSearch, + useHfTokenValidation, useInfiniteScroll, } from "@/hooks"; import { formatCompact } from "@/lib/utils"; @@ -102,10 +103,14 @@ export function DatasetSection() { isLoading, isLoadingMore, fetchMore, + error: hfSearchError, } = useHfDatasetSearch(debouncedQuery, { accessToken: hfToken || undefined, }); + const { error: tokenValidationError, isChecking: isCheckingToken } = + useHfTokenValidation(hfToken); + const resultIds = useMemo(() => { const ids = hfResults.map((r) => r.id); if (dataset && !ids.includes(dataset)) { @@ -249,6 +254,23 @@ export function DatasetSection() { + {(tokenValidationError ?? hfSearchError) && ( +

+ {tokenValidationError ?? hfSearchError} + {" — "} + + Get or update token + +

+ )} + {isCheckingToken && ( +

Checking token…

+ )} { const ids = hfResults.map((r) => r.id); if (selectedModel && !ids.includes(selectedModel)) { @@ -568,6 +573,23 @@ export function ModelSection() { onChange={(e) => setHfToken(e.target.value)} /> + {(tokenValidationError ?? hfSearchError) && ( +

+ {tokenValidationError ?? hfSearchError} + {" — "} + + Get or update token + +

+ )} + {isCheckingToken && ( +

Checking token…

+ )} diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index 2069da7a90..de82fb0fb5 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -20,7 +20,12 @@ import { TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { CONTEXT_LENGTHS, TARGET_MODULES } from "@/config/training"; +import { + CONTEXT_LENGTHS, + LR_SCHEDULER_OPTIONS, + OPTIMIZER_OPTIONS, + TARGET_MODULES, +} from "@/config/training"; import { useTrainingConfigStore } from "@/features/training"; import type { GradientCheckpointing } from "@/types/training"; import { @@ -508,6 +513,78 @@ export function ParamsSection(): ReactElement { value="optimization" className="mt-3 flex flex-col gap-3" > + + Optimization algorithm. 8-bit variants reduce memory usage. + Fused is recommended for vision models.{" "} + + Read more + + + } + > + + + + How the learning rate changes over training. Linear decays + steadily; cosine decays in a curve.{" "} + + Read more + + + } + > + + = { idle: "Idle", + downloading_model: "Downloading model", + downloading_dataset: "Downloading dataset", loading_model: "Loading model", loading_dataset: "Loading dataset", configuring: "Configuring", @@ -13,6 +15,10 @@ export const phaseLabel: Record = { export const phaseColors: Record = { idle: "bg-muted text-muted-foreground", + downloading_model: + "bg-sky-100 text-sky-700 dark:bg-sky-900 dark:text-sky-300", + downloading_dataset: + "bg-sky-100 text-sky-700 dark:bg-sky-900 dark:text-sky-300", loading_model: "bg-amber-100 text-amber-700 dark:bg-amber-900 dark:text-amber-300", loading_dataset: diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index 3752a0fb25..6e88c4c05a 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -35,6 +35,7 @@ import { Link, useNavigate } from "@tanstack/react-router"; import { useShallow } from "zustand/react/shallow"; import { useGpuUtilization } from "@/hooks"; import { setTrainingCompareHandoff } from "@/features/chat"; +import { OPTIMIZER_OPTIONS } from "@/config/training"; import { formatDuration, formatNumber, phaseColors, phaseLabel } from "./progress-section-lib"; export function ProgressSection(): ReactElement { @@ -71,6 +72,7 @@ export function ProgressSection(): ReactElement { maxSteps: state.maxSteps, contextLength: state.contextLength, warmupSteps: state.warmupSteps, + optimizerType: state.optimizerType, loraRank: state.loraRank, loraAlpha: state.loraAlpha, loraDropout: state.loraDropout, @@ -126,6 +128,10 @@ export function ProgressSection(): ReactElement { ? runtime.currentGradNorm : lastNonZeroValue(runtime.gradNormHistory) ?? runtime.currentGradNorm; + const optimizerLabel = + OPTIMIZER_OPTIONS.find((o) => o.value === config.optimizerType)?.label ?? + config.optimizerType; + const configItems = [ { section: "Hyperparams", @@ -133,6 +139,7 @@ export function ProgressSection(): ReactElement { ["Epochs", config.epochs], ["Batch size", config.batchSize], ["Learning rate", config.learningRate], + ["Optimizer", optimizerLabel], ["Max steps", config.maxSteps], ["Context length", config.contextLength], ["Warmup steps", config.warmupSteps], diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx index e9f01901b7..26353b9b7e 100644 --- a/studio/frontend/src/features/studio/studio-page.tsx +++ b/studio/frontend/src/features/studio/studio-page.tsx @@ -37,6 +37,9 @@ export function StudioPage(): ReactElement { const ensureModelDefaultsLoaded = useTrainingConfigStore( (s) => s.ensureModelDefaultsLoaded, ); + const ensureDatasetChecked = useTrainingConfigStore( + (s) => s.ensureDatasetChecked, + ); const dialogOpen = useDatasetPreviewDialogStore((s) => s.open); const dialogMode = useDatasetPreviewDialogStore((s) => s.mode); const dialogInitial = useDatasetPreviewDialogStore((s) => s.initialData); @@ -65,7 +68,8 @@ export function StudioPage(): ReactElement { useEffect(() => { ensureModelDefaultsLoaded(); - }, [selectedModel, ensureModelDefaultsLoaded]); + ensureDatasetChecked(); + }, [selectedModel, ensureModelDefaultsLoaded, ensureDatasetChecked]); return (
diff --git a/studio/frontend/src/features/studio/training-view.tsx b/studio/frontend/src/features/studio/training-view.tsx index 9e83a2f662..fe616ea88a 100644 --- a/studio/frontend/src/features/studio/training-view.tsx +++ b/studio/frontend/src/features/studio/training-view.tsx @@ -19,6 +19,8 @@ export function TrainingView(): ReactElement { ); const isPreparingPhase = + runtime.phase === "downloading_model" || + runtime.phase === "downloading_dataset" || runtime.phase === "loading_model" || runtime.phase === "loading_dataset" || runtime.phase === "configuring"; diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index 510cc542b8..ac93761b62 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -40,8 +40,8 @@ export function buildTrainingStartPayload( weight_decay: config.weightDecay, random_seed: config.randomSeed, packing: config.packing, - optim: "adamw_8bit", - lr_scheduler_type: "linear", + optim: config.optimizerType, + lr_scheduler_type: config.lrSchedulerType, use_lora: adapterMethod, lora_r: config.loraRank, lora_alpha: config.loraAlpha, diff --git a/studio/frontend/src/features/training/api/models-api.ts b/studio/frontend/src/features/training/api/models-api.ts index 6e6e3ff762..e22de5f581 100644 --- a/studio/frontend/src/features/training/api/models-api.ts +++ b/studio/frontend/src/features/training/api/models-api.ts @@ -9,6 +9,8 @@ interface BackendTrainingDefaults { max_seq_length?: number; num_epochs?: number; learning_rate?: number | string; + optim?: string; + lr_scheduler_type?: string; batch_size?: number; gradient_accumulation_steps?: number; warmup_steps?: number; diff --git a/studio/frontend/src/features/training/lib/model-defaults.ts b/studio/frontend/src/features/training/lib/model-defaults.ts index 07ffb2a422..35ce562dbf 100644 --- a/studio/frontend/src/features/training/lib/model-defaults.ts +++ b/studio/frontend/src/features/training/lib/model-defaults.ts @@ -7,6 +7,8 @@ type ModelDefaultsPatch = Partial< | "epochs" | "contextLength" | "learningRate" + | "optimizerType" + | "lrSchedulerType" | "loraRank" | "loraAlpha" | "loraDropout" @@ -86,6 +88,12 @@ export function mapBackendModelConfigToTrainingPatch( const learningRate = toNumber(training?.learning_rate); if (learningRate !== undefined) patch.learningRate = learningRate; + const optim = toStringValue(training?.optim); + if (optim !== undefined) patch.optimizerType = optim; + + const lrSchedulerType = toStringValue(training?.lr_scheduler_type); + if (lrSchedulerType !== undefined) patch.lrSchedulerType = lrSchedulerType; + const batchSize = toNumber(training?.batch_size); if (batchSize !== undefined) patch.batchSize = batchSize; diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index b98d2a3c7f..b6f9c01c42 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -43,12 +43,19 @@ let _datasetCheckController: AbortController | null = null; // AbortController for in-flight model default loads. let _modelConfigController: AbortController | null = null; +// Track whether the user has manually toggled trainOnCompletions +// since the last auto-set (model load or dataset change). +let _trainOnCompletionsManuallySet = false; + const NON_PERSISTED_STATE_KEYS: ReadonlySet = new Set([ "modelType", "isCheckingVision", "isLoadingModelDefaults", "modelDefaultsError", + "modelDefaultsAppliedFor", "isCheckingDataset", + "isDatasetMultimodal", + "trainOnCompletions", ]); function partializePersistedState( @@ -102,8 +109,17 @@ export const useTrainingConfigStore = create()( if (controller.signal.aborted) return; if (get().selectedModel !== modelName) return; + _trainOnCompletionsManuallySet = false; + const patch = mapBackendModelConfigToTrainingPatch(modelDetails.config); + + // If vision model + multimodal dataset already known, override + // trainOnCompletions to false regardless of backend default. + if (modelDetails.is_vision && get().isDatasetMultimodal === true) { + patch.trainOnCompletions = false; + } + set({ - ...mapBackendModelConfigToTrainingPatch(modelDetails.config), + ...patch, isVisionModel: modelDetails.is_vision, isLoadingModelDefaults: false, isCheckingVision: false, @@ -139,6 +155,40 @@ export const useTrainingConfigStore = create()( }); }; + const runDatasetCheck = (datasetName: string, split: string) => { + _datasetCheckController?.abort(); + const controller = new AbortController(); + _datasetCheckController = controller; + set({ isCheckingDataset: true }); + + const state = get(); + checkDatasetFormat({ + datasetName, + hfToken: state.hfToken.trim() || null, + subset: state.datasetSubset, + split, + }) + .then((res) => { + if (controller.signal.aborted) return; + const isMultimodal = !!res.is_multimodal; + const updates: Record = { + isDatasetMultimodal: isMultimodal, + isCheckingDataset: false, + }; + if (!_trainOnCompletionsManuallySet) { + const { isVisionModel } = get(); + if (isVisionModel && isMultimodal) { + updates.trainOnCompletions = false; + } + } + set(updates); + }) + .catch(() => { + if (controller.signal.aborted) return; + set({ isDatasetMultimodal: null, isCheckingDataset: false }); + }); + }; + return { ...initialState, setStep: (step) => set({ currentStep: step }), @@ -196,6 +246,7 @@ export const useTrainingConfigStore = create()( setDataset: (dataset) => { _datasetCheckController?.abort(); _datasetCheckController = null; + _trainOnCompletionsManuallySet = false; set({ dataset, datasetSubset: null, @@ -208,6 +259,7 @@ export const useTrainingConfigStore = create()( setDatasetSubset: (datasetSubset) => { _datasetCheckController?.abort(); _datasetCheckController = null; + _trainOnCompletionsManuallySet = false; set({ datasetSubset, datasetSplit: null, @@ -217,8 +269,6 @@ export const useTrainingConfigStore = create()( }); }, setDatasetSplit: (datasetSplit) => { - _datasetCheckController?.abort(); - _datasetCheckController = null; set({ datasetSplit, datasetManualMapping: emptyManualMapping(), @@ -233,27 +283,21 @@ export const useTrainingConfigStore = create()( : state.uploadedFile; if (!datasetName) return; - const controller = new AbortController(); - _datasetCheckController = controller; - set({ isCheckingDataset: true }); + runDatasetCheck(datasetName, datasetSplit || "train"); + }, + ensureDatasetChecked: () => { + const state = get(); + if (state.isCheckingDataset) return; + if (state.isDatasetMultimodal !== null) return; - checkDatasetFormat({ - datasetName, - hfToken: state.hfToken.trim() || null, - subset: state.datasetSubset, - split: datasetSplit || "train", - }) - .then((res) => { - if (controller.signal.aborted) return; - set({ - isDatasetMultimodal: !!res.is_multimodal, - isCheckingDataset: false, - }); - }) - .catch(() => { - if (controller.signal.aborted) return; - set({ isDatasetMultimodal: null, isCheckingDataset: false }); - }); + const datasetName = + state.datasetSource === "huggingface" + ? state.dataset + : state.uploadedFile; + if (!datasetName) return; + + const split = state.datasetSplit || "train"; + runDatasetCheck(datasetName, split); }, setDatasetManualMapping: (datasetManualMapping) => set({ datasetManualMapping }), @@ -261,6 +305,8 @@ export const useTrainingConfigStore = create()( setEpochs: (epochs) => set({ epochs }), setContextLength: (contextLength) => set({ contextLength }), setLearningRate: (learningRate) => set({ learningRate }), + setOptimizerType: (optimizerType) => set({ optimizerType }), + setLrSchedulerType: (lrSchedulerType) => set({ lrSchedulerType }), setLoraRank: (loraRank) => set({ loraRank }), setLoraAlpha: (loraAlpha) => set({ loraAlpha }), setLoraDropout: (loraDropout) => set({ loraDropout }), @@ -274,8 +320,10 @@ export const useTrainingConfigStore = create()( setSaveSteps: (saveSteps) => set({ saveSteps }), setEvalSteps: (evalSteps) => set({ evalSteps }), setPacking: (packing) => set({ packing }), - setTrainOnCompletions: (trainOnCompletions) => - set({ trainOnCompletions }), + setTrainOnCompletions: (trainOnCompletions) => { + _trainOnCompletionsManuallySet = true; + set({ trainOnCompletions }); + }, setGradientCheckpointing: (gradientCheckpointing) => set({ gradientCheckpointing }), setRandomSeed: (randomSeed) => set({ randomSeed }), @@ -300,7 +348,7 @@ export const useTrainingConfigStore = create()( }, { name: "unsloth_training_config_v1", - version: 3, + version: 5, migrate: (persisted, version) => { const s = persisted as Record; if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) { @@ -310,6 +358,12 @@ export const useTrainingConfigStore = create()( if (version < 3 && s.modelDefaultsAppliedFor == null) { s.modelDefaultsAppliedFor = null; } + if (version < 4 && s.optimizerType == null) { + s.optimizerType = DEFAULT_HYPERPARAMS.optimizerType; + } + if (version < 5 && s.lrSchedulerType == null) { + s.lrSchedulerType = DEFAULT_HYPERPARAMS.lrSchedulerType; + } return s as unknown as TrainingConfigStore; }, partialize: partializePersistedState, diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index 67f1d00edc..b268a08b79 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -30,6 +30,8 @@ export interface TrainingConfigState { epochs: number; contextLength: number; learningRate: number; + optimizerType: string; + lrSchedulerType: string; loraRank: number; loraAlpha: number; loraDropout: number; @@ -72,6 +74,7 @@ export interface TrainingConfigActions { setModelType: (type: ModelType) => void; setSelectedModel: (model: string | null) => void; ensureModelDefaultsLoaded: () => void; + ensureDatasetChecked: () => void; setTrainingMethod: (method: TrainingMethod) => void; setHfToken: (token: string) => void; setDatasetSource: (source: DatasetSource) => void; @@ -84,6 +87,8 @@ export interface TrainingConfigActions { setEpochs: (epochs: number) => void; setContextLength: (length: number) => void; setLearningRate: (rate: number) => void; + setOptimizerType: (value: string) => void; + setLrSchedulerType: (value: string) => void; setLoraRank: (rank: number) => void; setLoraAlpha: (alpha: number) => void; setLoraDropout: (dropout: number) => void; diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts index df24c020ca..7ebf09518d 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/studio/frontend/src/features/training/types/runtime.ts @@ -1,5 +1,7 @@ export type TrainingPhase = | "idle" + | "downloading_model" + | "downloading_dataset" | "loading_model" | "loading_dataset" | "configuring" diff --git a/studio/frontend/src/hooks/index.ts b/studio/frontend/src/hooks/index.ts index 16c2c82b75..b6c4fcad14 100644 --- a/studio/frontend/src/hooks/index.ts +++ b/studio/frontend/src/hooks/index.ts @@ -5,4 +5,5 @@ export { useHardwareInfo } from "./use-hardware-info"; export { useHfModelSearch } from "./use-hf-model-search"; export { useHfDatasetSearch } from "./use-hf-dataset-search"; export { useHfDatasetSplits } from "./use-hf-dataset-splits"; +export { useHfTokenValidation } from "./use-hf-token-validation"; export { useInfiniteScroll } from "./use-infinite-scroll"; diff --git a/studio/frontend/src/hooks/use-hf-token-validation.ts b/studio/frontend/src/hooks/use-hf-token-validation.ts new file mode 100644 index 0000000000..1882745eca --- /dev/null +++ b/studio/frontend/src/hooks/use-hf-token-validation.ts @@ -0,0 +1,59 @@ +import { whoAmI } from "@huggingface/hub"; +import { useCallback, useEffect, useRef, useState } from "react"; +import { useDebouncedValue } from "./use-debounced-value"; + +export interface HfTokenValidationState { + isValid: boolean | null; + error: string | null; + isChecking: boolean; +} + +const INITIAL: HfTokenValidationState = { + isValid: null, + error: null, + isChecking: false, +}; + +/** + * Validates the Hugging Face token by calling the whoami-v2 API. + * Debounces the token to avoid excessive requests while typing. + * Returns validation state: isValid (null = not checked), error message, and isChecking. + */ +export function useHfTokenValidation(token: string): HfTokenValidationState { + const debouncedToken = useDebouncedValue(token.trim(), 500); + const [state, setState] = useState(INITIAL); + const versionRef = useRef(0); + + const runCheck = useCallback(async (t: string) => { + if (!t) { + setState({ isValid: null, error: null, isChecking: false }); + return; + } + + const v = ++versionRef.current; + setState((prev) => ({ ...prev, isChecking: true, error: null })); + + try { + await whoAmI({ accessToken: t }); + if (versionRef.current !== v) return; + setState({ isValid: true, error: null, isChecking: false }); + } catch { + if (versionRef.current !== v) return; + setState({ + isValid: false, + error: "invalid or expired token", + isChecking: false, + }); + } + }, []); + + useEffect(() => { + if (!debouncedToken) { + setState(INITIAL); + return; + } + runCheck(debouncedToken); + }, [debouncedToken, runCheck]); + + return state; +} diff --git a/studio/frontend/src/lib/copy-to-clipboard.ts b/studio/frontend/src/lib/copy-to-clipboard.ts new file mode 100644 index 0000000000..3ef3df1177 --- /dev/null +++ b/studio/frontend/src/lib/copy-to-clipboard.ts @@ -0,0 +1,44 @@ +/** + * Copy text to clipboard in a way that works on Mac/Safari. + * Uses a synchronous textarea + execCommand fallback so the copy runs in the + * same user gesture as the click (required by Safari's clipboard security). + */ +export function copyToClipboard(text: string): boolean { + if (typeof text !== "string" || text.length === 0) { + return false; + } + + // Synchronous fallback: works in Safari/Mac when clipboard API fails + // because it runs entirely within the user gesture (click) stack. + if (document.queryCommandSupported?.("copy") !== false) { + const textarea = document.createElement("textarea"); + textarea.value = text; + textarea.style.position = "fixed"; + textarea.style.top = "0"; + textarea.style.left = "0"; + textarea.style.opacity = "0"; + textarea.setAttribute("aria-hidden", "true"); + document.body.appendChild(textarea); + textarea.focus({ preventScroll: true }); + textarea.select(); + try { + const ok = document.execCommand("copy"); + document.body.removeChild(textarea); + return ok; + } catch { + document.body.removeChild(textarea); + return false; + } + } + + // Modern API only when fallback not available (e.g. non-browser) + if (typeof navigator?.clipboard?.writeText === "function") { + navigator.clipboard.writeText(text).then( + () => {}, + () => {} + ); + return true; + } + + return false; +}