diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index a9fbe659b3..500bc9e706 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -217,6 +217,7 @@ class ExportOrchestrator: max_seq_length: int = 2048, load_in_4bit: bool = True, trust_remote_code: bool = False, + hf_token: Optional[str] = None, ) -> Tuple[bool, str]: """Load a checkpoint for export. @@ -227,6 +228,7 @@ class ExportOrchestrator: "max_seq_length": max_seq_length, "load_in_4bit": load_in_4bit, "trust_remote_code": trust_remote_code, + "hf_token": hf_token, } # Always kill existing subprocess and spawn fresh. diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index edf5b666a3..fd43fdb90d 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -3,6 +3,19 @@ import { SectionCard } from "@/components/section-card"; import { Button } from "@/components/ui/button"; +import { + Combobox, + ComboboxContent, + ComboboxEmpty, + ComboboxInput, + ComboboxItem, + ComboboxList, +} from "@/components/ui/combobox"; +import { + InputGroup, + InputGroupAddon, + InputGroupInput, +} from "@/components/ui/input-group"; import { Select, SelectContent, @@ -11,17 +24,34 @@ import { SelectValue, } from "@/components/ui/select"; import { Separator } from "@/components/ui/separator"; +import { Switch } from "@/components/ui/switch"; import { Spinner } from "@/components/ui/spinner"; import { Tooltip, TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { useTrainingConfigStore } from "@/features/training"; -import { AlertCircleIcon, InformationCircleIcon, PackageIcon } from "@hugeicons/core-free-icons"; +import { + listLocalModels, + type LocalModelInfo, + useTrainingConfigStore, +} from "@/features/training"; +import { + useDebouncedValue, + useHfModelSearch, + useHfTokenValidation, +} from "@/hooks"; +import { + AlertCircleIcon, + FolderSearchIcon, + InformationCircleIcon, + Key01Icon, + PackageIcon, + Search01Icon, +} from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { AnimatePresence, motion } from "motion/react"; -import { useCallback, useEffect, useMemo, useState } from "react"; +import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; import { collapseAnim } from "./anim"; import type { ModelCheckpoints } from "./api/export-api"; @@ -60,6 +90,21 @@ export function ExportPage() { const [selectedModelIdx, setSelectedModelIdx] = useState(null); const [checkpoint, setCheckpoint] = useState(null); + const [sourceMode, setSourceMode] = useState<"checkpoint" | "model">( + "checkpoint", + ); + const [modelSource, setModelSource] = useState<"hf" | "local">("hf"); + const [hfExportTrustRemoteCode, setHfExportTrustRemoteCode] = + useState(true); + const [modelInput, setModelInput] = useState(""); + const [selectedSourceModel, setSelectedSourceModel] = useState( + null, + ); + const [localModelInput, setLocalModelInput] = useState(""); + const [localModels, setLocalModels] = useState([]); + const [isLoadingLocalModels, setIsLoadingLocalModels] = useState(true); + const [localModelsError, setLocalModelsError] = useState(null); + const debouncedModelQuery = useDebouncedValue(modelInput); const [exportMethod, setExportMethod] = useState(null); const [quantLevels, setQuantLevels] = useState([]); @@ -74,6 +119,9 @@ export function ExportPage() { const [exportError, setExportError] = useState(null); const [exportSuccess, setExportSuccess] = useState(false); + const hfComboboxAnchorRef = useRef(null); + const localComboboxAnchorRef = useRef(null); + const tour = useGuidedTourController({ id: "export", steps: exportTourSteps, @@ -105,6 +153,27 @@ export function ExportPage() { }; }, []); + // ---- Fetch local models for direct export ---- + useEffect(() => { + const controller = new AbortController(); + void listLocalModels(controller.signal) + .then((models) => { + if (controller.signal.aborted) return; + setLocalModels(models); + }) + .catch((error) => { + if (controller.signal.aborted) return; + setLocalModelsError( + error instanceof Error ? error.message : "Failed to load local models", + ); + }) + .finally(() => { + if (controller.signal.aborted) return; + setIsLoadingLocalModels(false); + }); + return () => controller.abort(); + }, []); + // ---- Derived state ---- const selectedModelData = useMemo( () => @@ -127,6 +196,83 @@ export function ExportPage() { const trainingMethodLabel = selectedModelData?.peft_type ? "LoRA / QLoRA" : "Full Fine-tune"; + const sourceBaseModelName = sourceMode === "model" + ? selectedSourceModel ?? "—" + : baseModelName; + + const { + results: hfResults, + isLoading: isLoadingHfModels, + error: hfSearchError, + } = useHfModelSearch(debouncedModelQuery, { + accessToken: hfToken || undefined, + excludeGguf: true, + }); + const { error: tokenValidationError, isChecking: isCheckingToken } = + useHfTokenValidation(hfToken); + + const hfResultIds = useMemo(() => { + const ids = hfResults.map((r) => r.id); + if ( + selectedSourceModel && + modelSource === "hf" && + !ids.includes(selectedSourceModel) + ) { + ids.push(selectedSourceModel); + } + return ids; + }, [hfResults, modelSource, selectedSourceModel]); + + const exportableLocalModels = useMemo( + () => + localModels.filter((m) => { + if (m.path.endsWith(".gguf")) return false; + if (m.id.toLowerCase().includes("-gguf")) return false; + return true; + }), + [localModels], + ); + + const localMetaById = useMemo(() => { + const map = new Map(); + for (const model of exportableLocalModels) map.set(model.id, model); + return map; + }, [exportableLocalModels]); + + const localResultIds = useMemo(() => { + const ids = exportableLocalModels.map((model) => model.id); + const manual = localModelInput.trim(); + if (manual && !ids.includes(manual)) { + ids.unshift(manual); + } + return ids; + }, [exportableLocalModels, localModelInput]); + + const localFilteredIds = useMemo(() => { + const q = localModelInput.trim().toLowerCase(); + if (!q) return localResultIds; + return localResultIds.filter((id) => { + const meta = localMetaById.get(id); + if (id.toLowerCase().includes(q)) return true; + if (meta?.display_name.toLowerCase().includes(q)) return true; + if (meta?.path.toLowerCase().includes(q)) return true; + return false; + }); + }, [localMetaById, localModelInput, localResultIds]); + + const exportGuideSteps = useMemo( + () => + sourceMode === "model" + ? [ + "Select a Hugging Face or local model to export from", + "GGUF is used for non-finetuned model exports", + "Pick one or more GGUF quantization levels", + "Click Export and choose your destination", + "Test your model and compare outputs in Chat", + ] + : GUIDE_STEPS, + [sourceMode], + ); // Reset checkpoint when the selected model changes useEffect(() => { @@ -144,6 +290,25 @@ export function ExportPage() { } }, [isAdapter, isQuantized, exportMethod]); + const handleSourceModeSwitch = useCallback( + (next: "checkpoint" | "model") => { + setSourceMode(next); + if (next === "model") { + setExportMethod("gguf"); + } + setSelectedSourceModel(null); + setLocalModelInput(""); + setModelInput(""); + }, + [], + ); + + useEffect(() => { + setSelectedSourceModel(null); + setLocalModelInput(""); + setModelInput(""); + }, [modelSource]); + const handleMethodChange = (method: ExportMethod) => { setExportMethod(method); if (method !== "gguf") { @@ -152,19 +317,24 @@ export function ExportPage() { }; const estimatedSize = getEstimatedSize(exportMethod, quantLevels); - const canExport = - checkpoint && + const selectedExportSource = + sourceMode === "checkpoint" ? checkpoint : selectedSourceModel; + const canExport = !!( + selectedExportSource && exportMethod && - (exportMethod !== "gguf" || quantLevels.length > 0); + (exportMethod !== "gguf" || quantLevels.length > 0) + ); // ---- Export handler ---- const handleExport = useCallback(async () => { - if (!checkpoint) return; + const source = sourceMode === "checkpoint" ? checkpoint : selectedSourceModel; + if (!source) return; - const selectedCp = checkpointsForModel.find( - (cp) => cp.display_name === checkpoint, - ); - if (!selectedCp) return; + const selectedCp = sourceMode === "checkpoint" + ? checkpointsForModel.find((cp) => cp.display_name === checkpoint) + : null; + if (sourceMode === "checkpoint" && !selectedCp) return; + const checkpointPath = selectedCp?.path; setExporting(true); setExportError(null); @@ -174,7 +344,8 @@ export function ExportPage() { // For other formats, nest under training-run/checkpoint const saveDir = exportMethod === "gguf" - ? `${baseModelName.split("/").pop() ?? selectedModelIdx ?? "model"}-finetune-gguf` + ? `${(sourceBaseModelName.split("/").pop() ?? selectedModelIdx ?? "model") + .replace(/[^a-zA-Z0-9._-]/g, "-")}-gguf` : `${selectedModelIdx ?? "model"}/${checkpoint}`; const pushToHub = destination === "hub"; const repoId = pushToHub && hfUsername && modelName @@ -183,8 +354,18 @@ export function ExportPage() { const token = pushToHub && hfToken ? hfToken : undefined; try { - // 1. Load checkpoint - await loadCheckpoint({ checkpoint_path: selectedCp.path }); + // 1. Load model source + if (sourceMode === "checkpoint") { + if (!checkpointPath) return; + await loadCheckpoint({ checkpoint_path: checkpointPath }); + } else { + await loadCheckpoint({ + checkpoint_path: source, + load_in_4bit: false, + trust_remote_code: + modelSource === "hf" ? hfExportTrustRemoteCode : true, + }); + } // 2. Run export based on method if (exportMethod === "merged") { @@ -242,16 +423,21 @@ export function ExportPage() { }, [ checkpoint, checkpointsForModel, + sourceMode, + selectedSourceModel, selectedModelIdx, selectedModelData, exportMethod, isAdapter, + sourceBaseModelName, quantLevels, destination, hfUsername, modelName, hfToken, privateRepo, + modelSource, + hfExportTrustRemoteCode, ]); // ---- Render ---- @@ -265,14 +451,14 @@ export function ExportPage() { Export Model

- Export your fine-tuned model for deployment + Export fine-tuned or base models for deployment

} title="Export Configuration" - description="Select checkpoint, method, and quantization" + description="Select source, method, and quantization" accent="emerald" featured={true} className="shadow-border ring-1 ring-border" @@ -296,11 +482,10 @@ export function ExportPage() { <> {/* Top row: Dropdowns + metadata | Guide */}
-
- {/* Training run dropdown */} -
+
+
+ + ); + })} + + +
- {/* Checkpoint dropdown */} -
- + - - - - - {checkpointsForModel.map((cp) => ( - - - {cp.display_name} - {cp.loss != null && ( - - loss: {cp.loss.toFixed(4)} + + + {checkpointsForModel.map((cp) => ( + + + {cp.display_name} + {cp.loss != null && ( + + loss: {cp.loss.toFixed(4)} + + )} - )} - - - ))} - - -
+ + ))} + + +
+ + ) : ( + +
+ + +
-
- - Training Info - -
-
- Base Model - {baseModelName} -
-
- Method - - {trainingMethodLabel} - -
-
- Checkpoints - - {checkpointsForModel.length} - -
- {isAdapter && ( -
- LoRA Rank - {loraRank} + {modelSource === "hf" ? ( + <> +
+ +
+ { + setModelInput(val); + setSelectedSourceModel(null); + }} + itemToStringValue={(id) => id} + autoHighlight={true} + > + + + + + + + {isLoadingHfModels ? ( +
+ Searching… +
+ ) : ( + No models found + )} + + {(id: string) => ( + + + {id} + + + )} + + + +
+ {(tokenValidationError ?? hfSearchError) && ( +

+ {tokenValidationError ?? hfSearchError} +

+ )} +
+
+ + + + + + + + Loads custom Python from the repo if the model + needs it. Turn off if you do not trust the + source. + + +
+
+ + + + + + setHfToken(e.target.value)} + /> + + {isCheckingToken && ( +

Checking token…

+ )} +
+ + ) : ( +
+ +
+ { + const next = id ?? ""; + setLocalModelInput(next); + setSelectedSourceModel(next || null); + }} + onInputValueChange={setLocalModelInput} + itemToStringValue={(id) => id} + autoHighlight={true} + > + + setSelectedSourceModel(localModelInput.trim() || null) + } + onKeyDown={(event) => { + if (event.key !== "Enter") return; + event.preventDefault(); + setSelectedSourceModel(localModelInput.trim() || null); + }} + > + + + + + + {isLoadingLocalModels ? ( +
+ Scanning... +
+ ) : localModelsError ? ( +
+ {localModelsError} +
+ ) : ( + No local models found + )} + + {(id: string) => { + const model = localMetaById.get(id); + const source = + model?.source === "hf_cache" + ? "HF cache" + : "Local dir"; + return ( + + + {model?.display_name ?? id} + + + {source} + + + ); + }} + + + +
+ {isLoadingLocalModels ? ( +

+ Scanning local models... +

+ ) : localModelsError ? ( +

{localModelsError}

+ ) : ( +

+ {exportableLocalModels.length > 0 + ? `${exportableLocalModels.length} local/cached models found` + : "No local models found. Enter path manually."} +

+ )}
)} + +
+

+ Direct model exports currently support GGUF only. +

+
+ + )} + + + {sourceMode === "checkpoint" && ( +
+ + Training Info + +
+
+ Base Model + {baseModelName} +
+
+ Method + + {trainingMethodLabel} + +
+
+ Checkpoints + + {checkpointsForModel.length} + +
+ {isAdapter && ( +
+ LoRA Rank + {loraRank} +
+ )} +
-
+ )}
@@ -462,7 +915,7 @@ export function ExportPage() { Quick Guide
    - {GUIDE_STEPS.map((step, i) => ( + {exportGuideSteps.map((step, i) => (
  1. {exportMethod === "gguf" && ( - + )} @@ -530,12 +985,12 @@ export function ExportPage() {