pyright: Fix simple typing issues

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

@@ -126,6 +126,7 @@ def load_tokenizer(self) -> PreTrainedTokenizer:
 
         return tokenizer
 
+    @override
     def evaluate_strong_reject(self) -> dict[str, float]:
         """Evaluate attack on the `StrongRejectEvaluation` evaluator."""
         eval_cfg: StrongRejectEvaluationConfig = StrongRejectEvaluationConfig(

src/tamperbench/whitebox/evals/mbpp/mbpp.py

@@ -8,7 +8,7 @@
 """
 
 # pyright: reportUnknownArgumentType=false, reportUnknownMemberType=false, reportAny=false, reportMissingParameterType=false, reportUnknownParameterType=false, reportUnknownVariableType=false, reportMissingTypeArgument=false, reportArgumentType=false, reportCallIssue=false, reportOptionalMemberAccess=false
-# ruff: noqa: D401, B006, B905
+# ruff: noqa: D401, B905
 
 from __future__ import annotations
 
@@ -101,7 +101,7 @@ def postprocess_generation(generation: str, prompt: str) -> str:
 def compute_code_eval(
     predictions: list[list[str]],
     references: list[str],
-    k: list[int] = [1],
+    k: list[int] | None = None,
     num_workers: int = 4,
     timeout: float = 3.0,
 ) -> tuple[dict[str, float], dict]:
@@ -171,11 +171,10 @@ def check_correctness_fork(check_program, timeout, task_id, completion_id):
     total = np.array(total)
     correct = np.array(correct)
 
-    ks = k
-    if not isinstance(ks, list | tuple):
-        ks = [ks]
+    if k is None:
+        k = [1]
     pass_at_k = {
-        f"pass@{k_val}": estimate_pass_at_k(total, correct, k_val).mean() for k_val in ks if (total >= k_val).all()
+        f"pass@{k_val}": estimate_pass_at_k(total, correct, k_val).mean() for k_val in k if (total >= k_val).all()
     }
 
     return pass_at_k, results

src/tamperbench/whitebox/evals/mmlu_pro/mmlu_pro.py

@@ -65,7 +65,7 @@ def compute_inferences(self) -> DataFrame[InferenceSchema]:
         shot_pool = self._build_shot_pool(val_rows, test_rows)
         eval_rows = self._build_eval_items(val_rows, test_rows)
 
-        self._cached_eval_rows = list(eval_rows)
+        self._cached_eval_rows: list[dict] = list(eval_rows)
 
         payload: pl.DataFrame = run_in_isolation(
             target=compute_mmlu_inferences,

src/tamperbench/whitebox/evals/mt_bench/mt_bench.py

@@ -14,6 +14,7 @@
 from torch import Tensor
 from tqdm import tqdm
 from tqdm.contrib.concurrent import thread_map
+from transformers import PreTrainedModel, PreTrainedTokenizer
 from transformers.generation.utils import GenerateOutput
 from transformers.tokenization_utils_base import BatchEncoding
 from typing_extensions import override
@@ -99,10 +100,10 @@ def __init__(self, eval_config: MTBenchEvaluationConfig) -> None:
             eval_config: Configuration for the evaluation.
         """
         super().__init__(eval_config)
-        self.openai_client = OpenAI()
-        self.questions = load_mt_bench_questions_by_category(eval_config.categories)
+        self.openai_client: OpenAI = OpenAI()
+        self.questions: list[MTBenchQuestion] = load_mt_bench_questions_by_category(eval_config.categories)
 
-    def _format_prompt(self, messages: list[dict[str, str]], tokenizer) -> str:
+    def _format_prompt(self, messages: list[dict[str, str]], tokenizer: PreTrainedTokenizer) -> str:
         """Format messages into a prompt string using the tokenizer's chat template.
 
         Args:
@@ -113,7 +114,9 @@ def _format_prompt(self, messages: list[dict[str, str]], tokenizer) -> str:
             str: Formatted prompt string.
         """
         if hasattr(tokenizer, "apply_chat_template"):
-            return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            result = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            assert isinstance(result, str)
+            return result
         # Fallback for models without chat template
         prompt = ""
         for msg in messages:
@@ -126,7 +129,9 @@ def _format_prompt(self, messages: list[dict[str, str]], tokenizer) -> str:
         prompt += "Assistant: "
         return prompt
 
-    def _generate_response(self, prompt: str, model, tokenizer, max_new_tokens: int = 512) -> str:
+    def _generate_response(
+        self, prompt: str, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, max_new_tokens: int = 512
+    ) -> str:
         """Generate a response from the model.
 
         Args: