FEAT support any causal LM in GCG embedding helpers

pyrit/executor/promptgen/gcg/attack/base/attack_manager.py

@@ -23,10 +23,6 @@
 from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
 from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM
 from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
-from transformers.models.llama.modeling_llama import LlamaForCausalLM
-from transformers.models.mistral.modeling_mistral import MistralForCausalLM
-from transformers.models.mixtral.modeling_mixtral import MixtralForCausalLM
-from transformers.models.phi3.modeling_phi3 import Phi3ForCausalLM
 
 from pyrit.executor.promptgen.gcg.experiments.log import (
     log_gpu_memory,
@@ -195,69 +191,52 @@ def _get_worker_log_params(worker: ModelWorker) -> dict[str, Any]:
 
 def get_embedding_layer(model: Any) -> Any:
     """
-    Return the token embedding layer for a supported causal language model.
+    Return the token embedding layer for a causal language model.
+
+    Uses the ``PreTrainedModel.get_input_embeddings`` interface, so any model that
+    ``AutoModelForCausalLM`` can load is supported rather than a fixed set of
+    architectures.
+
+    Args:
+        model (Any): A loaded causal language model.
 
     Returns:
         Any: The model's token embedding layer.
-
-    Raises:
-        ValueError: If the model architecture is unsupported.
     """
-    if isinstance(model, (GPTJForCausalLM, GPT2LMHeadModel)):
-        return model.transformer.wte
-    if isinstance(model, LlamaForCausalLM):
-        return model.model.embed_tokens
-    if isinstance(model, GPTNeoXForCausalLM):
-        return model.base_model.embed_in
-    if isinstance(model, Phi3ForCausalLM):
-        return model.model.embed_tokens
-    raise ValueError(f"Unknown model type: {type(model)}")
+    return model.get_input_embeddings()
 
 
 def get_embedding_matrix(model: Any) -> Any:
     """
-    Return the token embedding matrix for a supported causal language model.
+    Return the token embedding matrix for a causal language model.
+
+    Args:
+        model (Any): A loaded causal language model.
 
     Returns:
         Any: The model's token embedding matrix.
-
-    Raises:
-        ValueError: If the model architecture is unsupported.
     """
-    if isinstance(model, (GPTJForCausalLM, GPT2LMHeadModel)):
-        return model.transformer.wte.weight
-    if isinstance(model, LlamaForCausalLM):
-        return model.model.embed_tokens.weight
-    if isinstance(model, GPTNeoXForCausalLM):
-        return model.base_model.embed_in.weight  # type: ignore[union-attr, unused-ignore]
-    if isinstance(model, (MixtralForCausalLM, MistralForCausalLM)):
-        return model.model.embed_tokens.weight
-    if isinstance(model, Phi3ForCausalLM):
-        return model.model.embed_tokens.weight
-    raise ValueError(f"Unknown model type: {type(model)}")
+    return model.get_input_embeddings().weight
 
 
 def get_embeddings(model: Any, input_ids: torch.Tensor) -> Any:
     """
-    Embed input token ids with a supported causal language model.
+    Embed input token ids with a causal language model.
+
+    Args:
+        model (Any): A loaded causal language model.
+        input_ids (torch.Tensor): Token ids to embed.
 
     Returns:
         Any: The embedded token tensor.
-
-    Raises:
-        ValueError: If the model architecture is unsupported.
     """
-    if isinstance(model, (GPTJForCausalLM, GPT2LMHeadModel)):
-        return model.transformer.wte(input_ids).half()
-    if isinstance(model, LlamaForCausalLM):
-        return model.model.embed_tokens(input_ids)
-    if isinstance(model, GPTNeoXForCausalLM):
-        return model.base_model.embed_in(input_ids).half()  # type: ignore[operator, unused-ignore]
-    if isinstance(model, (MixtralForCausalLM, MistralForCausalLM)):
-        return model.model.embed_tokens(input_ids)
-    if isinstance(model, Phi3ForCausalLM):
-        return model.model.embed_tokens(input_ids)
-    raise ValueError(f"Unknown model type: {type(model)}")
+    embeddings = model.get_input_embeddings()(input_ids)
+    # GPT-2, GPT-J and GPT-NeoX have always returned half precision here, while
+    # the other supported architectures return the embedding dtype unchanged.
+    # That asymmetry is preserved so this change stays a compatibility fix.
+    if isinstance(model, (GPTJForCausalLM, GPT2LMHeadModel, GPTNeoXForCausalLM)):
+        return embeddings.half()
+    return embeddings
 
 
 def get_nonascii_toks(tokenizer: Any, device: str = "cpu") -> torch.Tensor:

tests/unit/executor/promptgen/gcg/test_gcg_core.py

@@ -341,28 +341,63 @@ def test_non_ascii_filtering(self) -> None:
                 assert new_tok not in non_ascii_set, f"Candidate {i} position {pos}: sampled non-ASCII token {new_tok}"
 
 
+# Architectures built as tiny random models to exercise the embedding helpers.
+# The first three predate this generic path and must keep returning float16 from
+# get_embeddings; the rest were previously rejected outright.
+_HALF_PRECISION_ARCHITECTURES = ["gpt2", "gptj", "gpt_neox"]
+_OTHER_ARCHITECTURES = ["llama", "mistral", "mixtral", "phi3", "qwen3", "starcoder2"]
+
+_TINY_CONFIG = {
+    "hidden_size": 32,
+    "num_hidden_layers": 1,
+    "num_attention_heads": 4,
+    "num_key_value_heads": 4,
+    "intermediate_size": 64,
+    "vocab_size": 256,
+}
+_EXTRA_CONFIG = {
+    "phi3": {
+        "max_position_embeddings": 64,
+        "original_max_position_embeddings": 64,
+        "pad_token_id": 0,
+    },
+}
+
+
+def _tiny_model(model_type: str):
+    """Build a small randomly initialized model of the given architecture."""
+    transformers = pytest.importorskip("transformers", reason="transformers not installed")
+    config = transformers.AutoConfig.for_model(model_type, **_TINY_CONFIG, **_EXTRA_CONFIG.get(model_type, {}))
+    return transformers.AutoModelForCausalLM.from_config(config)
+
+
 class TestEmbeddingHelpers:
     """Tests for get_embedding_layer, get_embedding_matrix, get_embeddings."""
 
-    def test_get_embedding_layer_raises_for_unknown_model(self) -> None:
-        """Should raise ValueError for unsupported model types."""
-        mock_model = MagicMock()
-        # Ensure it doesn't match any isinstance checks
-        mock_model.__class__ = type("UnknownModel", (), {})
-        with pytest.raises(ValueError, match="Unknown model type"):
-            get_embedding_layer(mock_model)
-
-    def test_get_embedding_matrix_raises_for_unknown_model(self) -> None:
-        mock_model = MagicMock()
-        mock_model.__class__ = type("UnknownModel", (), {})
-        with pytest.raises(ValueError, match="Unknown model type"):
-            get_embedding_matrix(mock_model)
-
-    def test_get_embeddings_raises_for_unknown_model(self) -> None:
-        mock_model = MagicMock()
-        mock_model.__class__ = type("UnknownModel", (), {})
-        with pytest.raises(ValueError, match="Unknown model type"):
-            get_embeddings(mock_model, torch.tensor([1, 2, 3]))
+    @pytest.mark.parametrize("model_type", _HALF_PRECISION_ARCHITECTURES + _OTHER_ARCHITECTURES)
+    def test_helpers_resolve_embeddings_for_any_causal_model(self, model_type: str) -> None:
+        """Any model AutoModelForCausalLM can load should resolve through the helpers."""
+        model = _tiny_model(model_type)
+        expected = model.get_input_embeddings()
+
+        assert get_embedding_layer(model) is expected
+        assert get_embedding_matrix(model) is expected.weight
+
+        embedded = get_embeddings(model, torch.tensor([[1, 2, 3]]))
+        assert embedded.shape[-1] == model.config.hidden_size
+
+    @pytest.mark.parametrize("model_type", _HALF_PRECISION_ARCHITECTURES)
+    def test_get_embeddings_keeps_half_precision_for_legacy_architectures(self, model_type: str) -> None:
+        """GPT-2, GPT-J and GPT-NeoX returned float16 before this path existed."""
+        model = _tiny_model(model_type)
+        assert get_embeddings(model, torch.tensor([[1, 2, 3]])).dtype == torch.float16
+
+    @pytest.mark.parametrize("model_type", _OTHER_ARCHITECTURES)
+    def test_get_embeddings_keeps_embedding_dtype_for_other_architectures(self, model_type: str) -> None:
+        """Everything else keeps the embedding dtype rather than being downcast."""
+        model = _tiny_model(model_type)
+        expected_dtype = model.get_input_embeddings().weight.dtype
+        assert get_embeddings(model, torch.tensor([[1, 2, 3]])).dtype == expected_dtype
 
 
 class TestPromptManagerInit: