Make rollout completion tokens configurable

This commit is contained in:
hwq
2026-05-28 09:45:47 +00:00
parent 99212e3956
commit 786d57b5cf
20 changed files with 63 additions and 24 deletions
+5 -1
View File
@@ -27,10 +27,13 @@ class DocVQAAdapter(EnvAdapter):
edit_budget: int = 4,
seed: int = 42,
limit: int = 0,
image_detail: str = "auto", ) -> None:
image_detail: str = "auto",
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
@@ -75,6 +78,7 @@ class DocVQAAdapter(EnvAdapter):
exec_timeout=self.exec_timeout,
workers=self.workers,
image_detail=self.image_detail,
max_completion_tokens=self.max_completion_tokens,
diagnostic_mode=kwargs.get("diagnostic_mode", False),
diagnostic_instruction=kwargs.get("diagnostic_instruction", ""),
task_timeout=self.exec_timeout,
+5 -2
View File
@@ -134,6 +134,7 @@ def process_one(
max_turns: int = 1,
exec_timeout: int = 120,
image_detail: str = "auto",
max_completion_tokens: int = 16384,
diagnostic_mode: bool = False,
diagnostic_instruction: str = "",
) -> dict:
@@ -200,7 +201,7 @@ def process_one(
if turn == 0:
resp_text, _ = chat_target_messages(
messages=messages,
max_completion_tokens=768,
max_completion_tokens=max_completion_tokens,
retries=5,
stage="rollout",
timeout=exec_timeout,
@@ -214,7 +215,7 @@ def process_one(
]
resp_text, _ = chat_target_messages(
messages=refinement_messages,
max_completion_tokens=512,
max_completion_tokens=max_completion_tokens,
retries=5,
stage="rollout",
timeout=exec_timeout,
@@ -266,6 +267,7 @@ def run_batch(
exec_timeout: int = 120,
workers: int = 16,
image_detail: str = "auto",
max_completion_tokens: int = 16384,
diagnostic_mode: bool = False,
diagnostic_instruction: str = "",
task_timeout: int = 600,
@@ -325,6 +327,7 @@ def run_batch(
max_turns=max_turns,
exec_timeout=exec_timeout,
image_detail=image_detail,
max_completion_tokens=max_completion_tokens,
diagnostic_mode=diagnostic_mode,
diagnostic_instruction=diagnostic_instruction,
)