diff --git a/rdagent/log/ui/utils.py b/rdagent/log/ui/utils.py index 6c5a90e7..2f25f8b2 100644 --- a/rdagent/log/ui/utils.py +++ b/rdagent/log/ui/utils.py @@ -13,6 +13,7 @@ from rdagent.core.utils import cache_with_pickle from rdagent.log.ui.conf import UI_SETTING from rdagent.log.utils import extract_json from rdagent.oai.llm_utils import md5_hash +from rdagent.scenarios.kaggle.kaggle_crawler import get_metric_direction LITE = [ "aerial-cactus-identification", @@ -524,6 +525,7 @@ def compare( exp_list: list[str] = typer.Option(..., "--exp-list", help="List of experiment names.", show_default=False), output: str = typer.Option("merge_base_df.h5", help="Output summary file name."), hours: int | None = typer.Option(None, help="if None, use summary.pkl, else summary_{hours}h.pkl"), + select_best: bool = typer.Option(False, help="Select best experiment for each competition."), ): """ Generate summary and base dataframe for given experiment list, and save to a summary file. @@ -531,6 +533,23 @@ def compare( typer.secho(f"exp_list: {exp_list}", fg=typer.colors.GREEN) log_folders = [f"{UI_SETTING.amlt_path}/{exp}/combined_logs" for exp in exp_list] summary, base_df = get_summary_df(log_folders, hours=hours) + if select_best: + + def apply_func(cdf: pd.DataFrame): + cp = cdf["Competition"].values[0] + md = get_metric_direction(cp) + # If SOTA Exp Score (valid) column is empty, return the first index + if cdf["SOTA Exp Score (valid)"].dropna().empty: + return cdf.index[0] + if md: + best_idx = cdf["SOTA Exp Score (valid)"].idxmax() + else: + best_idx = cdf["SOTA Exp Score (valid)"].idxmin() + return best_idx + + best_idxs = base_df.groupby("Competition").apply(apply_func) + base_df = base_df[base_df.index.isin(best_idxs.values)] + summary = {k: v for k, v in summary.items() if k in best_idxs.values.tolist()} typer.secho(f"Summary keys: {list(summary.keys())}", fg=typer.colors.CYAN) typer.secho("Summary DataFrame:", fg=typer.colors.MAGENTA) typer.secho(str(base_df), fg=typer.colors.YELLOW)