pyright gcg: Fix simple typing issues

src/tamperbench/whitebox/attacks/gcg/implementation.py

@@ -46,10 +46,10 @@
 ]
 
 
-def get_nonascii_toks(tokenizer, device="cpu"):
+def get_nonascii_toks(tokenizer: transformers.PreTrainedTokenizer, device: str | torch.device = "cpu") -> Tensor:
     """Return tensor of non-ASCII token ids from the tokenizer vocabulary."""
 
-    def is_ascii(s):
+    def is_ascii(s: str) -> bool:
         return s.isascii() and s.isprintable()
 
     nonascii_toks = []
@@ -69,7 +69,7 @@ def is_ascii(s):
     return torch.tensor(nonascii_toks, device=device)
 
 
-def mellowmax(t: Tensor, alpha=1.0, dim=-1):
+def mellowmax(t: Tensor, alpha: float = 1.0, dim: int = -1) -> Tensor:
     """Compute the mellowmax of a tensor (smooth approximation to max)."""
     return (
         1.0
@@ -126,7 +126,7 @@ def find_executable_batch_size(function: callable = None, starting_batch_size: i
 
     batch_size = starting_batch_size
 
-    def decorator(*args, **kwargs):
+    def decorator(*args: object, **kwargs: object) -> object:
         nonlocal batch_size
         gc.collect()
         torch.cuda.empty_cache()
@@ -209,8 +209,8 @@ class AttackBuffer:
 
     def __init__(self, size: int):
         """Initialize the attack buffer with a given size."""
-        self.buffer = []  # elements are (loss: float, optim_ids: Tensor)
-        self.size = size
+        self.buffer: list[tuple[float, Tensor]] = []
+        self.size: int = size
 
     def add(self, loss: float, optim_ids: Tensor) -> None:
         """Add a candidate to the buffer."""
@@ -237,7 +237,7 @@ def get_highest_loss(self) -> float:
         """Return the highest loss in the buffer."""
         return self.buffer[-1][0]
 
-    def log_buffer(self, tokenizer):
+    def log_buffer(self, tokenizer: transformers.PreTrainedTokenizer) -> None:
         """Log the current buffer contents."""
         message = "buffer:"
         for loss, ids in self.buffer:
@@ -346,20 +346,33 @@ def __init__(
             tokenizer: The tokenizer for the model.
             config: Configuration for the GCG attack.
         """
-        self.model = model
-        self.tokenizer = tokenizer
-        self.config = config
+        self.model: transformers.PreTrainedModel = model
+        self.tokenizer: transformers.PreTrainedTokenizer = tokenizer
+        self.config: GCGConfig = config
 
-        self.embedding_layer = model.get_input_embeddings()
-        self.not_allowed_ids = None if config.allow_non_ascii else get_nonascii_toks(tokenizer, device=model.device)
-        self.prefix_cache = None
-        self.draft_prefix_cache = None
+        self.embedding_layer: torch.nn.Module = model.get_input_embeddings()
+        self.not_allowed_ids: Tensor | None = (
+            None if config.allow_non_ascii else get_nonascii_toks(tokenizer, device=model.device)
+        )
+        self.prefix_cache: list[tuple[Tensor, ...]] | None = None
+        self.draft_prefix_cache: list[tuple[Tensor, ...]] | None = None
+
+        self.stop_flag: bool = False
+
+        self.draft_model: transformers.PreTrainedModel | None = None
+        self.draft_tokenizer: transformers.PreTrainedTokenizer | None = None
+        self.draft_embedding_layer: torch.nn.Module | None = None
 
-        self.stop_flag = False
+        # Set during run()
+        self.target_ids: Tensor
+        self.before_embeds: Tensor
+        self.after_embeds: Tensor
+        self.target_embeds: Tensor
+        self.draft_target_ids: Tensor
+        self.draft_before_embeds: Tensor
+        self.draft_after_embeds: Tensor
+        self.draft_target_embeds: Tensor
 
-        self.draft_model = None
-        self.draft_tokenizer = None
-        self.draft_embedding_layer = None
         if self.config.probe_sampling_config:
             self.draft_model = self.config.probe_sampling_config.draft_model
             self.draft_tokenizer = self.config.probe_sampling_config.draft_tokenizer