Studio: use one shared Hugging Face token across Settings and training (#7152)
--------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
This commit is contained in:
Nilay
GitHub
parent
3be49070cc
commit
8c83478da0
5 changed files with 56 additions and 45 deletions
|
|
@ -2,7 +2,11 @@
|
|||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import type { RememberedLoadSettings } from "@/components/assistant-ui/model-selector/remembered-load-settings";
|
||||
import { cancelStagedModelDownload } from "@/features/hub";
|
||||
import {
|
||||
cancelStagedModelDownload,
|
||||
mirrorHfTokenInto,
|
||||
useHfTokenStore,
|
||||
} from "@/features/hub";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { create } from "zustand";
|
||||
import { isExternalModelId, parseExternalModelId } from "../external-providers";
|
||||
|
|
@ -25,8 +29,6 @@ import {
|
|||
import { useExternalProvidersStore } from "./external-providers-store";
|
||||
import { PLUS_MENU_PINS_STORAGE_KEY } from "./plus-menu-prefs-store";
|
||||
|
||||
const HF_TOKEN_KEY = "unsloth_hf_token";
|
||||
const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed";
|
||||
export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled";
|
||||
export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled";
|
||||
export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled";
|
||||
|
|
@ -495,17 +497,6 @@ export function saveSpeculativeType(value: string | null): void {
|
|||
}
|
||||
}
|
||||
|
||||
function notifyHfTokenChanged(value: string): void {
|
||||
if (!canUseStorage()) return;
|
||||
try {
|
||||
window.dispatchEvent(
|
||||
new CustomEvent(HF_TOKEN_CHANGED_EVENT, { detail: value }),
|
||||
);
|
||||
} catch {
|
||||
// ignore
|
||||
}
|
||||
}
|
||||
|
||||
/** A local model staged for a deferred load (see `pendingSelection`). Shape is
|
||||
* a subset of the load hook's `SelectedModelInput`, structurally assignable. */
|
||||
export type PendingModelSelection = {
|
||||
|
|
@ -1144,7 +1135,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
runningByThreadId: {},
|
||||
cancelByThreadId: {},
|
||||
autoTitle: false,
|
||||
hfToken: loadString(HF_TOKEN_KEY, ""),
|
||||
hfToken: useHfTokenStore.getState().token,
|
||||
modelsError: null,
|
||||
lastModelLoadError: null,
|
||||
activeGgufVariant: null,
|
||||
|
|
@ -1349,11 +1340,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
setScalarSettingVersion("autoTitle", autoTitle, state.autoTitle);
|
||||
return { autoTitle };
|
||||
}),
|
||||
setHfToken: (hfToken) => {
|
||||
saveString(HF_TOKEN_KEY, hfToken);
|
||||
set({ hfToken });
|
||||
notifyHfTokenChanged(hfToken);
|
||||
},
|
||||
setHfToken: (hfToken) => useHfTokenStore.getState().setToken(hfToken),
|
||||
setModelsError: (modelsError) => set({ modelsError }),
|
||||
setLastModelLoadError: (lastModelLoadError) => set({ lastModelLoadError }),
|
||||
setCheckpoint: (modelId, ggufVariant) =>
|
||||
|
|
@ -1837,6 +1824,12 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
setContextUsage: (contextUsage) => set({ contextUsage }),
|
||||
}));
|
||||
|
||||
// Mirror token edits made through the shared store (e.g. Studio's field).
|
||||
const unsubscribeHfTokenMirror = mirrorHfTokenInto(useChatRuntimeStore);
|
||||
if (import.meta.hot) {
|
||||
import.meta.hot.dispose(unsubscribeHfTokenMirror);
|
||||
}
|
||||
|
||||
export function resolveSpeculativeSettingsForLoad({
|
||||
usePersistedPreference = false,
|
||||
}: {
|
||||
|
|
|
|||
|
|
@ -2,3 +2,8 @@
|
|||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
export { cancelStagedModelDownload } from "./download-manager";
|
||||
export {
|
||||
getHfToken,
|
||||
mirrorHfTokenInto,
|
||||
useHfTokenStore,
|
||||
} from "./stores/hf-token-store";
|
||||
|
|
|
|||
|
|
@ -5,11 +5,9 @@ import { create } from "zustand";
|
|||
import { bumpInventoryVersion } from "./inventory-events";
|
||||
|
||||
const HF_TOKEN_KEY = "unsloth_hf_token";
|
||||
const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed";
|
||||
const LEGACY_TRAINING_KEY = "unsloth_training_config_v1";
|
||||
let storageSyncStarted = false;
|
||||
let storageSyncListener: ((event: StorageEvent) => void) | null = null;
|
||||
let tokenChangedListener: ((event: Event) => void) | null = null;
|
||||
|
||||
function canUseStorage(): boolean {
|
||||
return typeof window !== "undefined";
|
||||
|
|
@ -63,10 +61,6 @@ function stopStorageSync(): void {
|
|||
window.removeEventListener("storage", storageSyncListener);
|
||||
storageSyncListener = null;
|
||||
}
|
||||
if (tokenChangedListener !== null) {
|
||||
window.removeEventListener(HF_TOKEN_CHANGED_EVENT, tokenChangedListener);
|
||||
tokenChangedListener = null;
|
||||
}
|
||||
storageSyncStarted = false;
|
||||
}
|
||||
|
||||
|
|
@ -100,10 +94,6 @@ export const useHfTokenStore = create<HfTokenStore>((set) => {
|
|||
applyToken(event.newValue ?? "", false);
|
||||
};
|
||||
window.addEventListener("storage", storageSyncListener);
|
||||
tokenChangedListener = (event) => {
|
||||
applyToken((event as CustomEvent<string>).detail ?? "", false);
|
||||
};
|
||||
window.addEventListener(HF_TOKEN_CHANGED_EVENT, tokenChangedListener);
|
||||
}
|
||||
|
||||
return {
|
||||
|
|
@ -117,6 +107,21 @@ export function getHfToken(): string {
|
|||
return useHfTokenStore.getState().token;
|
||||
}
|
||||
|
||||
// Keep a plain zustand store's `hfToken` field in sync with the shared token:
|
||||
// seed the current value, then mirror later edits. Returns the unsubscribe so
|
||||
// callers can wire it to HMR disposal.
|
||||
export function mirrorHfTokenInto<T extends { hfToken: string }>(store: {
|
||||
getState: () => T;
|
||||
setState: (partial: Partial<T>) => void;
|
||||
}): () => void {
|
||||
store.setState({ hfToken: getHfToken() } as Partial<T>);
|
||||
return useHfTokenStore.subscribe((state) => {
|
||||
if (store.getState().hfToken !== state.token) {
|
||||
store.setState({ hfToken: state.token } as Partial<T>);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// HF's JS client throws on a non-empty token that isn't `hf_...` instead of
|
||||
// browsing anonymously, so treat anything malformed as no token.
|
||||
export function hfApiToken(
|
||||
|
|
|
|||
|
|
@ -54,7 +54,7 @@ import {
|
|||
// Imported directly from the store module rather than the "@/features/training"
|
||||
// barrel to avoid an import cycle (the barrel re-exports this section's siblings).
|
||||
import { hasSeparateStreamingEvalSplit } from "@/features/training/stores/training-config-store";
|
||||
import { useDebouncedValue, useHfTokenValidation } from "@/hooks";
|
||||
import { useDebouncedValue } from "@/hooks";
|
||||
import { translate, useT } from "@/i18n";
|
||||
import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
|
||||
import { toast } from "@/lib/toast";
|
||||
|
|
@ -398,9 +398,6 @@ export function DatasetSection() {
|
|||
enabled: pickerTab === "huggingface",
|
||||
});
|
||||
|
||||
const { error: tokenValidationError, isChecking: isCheckingToken } =
|
||||
useHfTokenValidation(hfToken);
|
||||
|
||||
const hfResultIds = useMemo(() => {
|
||||
const ids = hfResults.map((r) => r.id);
|
||||
if (dataset && !ids.includes(dataset)) {
|
||||
|
|
@ -1005,9 +1002,9 @@ export function DatasetSection() {
|
|||
</ComboboxContent>
|
||||
</Combobox>
|
||||
</div>
|
||||
{(tokenValidationError ?? hfSearchError) && (
|
||||
{hfSearchError && (
|
||||
<p className="text-xs text-destructive">
|
||||
{tokenValidationError ?? hfSearchError}
|
||||
{hfSearchError}
|
||||
{" — "}
|
||||
<a
|
||||
href="https://huggingface.co/settings/tokens"
|
||||
|
|
@ -1019,11 +1016,6 @@ export function DatasetSection() {
|
|||
</a>
|
||||
</p>
|
||||
)}
|
||||
{isCheckingToken && (
|
||||
<p className="text-xs text-muted-foreground">
|
||||
{t("studio.dataset.checkingToken")}
|
||||
</p>
|
||||
)}
|
||||
{pickerTab !== activeSourceTab && (
|
||||
<p className="text-[11px] text-muted-foreground">
|
||||
{t("studio.dataset.browsingSource", {
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
import { CPT_TARGET_MODULES, DEFAULT_HYPERPARAMS, LR_DEFAULT_CPT, LR_DEFAULT_FULL, LR_DEFAULT_LORA, STEPS, TARGET_MODULES } from "@/config/training";
|
||||
import { authFetch } from "@/features/auth";
|
||||
import { getHfToken, mirrorHfTokenInto, useHfTokenStore } from "@/features/hub";
|
||||
import { isAdapterMethod } from "@/types/training";
|
||||
import type { DatasetFormat } from "@/types/training";
|
||||
import type { ModelType, StepNumber, TrainingMethod } from "@/types/training";
|
||||
|
|
@ -117,7 +118,9 @@ let _datasetFormatAutoForcedByCpt = false;
|
|||
|
||||
// modelType / isVisionModel / isAudioModel persist so multimodal-only UI
|
||||
// paints right on reload; the model-config fetch still re-derives them.
|
||||
// hfToken mirrors the shared hf-token-store and is persisted there instead.
|
||||
const NON_PERSISTED_STATE_KEYS: ReadonlySet<keyof TrainingConfigState> = new Set([
|
||||
"hfToken",
|
||||
"isCheckingVision",
|
||||
"isEmbeddingModel",
|
||||
"isLoadingModelDefaults",
|
||||
|
|
@ -632,8 +635,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
),
|
||||
);
|
||||
},
|
||||
setHfToken: (hfToken) =>
|
||||
set({ hfToken: hfToken.trim().replace(/^["']+|["']+$/g, "") }),
|
||||
setHfToken: (hfToken) => useHfTokenStore.getState().setToken(hfToken),
|
||||
setDatasetSource: (datasetSource) => set({ datasetSource }),
|
||||
selectHfDataset: (dataset) => {
|
||||
_datasetCheckController?.abort();
|
||||
|
|
@ -923,7 +925,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
_learningRateManuallySet = false;
|
||||
_yamlLearningRate = undefined;
|
||||
clearCptDatasetFormatTracking();
|
||||
set(initialState);
|
||||
set({ ...initialState, hfToken: getHfToken() });
|
||||
},
|
||||
resetToModelDefaults: () => {
|
||||
const { selectedModel } = get();
|
||||
|
|
@ -947,7 +949,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
},
|
||||
{
|
||||
name: "unsloth_training_config_v1",
|
||||
version: 11,
|
||||
version: 12,
|
||||
migrate: (persisted, version) => {
|
||||
const s = persisted as Record<string, unknown>;
|
||||
if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) {
|
||||
|
|
@ -1000,6 +1002,15 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
// own version guard.
|
||||
s.datasetStreaming ??= false;
|
||||
}
|
||||
if (version < 12) {
|
||||
// hfToken moved to the shared hf-token-store; seed it once so an
|
||||
// existing Studio-only token isn't lost.
|
||||
const legacyToken = typeof s.hfToken === "string" ? s.hfToken.trim() : "";
|
||||
if (legacyToken && !getHfToken()) {
|
||||
useHfTokenStore.getState().setToken(legacyToken);
|
||||
}
|
||||
delete s.hfToken;
|
||||
}
|
||||
return s as unknown as TrainingConfigStore;
|
||||
},
|
||||
partialize: partializePersistedState,
|
||||
|
|
@ -1022,3 +1033,8 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
},
|
||||
),
|
||||
);
|
||||
|
||||
const unsubscribeHfTokenMirror = mirrorHfTokenInto(useTrainingConfigStore);
|
||||
if (import.meta.hot) {
|
||||
import.meta.hot.dispose(unsubscribeHfTokenMirror);
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue