Minor fixes

Author: matanbt

DESIGN.md

@@ -17,7 +17,7 @@ TROPT's design is guided by two technical principles:
 
 These two principles together address the *accessibility*, *adaptability*, *comparability*, and *extensibility* requirements motivated in the companion paper; this document is the technical complement, focusing on *how* the design realizes them rather than re-arguing *why* they matter.
 
-In the next segment I describe each of the ==four components==, starting from the lower-level model integration, through the loss modules, the optimizers that drive search, and finally the user-supplied inputs and targets. Crucially, one may abstract the internal design of these components, and merely compose attacks by combining different instances of them.
+In the next segment I describe each of the **four components**, starting from the lower-level model integration, through the loss modules, the optimizers that drive search, and finally the user-supplied inputs and targets. Crucially, one may abstract the internal design of these components, and merely compose attacks by combining different instances of them.
 
 In the final segment, I describe the glue: the **recipe**—an executable instantiation of all four components—along with the two existing interfaces (Recipe Hub and the config-driven runner) to run end-to-end optimization in the repo.
 

README.md

@@ -18,8 +18,8 @@
 
 <p align="center">
   <a href="https://pypi.org/project/tropt/"><img src="https://img.shields.io/pypi/v/tropt?logo=python&logoColor=white&color=3776ab" alt="PyPI"></a>
-  <a href="https://github.com/matanbt/tropt"><img src="https://img.shields.io/github/stars/matanbt/tropt?style=flat&logo=github&color=181717" alt="GitHub stars"></a>
-  <a href="https://github.com/matanbt/tropt/actions/workflows/test.yml"><img src="https://img.shields.io/github/actions/workflow/status/matanbt/tropt/test.yml?branch=main&label=tests" alt="Tests"></a>
+  <a href="https://github.com/matanbt/TROPT"><img src="https://img.shields.io/github/stars/matanbt/TROPT?style=flat&logo=github&color=181717" alt="GitHub stars"></a>
+  <a href="https://github.com/matanbt/TROPT/actions/workflows/test.yml"><img src="https://img.shields.io/github/actions/workflow/status/matanbt/TROPT/test.yml?branch=main&label=tests" alt="Tests"></a>
   <a href="https://opensource.org/licenses/MIT"><img src="https://img.shields.io/badge/License-MIT-green?style=flat" alt="License"></a>
 </p>
 
@@ -61,7 +61,7 @@ TROPT lets you run an optimization scheme — a _recipe_ — in a single functio
 For instance, you can import the GCG jailbreak ([Zou et al., 2023](https://arxiv.org/abs/2307.15043)) and reproduce it on `gemma-3-1b-it` in a few lines:
 
 ```python
-from tropt.recipe_hub.GCG import gcg__zou2023
+from tropt.recipe_hub import gcg__zou2023
 
 result = gcg__zou2023(
     model_name="google/gemma-3-1b-it",
@@ -142,6 +142,6 @@ If you find this package useful, please cite our paper as follows:
 @misc{tropt2026,
   title        = {TROPT: An Open Framework for Unifying and Advancing Discrete Text Optimization},
   year         = {2026},
-  howpublished = {\url{https://github.com/matanbt/tropt}},
+  howpublished = {\url{https://github.com/matanbt/TROPT}},
 }
 ```

tropt/common.py

@@ -86,7 +86,7 @@ class MessageTargets(pydantic.BaseModel):
 
     target_gradient: Optional[Float[Tensor, "n_params"]] = None
     """Target weight-gradient to align with, flattened over the trainable params.
-    Precompute via ``tropt.loss.gradient_matching.compute_weight_gradient``.
+    Precompute the target weight-gradient externally and pass it here.
     """
 
     # ── Classifier targets ─────────────────────────────────────────────────
@@ -172,9 +172,8 @@ class Targets(pydantic.BaseModel):
     """Target weight-gradients, one per template (flattened over the trainable params).
 
     Shape: (n_templates, n_params).
-    Used by: gradient-matching losses. Precompute via
-    ``tropt.loss.gradient_matching.compute_weight_gradient`` and stack/replicate
-    across templates.
+    Used by: gradient-matching losses. Precompute the per-template target
+    weight-gradients externally and stack across templates.
     """
 
     # ── Classifier targets ─────────────────────────────────────────────────

tropt/loss/losses.py

@@ -468,7 +468,7 @@ class SteeringActivationLoss(HiddenStateBasedLoss):
 
     targeted_layers: slice = slice(None)
     steer_away: bool = False
-    slc_name: str = SliceKey.INPUT_LAST_TOKEN
+    slc_name: SliceKey = SliceKey.INPUT_LAST_TOKEN
     do_cosine_sim: bool = False
     apply_square: bool = False
     apply_abs: bool = False

tropt/model/huggingface/base.py

@@ -972,10 +972,10 @@ def compute_loss_from_tokens(
         matching, whose value is itself a weight-gradient).
 
         Args:
-            search_batch_size : int
-                the number of candidate sequences to evaluate in a given batch
-            inputs_embeds : Tensor, shape = (search_width, seq_len, embd_dim)
-                the embeddings of the `search_width` candidate sequences to evaluate
+            candidate_trigger_ids : Tensor, shape = (n_candidates, trigger_seq_len)
+                candidate trigger token ids to evaluate
+            loss_func : BaseLoss
+                the loss to compute for each candidate
             keep_message_dim : bool
                 whether to return the loss per message (shape = (n_templates, n_candidates))
         Returns:
@@ -1015,7 +1015,7 @@ def _compute_candidates_loss__batched(
                 logger.debug(f"from loss [msg={template_idx}]: {(cand_idx_end - cand_idx)}")
 
                 # Only enable gradient is it's required by the loss (e.g. for gradient matching losses); mostly false.
-                with torch.set_grad_enabled(loss_func.require_gradients):                
+                with torch.set_grad_enabled(loss_func.require_gradients):
                     model_input = input_manager.get_triggered_inputs(
                         chosen_template_idx=template_idx,
                         trigger_ids=batch_candidate_trigger_ids,

tropt/model/litellm_proxy/lm.py

@@ -1,8 +1,6 @@
 import logging
 from typing import Dict, List, Optional
 
-import litellm
-
 from tropt.common import ModelOutput
 from tropt.model import BaseTokenizer, LMBaseModel, LossTextAccessMixin
 from tropt.model.openai.encoder import OpenAITokenizer
@@ -134,6 +132,9 @@ def invoke_from_texts(
             generation_kwargs["logprobs"] = True
             generation_kwargs["top_logprobs"] = _MAX_TOP_LOGPROBS
 
+        # Import litellm lazily: optional dependency (`tropt[litellm]`).
+        import litellm
+
         outputs = litellm.batch_completion(
             model=self.model_name,
             messages=prompts,

tropt/model/openai/encoder.py

@@ -1,9 +1,7 @@
 from typing import Annotated, List, Literal, Optional
 
 import numpy as np
-import tiktoken
 import torch
-from openai import OpenAI
 from tenacity import retry, stop_after_attempt, wait_exponential
 from transformers import BatchEncoding
 
@@ -31,6 +29,9 @@ class OpenAITokenizer(BaseTokenizer):
     A wrapper around OpenAI's tokenizer that mimics the HuggingFace interface.
     """
     def __init__(self, model_name: str) -> None:
+        # Import tiktoken lazily: optional dependency (`tropt[openai]`).
+        import tiktoken
+
         # Get the tokeniser corresponding to a specific model in the OpenAI API
         try:
             self._encoding = tiktoken.encoding_for_model(model_name)
@@ -158,6 +159,8 @@ def __init__(
             base_url: Optional base URL for the OpenAI client. If None, the default OpenAI API URL is used.
         """
         # Import openai only when instantiating (optional dependency)
+        from openai import OpenAI
+
         self._client = OpenAI(api_key=api_key, base_url=base_url)
         self.model_name = model_name
 

tropt/optimizer/gasliteplus_optimizer.py

@@ -216,7 +216,7 @@ def optimize_trigger(
             losses=[losses[i].item() for i in range(self.buffer_size)],
         )
 
-        trigger_str = tokenizer.decode(buffer.get_best_trigger(), skip_special_tokens=True)
+        trigger_str = tokenizer.decode_trigger(buffer.get_best_trigger())
         self.log(loss=buffer.get_lowest_loss(), trigger_str=trigger_str)
 
         for step in self.track_steps(range(self.num_steps), desc="Optimizing with GASLITE..."):

tropt/optimizer/gcg_optimizer.py

@@ -118,7 +118,7 @@ def optimize_trigger(
             # Compute loss on all candidate sequences
             losses = self.model.compute_loss_from_tokens(
                 candidate_trigger_ids, loss_func=self.loss_func
-            )  # shape: (n_templates, n_candidates)
+            )  # shape: (n_candidates,)
             current_loss = losses.min().item()
             trigger_ids = candidate_trigger_ids[losses.argmin()]
             trigger_str = tokenizer.decode_trigger(trigger_ids)

tropt/optimizer/utils/retokenization.py

@@ -158,9 +158,6 @@ def full_messages_retokenize_filtering(
                 [triggered_template], add_special_tokens=False
             ).input_ids[0]
 
-            print("old texts:", tokenizer.decode(triggered_template_ids))
-            print("new texts:", tokenizer.decode(triggered_template_new_ids))
-
             # 3. We want the original to match the (re)tokenization:
             if triggered_template_ids != triggered_template_new_ids:
                 is_curr_trigger_valid = False