From f2831e7442510668b0ca75953b3359894803ef3c Mon Sep 17 00:00:00 2001
From: Way2Learn <118058822+Xisen-Wang@users.noreply.github.com>
Date: Fri, 20 Sep 2024 16:40:47 +0800
Subject: [PATCH] feat: Model context for tuning and selection (#284)

* Revised for tuning

* Fixed a bug

* fix a ci bug

---------

Co-authored-by: WinstonLiye <1957922024@qq.com>
---
 .../model_coder/CoSTEER/evolving_strategy.py  | 38 +++++++++++++++++--
 .../components/coder/model_coder/prompts.yaml |  5 +++
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/rdagent/components/coder/model_coder/CoSTEER/evolving_strategy.py b/rdagent/components/coder/model_coder/CoSTEER/evolving_strategy.py
index f824c2ba..b67f2cc6 100644
--- a/rdagent/components/coder/model_coder/CoSTEER/evolving_strategy.py
+++ b/rdagent/components/coder/model_coder/CoSTEER/evolving_strategy.py
@@ -11,7 +11,11 @@
 from rdagent.components.coder.model_coder.CoSTEER.knowledge_management import (
     ModelQueriedKnowledge,
 )
-from rdagent.components.coder.model_coder.model import ModelFBWorkspace, ModelTask
+from rdagent.components.coder.model_coder.model import (
+    ModelExperiment,
+    ModelFBWorkspace,
+    ModelTask,
+)
 from rdagent.core.conf import RD_AGENT_SETTINGS
 from rdagent.core.evolving_framework import EvolvingStrategy
 from rdagent.core.prompts import Prompts
@@ -26,8 +30,23 @@ def implement_one_model(
         self,
         target_task: ModelTask,
         queried_knowledge: ModelQueriedKnowledge = None,
+        exp: ModelExperiment = None,  # Add this parameter
     ) -> str:
         model_information_str = target_task.get_task_information()
+        model_type = target_task.model_type
+
+        # Get the current code from the experiment using build_from_SOTA
+        current_code = ""
+        if exp is not None:
+            self.build_from_SOTA(exp)
+            model_file_mapping = {
+                "XGBoost": "model_xgb.py",
+                "RandomForest": "model_rf.py",
+                "LightGBM": "model_lgb.py",
+                "NN": "model_nn.py",
+            }
+            if model_type in model_file_mapping:
+                current_code = exp.experiment_workspace.code_dict.get(model_file_mapping[model_type], "")
 
         if queried_knowledge is not None and model_information_str in queried_knowledge.success_task_to_knowledge_dict:
             return queried_knowledge.success_task_to_knowledge_dict[model_information_str].implementation
@@ -55,6 +74,7 @@ def implement_one_model(
                 .render(
                     scenario=self.scen.get_scenario_all_desc(),
                     queried_former_failed_knowledge=queried_former_failed_knowledge_to_render,
+                    current_code=current_code,  # Add this line
                 )
             )
 
@@ -67,6 +87,7 @@ def implement_one_model(
                     )
                     .render(
                         model_information_str=model_information_str,
+                        model_type=model_type,  # Add model type to the prompt
                         queried_similar_successful_knowledge=queried_similar_successful_knowledge_to_render,
                         queried_former_failed_knowledge=queried_former_failed_knowledge_to_render,
                     )
@@ -103,7 +124,7 @@ def evolve(
         queried_knowledge: ModelQueriedKnowledge | None = None,
         **kwargs,
     ) -> ModelEvolvingItem:
-        # 1.找出需要evolve的model
+        # 1. Find the models that need to be evolved
         to_be_finished_task_index = []
         for index, target_model_task in enumerate(evo.sub_tasks):
             target_model_task_desc = target_model_task.get_task_information()
@@ -125,10 +146,21 @@ def evolve(
             n=RD_AGENT_SETTINGS.multi_proc_n,
         )
 
+        model_file_mapping = {
+            "XGBoost": "model_xgb.py",
+            "RandomForest": "model_rf.py",
+            "LightGBM": "model_lgb.py",
+            "NN": "model_nn.py",
+        }
+
         for index, target_index in enumerate(to_be_finished_task_index):
             if evo.sub_workspace_list[target_index] is None:
                 evo.sub_workspace_list[target_index] = ModelFBWorkspace(target_task=evo.sub_tasks[target_index])
-            evo.sub_workspace_list[target_index].inject_code(**{"model.py": result[index]})
+            model_type = evo.sub_tasks[target_index].model_type
+            if model_type in model_file_mapping:
+                evo.sub_workspace_list[target_index].inject_code(**{model_file_mapping[model_type]: result[index]})
+            else:
+                raise ValueError(f"Unsupported model type: {model_type}")
 
         evo.corresponding_selection = to_be_finished_task_index
 
diff --git a/rdagent/components/coder/model_coder/prompts.yaml b/rdagent/components/coder/model_coder/prompts.yaml
index b25fe64d..b58bb92d 100644
--- a/rdagent/components/coder/model_coder/prompts.yaml
+++ b/rdagent/components/coder/model_coder/prompts.yaml
@@ -52,6 +52,11 @@ evolving_strategy_model_coder:
 
         Your must write your code based on your former latest attempt below which consists of your former code and code feedback, you should read the former attempt carefully and must not modify the right part of your former code.
 
+        {% if current_code %}
+        --------------Current code in the workspace:--------------- You need to tune the model based on this! If it is not None, do not write from scratch. 
+        {{ current_code }}
+        {% endif %}
+
         {% if queried_former_failed_knowledge|length != 0 %}
         --------------Your former latest attempt:---------------
         =====Code to the former implementation=====