Make rollout completion tokens configurable

skillopt/envs/livemathematicianbench/adapter.py

@@ -60,10 +60,13 @@ class LiveMathematicianBenchAdapter(EnvAdapter):
         limit: int = 0,
         shuffle_choices: bool = True,
         use_theorem: bool = False,
-        use_sketch: bool = False,    ) -> None:
+        use_sketch: bool = False,
+        max_completion_tokens: int = 16384,
+    ) -> None:
         self.max_turns = max_turns
         self.exec_timeout = exec_timeout
         self.workers = workers
+        self.max_completion_tokens = int(max_completion_tokens)
         self.analyst_workers = analyst_workers
         self.failure_only = failure_only
         self.minibatch_size = minibatch_size
@@ -115,6 +118,7 @@ class LiveMathematicianBenchAdapter(EnvAdapter):
             max_turns=self.max_turns,
             exec_timeout=self.exec_timeout,
             workers=self.workers,
+            max_completion_tokens=self.max_completion_tokens,
             use_theorem=self.use_theorem,
             use_sketch=self.use_sketch,
             diagnostic_mode=kwargs.get("diagnostic_mode", False),

skillopt/envs/livemathematicianbench/rollout.py

@@ -120,6 +120,7 @@ def process_one(
     diagnostic_instruction: str = "",
     diagnostic_trace_context: str = "",
     exec_timeout: int = 300,
+    max_completion_tokens: int = 16384,
 ) -> dict:
     item_id = str(item["id"])
     result = {
@@ -219,7 +220,7 @@ def process_one(
                 resp_text, _ = chat_target(
                     system=system,
                     user=user,
-                    max_completion_tokens=16384,
+                    max_completion_tokens=max_completion_tokens,
                     retries=5,
                     stage="rollout",
                     timeout=exec_timeout,
@@ -233,7 +234,7 @@ def process_one(
                 resp_text, _ = chat_target(
                     system=system,
                     user=refinement,
-                    max_completion_tokens=16384,
+                    max_completion_tokens=max_completion_tokens,
                     retries=5,
                     stage="rollout",
                     timeout=exec_timeout,
@@ -293,6 +294,7 @@ def run_batch(
     max_turns: int = 1,
     exec_timeout: int = 300,
     workers: int = 64,
+    max_completion_tokens: int = 16384,
     use_theorem: bool = False,
     use_sketch: bool = False,
     diagnostic_mode: bool = False,
@@ -338,6 +340,7 @@ def run_batch(
             skill_content,
             max_turns=max_turns,
             exec_timeout=exec_timeout,
+            max_completion_tokens=max_completion_tokens,
             use_theorem=use_theorem,
             use_sketch=use_sketch,
             diagnostic_mode=diagnostic_mode,