[WIP] Ponytail's code cleanup

CLAUDE.md

@@ -65,10 +65,9 @@ Naming: `{Value}{InputType}AccessMixin` (e.g. `LossTokenAccessMixin`, `GradientT
 
 Canonical definitions and full field lists live in `tropt/common.py`. Don't enumerate them here.
 
-### Glue: Recipe Hub & Config Runner
+### Glue: Recipe Hub
 
 - **Recipe Hub** (`tropt/recipe_hub/`): pre-configured Model + Loss + Optimizer + Inputs/Targets wirings, each callable as one function. Enumerate via `list_recipes()`; full registry in `tropt/recipe_hub/__init__.py`.
-- **Config Runner** (`runner/main.py`): YAML-driven runner via Hydra. (Hydra config support is incomplete — see Known Limitations.)
 
 ## Component Interactions
 
@@ -240,8 +239,6 @@ Auto-generated docs (`docs/api/`, `docs/guides/compatibility_matrix.md`) are pro
 
 For contributions back to the package (file placement, exports, tests, Recipe Hub naming convention), see `CONTRIBUTING.md`.
 
-For testing conventions (mirror layout, fixtures, numerical tolerances, what to test per component), see `TESTING.md`.
-
 ## Important Notes
 
 ### Scripts: separate repo
@@ -257,13 +254,12 @@ This codebase is explicitly designed for adversarial robustness research and red
 ### Known Limitations
 - Multiple-message prefix caching currently disabled due to edge cases
 - Tracker should be initialized per RUN, not per optimizer instance
-- Hydra config-runner support is incomplete
 
 ## Dependencies
 
-Core: PyTorch, Transformers, Accelerate, Hydra, Pydantic
-Models: HuggingFace, SentenceTransformers, OpenAI, LiteLLM
-Tracking: Weights & Biases, LiveLossPlot
+Core: PyTorch, Transformers, Accelerate, SentenceTransformers, Pydantic
+Optional extras: OpenAI, Google, Voyage, LiteLLM, vision, tracking, notebooks
+Tracking: Weights & Biases, Trackio, LiveLossPlot
 Dev: pytest, ruff, ty, pre-commit
 
 ## Repository Structure
@@ -279,7 +275,6 @@ tropt/
 └── utils/           # Shared utilities
 
 tests/               # Test suite mirroring tropt/ structure
-runner/              # Hydra-driven config runner
 scripts/             # Separate repo — see "Scripts: separate repo" above
 docs/                # Sphinx documentation
 ├── api/             # Auto-generated API reference (rst)
@@ -288,6 +283,5 @@ docs/                # Sphinx documentation
 └── conf.py          # Sphinx config
 quickstart.ipynb     # End-to-end notebook
 DESIGN.md            # Design philosophy and rationale
-TESTING.md           # Testing guidelines and conventions
 CONTRIBUTING.md      # Contribution workflow (file placement, exports, tests)
 ```

DESIGN.md

@@ -156,7 +156,7 @@ return resolve_and_compute_loss(model_output, model_input, loss_func)  # the los
 This keeps loss logic in one location, makes adding new loss types straightforward, and means new models only provide data---not loss implementation.
 
 **Convention.** As a good practice, we divide the losses with superclasses according to the type of input that the loss accepts (which is, in turn, mostly the type of output of the model). For instance, Cross-Entropy-based losses utilize the logit outputs, and thus they will inherit from `PrefillBasedLoss`.
-In this way, from the model end, we would be disable unneeded calculation: for instnace, if the loss does not requrie hidden states (i.e., the loss is not subclass of `HiddenStateBasedLoss`), we can know in advance to avoid saving them for efficiency.
+In this way, from the model end, we would be disable unneeded calculation: for instnace, if the loss does not requrie hidden states (i.e., the loss does not set `require_hidden_states`), we can know in advance to avoid saving them for efficiency.
 
 ## Component 3: Optimizers
 
@@ -275,7 +275,6 @@ The repository exposes two interfaces for managing recipes:
 
 * **Recipe Hub [`tropt/recipe_hub/`].** Python modules that bind the four components to reproduce existing attacks. E.g., the `gcg__zou2023` module wires `LMHFModel`, `PrefillCELoss`, `GCGOptimizer`, and a standard suffix template to reproduce the GCG attack. Useful for researchers who want to quickly run published attacks, benchmark them, or fork-and-modify.
 
-<!-- * **Config Runner [`runner/main.py`].** A flexible runner that constructs a recipe from a YAML configuration file. The runner uses [Hydra](https://hydra.cc/) to manage configurations, allowing users to specify the model, loss, optimizer, and their parameters in a structured way without writing new code. -->
 
 ## Summary
 

docs/api/loss.rst

@@ -27,17 +27,7 @@ Loss Classes Interfaces
    :undoc-members:
    :show-inheritance:
 
-.. autoclass:: TriggerLogitBasedLoss
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-.. autoclass:: AttentionBasedLoss
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-.. autoclass:: EmbeddingBasedLoss
+.. autoclass:: PrefillBasedLoss
    :members:
    :undoc-members:
    :show-inheritance:
@@ -66,7 +56,7 @@ Loss Implementations
 
 .. automodule:: tropt.loss
    :members:
-   :exclude-members: BaseLoss, TriggerLogitBasedLoss, AttentionBasedLoss, EmbeddingBasedLoss, TextBasedLoss, SteeringActivationLoss, CombinedLoss
+   :exclude-members: BaseLoss, PrefillBasedLoss, TextBasedLoss, SteeringActivationLoss, CombinedLoss
    :undoc-members:
    :show-inheritance:
    :imported-members:

docs/api/models.rst

@@ -98,7 +98,7 @@ Model Implementations
 
 .. automodule:: tropt.model
    :members:
-   :exclude-members: BaseModel, LMBaseModel, EncoderBaseModel, BaseTokenizer, TokenAccessMixin, LossTokenAccessMixin, LogitsTokenAccessMixin, GradientTokenAccessMixin, TextAccessMixin, LossTextAccessMixin, InputsManager, TextInputManager, TokenInputManager, DefaultTokenInputManager
+   :exclude-members: BaseModel, LMBaseModel, EncoderBaseModel, BaseTokenizer, TokenAccessMixin, LossTokenAccessMixin, LogitsTokenAccessMixin, GradientTokenAccessMixin, TextAccessMixin, LossTextAccessMixin, TextInputManager, TokenInputManager, DefaultTokenInputManager
    :undoc-members:
    :show-inheritance:
    :imported-members:

docs/build_docs.py

@@ -3,41 +3,36 @@
 import shutil
 import stat
 import subprocess
+import sys
 import time
 import webbrowser
 
 
-def make_writable(path):
-    """Force a file/directory to be writable."""
-    try:
-        os.chmod(path, stat.S_IWRITE)
-    except Exception:
-        pass
+def _force_writable(func, path, _exc):
+    """rmtree error hook: clear the read-only bit (git/GitHub files) and retry."""
+    os.chmod(path, stat.S_IWRITE)
+    func(path)
+
+
+# `onexc` replaced `onerror` in Python 3.12; the project supports 3.10+.
+_RMTREE_HOOK = (
+    {"onexc": _force_writable} if sys.version_info >= (3, 12)
+    else {"onerror": _force_writable}
+)
+
 
 def robust_cleanup(path):
-    """
-    Aggressively cleans up a directory.
-    1. Walks tree to fix permissions (handling read-only Git/GitHub files).
-    2. Tries to delete.
-    3. Returns True if successful, False if locked.
+    """Delete `path`, retrying while Google Drive / Windows holds a lock.
+
+    Returns True if the tree is gone, False if it stayed locked.
     """
     if not os.path.exists(path):
         return True
 
     print(f"Cleaning up {path}...")
-
-    # 1. Force permissions first (Pre-emptive strike)
-    for root, dirs, files in os.walk(path):
-        for d in dirs:
-            make_writable(os.path.join(root, d))
-        for f in files:
-            make_writable(os.path.join(root, f))
-    make_writable(path)
-
-    # 2. Try deletion with retries
-    for i in range(5):
+    for _ in range(5):
         try:
-            shutil.rmtree(path)
+            shutil.rmtree(path, **_RMTREE_HOOK)
             return True
         except OSError:
             time.sleep(0.5)  # Wait for Google Drive/Windows to release lock
@@ -84,41 +79,10 @@ def build_docs():
     # ---------------------------------------------------------
     # 2. Inject 'from __future__ import annotations'
     # ---------------------------------------------------------
+    # Same script CI runs (deploy_docs.yml), pointed at the throwaway copy.
     print("Injecting future annotations...")
-    import ast
-    for root, _, files in os.walk(src_copy):
-        for file in files:
-            # common.py is excluded to match CI (deploy_docs.yml): future
-            # annotations stringify its pydantic/jaxtyping runtime field types
-            # and break Sphinx autodoc (sphinx-doc/sphinx#11211).
-            if file.endswith(".py") and file != "common.py":
-                path = os.path.join(root, file)
-                try:
-                    with open(path, "r", encoding="utf-8") as f:
-                        content = f.read()
-
-                    if "from __future__ import annotations" not in content:
-                        # Insert AFTER a leading module docstring so the docstring
-                        # isn't demoted to a bare string expression (which would
-                        # strip it from the autodoc-rendered module page).
-                        insert_at = 0
-                        try:
-                            mod = ast.parse(content)
-                            first = mod.body[0] if mod.body else None
-                            if (
-                                isinstance(first, ast.Expr)
-                                and isinstance(getattr(first, "value", None), ast.Constant)
-                                and isinstance(first.value.value, str)
-                            ):
-                                insert_at = first.end_lineno
-                        except SyntaxError:
-                            insert_at = 0
-                        lines = content.splitlines(keepends=True)
-                        lines.insert(insert_at, "from __future__ import annotations\n")
-                        with open(path, "w", encoding="utf-8") as f:
-                            f.write("".join(lines))
-                except Exception as e:
-                    print(f"  Skipping {file}: {e}")
+    from docs.scripts.inject_annotations import inject
+    print(f"  Injected into {inject(src_copy)} module(s).")
 
     # NOTE: inlining + sanitizing tropt/recipe_hub/README.md now happens in
     # conf.py (the `include-read` event), so it runs for every build path.

docs/conf.py

@@ -56,7 +56,7 @@
 myst_enable_extensions = ["colon_fence", "deflist", "attrs_inline"]
 myst_heading_anchors = 3
 
-autodoc_mock_imports = ['runner', 'sentence_transformers', 'wandb', 'livelossplot', 'openai', 'litellm', 'IPython', 'tqdm', 'transformers', 'accelerate', 'hydra', 'omegaconf', 'huggingface_hub', 'datasets', 'PIL', 'diffusers', 'trackio']
+autodoc_mock_imports = ['sentence_transformers', 'wandb', 'livelossplot', 'openai', 'litellm', 'IPython', 'tqdm', 'transformers', 'accelerate', 'huggingface_hub', 'datasets', 'PIL', 'diffusers', 'trackio']
 autodoc_typehints = "description"
 
 

docs/guides/adding_a_loss.md

@@ -130,12 +130,12 @@ Breaking down the additions:
 **Convention: Hyperparameters as Fields.** Anything you want callers to tweak at construction time — temperatures, margins, layer ranges, mode flags — goes here. Crucially, *per-template* data (target tokens, target vectors, target classes) does **not** belong here; it belongs on {py:class}`~tropt.common.Targets`, and the loss pulls it in by parameter name (`target_response_toks` above). This keeps the loss instance stateless w.r.t. the attack — the optimizer can resample or subsample templates without telling the loss.
 
 
-**Convention: Inheriting from a category base class.** {py:mod}`tropt.loss` defines abstract bases like {py:class}`~tropt.loss.PrefillBasedLoss`, {py:class}`~tropt.loss.EmbeddingBasedLoss`, {py:class}`~tropt.loss.HiddenStateBasedLoss`, {py:class}`~tropt.loss.AttentionBasedLoss`, {py:class}`~tropt.loss.ClassificationBasedLoss`, and {py:class}`~tropt.loss.TextBasedLoss`. Each sets the right `require_*` flag and pins a typed `__call__` signature for its category — inheriting from the right base saves boilerplate and signals intent. The categorization is a **convention for readability**, not a hard requirement: the resolver dispatches by parameter names, not by base class. If your loss doesn't fit any existing category, inheriting from `BaseLoss` directly and setting the flags yourself is equally valid.
+**Convention: inherit from `BaseLoss`, declare your own flags.** The resolver dispatches on **parameter names**, never on base class — so subclassing `BaseLoss` directly and setting the `require_*` flags you need is the normal case. Two shared bases exist only where they carry something real: {py:class}`~tropt.loss.PrefillBasedLoss` (sets `require_target_prefill=True` and pins the prefill `__call__` signature shared by four losses) and {py:class}`~tropt.loss.TextBasedLoss` (sets `is_differentiable=False` for text-scoring losses). Don't add a new base class for a single loss.
 
 
 ## Practical Example: Activation Steering
 
-Up to here we asked the model for logits over a target response. The same pattern works just as well for *any* model artifact — swap the base class, swap the `require_*` flag, and swap the `__call__` parameter names, and you have a loss over a completely different signal.
+Up to here we asked the model for logits over a target response. The same pattern works just as well for *any* model artifact — swap the `require_*` flag and swap the `__call__` parameter names, and you have a loss over a completely different signal.
 
 To demonstrate, we rewrite the loss as **activation steering** ([Arditi et al., 2024](https://arxiv.org/abs/2406.11717)): encouraging the model's hidden activations at chosen layers/positions to align with (or away from) a target direction in activation space. This has been used both to suppress refusal in jailbreaks and to probe internal representations. It's what {py:class}`~tropt.loss.SteeringActivationLoss` does in TROPT.
 
@@ -148,19 +148,18 @@ import torch.nn.functional as F
 from jaxtyping import Float
 from torch import Tensor
 
-from tropt.loss import HiddenStateBasedLoss
+from tropt.loss import BaseLoss
 
 
 @dataclass
-class MyLoss(HiddenStateBasedLoss):
+class MyLoss(BaseLoss):
     """Steers hidden activations along (or away from) a target direction."""
 
     # Hyperparameters
     targeted_layers: slice = slice(None)
     steer_away: bool = False  # if True, push activations *away* from the direction
 
     # Ask the model to return all hidden states.
-    # (HiddenStateBasedLoss already declares this; shown here for emphasis.)
     require_hidden_states: ClassVar[bool] = True
 
     def __call__(
@@ -177,11 +176,11 @@ class MyLoss(HiddenStateBasedLoss):
 
 Structurally nothing is new — same `BaseLoss` lineage, same dataclass-fields-as-hyperparameters, same `require_*` declaration, same per-sample return shape. The differences are entirely in *which* names appear:
 
-- **Base class & flag:** {py:class}`~tropt.loss.HiddenStateBasedLoss` (sets `require_hidden_states=True`), replacing `PrefillBasedLoss` / `require_target_prefill`.
+- **Flag:** `require_hidden_states=True` replaces `require_target_prefill`.
 - **Model output:** `full_hidden_states` replaces `prefill_response_logits`.
 - **Per-template target:** `target_directions` (provided via `Targets(target_directions=...)`) replaces `target_response_toks`.
 
-The same one-knob swap also produces attention-based losses (inherit from `AttentionBasedLoss`, name `full_attentions`) and classifier losses (inherit from `ClassificationBasedLoss`, name `output_class_logits`) — see the existing implementations of {py:class}`~tropt.loss.AttentionEnhLoss` and {py:class}`~tropt.loss.MisclassCELoss` for those variants.
+The same one-knob swap also produces attention-based losses (`require_attentions=True`, name `full_attentions`) and classifier losses (no flag, name `output_class_logits`) — see the existing implementations of {py:class}`~tropt.loss.AttentionEnhLoss` and {py:class}`~tropt.loss.MisclassCELoss` for those variants.
 
 
 ## Going Non-Differentiable: Scoring Triggered Text

docs/guides/adding_a_model.md

@@ -92,7 +92,7 @@ class ModelOutput:
     full_ids: ...
 ```
 
-The fields you populate determine which loss types are compatible with your model. For example, `output_embeddings` enables `EmbeddingBasedLoss` (e.g., `SimilarityLoss`), while `generated_response_strs` enables `TextBasedLoss` (e.g., `GeneratedResponseBasedLoss`). The {py:func}`loss resolution system <tropt.loss.resolve_and_compute_loss>` validates this at runtime and raises clear errors if a required field is missing.
+The fields you populate determine which losses are compatible with your model. For example, `output_embeddings` enables embedding losses (e.g., `SimilarityLoss`), while `generated_response_strs` enables generation-scoring losses (e.g., `ResponseHarmfulnessLoss`). The {py:func}`loss resolution system <tropt.loss.resolve_and_compute_loss>` validates this at runtime and raises clear errors if a required field is missing.
 
 ### Model compatibility
 

docs/guides/adding_an_optimizer.md

@@ -211,7 +211,7 @@ optimizer.optimize_trigger(...)
 
 **Trigger initialization.** Callers can pass an explicit `initial_trigger` (the `DEFAULT_INIT_TRIGGER` is `"! ! ! ..."`-style). For something smarter, sample from the constrained vocabulary itself via {py:func}`~tropt.optimizer.utils.token_initializers.get_printable_random_trigger`.
 
-The remaining building blocks in [optimizer utilities](../api/optimizer_utils) worth knowing about: `retokenize_filtering` (drop candidates that don't survive a decode → encode round-trip), `TriggerBuffer` (best-K pool instead of a single best), `NFlipScheduler` (control how many positions to mutate per step). Pull them in only when your search actually needs them — they're not boilerplate.
+The remaining building blocks in [optimizer utilities](../api/optimizer_utils) worth knowing about: `retokenize_filtering` (drop candidates that don't survive a decode → encode round-trip), `TriggerBuffer` (best-K pool instead of a single best), `LinearScheduler` (a `(step) -> n_flip` callable controlling how many positions to mutate per step; any callable works). Pull them in only when your search actually needs them — they're not boilerplate.
 
 
 ## Going White-Box: Gradient-Guided Search

docs/scripts/inject_annotations.py

@@ -0,0 +1,51 @@
+"""Insert `from __future__ import annotations` into every `tropt/*.py`.
+
+Sphinx autodoc can't resolve TROPT's jaxtyping/pydantic forward refs without it
+(https://github.com/sphinx-doc/sphinx/issues/11211). `common.py` is excluded:
+its pydantic models use runtime jaxtyping field types that break when
+stringified. The import goes *after* any module docstring so the docstring still
+renders on the autodoc module pages.
+
+Run against a throwaway copy of the tree — it rewrites files in place.
+
+Usage: python docs/scripts/inject_annotations.py [package_dir]
+"""
+
+import ast
+import pathlib
+import sys
+
+EXCLUDED = {"common.py"}
+
+
+def inject(package_dir: str | pathlib.Path = "tropt") -> int:
+    """Rewrite every eligible module under `package_dir`. Returns the count."""
+    n = 0
+    for path in pathlib.Path(package_dir).rglob("*.py"):
+        if path.name in EXCLUDED:
+            continue
+        src = path.read_text(encoding="utf-8")
+        if "from __future__ import annotations" in src:
+            continue
+
+        insert_at = 0
+        try:
+            first = (ast.parse(src).body or [None])[0]
+            if (
+                isinstance(first, ast.Expr)
+                and isinstance(getattr(first, "value", None), ast.Constant)
+                and isinstance(first.value.value, str)
+            ):
+                insert_at = first.end_lineno
+        except SyntaxError:
+            insert_at = 0
+
+        lines = src.splitlines(keepends=True)
+        lines.insert(insert_at, "from __future__ import annotations\n")
+        path.write_text("".join(lines), encoding="utf-8")
+        n += 1
+    return n
+
+
+if __name__ == "__main__":
+    print(f"Injected future annotations into {inject(*sys.argv[1:])} module(s).")