Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyrit/executor/promptgen/gcg/attack/base/attack_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -1819,7 +1819,7 @@ class ModelWorker:
def __init__(
self,
model_path: str,
token: str,
token: str | None,
model_kwargs: dict[str, Any],
tokenizer: Any,
device: str,
Expand Down
5 changes: 4 additions & 1 deletion pyrit/executor/promptgen/gcg/generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -410,7 +410,10 @@ def _to_attack_params(self, *, context: GCGContext) -> Any:

all_models = self._models + self._test_models
return SimpleNamespace(
token=self._hf_token or "",
# None means anonymous access. An empty string makes huggingface_hub send
# "Authorization: Bearer " with no credential, which httpx rejects as an
# illegal header value, so public models would fail to load.
token=self._hf_token,
tokenizer_paths=[m.name for m in all_models],
tokenizer_kwargs=[m.tokenizer_kwargs for m in all_models],
model_paths=[m.name for m in all_models],
Expand Down
21 changes: 21 additions & 0 deletions tests/unit/executor/promptgen/gcg/test_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,27 @@ def test_test_length_mismatch_raises(self) -> None:
)


class TestToAttackParamsToken:
"""The HuggingFace token forwarded to model loading."""

def test_token_is_none_when_not_supplied(self) -> None:
# An empty string makes huggingface_hub send "Authorization: Bearer " with no
# credential, which httpx rejects as an illegal header value. None means anonymous,
# which is what loading a public model needs.
gen = GCGGenerator(models=[GCGModelConfig(name=_LLAMA_2)])

params = gen._to_attack_params(context=GCGContext(goals=["g"], targets=["t"]))

assert params.token is None

def test_token_is_forwarded_when_supplied(self) -> None:
gen = GCGGenerator(models=[GCGModelConfig(name=_LLAMA_2)], hf_token="hf_example")

params = gen._to_attack_params(context=GCGContext(goals=["g"], targets=["t"]))

assert params.token == "hf_example"


@pytest.fixture
def patched_get_workers():
"""Patch the heavy worker spawn so lifecycle tests don't try to load real models."""
Expand Down