diff --git a/src/tamperbench/whitebox/attacks/gcg/gcg.py b/src/tamperbench/whitebox/attacks/gcg/gcg.py index d84fb51d..8a0ec3c1 100644 --- a/src/tamperbench/whitebox/attacks/gcg/gcg.py +++ b/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( diff --git a/src/tamperbench/whitebox/evals/mbpp/mbpp.py b/src/tamperbench/whitebox/evals/mbpp/mbpp.py index 07f236ab..c5a4ff79 100644 --- a/src/tamperbench/whitebox/evals/mbpp/mbpp.py +++ b/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 diff --git a/src/tamperbench/whitebox/evals/mmlu_pro/mmlu_pro.py b/src/tamperbench/whitebox/evals/mmlu_pro/mmlu_pro.py index 61507476..92aa2e92 100644 --- a/src/tamperbench/whitebox/evals/mmlu_pro/mmlu_pro.py +++ b/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, diff --git a/src/tamperbench/whitebox/evals/mt_bench/mt_bench.py b/src/tamperbench/whitebox/evals/mt_bench/mt_bench.py index cc760de9..e8708bb3 100644 --- a/src/tamperbench/whitebox/evals/mt_bench/mt_bench.py +++ b/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: