feat: Model context for tuning and selection (#284)

* Revised for tuning * Fixed a bug * fix a ci bug --------- Co-authored-by: WinstonLiye <[email protected]>
microsoft · Sep 20, 2024 · f2831e7 · f2831e7
1 parent 0047641
commit f2831e7
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 3 deletions.
diff --git 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
@@ -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=====