From 214da4826fac46858f63282a489cc2e1e9af4e8c Mon Sep 17 00:00:00 2001 From: Roland Minrui <114476598+RolandMinrui@users.noreply.github.com> Date: Mon, 14 Apr 2025 11:02:45 +0800 Subject: [PATCH] fix: add wait_retry to exp_gen v2 (#783) * add wait retry to v2 * format * fix a bug --------- Co-authored-by: Xu Co-authored-by: yuanteli <1957922024@qq.com> --- .../data_science/proposal/exp_gen/proposal.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py b/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py index f9c37874..d7579cde 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py @@ -242,13 +242,7 @@ class DSProposalV2ExpGen(ExpGen): ) return json.loads(response) - def identify_feedback_problem( - self, - scenario_desc: str, - exp_feedback_list_desc: str, - sota_exp_desc: str, - pipeline: bool, - ) -> Dict: + def identify_feedback_problem(self, scenario_desc: str, exp_feedback_list_desc: str, sota_exp_desc: str) -> Dict: sys_prompt = T(".prompts_v2:scenario_problem.system").r( problem_spec=T(".prompts_v2:specification.problem").r(), problem_output_format=T(".prompts_v2:output_format.problem").r(), @@ -266,6 +260,13 @@ class DSProposalV2ExpGen(ExpGen): ) return json.loads(response) + def _append_retry(args: tuple, kwargs: dict) -> tuple[tuple, dict]: + # Only modify the user_prompt on retries (i > 0) + user_prompt = args[0] + user_prompt += "\n\nretrying..." + return (user_prompt,), kwargs + + @wait_retry(retry_n=5, transform_args_fn=_append_retry) def hypothesis_gen( self, component_desc: str, @@ -293,7 +294,13 @@ class DSProposalV2ExpGen(ExpGen): json_mode=True, json_target_type=Dict[str, Dict[str, str | Dict[str, str | int]]], ) - return json.loads(response) + resp_dict = json.loads(response) + for key, value in resp_dict.items(): + assert "reason" in value, "Reason not provided." + assert "component" in value, "Component not provided." + assert "hypothesis" in value, "Hypothesis not provided." + assert "evaluation" in value, "Evaluation not provided." + return resp_dict def hypothesis_rank(self, hypothesis_dict: dict, problem_dict: dict, pipeline: bool) -> DSHypothesis: weights = { @@ -361,7 +368,6 @@ class DSProposalV2ExpGen(ExpGen): component_desc=component_desc, workflow_check=not pipeline and hypothesis.component != "Workflow", ) - user_prompt = T(".prompts_v2:task_gen.user").r( scenario_desc=scenario_desc, sota_exp_desc=sota_exp_desc, @@ -394,7 +400,6 @@ class DSProposalV2ExpGen(ExpGen): # exp.experiment_workspace.inject_code_from_folder(sota_exp.experiment_workspace.workspace_path) if sota_exp is not None: exp.experiment_workspace.inject_code_from_file_dict(sota_exp.experiment_workspace) - if not pipeline and new_workflow_desc != "No update needed": workflow_task = WorkflowTask( name="Workflow", @@ -442,7 +447,6 @@ class DSProposalV2ExpGen(ExpGen): scenario_desc=scenario_desc, exp_feedback_list_desc=exp_feedback_list_desc, sota_exp_desc=sota_exp_desc, - pipeline=pipeline, ) all_problems = {**scen_problems, **fb_problems}