Make rollout completion tokens configurable
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user