small fixes to turn completer integration in pb

Author: thesofakillers

project/paperbench/paperbench/judge/simple.py

@@ -445,9 +445,7 @@ async def _prepare_relevant_files(
 
         file_content_tasks = [
             read_file_content(
-                (self.submission_dir / rel_path.strip().strip("/")).relative_to(
-                    self.submission_dir
-                ),
+                self.submission_dir / rel_path.strip().strip("/"),
                 self.computer,
             )
             for rel_path in selected_files.split("\n")[: max_files or None]

project/paperbench/tests/unit/test_judge.py

@@ -1,21 +1,29 @@
 import json
 import math
+import os
 import shutil
 from pathlib import Path
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 from typing import Callable, Generator
 
 import pytest
+from dotenv import load_dotenv
+from preparedness_turn_completer.oai_turn_completer import OpenAITurnCompleter
 
 from paperbench.judge.base import Judge
 from paperbench.judge.dummyrandom import DummyJudge
 from paperbench.judge.simple import SimpleJudge
 from paperbench.rubric.tasks import TaskNode
-from paperbench.utils import in_ci
+from paperbench.utils import find_dotenv, in_ci
+
+load_dotenv(find_dotenv())
 
 non_dummy_judges = [SimpleJudge]
 
 
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+
+
 def get_ancestor(name: str) -> Path:
     """Returns the path to an ancestor directory with `name`, starting from the current file."""
 
@@ -119,7 +127,7 @@ async def test_all_gold_submissions_achieve_a_perfect_score_on_a_trivial_rubric(
         addendum=None,
         judge_addendum=None,
         submission_dir=gold_submission,
-        model="gpt-4o",
+        completer_config=OpenAITurnCompleter.Config(model="gpt-4o"),
         paper_md=empty_markdown,
     )
 
@@ -134,7 +142,8 @@ async def test_all_gold_submissions_achieve_a_perfect_score_on_a_trivial_rubric(
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("JudgeClass", non_dummy_judges)  # Skip DummyJudge
-@pytest.mark.skipif(in_ci(), reason="Test does not apply to DummyJudge")
+@pytest.mark.skipif(in_ci(), reason="Skip non-dummy judge in CI.")
+@pytest.mark.skipif(OPENAI_API_KEY is None, reason="OPENAI_API_KEY is not set.")
 @pytest.mark.parametrize("task", ["empty", "hex_flag", "hex_flags", "nested_hex_flags"])
 async def test_all_gold_submissions_achieve_a_null_score_on_an_impossible_rubric(
     task: str,
@@ -153,7 +162,7 @@ async def test_all_gold_submissions_achieve_a_null_score_on_an_impossible_rubric
         addendum=None,
         judge_addendum=None,
         submission_dir=gold_submission,
-        model="gpt-4o",
+        completer_config=OpenAITurnCompleter.Config(model="gpt-4o"),
         paper_md=empty_markdown,
     )
 
@@ -168,7 +177,8 @@ async def test_all_gold_submissions_achieve_a_null_score_on_an_impossible_rubric
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("JudgeClass", non_dummy_judges)  # Skip DummyJudge
-@pytest.mark.skipif(in_ci(), reason="Test does not apply to DummyJudge")
+@pytest.mark.skipif(in_ci(), reason="Skip non-dummy judge in CI.")
+@pytest.mark.skipif(OPENAI_API_KEY is None, reason="OPENAI_API_KEY is not set.")
 @pytest.mark.parametrize("task", ["empty", "hex_flag", "hex_flags", "nested_hex_flags"])
 async def test_all_gold_submissions_achieve_a_perfect_score_on_their_corresponding_rubric(
     task: str,
@@ -187,7 +197,7 @@ async def test_all_gold_submissions_achieve_a_perfect_score_on_their_correspondi
         addendum=None,
         judge_addendum=None,
         submission_dir=gold_submission,
-        model="gpt-4o",
+        completer_config=OpenAITurnCompleter.Config(model="gpt-4o"),
         paper_md=empty_markdown,
     )
 
@@ -202,7 +212,8 @@ async def test_all_gold_submissions_achieve_a_perfect_score_on_their_correspondi
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("JudgeClass", non_dummy_judges)  # Skip DummyJudge
-@pytest.mark.skipif(in_ci(), reason="Test does not apply to DummyJudge")
+@pytest.mark.skipif(in_ci(), reason="Skip non-dummy judge in CI.")
+@pytest.mark.skipif(OPENAI_API_KEY is None, reason="OPENAI_API_KEY is not set.")
 @pytest.mark.parametrize("task", ["hex_flag", "hex_flags", "nested_hex_flags"])
 async def test_empty_submission_achieves_a_null_score_on_all_non_trvial_rubrics(
     task: str,
@@ -221,7 +232,7 @@ async def test_empty_submission_achieves_a_null_score_on_all_non_trvial_rubrics(
         addendum=None,
         judge_addendum=None,
         submission_dir=empty_submission,
-        model="gpt-4o",
+        completer_config=OpenAITurnCompleter.Config(model="gpt-4o"),
         paper_md=empty_markdown,
     )
 
@@ -236,7 +247,8 @@ async def test_empty_submission_achieves_a_null_score_on_all_non_trvial_rubrics(
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("JudgeClass", non_dummy_judges)  # Skip DummyJudge
-@pytest.mark.skipif(in_ci(), reason="Test does not apply to DummyJudge")
+@pytest.mark.skipif(in_ci(), reason="Skip non-dummy judge in CI.")
+@pytest.mark.skipif(OPENAI_API_KEY is None, reason="OPENAI_API_KEY is not set.")
 @pytest.mark.parametrize(
     "n_missing",
     [
@@ -268,7 +280,7 @@ async def test_submission_with_n_missing_files_to_the_hex_flags_task_achieves_a_
         addendum=None,
         judge_addendum=None,
         submission_dir=submission,
-        model="gpt-4o",
+        completer_config=OpenAITurnCompleter.Config(model="gpt-4o"),
         paper_md=empty_markdown,
     )
 
@@ -284,7 +296,8 @@ async def test_submission_with_n_missing_files_to_the_hex_flags_task_achieves_a_
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("JudgeClass", non_dummy_judges)  # Skip DummyJudge
-@pytest.mark.skipif(in_ci(), reason="Test does not apply to DummyJudge")
+@pytest.mark.skipif(in_ci(), reason="Skip non-dummy judge in CI.")
+@pytest.mark.skipif(OPENAI_API_KEY is None, reason="OPENAI_API_KEY is not set.")
 async def test_nested_context_preserved_in_grading(
     JudgeClass: type[Judge],
     empty_pdf: Path,
@@ -308,7 +321,7 @@ async def test_nested_context_preserved_in_grading(
         addendum=None,
         judge_addendum=None,
         submission_dir=submission,
-        model="gpt-4o",
+        completer_config=OpenAITurnCompleter.Config(model="gpt-4o"),
         paper_md=empty_markdown,
     )
 

project/preparedness_turn_completer/preparedness_turn_completer/oai_turn_completer.py

@@ -30,7 +30,12 @@ class OpenAITurnCompleter(TurnCompleter):
     def __init__(self, model: str, reasoning_effort: str | None = None):
         self.model: str = model
         self.reasoning_effort: str | None = reasoning_effort
-        self.encoding_name: str = tiktoken.encoding_name_for_model(model)
+        try:
+            self.encoding_name: str = tiktoken.encoding_name_for_model(model)
+        except KeyError:
+            # Fallback to o200k_base
+            logger.warning(f"Model {model} not found in tiktoken, using o200k_base")
+            self.encoding_name: str = "o200k_base"
         self.n_ctx: int = get_model_context_window_length(model)
 
     class Config(TurnCompleter.Config):
@@ -180,9 +185,11 @@ def get_model_context_window_length(model: str | None) -> int:
         "o1-mini-2024-09-12": 128000,
         "o1": 200000,
         "o1-2024-12-17": 200000,
+        "o3": 200000,
         "o3-mini-2024-12-17": 128000,
         "o3-mini-2025-01-31": 200000,
         "o3-mini": 200000,
+        "o4-mini": 200000,
         "o1-preview": 128000,
         "gpt-4-turbo": 128000,
     }