diff --git a/rdagent/scenarios/data_science/proposal/exp_gen.py b/rdagent/scenarios/data_science/proposal/exp_gen.py index 054ca3d8..25a6cfbc 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen.py @@ -1,4 +1,7 @@ import json +from typing import Literal + +import pandas as pd from rdagent.components.coder.data_science.ensemble.exp import EnsembleTask from rdagent.components.coder.data_science.feature.exp import FeatureTask @@ -110,14 +113,33 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]): return False def experiment_and_feedback_list_after_init( - self, return_success=True + self, return_type: Literal["sota", "failed", "all"] ) -> list[tuple[DSExperiment, ExperimentFeedback]]: + """ + Retrieve a list of experiments and feedbacks based on the return_type. + + Parameters + ---------- + return_type : str + One of "sota", "failed", "all". + + Returns + ------- + list[tuple[DSExperiment, ExperimentFeedback]] + List of experiments and feedbacks. + """ + final_component = self.COMPLETE_ORDER[-1] has_final_component = False exp_and_feedback_list = [] for exp, fb in self.hist: - if has_final_component and bool(fb) == return_success: - exp_and_feedback_list.append((exp, fb)) + if has_final_component: + if return_type == "all": + exp_and_feedback_list.append((exp, fb)) + elif return_type == "failed" and not fb.decision: + exp_and_feedback_list.append((exp, fb)) + elif return_type == "sota" and fb.decision: + exp_and_feedback_list.append((exp, fb)) if exp.hypothesis.component == final_component and fb: has_final_component = True return exp_and_feedback_list @@ -307,10 +329,18 @@ class DSExpGen(ExpGen): ) ) # we use file_dict for hitting the cache when replicate the experiment in another machine. - sota_exp_feedback_list = trace.experiment_and_feedback_list_after_init(return_success=True) - failed_exp_feedback_list = trace.experiment_and_feedback_list_after_init(return_success=False)[ + sota_exp_feedback_list = trace.experiment_and_feedback_list_after_init(return_type="sota") + failed_exp_feedback_list = trace.experiment_and_feedback_list_after_init(return_type="failed")[ -self.max_trace_hist : ] + all_exp_feedback_list = trace.experiment_and_feedback_list_after_init(return_type="all") + trace_component_to_feedback_df = pd.DataFrame(columns=["component", "hypothesis", "decision"]) + for index, (exp, fb) in enumerate(all_exp_feedback_list): + trace_component_to_feedback_df.loc[f"trial {index + 1}"] = [ + exp.hypothesis.component, + exp.hypothesis.hypothesis, + fb.decision, + ] sota_exp_feedback_list_desc = T("scenarios.data_science.share:describe.trace").r( exp_and_feedback_list=sota_exp_feedback_list, @@ -332,6 +362,11 @@ class DSExpGen(ExpGen): component_user_prompt = T(".prompts:component_gen.user").r( sota_exp_and_feedback_list_desc=sota_exp_feedback_list_desc, failed_exp_and_feedback_list_desc=failed_exp_feedback_list_desc, + component_and_feedback_df=( + trace_component_to_feedback_df.to_string() + if len(trace_component_to_feedback_df) > 0 + else "No experiment and feedback provided" + ), ) resp_dict_component: dict = json.loads( @@ -341,6 +376,7 @@ class DSExpGen(ExpGen): ) component = resp_dict_component.get("component", "Component not provided") + component_reason = resp_dict_component.get("reason", "Reason not provided") sota_exp_model_file_count = len( [ k @@ -396,7 +432,7 @@ class DSExpGen(ExpGen): hypothesis = DSHypothesis( component=component, hypothesis=hypothesis_proposal.get("hypothesis", ""), - reason=hypothesis_proposal.get("reason", ""), + reason=component_reason + "\n" + hypothesis_proposal.get("reason", ""), concise_reason=hypothesis_proposal.get("concise_reason", ""), concise_observation=hypothesis_proposal.get("concise_observation", ""), concise_justification=hypothesis_proposal.get("concise_justification", ""), diff --git a/rdagent/scenarios/data_science/proposal/prompts.yaml b/rdagent/scenarios/data_science/proposal/prompts.yaml index b4d4db4a..43dda637 100644 --- a/rdagent/scenarios/data_science/proposal/prompts.yaml +++ b/rdagent/scenarios/data_science/proposal/prompts.yaml @@ -277,6 +277,11 @@ component_gen: Also, here's the former failed experiments and their feedbacks: {{ failed_exp_and_feedback_list_desc }} + All former trials and their feedbacks are provided in pandas DataFrame format. The user has already made several hypothesis on this scenario and did several evaluation on them: + {{ component_and_feedback_df }} + + Please choose the most proper component to focus on based on the information above. Please balance the exploration and exploitation. + exp_and_feedback: |- {% for experiment, feedback in trace.hist[-10:] %} @@ -298,6 +303,7 @@ hypothesis_specification: |- output_format: component: |- { + "reason": "The reason why you choose this component. Based on the current status and former trials, why this component is the most promising one to focus on.", "component": "The component you suggest to focus on. It must be one of ['DataLoadSpec', 'FeatureEng', 'Model', 'Ensemble', 'Workflow']." } hypothesis: |-