chore: merge with more trace (#982)

* merge support more traces
* use feedback from all traces
This commit is contained in:
Tim
2025-06-24 10:47:45 +08:00
committed by GitHub
parent 46367e15f5
commit ed84b6b0a3
@@ -119,15 +119,26 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
resp_dict = json.loads(response)
return resp_dict
def get_exp_index(self, trace: DSTrace) -> int:
leaves: list[int] = trace.get_leaves()
if trace.sota_exp_to_submit is not None:
sota_submit_value = trace.sota_exp_to_submit.result.loc["ensemble"].iloc[0]
trace_scores = []
for i, leaf in enumerate(leaves):
if leaf == trace.current_selection[0]:
continue
fb = trace.sota_experiment_fb(selection=(leaf,))
if fb is None:
continue
final_score = fb[0].result.loc["ensemble"].iloc[0]
trace_scores.append((i, abs(final_score - sota_submit_value)))
if trace_scores:
return min(trace_scores, key=lambda item: item[1])[0]
return next((i for i, leaf in enumerate(leaves) if leaf != trace.current_selection[0]))
def gen(self, trace: DSTrace) -> DSExperiment:
# Ignore the selection argument and use all leaves instead.
leaves: list[int] = trace.get_leaves()
sota_exp_fb = trace.sota_experiment_fb(selection=trace.current_selection)
exp_index = leaves[1] if trace.current_selection[0] == leaves[0] else leaves[0]
exp_to_merge_fb = trace.sota_experiment_fb(selection=(exp_index,))
if exp_to_merge_fb is None:
exp_to_merge_fb = trace.hist[exp_index]
if sota_exp_fb:
sota_exp_desc = T("scenarios.data_science.share:describe.exp").r(
@@ -139,9 +150,30 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
sota_exp_desc = ""
eda_output = None
success_fb_list = trace.experiment_and_feedback_list_after_init(
return_type="sota", search_type="ancestors", selection=(exp_index,)
)
trace_fbs = []
# find the best exp to merge
leaves: list[int] = trace.get_leaves()
for leaf in leaves:
if leaf == trace.current_selection[0]:
continue
trace_fbs.append(
trace.experiment_and_feedback_list_after_init(
return_type="sota",
search_type="ancestors",
selection=(leaf,),
)
)
num_to_slice = 20
if sum(len(fb_list) for fb_list in trace_fbs) > num_to_slice:
success_fb_trace_count = sum(1 for fb_list in trace_fbs if fb_list)
success_fb_list = [
fb for fb_list in trace_fbs for fb in fb_list[-(num_to_slice // success_fb_trace_count) :]
]
else:
success_fb_list = [fb for fb_list in trace_fbs for fb in fb_list]
if len(success_fb_list) > 0:
exp_to_merge_fb_desc = T("scenarios.data_science.proposal.exp_gen.merge:trace").r(
exp_and_feedback_list=success_fb_list,
@@ -151,6 +183,11 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
pipeline=DS_RD_SETTING.coder_on_whole_pipeline,
)
else:
exp_index = self.get_exp_index(trace)
exp_to_merge_fb = trace.sota_experiment_fb(selection=(exp_index,))
if exp_to_merge_fb is None:
exp_to_merge_fb = trace.hist[exp_index]
exp_to_merge_fb_desc = T("scenarios.data_science.share:describe.feedback").r(
exp_and_feedback=exp_to_merge_fb,
heading="The feedback for the solution to be merged",
@@ -188,6 +225,7 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
class ExpGen2TraceAndMerge(ExpGen):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.merge_exp_gen = MergeExpGen(self.scen)
@@ -396,7 +434,9 @@ class ExpGen2TraceAndMergeV3(ExpGen):
else:
selection = (leaves[0],)
if trace.sota_exp_to_submit is not None:
if trace.is_parent(trace.exp2idx(trace.sota_exp_to_submit), leaves[1]):
selection = (leaves[1],)
for i in range(1, len(leaves)):
if trace.is_parent(trace.exp2idx(trace.sota_exp_to_submit), leaves[i]):
selection = (leaves[i],)
break
trace.set_current_selection(selection)
return self.merge_exp_gen.gen(trace)