From 9e216210715808bb12e1785eb5a5192abef4d3c3 Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:31:16 -0400 Subject: [PATCH] fix(colab): verify model access before loading the pipeline --- ci/pr-test-policy.json | 3 +- notebooks/abliterate.ipynb | 84 +++++++++++---- tests/test_notebook_contract.py | 186 ++++++++++++++++++++++++++++++++ 3 files changed, 252 insertions(+), 21 deletions(-) diff --git a/ci/pr-test-policy.json b/ci/pr-test-policy.json index c780f11..48184e6 100644 --- a/ci/pr-test-policy.json +++ b/ci/pr-test-policy.json @@ -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/**", diff --git a/notebooks/abliterate.ipynb b/notebooks/abliterate.ipynb index a89c6e5..2739c03 100644 --- a/notebooks/abliterate.ipynb +++ b/notebooks/abliterate.ipynb @@ -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" ] }, { diff --git a/tests/test_notebook_contract.py b/tests/test_notebook_contract.py index 6080b30..7b8eeb1 100644 --- a/tests/test_notebook_contract.py +++ b/tests/test_notebook_contract.py @@ -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()