feat(secrets): add provider-neutral credential hooks

This commit is contained in:
Joseph Magly
2026-08-22 10:02:49 -04:00
parent 2c7f6d78f1
commit 3390fe551b
12 changed files with 530 additions and 24 deletions
+10 -12
View File
@@ -58,6 +58,7 @@ if "HF_HOME" not in os.environ:
import gradio as gr
import torch
from obliteratus import device as dev
from obliteratus.credential_sources import resolve_first, resolve_secret, secret_available
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
# ── ZeroGPU support ─────────────────────────────────────────────────
@@ -347,7 +348,6 @@ METHODS = {
# Shared org + token so users can auto-push without their own HF_TOKEN.
# Set OBLITERATUS_HUB_TOKEN as a Space secret with write access to the org.
_HUB_COMMUNITY_ORG = os.environ.get("OBLITERATUS_HUB_ORG", "OBLITERATUS")
_HUB_COMMUNITY_TOKEN = os.environ.get("OBLITERATUS_HUB_TOKEN")
# Import preset configs for Advanced Settings defaults
from obliteratus.abliterate import METHODS as _PRESET_CONFIGS # noqa: E402
@@ -466,7 +466,6 @@ def _on_dataset_change(dataset_label: str):
def _validate_hub_repo(hub_repo: str) -> str:
"""Validate Hub repo ID format and check HF_TOKEN. Returns warning HTML or empty string."""
import os
import re
repo = hub_repo.strip() if hub_repo else ""
if not repo:
@@ -477,7 +476,7 @@ def _validate_hub_repo(hub_repo: str) -> str:
"Invalid repo format — use `username/model-name` "
"(letters, numbers, hyphens, dots only)"
)
if not os.environ.get("HF_TOKEN") and not os.environ.get("HF_PUSH_TOKEN") and not _HUB_COMMUNITY_TOKEN:
if not secret_available("HF_TOKEN", "HF_PUSH_TOKEN", "OBLITERATUS_HUB_TOKEN"):
warnings.append(
"No Hub token available — push will fail. "
"Set HF_PUSH_TOKEN, HF_TOKEN, or OBLITERATUS_HUB_TOKEN."
@@ -591,7 +590,6 @@ def push_session_to_hub(
progress=gr.Progress(),
):
"""Push a session model to HuggingFace Hub, with optional refinement."""
import os
import re
if not session_label or session_label.startswith("("):
@@ -622,7 +620,7 @@ def push_session_to_hub(
# Resolve token
token = hub_token_input.strip() if hub_token_input else None
if not token:
token = os.environ.get("HF_PUSH_TOKEN") or os.environ.get("HF_TOKEN") or _HUB_COMMUNITY_TOKEN
token = resolve_first("HF_PUSH_TOKEN", "HF_TOKEN", "OBLITERATUS_HUB_TOKEN")
if not token:
yield (
"**Error:** No Hub token available. Enter a token above, "
@@ -720,7 +718,7 @@ def _should_quantize(model_id: str, is_preset: bool = False) -> str | None:
try:
from obliteratus.models.loader import _estimate_model_memory_gb, _available_gpu_memory_gb
from transformers import AutoConfig
token = os.environ.get("HF_TOKEN") or None
token = resolve_secret("HF_TOKEN")
config = AutoConfig.from_pretrained(model_id, trust_remote_code=is_preset, token=token)
# Skip if model already ships with native quantization (e.g. Mxfp4Config)
if getattr(config, "quantization_config", None) is not None:
@@ -1835,8 +1833,6 @@ def obliterate(model_choice: str, method_choice: str,
5 minutes). The @spaces.GPU decorator allocates a GPU at call time and
releases it when the function returns.
"""
import os
model_id = MODELS.get(model_choice, model_choice)
is_preset = model_choice in MODELS
method = METHODS.get(method_choice, "advanced")
@@ -1874,7 +1870,7 @@ def obliterate(model_choice: str, method_choice: str,
# Early validation: gated model access
from obliteratus.presets import is_gated
if is_gated(model_id) and not (os.environ.get("HF_TOKEN") or os.environ.get("HF_PUSH_TOKEN")):
if is_gated(model_id) and not secret_available("HF_TOKEN", "HF_PUSH_TOKEN"):
yield (
f"**Error: Gated model requires authentication.**\n\n"
f"`{model_id}` is a gated HuggingFace repo. To use it:\n\n"
@@ -2969,7 +2965,7 @@ def ab_chat_respond(message: str, history_left: list[dict], history_right: list[
model_id, torch_dtype=torch.float16,
trust_remote_code=is_preset,
low_cpu_mem_usage=True,
token=os.environ.get("HF_TOKEN") or None,
token=resolve_secret("HF_TOKEN"),
)
streamer_orig = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=stream_timeout)
@@ -4991,8 +4987,10 @@ To opt out, set the environment variable `OBLITERATUS_TELEMETRY=0` before launch
diag.append(f"- Telemetry enabled: `{is_enabled()}`")
diag.append(f"- On HF Spaces: `{_ON_HF_SPACES}`")
diag.append(f"- Repo: `{_TELEMETRY_REPO or '(not set)'}`")
diag.append(f"- HF_TOKEN set: `{bool(os.environ.get('HF_TOKEN'))}`")
diag.append(f"- HF_PUSH_TOKEN set: `{bool(os.environ.get('HF_PUSH_TOKEN'))}`")
diag.append(f"- HF token available: `{secret_available('HF_TOKEN')}`")
diag.append(
f"- HF push token available: `{secret_available('HF_PUSH_TOKEN')}`"
)
diag.append(f"- Local file: `{TELEMETRY_FILE}`")
diag.append(f"- Local file exists: `{TELEMETRY_FILE.exists()}`")
n_records = len(read_telemetry()) if TELEMETRY_FILE.exists() else 0