Merge branch 'nightly' into feat/embedding-models

This commit is contained in:
Manan Shah 2026-03-10 14:16:05 -07:00 committed by GitHub
commit f696ef81e8
9 changed files with 182 additions and 26 deletions

View file

@ -249,17 +249,42 @@ class LlamaCppBackend:
)
# Determine the filename from the variant (e.g., "Q4_K_M" -> find matching file)
# For split GGUFs (e.g., *-00001-of-00003.gguf) we must download ALL shards.
gguf_filename = None
gguf_extra_shards: list[str] = []
if hf_variant:
# Try common naming patterns
try:
import re
from huggingface_hub import list_repo_files
files = list_repo_files(hf_repo, token=hf_token)
variant_lower = hf_variant.lower()
for f in files:
if f.endswith(".gguf") and variant_lower in f.lower():
gguf_filename = f
break
# Use word-boundary matching so "Q8_0" doesn't also
# match "IQ8_0" or other superset variant names.
boundary = re.compile(
r'(?<![a-zA-Z0-9])' + re.escape(variant_lower) + r'(?![a-zA-Z0-9])'
)
gguf_files = sorted(
f for f in files
if f.endswith(".gguf") and boundary.search(f.lower())
)
if gguf_files:
gguf_filename = gguf_files[0]
# For split GGUFs (e.g. model-Q8_0-00001-of-00003.gguf)
# discover siblings by exact basename + total match
# so "model-Q8_0-v2-*" isn't pulled in as a sibling.
shard_pat = re.compile(r'^(.*)-\d{5}-of-(\d{5})\.gguf$')
m = shard_pat.match(gguf_filename)
if m:
prefix = m.group(1)
total = m.group(2)
sibling_pat = re.compile(
r'^' + re.escape(prefix) + r'-\d{5}-of-' + re.escape(total) + r'\.gguf$'
)
gguf_extra_shards = [
f for f in gguf_files[1:]
if sibling_pat.match(f)
]
except Exception as e:
logger.warning(f"Could not list repo files: {e}")
@ -269,13 +294,23 @@ class LlamaCppBackend:
repo_name = hf_repo.split("/")[-1].replace("-GGUF", "")
gguf_filename = f"{repo_name}-{hf_variant}.gguf"
logger.info(f"Downloading GGUF: {hf_repo}/{gguf_filename}")
logger.info(f"Downloading GGUF: {hf_repo}/{gguf_filename}"
+ (f" (+{len(gguf_extra_shards)} shards)" if gguf_extra_shards else ""))
try:
local_path = hf_hub_download(
repo_id=hf_repo,
filename=gguf_filename,
token=hf_token,
)
# Download remaining shards for split GGUFs — llama-server
# auto-discovers them when they are in the same directory.
for shard in gguf_extra_shards:
logger.info(f"Downloading GGUF shard: {shard}")
hf_hub_download(
repo_id=hf_repo,
filename=shard,
token=hf_token,
)
except Exception as e:
raise RuntimeError(
f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}"

View file

@ -1866,10 +1866,34 @@ class UnslothTrainer:
elif dataset_source:
# Load from Hugging Face
load_kwargs = {"path": dataset_source, "split": train_split or "train"}
split_name = train_split or "train"
load_kwargs = {"path": dataset_source, "split": split_name}
if subset:
load_kwargs["name"] = subset
dataset = load_dataset(**load_kwargs)
_slice_start = dataset_slice_start or 0
if (dataset_slice_end is not None
and dataset_slice_end >= 0
and dataset_slice_end >= _slice_start):
# Manual slice — stream only the rows we need instead of
# downloading the entire dataset.
rows_to_stream = dataset_slice_end + 1
print(
f"[dataset-slice] Manual slice specified "
f"(start={dataset_slice_start}, end={dataset_slice_end}), "
f"streaming {rows_to_stream} rows\n"
)
stream = load_dataset(**load_kwargs, streaming=True)
dataset = Dataset.from_list(list(stream.take(rows_to_stream)))
print(
f"[dataset-slice] Downloaded {len(dataset)} rows "
f"(requested {rows_to_stream})\n"
)
self._update_progress(
status_message=f"Streamed {len(dataset)} rows from HuggingFace"
)
else:
dataset = load_dataset(**load_kwargs)
# Check if stopped during dataset loading
if self.should_stop:
@ -1877,7 +1901,7 @@ class UnslothTrainer:
return None
self._update_progress(status_message=f"Loaded dataset from HuggingFace: {dataset_source}")
print(f"Loaded dataset from Hugging Face: {dataset_source}\n")
print(f"Loaded dataset from Hugging Face: {dataset_source} ({len(dataset)} rows)\n")
# Resolve eval split from a separate HF split (explicit or auto-detected)
if eval_enabled:

View file

@ -1,6 +1,6 @@
// SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0
// Copyright © 2025 Unsloth AI
// SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0
// Copyright © 2025 Unsloth AI
import { Tooltip as TooltipPrimitive } from "radix-ui";
import type * as React from "react";
@ -46,14 +46,14 @@ function TooltipContent({
<TooltipPrimitive.Content
data-slot="tooltip-content"
sideOffset={sideOffset}
className={cn(
"data-open:animate-in data-open:fade-in-0 data-open:zoom-in-95 data-[state=delayed-open]:animate-in data-[state=delayed-open]:fade-in-0 data-[state=delayed-open]:zoom-in-95 data-closed:animate-out data-closed:fade-out-0 data-closed:zoom-out-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 rounded-2xl corner-squircle px-3 py-1.5 text-xs **:data-[slot=kbd]:rounded-4xl bg-foreground text-background border border-foreground/40 shadow-lg z-50 w-fit max-w-xs origin-(--radix-tooltip-content-transform-origin)",
className,
)}
className={cn(
"data-open:animate-in data-open:fade-in-0 data-open:zoom-in-95 data-[state=delayed-open]:animate-in data-[state=delayed-open]:fade-in-0 data-[state=delayed-open]:zoom-in-95 data-closed:animate-out data-closed:fade-out-0 data-closed:zoom-out-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 rounded-2xl corner-squircle px-3 py-1.5 text-xs **:data-[slot=kbd]:rounded-4xl bg-foreground text-background border border-foreground/40 shadow-lg z-[999999] w-fit max-w-xs origin-(--radix-tooltip-content-transform-origin)",
className,
)}
{...props}
>
{children}
<TooltipPrimitive.Arrow className="size-2.5 translate-y-[calc(-50%_-_2px)] rotate-45 rounded-[2px] data-[side=left]:translate-x-[-1.5px] data-[side=right]:translate-x-[1.5px] bg-foreground fill-foreground z-50 translate-y-[calc(-50%_-_2px)]" />
<TooltipPrimitive.Arrow className="size-2.5 translate-y-[calc(-50%_-_2px)] rotate-45 rounded-[2px] data-[side=left]:translate-x-[-1.5px] data-[side=right]:translate-x-[1.5px] bg-foreground fill-foreground z-[999999] translate-y-[calc(-50%_-_2px)]" />
</TooltipPrimitive.Content>
</TooltipPrimitive.Portal>
);

View file

@ -89,6 +89,13 @@ function formatUpdatedDate(timestamp: number | null): string {
return new Date(timestamp * 1000).toLocaleDateString();
}
function normalizeSliceInput(value: string): string | null {
const trimmed = value.trim();
if (!trimmed) return null;
if (!/^\d+$/.test(trimmed)) return null;
return trimmed;
}
export function DatasetSection() {
const {
dataset,
@ -641,6 +648,19 @@ export function DatasetSection() {
datasetEvalSplit={datasetEvalSplit}
setDatasetEvalSplit={setDatasetEvalSplit}
/>
) : !selectedDatasetName ? (
<HfDatasetSubsetSplitSelectors
variant="studio"
enabled={false}
datasetName={null}
accessToken={hfToken || undefined}
datasetSubset={datasetSubset}
setDatasetSubset={setDatasetSubset}
datasetSplit={datasetSplit}
setDatasetSplit={setDatasetSplit}
datasetEvalSplit={datasetEvalSplit}
setDatasetEvalSplit={setDatasetEvalSplit}
/>
) : datasetSource === "upload" && selectedLocalDataset ? (
<div className="rounded-lg border bg-muted/20 px-3.5 py-3">
<div className="mb-3 flex items-center justify-between gap-3">
@ -770,11 +790,14 @@ export function DatasetSection() {
</Tooltip>
</span>
<Input
type="number"
inputMode="numeric"
min={0}
step={1}
placeholder="0"
value={datasetSliceStart ?? ""}
onChange={(e) =>
setDatasetSliceStart(e.target.value || null)
setDatasetSliceStart(normalizeSliceInput(e.target.value))
}
/>
</div>
@ -802,11 +825,14 @@ export function DatasetSection() {
</Tooltip>
</span>
<Input
type="number"
inputMode="numeric"
min={0}
step={1}
placeholder="End"
value={datasetSliceEnd ?? ""}
onChange={(e) =>
setDatasetSliceEnd(e.target.value || null)
setDatasetSliceEnd(normalizeSliceInput(e.target.value))
}
/>
</div>

View file

@ -12,7 +12,7 @@ function parseSliceValue(value: string | null): number | null {
const trimmed = value.trim();
if (!trimmed) return null;
const num = Number(trimmed);
if (!Number.isFinite(num) || !Number.isInteger(num)) return null;
if (!Number.isFinite(num) || !Number.isInteger(num) || num < 0) return null;
return num;
}

View file

@ -67,8 +67,8 @@ export interface ModelConfigResponse {
config?: BackendModelConfig | null;
is_vision: boolean;
is_embedding?: boolean;
is_audio: boolean;
is_lora: boolean;
is_audio?: boolean;
base_model?: string | null;
model_type?: "text" | "vision" | "audio" | "embeddings" | null;
}

View file

@ -56,6 +56,8 @@ export function HfDatasetSubsetSplitSelectors({
} = useHfDatasetSplits(enabled ? datasetName : null, datasetSubset, {
accessToken,
});
const showPlaceholderDropdowns =
variant === "studio" && !enabled && !datasetName;
// Auto-select subset and split in one pass to avoid racing effects
useEffect(() => {
@ -83,12 +85,48 @@ export function HfDatasetSubsetSplitSelectors({
setDatasetSplit,
]);
if (!enabled || !datasetName) return null;
const showDropdowns = !isLoading && !error && hfSubsets.length > 0;
return (
<>
{showPlaceholderDropdowns && (
<>
<div className="grid gap-3 sm:grid-cols-2">
<SelectorDropdown
variant={variant}
label="Subset"
tooltip="Select which subset (config) of the dataset to use."
value={null}
onChange={setDatasetSubset}
options={[]}
placeholder="Select a subset..."
disabled={true}
/>
<SelectorDropdown
variant={variant}
label="Train Split"
tooltip="Select which split to use for training."
value={null}
onChange={setDatasetSplit}
options={[]}
placeholder="Select a split..."
disabled={true}
/>
</div>
<SelectorDropdown
variant={variant}
label="Eval Split"
tooltip="Select which split to use for evaluation. None means no evaluation during training."
value={null}
onChange={setDatasetEvalSplit}
options={[]}
placeholder="None"
allowNone
disabled={true}
/>
</>
)}
{isLoading && (
<div
className={
@ -184,6 +222,7 @@ function SelectorDropdown({
options,
placeholder,
allowNone = false,
disabled = false,
}: {
variant: "wizard" | "studio";
label: string;
@ -193,7 +232,11 @@ function SelectorDropdown({
options: string[];
placeholder: string;
allowNone?: boolean;
disabled?: boolean;
}) {
const selectValue =
value ?? (allowNone && !disabled ? "_none" : undefined);
if (variant === "wizard") {
return (
<Field>
@ -217,8 +260,9 @@ function SelectorDropdown({
</Tooltip>
</FieldLabel>
<Select
value={value ?? "_none"}
value={selectValue}
onValueChange={(v) => onChange(v === "_none" ? null : v)}
disabled={disabled}
>
<SelectTrigger className="w-full">
<SelectValue placeholder={placeholder} />
@ -260,8 +304,9 @@ function SelectorDropdown({
</Tooltip>
</span>
<Select
value={value ?? "_none"}
value={selectValue}
onValueChange={(v) => onChange(v === "_none" ? null : v)}
disabled={disabled}
>
<SelectTrigger className="w-full">
<SelectValue placeholder={placeholder} />

View file

@ -37,6 +37,7 @@ const initialState: TrainingConfigState = {
isCheckingVision: false,
isVisionModel: false,
isEmbeddingModel: false,
isAudioModel: false,
isLoadingModelDefaults: false,
modelDefaultsError: null,
modelDefaultsAppliedFor: null,
@ -60,6 +61,7 @@ const NON_PERSISTED_STATE_KEYS: ReadonlySet<keyof TrainingConfigState> = new Set
"modelType",
"isCheckingVision",
"isEmbeddingModel",
"isAudioModel",
"isLoadingModelDefaults",
"modelDefaultsError",
"modelDefaultsAppliedFor",
@ -129,6 +131,16 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
patch.trainOnCompletions = false;
}
const isAudio = !!modelDetails.is_audio;
// Pure audio model → always uncheck trainOnCompletions.
if (isAudio && !modelDetails.is_vision) {
patch.trainOnCompletions = false;
}
// Audio-capable vision model (e.g. gemma3n) + audio dataset → uncheck.
if (isAudio && modelDetails.is_vision && get().isDatasetAudio) {
patch.trainOnCompletions = false;
}
// Use backend-provided model_type when available, otherwise
// infer from capability flags.
const isEmbedding = !!modelDetails.is_embedding;
@ -140,6 +152,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
modelType: inferredModelType,
isVisionModel: modelDetails.is_vision,
isEmbeddingModel: isEmbedding,
isAudioModel: isAudio,
isLoadingModelDefaults: false,
isCheckingVision: false,
modelDefaultsError: null,
@ -152,6 +165,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
set({
isLoadingModelDefaults: false,
isAudioModel: false,
modelDefaultsError:
error instanceof Error
? error.message
@ -165,12 +179,13 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
set({
modelType: isVision ? "vision" : "text",
isVisionModel: isVision,
isAudioModel: false,
isCheckingVision: false,
});
})
.catch(() => {
if (get().selectedModel !== modelName) return;
set({ isCheckingVision: false });
set({ isCheckingVision: false, isAudioModel: false });
});
});
};
@ -198,10 +213,18 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
isCheckingDataset: false,
};
if (!_trainOnCompletionsManuallySet) {
const { isVisionModel } = get();
const { isVisionModel, isAudioModel } = get();
if (isVisionModel && isImage) {
updates.trainOnCompletions = false;
}
// Pure audio model → always uncheck regardless of dataset.
if (isAudioModel && !isVisionModel) {
updates.trainOnCompletions = false;
}
// Audio-capable vision model (e.g. gemma3n) + audio dataset → uncheck.
if (isAudioModel && isVisionModel && isAudio) {
updates.trainOnCompletions = false;
}
}
set(updates);
})
@ -238,6 +261,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
isCheckingVision: false,
isVisionModel: false,
isEmbeddingModel: false,
isAudioModel: false,
isDatasetAudio: false,
isLoadingModelDefaults: false,
modelDefaultsError: null,
@ -255,6 +279,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
isCheckingVision: false,
isVisionModel: false,
isEmbeddingModel: false,
isAudioModel: false,
isDatasetAudio: false,
isLoadingModelDefaults: false,
modelDefaultsError: null,

View file

@ -61,6 +61,7 @@ export interface TrainingConfigState {
isCheckingVision: boolean;
isVisionModel: boolean;
isEmbeddingModel: boolean;
isAudioModel: boolean;
isLoadingModelDefaults: boolean;
modelDefaultsError: string | null;
modelDefaultsAppliedFor: string | null;