diff --git a/metagpt/ext/sela/data/custom_task.py b/metagpt/ext/sela/data/custom_task.py index 08a7cbabb..69abaff5e 100644 --- a/metagpt/ext/sela/data/custom_task.py +++ b/metagpt/ext/sela/data/custom_task.py @@ -6,6 +6,7 @@ from metagpt.ext.sela.runner.mle_bench.instructions import ( INSTRUCTIONS, INSTRUCTIONS_OBFUSCATED, ) +from metagpt.ext.sela.utils import mcts_logger MLE_BENCH_FILES = ["description.md", "description_obfuscated.md"] @@ -70,5 +71,5 @@ def get_mle_bench_requirements(dataset_dir, data_config, special_instruction, ob output_dir=output_dir, special_instruction=special_instruction, ) - print(mle_requirement) + mcts_logger.info(mle_requirement) return mle_requirement diff --git a/metagpt/ext/sela/data/dataset.py b/metagpt/ext/sela/data/dataset.py index ef4179011..79c67b2b5 100644 --- a/metagpt/ext/sela/data/dataset.py +++ b/metagpt/ext/sela/data/dataset.py @@ -10,7 +10,7 @@ import yaml from sklearn.model_selection import train_test_split from metagpt.ext.sela.insights.solution_designer import SolutionDesigner -from metagpt.ext.sela.utils import DATA_CONFIG +from metagpt.ext.sela.utils import DATA_CONFIG, mcts_logger BASE_USER_REQUIREMENT = """ This is a {datasetname} dataset. Your goal is to predict the target column `{target_col}`. @@ -191,7 +191,7 @@ def generate_task_requirement(task_name, data_config, is_di=True, special_instru additional_instruction=additional_instruction, data_info_path=data_info_path, ) - print(user_requirement) + mcts_logger.info(user_requirement) return user_requirement @@ -286,16 +286,16 @@ class ExpDataset: def save_dataset(self, target_col): df, test_df = self.get_raw_dataset() if not self.check_dataset_exists() or self.force_update: - print(f"Saving Dataset {self.name} in {self.dataset_dir}") + mcts_logger.info(f"Saving Dataset {self.name} in {self.dataset_dir}") self.split_and_save(df, target_col, test_df=test_df) else: - print(f"Dataset {self.name} already exists") + mcts_logger.info(f"Dataset {self.name} already exists") if not self.check_datasetinfo_exists() or self.force_update: - print(f"Saving Dataset info for {self.name}") + mcts_logger.info(f"Saving Dataset info for {self.name}") dataset_info = self.get_dataset_info() self.save_datasetinfo(dataset_info) else: - print(f"Dataset info for {self.name} already exists") + mcts_logger.info(f"Dataset info for {self.name} already exists") def save_datasetinfo(self, dataset_info): with open(Path(self.dataset_dir, self.name, "dataset_info.json"), "w", encoding="utf-8") as file: diff --git a/metagpt/ext/sela/evaluation/visualize_mcts.py b/metagpt/ext/sela/evaluation/visualize_mcts.py index 6f803a91c..266766cae 100644 --- a/metagpt/ext/sela/evaluation/visualize_mcts.py +++ b/metagpt/ext/sela/evaluation/visualize_mcts.py @@ -4,6 +4,7 @@ import matplotlib.pyplot as plt import networkx as nx from metagpt.ext.sela.search.tree_search import Node +from metagpt.ext.sela.utils import mcts_logger NODE_TEMPLATE = """\ [Node {id}] @@ -139,7 +140,7 @@ def build_tree_recursive(graph, parent_id, node, node_order, start_task_id=2): instruction = "\n\n".join([role.planner.plan.tasks[i].instruction for i in range(start_task_id)]) else: instruction = role.planner.plan.tasks[depth + start_task_id - 1].instruction - print(instruction) + mcts_logger.info(instruction) # Add the current node with attributes to the graph dev_score = node.raw_reward.get("dev_score", 0) * 100 avg_score = node.avg_value() * 100 diff --git a/metagpt/ext/sela/experimenter.py b/metagpt/ext/sela/experimenter.py index b05ea2fc3..2f8bef96d 100644 --- a/metagpt/ext/sela/experimenter.py +++ b/metagpt/ext/sela/experimenter.py @@ -133,9 +133,9 @@ class Experimenter(DataInterpreter): if self.planner.plan.goal != "": self.set_actions([WriteAnalysisCode]) self._set_state(0) - print("Plan already exists, skipping initialization.") + mcts_logger.info("Plan already exists, skipping initialization.") return self - print("Initializing plan and tool...") + mcts_logger.info("Initializing plan and tool...") return super().set_plan_and_tool() async def _act_on_task(self, current_task: Task) -> TaskResult: diff --git a/metagpt/ext/sela/runner/aide.py b/metagpt/ext/sela/runner/aide.py index 50fae94c1..6fd17b87c 100644 --- a/metagpt/ext/sela/runner/aide.py +++ b/metagpt/ext/sela/runner/aide.py @@ -3,6 +3,8 @@ import time import aide +from metagpt.ext.sela.utils import mcts_logger + os.environ["OPENAI_API_KEY"] = "sk-xxx" os.environ["OPENAI_BASE_URL"] = "your url" @@ -27,9 +29,9 @@ exp = aide.Experiment( best_solution = exp.run(steps=10) -print(f"Best solution has validation metric: {best_solution.valid_metric}") -print(f"Best solution code: {best_solution.code}") +mcts_logger.info(f"Best solution has validation metric: {best_solution.valid_metric}") +mcts_logger.info(f"Best solution code: {best_solution.code}") end_time = time.time() execution_time = end_time - start_time -print(f"run time : {execution_time} seconds") +mcts_logger.info(f"run time : {execution_time} seconds") diff --git a/metagpt/ext/sela/runner/mcts.py b/metagpt/ext/sela/runner/mcts.py index 8b6c14100..bebd2400c 100644 --- a/metagpt/ext/sela/runner/mcts.py +++ b/metagpt/ext/sela/runner/mcts.py @@ -7,6 +7,7 @@ from metagpt.ext.sela.evaluation.evaluation import ( from metagpt.ext.sela.evaluation.visualize_mcts import get_tree_text from metagpt.ext.sela.runner.runner import Runner from metagpt.ext.sela.search.search_algorithm import MCTS, Greedy, Random +from metagpt.ext.sela.utils import mcts_logger class MCTSRunner(Runner): @@ -46,7 +47,7 @@ class MCTSRunner(Runner): text += f"Best node: {best_node.id}, score: {best_node.raw_reward}\n" text += f"Dev best node: {dev_best_node.id}, score: {dev_best_node.raw_reward}\n" text += f"Grader score: {additional_scores['grader']}\n" - print(text) + mcts_logger.info(text) results = [ { "best_node": best_node.id, diff --git a/metagpt/ext/sela/runner/random_search.py b/metagpt/ext/sela/runner/random_search.py index b1f43ac0c..a9424cbe5 100644 --- a/metagpt/ext/sela/runner/random_search.py +++ b/metagpt/ext/sela/runner/random_search.py @@ -1,7 +1,7 @@ from metagpt.ext.sela.experimenter import Experimenter from metagpt.ext.sela.insights.instruction_generator import InstructionGenerator from metagpt.ext.sela.runner.runner import Runner -from metagpt.ext.sela.utils import get_exp_pool_path +from metagpt.ext.sela.utils import get_exp_pool_path, mcts_logger EXPS_PROMPT = """ When doing the tasks, you can refer to the insights below: @@ -37,7 +37,7 @@ class RandomSearchRunner(Runner): di = Experimenter(node_id=str(i), use_reflection=self.args.reflection, role_timeout=self.args.role_timeout) di.role_dir = f"{di.role_dir}_{self.args.task}" requirement = user_requirement + EXPS_PROMPT.format(experience=exps[i]) - print(requirement) + mcts_logger.info(requirement) score_dict = await self.run_di(di, requirement, run_idx=i) results.append( { diff --git a/metagpt/ext/sela/runner/runner.py b/metagpt/ext/sela/runner/runner.py index 4b5504e09..f10aa29f3 100644 --- a/metagpt/ext/sela/runner/runner.py +++ b/metagpt/ext/sela/runner/runner.py @@ -8,7 +8,7 @@ import pandas as pd from metagpt.ext.sela.evaluation.evaluation import evaluate_score from metagpt.ext.sela.experimenter import Experimenter from metagpt.ext.sela.search.tree_search import create_initial_state -from metagpt.ext.sela.utils import DATA_CONFIG, save_notebook +from metagpt.ext.sela.utils import DATA_CONFIG, mcts_logger, save_notebook class Runner: @@ -38,7 +38,7 @@ class Runner: score_dict = self.evaluate(score_dict, self.state) run_finished = True except Exception as e: - print(f"Error: {e}") + mcts_logger.info(f"Error: {e}") num_runs += 1 # save_notebook(role=di, save_dir=self.result_path, name=f"{self.args.task}_{self.start_time}_{run_idx}") save_name = self.get_save_name() diff --git a/metagpt/ext/sela/search/tree_search.py b/metagpt/ext/sela/search/tree_search.py index eac26c86c..94d685080 100644 --- a/metagpt/ext/sela/search/tree_search.py +++ b/metagpt/ext/sela/search/tree_search.py @@ -374,7 +374,7 @@ class BaseTreeSearch: return best_score, best_child for child in self.children[node]: score = child.normalized_reward[split] - print(child.id, split, score) + mcts_logger.info(f"{child.id} {split} {score}") if score > best_score: best_score = score best_child = child diff --git a/metagpt/ext/sela/utils.py b/metagpt/ext/sela/utils.py index 21b311e7f..eb112dc09 100644 --- a/metagpt/ext/sela/utils.py +++ b/metagpt/ext/sela/utils.py @@ -57,7 +57,7 @@ def get_exp_pool_path(task_name, data_config, pool_name="analysis_pool"): def change_plan(role, plan): - print(f"Change next plan to: {plan}") + mcts_logger.info(f"Change next plan to: {plan}") tasks = role.planner.plan.tasks finished = True for i, task in enumerate(tasks): @@ -115,8 +115,8 @@ async def load_execute_notebook(role): # await executor.build() for code in codes: outputs, success = await executor.run(code) - print(f"Execution success: {success}, Output: {outputs}") - print("Finish executing the loaded notebook") + mcts_logger.info(f"Execution success: {success}, Output: {outputs}") + mcts_logger.info("Finish executing the loaded notebook") return executor