Minor fixes round

Author: matanbt

CLAUDE.md

@@ -25,7 +25,6 @@ TROPT (Textual Trigger Optimization Toolbox) is a research platform for optimizi
 
 ```bash
 uv sync --all-extras    # install in dev mode with all extras
-pre-commit install      # install pre-commit hooks
 ```
 
 Always invoke tools via `uv run` (e.g. `uv run ruff check`, `uv run ty check`, `uv run pytest`).

README.md

@@ -30,10 +30,10 @@
 
 
 - ⚔️ **Red-team LLMs out of the box:** Craft jailbreaks and other LLM attacks with **30+ ready-to-run recipes** — spanning white- and black-box methods (GCG, BEAST, MAC, GASLITE, …) — each invocable in a single call, to evaluate model and defense robustness.
-- 🔁 **Extend to any NLP model:** Seamlessly port existing optimization schemes (e.g., LLM jailreabks) to any model (e.g., retrievers, classifiers, multimodal systems), or to novel tasks (e.g., new attack vectors, interpretability research).
+- 🔁 **Extend to any NLP model:** Seamlessly port existing optimization schemes (e.g., LLM jailbreaks) to any model (e.g., retrievers, classifiers, multimodal systems), or to novel tasks (e.g., new attack vectors, interpretability research).
 - 🧩 **Compose new optimization recipes:** Mix and match any optimizer (gradient-based, continuous-relaxation, black-box) with any loss (logits, embeddings, attention, activations, LM-as-judge) to create adaptive and novel optimization recipes in new domains.
-- 🔬 **Build new optimizers and losses:** Build **new optimizers** leveraging TROPT's standardized, lightweight optimizer implementation and its exetensive toolkit. Or, **customize loss** by only defining its core logic. TROPT **automatcally integrate** new optimizers and losses to with any model and recipe (including batching, trigger combination, gradients), avoiding annoying yet subtle boierplate.
-- 🛡️ **Reliable Benchmarking:** Run fair, reproducible comparisons of optimizers and their enhancements on shared infrastructure and a rich bank of optimziers, losses, etc.
+- 🔬 **Build new optimizers and losses:** Build **new optimizers** leveraging TROPT's standardized, lightweight optimizer implementation and its extensive toolkit. Or, **customize loss** by only defining its core logic. TROPT **automatically integrates** new optimizers and losses with any model and recipe (including batching, trigger combination, gradients), avoiding annoying yet subtle boilerplate.
+- 🛡️ **Reliable Benchmarking:** Run fair, reproducible comparisons of optimizers and their enhancements on shared infrastructure and a rich bank of optimizers, losses, etc.
 
 
 ## 🚀 Getting Started
