mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-09-21 08:50:42 +02:00
fix(colab): verify model access before loading the pipeline
This commit is contained in:
Joseph Magly
@@ -6,7 +6,8 @@
|
||||
"tests/test_cli_boundaries.py",
|
||||
"tests/test_runtime_contracts.py",
|
||||
"tests/test_abliterate.py",
|
||||
"tests/test_strategies.py"
|
||||
"tests/test_strategies.py",
|
||||
"tests/test_notebook_contract.py"
|
||||
],
|
||||
"infrastructure_paths": [
|
||||
".aiwg/**",
|
||||
|
||||
+64
-20
@@ -33,8 +33,9 @@
|
||||
"**How to use:**\n",
|
||||
"1. Make sure GPU runtime is enabled: `Runtime > Change runtime type > T4 GPU`\n",
|
||||
"2. Set your model and method in the config cell below\n",
|
||||
"3. Run All (`Runtime > Run all` or `Ctrl+F9`)\n",
|
||||
"4. Download the abliterated model from the output"
|
||||
"3. For gated models, complete the Hugging Face access setup in section 2 before running.\n",
|
||||
"4. Run All (`Runtime > Run all` or `Ctrl+F9`)\n",
|
||||
"5. Download the abliterated model from the output"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -61,9 +62,15 @@
|
||||
"id": "config-header"
|
||||
},
|
||||
"source": [
|
||||
"## 2. Configure\n",
|
||||
"## 2. Configure and prepare model access\n",
|
||||
"\n",
|
||||
"Edit the cell below to set your target model and abliteration method."
|
||||
"Edit the cell below to set your target model and abliteration method.\n",
|
||||
"\n",
|
||||
"**The default Llama model is gated.** Open its [model page](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct), accept the license/request access, and wait for approval on your Hugging Face account. A token alone does not grant access.\n",
|
||||
"\n",
|
||||
"Create a **read** token for that account (or a fine-grained token permitted to read the selected model). In Colab, open **Secrets** (the key icon), add a secret named `HF_TOKEN`, and enable **Notebook access**. Never paste a token into a code cell or saved output. Outside Colab, an existing `hf auth login` session or an `HF_TOKEN` environment variable also works.\n",
|
||||
"\n",
|
||||
"Ungated public models such as `Qwen/Qwen2.5-7B-Instruct` or `openai-community/gpt2` need no login. Step 3 checks access to the selected model's small `config.json` before constructing the pipeline or downloading weights; if access cannot be confirmed, it stops with setup instructions. After fixing access or changing the model, rerun the configuration and step 3. The check runs again even when step 3 is executed directly.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -115,6 +122,38 @@
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"from huggingface_hub import get_token, hf_hub_download\n",
|
||||
"from obliteratus.credential_sources import resolve_secret\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def check_model_access(model_name):\n",
|
||||
" # Resolve the same configured sources as the loader, then Hub/Colab login.\n",
|
||||
" # Only the small config is fetched; force a live check even if it is cached.\n",
|
||||
" try:\n",
|
||||
" token = resolve_secret(\"HF_TOKEN\") or get_token()\n",
|
||||
" hf_hub_download(\n",
|
||||
" repo_id=model_name, filename=\"config.json\",\n",
|
||||
" token=token or False, force_download=True,\n",
|
||||
" )\n",
|
||||
" except Exception:\n",
|
||||
" # Hub/credential exceptions may contain sensitive request details.\n",
|
||||
" raise RuntimeError(\n",
|
||||
" \"Model access could not be confirmed. Check the model ID and your \"\n",
|
||||
" \"connection. For gated/private models, accept the license/request \"\n",
|
||||
" \"access on the model page and wait for approval; use a read token \"\n",
|
||||
" \"from that account. In Colab, add HF_TOKEN in Secrets and enable \"\n",
|
||||
" \"Notebook access (outside Colab use hf auth login or HF_TOKEN). \"\n",
|
||||
" \"Then rerun this cell, or choose an ungated public model. \"\n",
|
||||
" \"No model weights have been loaded.\"\n",
|
||||
" ) from None\n",
|
||||
" if token:\n",
|
||||
" # Runtime only: keep the loader on the credential that passed the check.\n",
|
||||
" os.environ[\"HF_TOKEN\"] = token\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"check_model_access(MODEL)\n",
|
||||
"\n",
|
||||
"from obliteratus.abliterate import AbliterationPipeline\n",
|
||||
"\n",
|
||||
"# Build kwargs, only pass overrides if non-zero\n",
|
||||
@@ -217,25 +256,30 @@
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#@title Option B: Push to HuggingFace Hub\n",
|
||||
"#@markdown Set your HF repo name. You'll need to be logged in (`huggingface-cli login`).\n",
|
||||
"#@title Option B: Push to HuggingFace Hub (opt-in)\n",
|
||||
"#@markdown Enable only when ready to upload. Publishing requires a write token with permission for the destination repo; read access to the source model is separate.\n",
|
||||
"UPLOAD_TO_HUB = False #@param {type: \"boolean\"}\n",
|
||||
"HF_REPO = \"your-username/model-name-abliterated\" #@param {type: \"string\"}\n",
|
||||
"\n",
|
||||
"from huggingface_hub import HfApi\n",
|
||||
"api = HfApi()\n",
|
||||
"if UPLOAD_TO_HUB:\n",
|
||||
" from huggingface_hub import HfApi, get_token\n",
|
||||
" from obliteratus.credential_sources import resolve_secret\n",
|
||||
"\n",
|
||||
"# Login if needed\n",
|
||||
"from huggingface_hub import notebook_login\n",
|
||||
"notebook_login()\n",
|
||||
"\n",
|
||||
"# Upload\n",
|
||||
"api.create_repo(HF_REPO, exist_ok=True)\n",
|
||||
"api.upload_folder(\n",
|
||||
" folder_path=str(model_dir),\n",
|
||||
" repo_id=HF_REPO,\n",
|
||||
" repo_type=\"model\",\n",
|
||||
")\n",
|
||||
"print(f\"\\nUploaded to: https://huggingface.co/{HF_REPO}\")"
|
||||
" if HF_REPO == \"your-username/model-name-abliterated\":\n",
|
||||
" raise ValueError(\"Set HF_REPO to your destination repository before uploading.\")\n",
|
||||
" upload_token = resolve_secret(\"HF_TOKEN\") or get_token()\n",
|
||||
" if not upload_token:\n",
|
||||
" raise RuntimeError(\"Set a write token in HF_TOKEN Secrets before uploading.\")\n",
|
||||
" api = HfApi(token=upload_token)\n",
|
||||
" api.create_repo(HF_REPO, exist_ok=True)\n",
|
||||
" api.upload_folder(\n",
|
||||
" folder_path=str(model_dir),\n",
|
||||
" repo_id=HF_REPO,\n",
|
||||
" repo_type=\"model\",\n",
|
||||
" )\n",
|
||||
" print(f\"\\nUploaded to: https://huggingface.co/{HF_REPO}\")\n",
|
||||
"else:\n",
|
||||
" print(\"Hub upload skipped. Enable UPLOAD_TO_HUB to publish explicitly.\")\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -37,3 +37,189 @@ def test_abliterate_notebook_stage_callback_uses_stage_result_contract(capsys):
|
||||
|
||||
assert "STAGE: PROBE" in output
|
||||
assert "loading model" in output
|
||||
|
||||
|
||||
def _notebook_code(marker):
|
||||
notebook = json.loads(Path("notebooks/abliterate.ipynb").read_text())
|
||||
return next(
|
||||
"".join(cell["source"])
|
||||
for cell in notebook["cells"]
|
||||
if cell["cell_type"] == "code" and marker in "".join(cell["source"])
|
||||
)
|
||||
|
||||
|
||||
def _run_harness(monkeypatch, *, token=None, failure=None):
|
||||
"""Execute real notebook cells with an offline Hub and observable pipeline."""
|
||||
import sys
|
||||
from unittest.mock import Mock
|
||||
|
||||
monkeypatch.delenv("HF_TOKEN", raising=False)
|
||||
events = []
|
||||
|
||||
def download(**kwargs):
|
||||
events.append(("access", kwargs))
|
||||
if failure:
|
||||
raise failure
|
||||
return "offline-config.json"
|
||||
|
||||
def pipeline(**kwargs):
|
||||
import os
|
||||
events.append(("construct", kwargs, os.environ.get("HF_TOKEN")))
|
||||
return SimpleNamespace(run=lambda: events.append(("run",)) or "output")
|
||||
|
||||
hub = SimpleNamespace(
|
||||
get_token=Mock(return_value=token), hf_hub_download=download,
|
||||
HfApi=Mock(side_effect=AssertionError("Unexpected upload")),
|
||||
notebook_login=Mock(side_effect=AssertionError("Unexpected login widget")),
|
||||
)
|
||||
secret = Mock(return_value=None)
|
||||
monkeypatch.setitem(sys.modules, "huggingface_hub", hub)
|
||||
monkeypatch.setitem(sys.modules, "obliteratus.credential_sources", SimpleNamespace(resolve_secret=secret))
|
||||
monkeypatch.setitem(sys.modules, "obliteratus.abliterate", SimpleNamespace(AbliterationPipeline=pipeline))
|
||||
namespace = {}
|
||||
exec(_notebook_code("#@title Abliteration Config"), namespace)
|
||||
return namespace, events, hub, secret
|
||||
|
||||
|
||||
def test_default_gated_model_stops_before_pipeline_without_access(monkeypatch, capsys):
|
||||
import traceback
|
||||
import pytest
|
||||
|
||||
namespace, events, hub, _ = _run_harness(
|
||||
monkeypatch, failure=PermissionError("private request details fake-secret-value"),
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="HF_TOKEN") as caught:
|
||||
exec(_notebook_code("def check_model_access"), namespace)
|
||||
assert [event[0] for event in events] == ["access"]
|
||||
assert events[0][1] == {
|
||||
"repo_id": "meta-llama/Llama-3.1-8B-Instruct", "filename": "config.json",
|
||||
"token": False, "force_download": True,
|
||||
}
|
||||
rendered = "".join(traceback.format_exception(caught.type, caught.value, caught.tb))
|
||||
assert "fake-secret-value" not in rendered + capsys.readouterr().out
|
||||
assert "wait for approval" in str(caught.value)
|
||||
hub.notebook_login.assert_not_called()
|
||||
|
||||
|
||||
def test_authorized_colab_token_reaches_loader_without_output(monkeypatch, capsys):
|
||||
namespace, events, _, _ = _run_harness(monkeypatch, token="fake-private-read-token")
|
||||
exec(_notebook_code("def check_model_access"), namespace)
|
||||
assert [event[0] for event in events] == ["access", "construct", "run"]
|
||||
assert events[0][1]["token"] == "fake-private-read-token"
|
||||
assert events[1][2] == "fake-private-read-token"
|
||||
assert "fake-private-read-token" not in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_ungated_model_runs_anonymously_and_upload_is_opt_in(monkeypatch):
|
||||
namespace, events, hub, _ = _run_harness(monkeypatch)
|
||||
namespace["MODEL"] = "Qwen/Qwen2.5-7B-Instruct"
|
||||
exec(_notebook_code("def check_model_access"), namespace)
|
||||
exec(_notebook_code("UPLOAD_TO_HUB ="), namespace)
|
||||
assert [event[0] for event in events] == ["access", "construct", "run"]
|
||||
assert events[0][1]["token"] is False
|
||||
assert events[1][2] is None
|
||||
hub.HfApi.assert_not_called()
|
||||
hub.notebook_login.assert_not_called()
|
||||
|
||||
|
||||
def test_direct_rerun_rechecks_changed_model_and_blocks_expired_access(monkeypatch):
|
||||
import pytest
|
||||
|
||||
namespace, events, hub, _ = _run_harness(monkeypatch, token="fake-token")
|
||||
source = _notebook_code("def check_model_access")
|
||||
exec(source, namespace)
|
||||
namespace["MODEL"] = "other/private-model"
|
||||
|
||||
def denied(**kwargs):
|
||||
events.append(("access", kwargs))
|
||||
raise PermissionError("expired")
|
||||
|
||||
hub.hf_hub_download = denied
|
||||
with pytest.raises(RuntimeError, match="Model access"):
|
||||
exec(source, namespace)
|
||||
assert [event[0] for event in events] == ["access", "construct", "run", "access"]
|
||||
assert events[-1][1]["repo_id"] == "other/private-model"
|
||||
|
||||
|
||||
def test_configured_secret_precedes_hub_cached_token(monkeypatch):
|
||||
namespace, events, hub, secret = _run_harness(monkeypatch, token="fake-other-account")
|
||||
secret.return_value = "fake-configured-token"
|
||||
exec(_notebook_code("def check_model_access"), namespace)
|
||||
assert events[0][1]["token"] == "fake-configured-token"
|
||||
hub.get_token.assert_not_called()
|
||||
|
||||
|
||||
def test_secret_resolution_failure_cannot_fall_back_or_run(monkeypatch):
|
||||
import pytest
|
||||
|
||||
namespace, events, hub, secret = _run_harness(monkeypatch, token="fake-cached-token")
|
||||
secret.side_effect = RuntimeError("sensitive credential source")
|
||||
with pytest.raises(RuntimeError, match="Model access"):
|
||||
exec(_notebook_code("def check_model_access"), namespace)
|
||||
assert events == []
|
||||
hub.get_token.assert_not_called()
|
||||
|
||||
|
||||
def test_notebook_preflight_network_failure_stops_before_loading(monkeypatch):
|
||||
import pytest
|
||||
|
||||
namespace, events, _, _ = _run_harness(monkeypatch, failure=ConnectionError("offline"))
|
||||
with pytest.raises(RuntimeError, match="connection"):
|
||||
exec(_notebook_code("def check_model_access"), namespace)
|
||||
assert [event[0] for event in events] == ["access"]
|
||||
|
||||
|
||||
def test_notebook_has_no_persisted_execution_outputs():
|
||||
notebook = json.loads(Path("notebooks/abliterate.ipynb").read_text())
|
||||
for cell in notebook["cells"]:
|
||||
if cell["cell_type"] == "code":
|
||||
assert cell.get("outputs", []) == []
|
||||
assert cell.get("execution_count") is None
|
||||
|
||||
|
||||
def test_token_without_gated_approval_cannot_run(monkeypatch):
|
||||
import pytest
|
||||
|
||||
namespace, events, _, _ = _run_harness(
|
||||
monkeypatch, token="fake-unapproved-token", failure=PermissionError("403"),
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="wait for approval"):
|
||||
exec(_notebook_code("def check_model_access"), namespace)
|
||||
assert [event[0] for event in events] == ["access"]
|
||||
|
||||
|
||||
def test_explicit_upload_uses_resolved_token_and_destination(monkeypatch):
|
||||
from unittest.mock import Mock
|
||||
|
||||
namespace, _, hub, secret = _run_harness(monkeypatch)
|
||||
secret.return_value = "fake-write-token"
|
||||
api = Mock()
|
||||
hub.HfApi = Mock(return_value=api)
|
||||
namespace["model_dir"] = Path("saved-model")
|
||||
source = _notebook_code("UPLOAD_TO_HUB =").replace(
|
||||
"UPLOAD_TO_HUB = False", "UPLOAD_TO_HUB = True",
|
||||
).replace(
|
||||
'HF_REPO = "your-username/model-name-abliterated"',
|
||||
'HF_REPO = "researcher/output"',
|
||||
)
|
||||
exec(source, namespace)
|
||||
hub.HfApi.assert_called_once_with(token="fake-write-token")
|
||||
api.create_repo.assert_called_once_with("researcher/output", exist_ok=True)
|
||||
api.upload_folder.assert_called_once_with(
|
||||
folder_path="saved-model", repo_id="researcher/output", repo_type="model",
|
||||
)
|
||||
|
||||
|
||||
def test_explicit_upload_stops_without_credentials(monkeypatch):
|
||||
import pytest
|
||||
|
||||
namespace, _, hub, _ = _run_harness(monkeypatch)
|
||||
source = _notebook_code("UPLOAD_TO_HUB =").replace(
|
||||
"UPLOAD_TO_HUB = False", "UPLOAD_TO_HUB = True",
|
||||
).replace(
|
||||
'HF_REPO = "your-username/model-name-abliterated"',
|
||||
'HF_REPO = "researcher/output"',
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="write token"):
|
||||
exec(source, namespace)
|
||||
hub.HfApi.assert_not_called()
|
||||
|
||||
Reference in New Issue
Block a user