@@ -48,10 +48,9 @@ pip install tropt[all]      # all optional extras (OpenAI, LiteLLM, tracking, ..
 For development, we use [uv](https://docs.astral.sh/uv/):
 
 ```bash
-git clone https://github.com/matanbt/tropt.git
+git clone https://github.com/matanbt/TROPT.git
 cd tropt
 uv sync --extra dev
-pre-commit install
 ```
 
 ### Quick Start: Run a Recipe 🥗
@@ -130,7 +129,7 @@ You can help improve TROPT in the following two ways:
 ## Intended Use
 
 TROPT is built for defensive research:
-auditing, interpretabiliy, robustness evaluation,
+auditing, interpretability, robustness evaluation,
 and authorized red-teaming of NLP models.
 **Do not use TROPT to attack systems you don't own or to elicit harmful behaviors from deployed models in the wild.**
 
@@ -143,6 +142,7 @@ If you find this package useful, please cite our paper as follows:
 ```bibtex
 @misc{tropt2026,
   title        = {TROPT: An Open Framework for Unifying and Advancing Discrete Text Optimization},
+  author       = {Ben-Tov, Matan and Sharif, Mahmood},
   year         = {2026},
   howpublished = {\url{https://github.com/matanbt/TROPT}},
 }

skills/tropt/SKILL.md

@@ -37,11 +37,11 @@ optimizer        (declares model_requirements, drives the search)──┘
 Two design principles to keep in mind: *modularity* (each component swaps largely independently) and *backend–frontend separation* (tokenization, batching, gradient computation live in the model "backend"; losses and optimizers stay lightweight "frontend").
 
 - **Model** (`tropt/model/`) — wraps the target model and exposes capabilities via **access mixins**. The optimizer declares which mixins it needs; `BaseOptimizer.__init__` validates the model satisfies them. Common mixins: `LossTokenAccessMixin` (grey-box loss from tokens), `GradientTokenAccessMixin` (white-box gradients), `LossTextAccessMixin` (black-box, text-only), `LogitsTokenAccessMixin`, `GradientEmbedAccessMixin`. Backends: `LMHFModel` (HF causal LMs — most mixins), `EncoderHFModel`, `EncoderOpenAIModel`, `LiteLLMModel`, `EncoderGeminiModel`.
-- **Loss** (`tropt/loss/`) — pure objective. Receives fields from `ModelOutput` / `ModelInput` / `MessageTargets` by **parameter-name matching** in its `__call__` signature (the one rule). Categories: `LogitBasedLoss`, `EmbeddingBasedLoss`, `TextBasedLoss`, `AttentionBasedLoss`, `HiddenStateBasedLoss`, `CombinedLoss`.
+- **Loss** (`tropt/loss/`) — pure objective. Receives fields from `ModelOutput` / `ModelInput` / `MessageTargets` by **parameter-name matching** in its `__call__` signature (the one rule). Categories: `PrefillBasedLoss`, `TriggerLogitBasedLoss`, `EmbeddingBasedLoss`, `TextBasedLoss`, `AttentionBasedLoss`, `HiddenStateBasedLoss`, `ClassificationBasedLoss`, `CombinedLoss`.
 - **Optimizer** (`tropt/optimizer/`) — the search algorithm. Self-contained, one file per optimizer (HuggingFace "Repeat Yourself"). Declares `model_requirements = (Mixin1, Mixin2)` at class level.
 - **Inputs & targets** — the user-supplied input template(s) carrying the `{{OPTIMIZED_TRIGGER}}` placeholder, plus a `Targets` dataclass holding per-template targets (`target_response_strs`, `target_vectors`, `target_directions`, `target_class_idx`, …). Each set field has length `n_templates`. Templates and targets are passed to `optimizer.optimize_trigger(templates=..., targets=...)`; the framework handles trigger/template combination internally.
 
-The **Recipe Hub** (`tropt/recipe_hub/`) ships ~30 pre-wired (Model, Loss, Optimizer, defaults) combinations as one-call functions. That's almost always the right starting point.
+The **Recipe Hub** (`tropt/recipe_hub/`) ships ~38 pre-wired (Model, Loss, Optimizer, defaults) combinations as one-call functions. That's almost always the right starting point.
 
 ## 3. Install or update TROPT
 
@@ -185,17 +185,17 @@ These apply across optimizers and recipes — surface them when relevant rather
 - **Missing `{{OPTIMIZED_TRIGGER}}`** in templates → optimizer raises. Every template must contain the placeholder string (exported as `tropt.common.OPTIMIZED_TRIGGER_PLACEHOLDER`).
 - **Mixin mismatch** — pairing a black-box `LiteLLMModel` with `GCGOptimizer` (which needs `GradientTokenAccessMixin`) fails at construction time, not silently. If the user wants a gradient-free attack on a black-box model, point them at `prs__andriushchenko2024` / `rs_emb` / `gcgp_blackbox__hayase2024`.
 - **`Targets` field length mismatch** — every set field must have length `n_templates`. Universal triggers across N templates need N targets (often repeated: `["Sure, here is"] * N`).
-- **Thinking-model target misalignment (Qwen3, etc.)** — these models prefix every emits a `<think>...</think>` block, so the affirmative target the optimizer chases (e.g. `"Sure, here's how:"`) never appears at position 0 of the generation. Prepend the empty thinking block to the target: `target = "<think>\n\n</think>\n\n" + "Sure, here's how:"`. Without this fix, prefill-CE losses chase a position the model will never write to.
+- **Thinking-model target misalignment (Qwen3, etc.)** — these models always emit a `<think>...</think>` block first, so the affirmative target the optimizer chases (e.g. `"Sure, here's how:"`) never appears at position 0 of the generation. Prepend the empty thinking block to the target: `target = "<think>\n\n</think>\n\n" + "Sure, here's how:"`. Without this fix, prefill-CE losses chase a position the model will never write to.
 - **Some losses impose model-construction requirements** beyond just the access mixins. Check the loss's docstring before pairing it with a recipe:
   - `AttentionBasedLoss` family (e.g. `AttentionEnhLoss`) needs the model loaded with `use_eager_attention=True` — FlashAttention/SDPA paths don't return attention weights. `LMHFModel(..., use_eager_attention=True)`.
-  - `AttentionEnhLoss` with `dst=INPUT_AFTER` (and other suffix-position-dependent losses) requires the trigger to be a **true suffix** of the user message — strip any trailing punctuation after `{{OPTIMIZED_TRIGGER}}` in the template (`"... harmful thing. {{OPTIMIZED_TRIGGER}}"` → `"... harmful thing {{OPTIMIZED_TRIGGER}}"`).
+  - `AttentionEnhLoss` with `dst_slc_name=SliceKey.INPUT_AFTER` (and other suffix-position-dependent losses) requires the trigger to be a **true suffix** of the user message — strip any trailing punctuation after `{{OPTIMIZED_TRIGGER}}` in the template (`"... harmful thing. {{OPTIMIZED_TRIGGER}}"` → `"... harmful thing {{OPTIMIZED_TRIGGER}}"`).
   - `TriggerPerplexityLoss` and other trigger-slice-dependent losses require `use_prefix_cache=False` on the model — prefix caching shifts the trigger slice's start position.
   - Hidden-state / activation losses (`SteeringActivationLoss`, etc.) often need pre-computed direction vectors (e.g. `compute_refusal_directions(model, n_samples=128)`) passed via `Targets(target_directions=...)`.
 - **Stop by budget, not just by `num_steps`** — TROPT optimizers support `optimizer.set_budget(BUDGET, metric="total_flops")` (white-box) and `set_budget(BUDGET, metric="total_tokens", scope="target")` (black-box / API), letting you cap compute or API cost directly instead of guessing a step count. Pair with `model.set_flop_counting("manual")` on the model. Far more reliable across optimizers with different per-step costs.
-- **Initial trigger init** — prefer `tropt.optimizer.utils.token_initializers.get_printable_random_trigger(trigger_len, tokenizer, blacklist_ids=TokenConstraints().get_blacklist_ids(tokenizer))` over hand-crafted `"! ! ! ! !"` strings. The hand-crafted form depends on the tokenizer treating `!` as a single token (not universal) and gives every position the same starting gradient, which hurts on some models.
+- **Initial trigger init** — prefer `tropt.optimizer.utils.token_initializers.get_printable_random_trigger(trigger_len, tokenizer=tokenizer, blacklist_ids=TokenConstraints().get_blacklist_ids(tokenizer))` over hand-crafted `"! ! ! ! !"` strings. The hand-crafted form depends on the tokenizer treating `!` as a single token (not universal) and gives every position the same starting gradient, which hurts on some models.
 - **Pass `TokenConstraints` to the optimizer, not just to trigger init** — without `token_constraints=`, optimizers may emit triggers with non-ASCII Unicode, model special tokens (`<|im_start|>`, `<bos>`, etc.), or other artifacts that break downstream evaluation, leak chat-template structure, or are trivially filterable. For red-teaming / robustness work the safe default is `TokenConstraints(disallow_non_ascii=True, disallow_special_tokens=True)`, passed both to the optimizer and to `get_blacklist_ids(tokenizer)` for trigger init. Recipe Hub recipes already do this; raw `Pattern B` compositions in user scripts often forget.
 - **Multi-model OOM** — when a recipe needs a *second* model (teacher for IRIS / FLRT distillation, proxy for PAL/RAL/QCG in black-box mode, utility LM for AdvDecoding/BeamSearch), do not let it co-reside with the victim. Standard pattern: load aux → generate / extract what you need → `del aux._model; del aux; gc.collect(); torch.cuda.empty_cache(); torch.cuda.synchronize()` → THEN load victim. Holding both on one GPU blows memory on most consumer setups.
-- **Black-box loss is not the same as white-box loss** — `PrefillCELoss` reads `full_logits`, which an API/LiteLLM model doesn't expose. For API targets swap to `FirstTokenNLLLoss(target_token="Sure")` (or another text/first-token-based loss). Optimizer-side, the model itself also flips to `LiteLLMModel`, and only optimizers declaring `LossTextAccessMixin` in their `model_requirements` will accept it.
+- **Black-box loss is not the same as white-box loss** — `PrefillCELoss` reads `prefill_response_logits` (target prefill), which an API/LiteLLM model can't provide. For API targets swap to `FirstTokenNLLLoss(target_token="Sure")` (or another text/first-token-based loss). Optimizer-side, the model itself also flips to `LiteLLMModel`, and only optimizers declaring `LossTextAccessMixin` in their `model_requirements` will accept it.
 - **Large models on a small machine** — don't suggest `meta-llama/Llama-3.1-8B-Instruct` for someone debugging locally on a laptop. Default to `google/gemma-3-270m-it` or `google/gemma-3-1b-it` for smoke tests; the recipe still reproduces, just with weaker results.
 - **Recipe naming** — `foo__year2024` means it reproduces the paper's algorithm + hyperparameters. Bare `foo` means it's a variant or a setting-port. When recommending a recipe for paper reproduction, prefer the tagged form.
 - **WandB tracking is opt-in** — instantiate `WandbTracker()` and pass it to the optimizer; otherwise nothing logs. `LiveLossPlot` is the lighter local alternative.

tropt/loss/base.py

@@ -39,7 +39,7 @@ class BaseLoss(ABC):
     """Whether this loss requires autoregressive generation."""
 
     require_hidden_states: ClassVar[bool] = False
-    """Whether this loss requires the model to provid the forward pass's hidden states."""
+    """Whether this loss requires the model to provide the forward pass's hidden states."""
 
     require_attentions: ClassVar[bool] = False
     """Whether this loss requires the model to return attention weights."""
@@ -128,7 +128,7 @@ def __call__(self, losses: Float[Tensor, "n_losses bsz"]) -> Float[Tensor, "bsz"
 
         weights = self.weights.to(losses).unsqueeze(-1)  # shape: (n_losses, 1)
         loss = losses * weights
-        loss = loss.sum(dim=0)  # recude over n_losses
+        loss = loss.sum(dim=0)  # reduce over n_losses
 
         return loss  # shape: (bsz,)
 

tropt/loss/losses.py

@@ -353,8 +353,7 @@ class AttentionEnhLoss(AttentionBasedLoss):
     Enable to instantiate the (different) losses from:
     https://arxiv.org/abs/2506.12880, https://arxiv.org/abs/2410.09040
 
-    Note that it requires setting `use_eager_attention=True` when loading the model (for explicit attention computations); also it
-    is some slices are not supported when LM prefix caching is enabled, so it should be set to `use_prefix_cache=False` when loading the model.
+    Note that it requires setting `use_eager_attention=True` when loading the model (for explicit attention computations); also, some slices are not supported when LM prefix caching is enabled, so set `use_prefix_cache=False` when loading the model.
     """
 
     targeted_layers: slice = slice(None)
@@ -367,7 +366,7 @@ def __call__(
         input_slices: dict[SliceKey, slice],
     ) -> Float[Tensor, "bsz"]:
         if SliceKey.INPUT_AFTER in (self.src_slc_name, self.dst_slc_name):
-            logger.debug("Note: `INPUT_AFTER` slice is currently only correct for LMs and on suffix attacks. If the usage is different, somethings may break, or worse -- be wrong.")
+            logger.debug("Note: `INPUT_AFTER` slice is currently only correct for LMs and on suffix attacks. If the usage is different, something may break, or worse -- be wrong.")
         slc_src = input_slices.get(self.src_slc_name, slice(None))
         slc_dst = input_slices.get(self.dst_slc_name, slice(None))
 
@@ -461,9 +460,10 @@ class SteeringActivationLoss(HiddenStateBasedLoss):
     Args:
         targeted_layers: Which layers to apply steering on (default: all layers)
         steer_away: Whether to minimize alignment instead of maximizing (default: False = steer towards)
-        slc_name: Which token positions to apply steering on (default: "last_input_token")
+        slc_name: Which token positions to apply steering on (default: "input_last_token")
         do_cosine_sim: Whether to use cosine similarity instead of dot product (default: False)
         apply_square: Whether to square the similarity scores (default: False)
+        apply_abs: Whether to take the absolute value of the similarity scores (default: False).
     """
 
     targeted_layers: slice = slice(None)
@@ -483,7 +483,7 @@ def __call__(
         Compute steering loss by measuring cosine similarity between hidden states and target directions.
 
         Args:
-            output_hidden_states: Model hidden states from all layers and positions (bsz, n_layers, seq_len, d_model)
+            full_hidden_states: Model hidden states from all layers and positions (bsz, n_layers, seq_len, d_model)
             target_directions: Direction vectors to align with (, d_model)
             input_slices: Position slices reflecting the input tokens (dict mapping slice names to slices)
 

tropt/loss/resolution.py

@@ -73,15 +73,14 @@ def resolve_and_compute_loss(
 
     Examples:
         >>> from tropt.common import MessageTargets
-        >>> # Encoder model with SimilarityLoss(output_embeddings, target_embeddings)
+        >>> # Encoder model with SimilarityLoss(output_embeddings, target_vectors)
         >>> output = ModelOutput(output_embeddings=torch.randn(4, 768))
         >>> input_data = ModelInput(message_targets=MessageTargets(target_vectors=target_vecs))
         >>> loss = resolve_and_compute_loss(output, input_data, SimilarityLoss())
 
         >>> # Language model with PrefillCELoss(prefill_response_logits, message_targets)
         >>> output = ModelOutput(prefill_response_logits=torch.randn(2, 50, 32000))
         >>> input_data = ModelInput(
-        ...     input_slices=[{SliceKey.APPENDED: slice(40, 50)}] * 2,
         ...     message_targets=MessageTargets(target_response_toks=target_ids)
         ... )
         >>> loss = resolve_and_compute_loss(output, input_data, PrefillCELoss())

tropt/loss/text_losses.py

@@ -1,5 +1,5 @@
 """
-General loss functions.
+Text-based (black-box / non-differentiable) loss functions.
 
 Important note: The losses arguments must match the fields in ModelOutput and ModelInput
 for unified loss resolution to work properly.

tropt/model/flop_counter.py

@@ -21,11 +21,11 @@
 Usage::
 
     model = LMHFModel("meta-llama/Llama-3.2-1B", ...)
-    model.set_flop_counting("manual")   # False to disable
+    model.set_flop_counting("manual")   # "none" to disable
     model.reset_usage_stats()
     # ... run optimization ...
     stats = model.get_usage_stats()
-    print(stats["usage/total_flops"])
+    print(stats["total_flops"])
 """
 
 import logging

tropt/model/google/encoder.py

@@ -74,7 +74,7 @@ def __init__(
 
         Args:
             model_name: The name of the Gemini embedding model to use.
-            d_model: The dimensionality of the embeddings (e.g., 768, 3072).
+            d_model: The dimensionality of the embeddings (e.g., 768, 1536, 3072).
             use_vertex: Embed via the Vertex AI backend (ADC) instead of AI Studio.
                 Needed for Vertex-only models such as ``text-embedding-005``.
             project: Vertex project (only used when ``use_vertex``; falls back to the

tropt/model/huggingface/clip_encoder.py

@@ -105,6 +105,8 @@ def set_inputs_from_tokens(
     ) -> None:
         """Prepare and store the given templates in the inputs manager."""
         assert isinstance(templates, list)
+        if targets is None:
+            targets = Targets()
 
         tok_ids = self._tokenizer(templates, add_special_tokens=True)["input_ids"]
         self._token_input_manager = HuggingFaceTokenInputManager(