Data mining (#103)

* scen

* scen2

* app

* fix

* Simplify workflow

* We can share more code in new scenarios

* rename model to rd loop

* Optimize data path

* Update rdagent/app/data_mining/model.py

* Add TODO

* Support GPU

* gpu

---------

Co-authored-by: SH-Src <suhan.c@outlook.com>

TODO.md

@@ -0,0 +1,10 @@
+We encourage to set the TODOs in code. But some TODOs are more global.
+So we place it here.
+
+
+- [ ] Aligning the naming of files in components & scenarios.
+  - We would like to have the same logic for naming convention in components(reusable components for all scenarios) and scenarios (componets for specific scenario).
+  - But now we have following mismatch
+    - `coder` in `components` & `developer` in `components`
+- [ ] The name of the folders mismatch with the content in them.
+  - Why are scenarios in experiments?

rdagent/app/data mining/README.md

@@ -1,5 +0,0 @@
-
-
-
-# TODO
-If we have more efforts, include more scenario.

rdagent/app/data_mining/conf.py

@@ -0,0 +1,29 @@
+from pathlib import Path
+
+
+from rdagent.components.workflow.conf import BasePropSetting
+
+
+class PropSetting(BasePropSetting):
+    class Config:
+        env_prefix = "DM_"  # Use MODEL_CODER_ as prefix for environment variables
+        protected_namespaces = ()  # Add 'model_' to the protected namespaces
+
+    # 1) overriding the default
+    scen: str = "rdagent.scenarios.data_mining.experiment.model_experiment.DMModelScenario"
+    hypothesis_gen: str = "rdagent.scenarios.data_mining.proposal.model_proposal.DMModelHypothesisGen"
+    hypothesis2experiment: str = "rdagent.scenarios.data_mining.proposal.model_proposal.DMModelHypothesis2Experiment"
+    coder: str = "rdagent.scenarios.data_mining.developer.model_coder.DMModelCoSTEER"
+    runner: str = "rdagent.scenarios.data_mining.developer.model_runner.DMModelRunner"
+    summarizer: str = "rdagent.scenarios.data_mining.developer.feedback.DMModelHypothesisExperiment2Feedback"
+
+    evolving_n: int = 10
+
+    # 2) Extra config for the scenario 
+    # physionet account
+    # NOTE: You should apply the account in https://physionet.org/
+    username: str = ''
+    password: str = ''
+
+
+PROP_SETTING = PropSetting()

rdagent/app/data_mining/model.py

@@ -0,0 +1,27 @@
+import fire
+from rdagent.app.data_mining.conf import PROP_SETTING
+from rdagent.components.workflow.rd_loop import RDLoop
+from rdagent.core.exception import ModelEmptyError
+
+class ModelRDLoop(RDLoop):
+    skip_loop_error = (ModelEmptyError,)
+
+
+def main(path=None, step_n=None):
+    """
+    You can continue running session by
+
+    .. code-block:: python
+
+        dotenv run -- python rdagent/app/data_mining/model.py $LOG_PATH/__session__/1/0_propose  --step_n 1   # `step_n` is a optional paramter
+
+    """
+    if path is None:
+        model_loop = ModelRDLoop(PROP_SETTING)
+    else:
+        model_loop = ModelRDLoop.load(path)
+    model_loop.run(step_n=step_n)
+
+
+if __name__ == "__main__":
+    fire.Fire(main)

rdagent/components/proposal/model_proposal.py

@@ -15,12 +15,14 @@ from rdagent.core.proposal import (
 )
 from rdagent.oai.llm_utils import APIBackend
 
-prompt_dict = Prompts(file_path=Path(__file__).parent / "prompts.yaml")
 
 ModelHypothesis = Hypothesis
 
+prompt_dict = Prompts(file_path=Path(__file__).parent / "prompts.yaml")
 
 class ModelHypothesisGen(HypothesisGen):
+    prompts: Prompts = prompt_dict
+
     # The following methods are scenario related so they should be implemented in the subclass
     @abstractmethod
     def prepare_context(self, trace: Trace) -> Tuple[dict, bool]:
@@ -35,7 +37,7 @@ class ModelHypothesisGen(HypothesisGen):
 
         system_prompt = (
             Environment(undefined=StrictUndefined)
-            .from_string(prompt_dict["hypothesis_gen"]["system_prompt"])
+            .from_string(ModelHypothesisGen.prompts["hypothesis_gen"]["system_prompt"])
             .render(
                 targets="model",
                 scenario=self.scen.get_scenario_all_desc(),
@@ -45,7 +47,7 @@ class ModelHypothesisGen(HypothesisGen):
         )
         user_prompt = (
             Environment(undefined=StrictUndefined)
-            .from_string(prompt_dict["hypothesis_gen"]["user_prompt"])
+            .from_string(ModelHypothesisGen.prompts["hypothesis_gen"]["user_prompt"])
             .render(
                 targets="model",
                 hypothesis_and_feedback=context_dict["hypothesis_and_feedback"],
@@ -61,6 +63,8 @@ class ModelHypothesisGen(HypothesisGen):
 
 
 class ModelHypothesis2Experiment(Hypothesis2Experiment[ModelExperiment]):
+    prompts: Prompts = prompt_dict
+
     def __init__(self) -> None:
         super().__init__()
 
@@ -76,7 +80,7 @@ class ModelHypothesis2Experiment(Hypothesis2Experiment[ModelExperiment]):
         context, json_flag = self.prepare_context(hypothesis, trace)
         system_prompt = (
             Environment(undefined=StrictUndefined)
-            .from_string(prompt_dict["hypothesis2experiment"]["system_prompt"])
+            .from_string(ModelHypothesis2Experiment.prompts["hypothesis2experiment"]["system_prompt"])
             .render(
                 targets="model",
                 scenario=trace.scen.get_scenario_all_desc(),
@@ -85,7 +89,7 @@ class ModelHypothesis2Experiment(Hypothesis2Experiment[ModelExperiment]):
         )
         user_prompt = (
             Environment(undefined=StrictUndefined)
-            .from_string(prompt_dict["hypothesis2experiment"]["user_prompt"])
+            .from_string(ModelHypothesis2Experiment.prompts["hypothesis2experiment"]["user_prompt"])
             .render(
                 targets="model",
                 target_hypothesis=context["target_hypothesis"],

rdagent/components/workflow/conf.py

@@ -0,0 +1,21 @@
+from pydantic_settings import BaseSettings
+
+class BasePropSetting(BaseSettings):
+    """
+    The common part of the config for RD Loop to propose and developement
+    You can add following config in the subclass to distinguish the environment variables.
+
+    .. code-block:: python
+
+        class Config:
+            env_prefix = "DM_MODEL_"  # Use MODEL_CODER_ as prefix for environment variables
+            protected_namespaces = ()  # Add 'model_' to the protected namespaces
+    """
+    scen: str = ""
+    hypothesis_gen: str = ""
+    hypothesis2experiment: str = ""
+    coder: str = ""
+    runner: str = ""
+    summarizer: str = ""
+
+    evolving_n: int = 10

rdagent/components/workflow/rd_loop.py

@@ -0,0 +1,64 @@
+"""
+Model workflow with session control
+It is from `rdagent/app/qlib_rd_loop/model.py` and try to replace `rdagent/app/qlib_rd_loop/RDAgent.py`
+"""
+
+from typing import Any
+from rdagent.components.workflow.conf import BasePropSetting
+from rdagent.core.developer import Developer
+from rdagent.core.proposal import (
+    Hypothesis2Experiment,
+    HypothesisExperiment2Feedback,
+    HypothesisGen,
+    Trace,
+)
+from rdagent.core.scenario import Scenario
+from rdagent.core.utils import import_class
+from rdagent.log import rdagent_logger as logger
+
+from rdagent.utils.workflow import LoopMeta, LoopBase
+
+class RDLoop(LoopBase, metaclass=LoopMeta):
+
+    def __init__(self, PROP_SETTING: BasePropSetting):
+        scen: Scenario = import_class(PROP_SETTING.scen)()
+
+        self.hypothesis_gen: HypothesisGen = import_class(PROP_SETTING.hypothesis_gen)(scen)
+
+        self.hypothesis2experiment: Hypothesis2Experiment = import_class(PROP_SETTING.hypothesis2experiment)()
+
+        self.coder: Developer = import_class(PROP_SETTING.coder)(scen)
+        self.runner: Developer = import_class(PROP_SETTING.runner)(scen)
+
+        self.summarizer: HypothesisExperiment2Feedback = import_class(PROP_SETTING.summarizer)(scen)
+        self.trace = Trace(scen=scen)
+        super().__init__()
+
+    def propose(self, prev_out: dict[str, Any]):
+        with logger.tag("r"):  # research
+            hypothesis = self.hypothesis_gen.gen(self.trace)
+            logger.log_object(hypothesis, tag="hypothesis generation")
+        return hypothesis
+
+    def exp_gen(self, prev_out: dict[str, Any]):
+        with logger.tag("r"):  # research
+            exp = self.hypothesis2experiment.convert(prev_out["propose"], self.trace)
+            logger.log_object(exp.sub_tasks, tag="experiment generation")
+        return exp
+
+    def coding(self, prev_out: dict[str, Any]):
+        with logger.tag("d"):  # develop
+            exp = self.coder.develop(prev_out["exp_gen"])
+            logger.log_object(exp.sub_workspace_list, tag="coder result")
+        return exp
+
+    def running(self, prev_out: dict[str, Any]):
+        with logger.tag("ef"):  # evaluate and feedback
+            exp = self.runner.develop(prev_out["coding"])
+            logger.log_object(exp, tag="runner result")
+        return exp
+
+    def feedback(self, prev_out: dict[str, Any]):
+        feedback = self.summarizer.generate_feedback(prev_out["running"], prev_out["propose"], self.trace)
+        logger.log_object(feedback, tag="feedback")
+        self.trace.hist.append((prev_out["propose"],prev_out["running"] , feedback))

rdagent/core/proposal.py

@@ -9,6 +9,7 @@ from typing import Generic, TypeVar
 
 from rdagent.core.evaluation import Feedback
 from rdagent.core.experiment import ASpecificExp, Experiment
+from rdagent.core.prompts import Prompts
 from rdagent.core.scenario import Scenario
 
 # class data_ana: XXX
@@ -81,6 +82,12 @@ class Trace(Generic[ASpecificScen]):
 
 
 class HypothesisGen(ABC):
+    # NOTE: the design is a little wierd
+    # - Sometimes we want accurate access the prompts in a specific level
+    #   - It renders the prompt to a specific abstract level
+    # - Sometimes we want to access the most recent level prompts
+    prompts: Prompts  # this is a class level prompt.
+
     def __init__(self, scen: Scenario) -> None:
         self.scen = scen
 

rdagent/scenarios/data_mining/developer/feedback.py

@@ -0,0 +1,70 @@
+# TODO:
+# Implement to feedback.
+
+import json
+from pathlib import Path
+
+from jinja2 import Environment, StrictUndefined
+
+from rdagent.core.experiment import Experiment
+from rdagent.core.prompts import Prompts
+from rdagent.core.proposal import (
+    Hypothesis,
+    HypothesisExperiment2Feedback,
+    HypothesisFeedback,
+    Trace,
+)
+from rdagent.log import rdagent_logger as logger
+from rdagent.oai.llm_utils import APIBackend
+from rdagent.utils import convert2bool
+
+feedback_prompts = Prompts(file_path=Path(__file__).parent.parent.parent/ "qlib" / "prompts.yaml")
+DIRNAME = Path(__file__).absolute().resolve().parent
+
+
+class DMModelHypothesisExperiment2Feedback(HypothesisExperiment2Feedback):
+    """Generated feedbacks on the hypothesis from **Executed** Implementations of different tasks & their comparisons with previous performances"""
+
+    def generate_feedback(self, exp: Experiment, hypothesis: Hypothesis, trace: Trace) -> HypothesisFeedback:
+        """
+        The `ti` should be executed and the results should be included, as well as the comparison between previous results (done by LLM).
+        For example: `mlflow` of Qlib will be included.
+        """
+
+        logger.info("Generating feedback...")
+        # Define the system prompt for hypothesis feedback
+        system_prompt = feedback_prompts["model_feedback_generation"]["system"]
+
+        # Define the user prompt for hypothesis feedback
+        context = trace.scen
+        SOTA_hypothesis, SOTA_experiment = trace.get_sota_hypothesis_and_experiment()
+
+        user_prompt = (
+            Environment(undefined=StrictUndefined)
+            .from_string(feedback_prompts["model_feedback_generation"]["user"])
+            .render(
+                context=context,
+                last_hypothesis=SOTA_hypothesis,
+                last_task=SOTA_experiment.sub_tasks[0].get_task_information() if SOTA_hypothesis else None,
+                last_result=SOTA_experiment.result if SOTA_hypothesis else None,
+                hypothesis=hypothesis,
+                exp=exp,
+            )
+        )
+
+        # Call the APIBackend to generate the response for hypothesis feedback
+        response_hypothesis = APIBackend().build_messages_and_create_chat_completion(
+            user_prompt=user_prompt,
+            system_prompt=system_prompt,
+            json_mode=True,
+        )
+
+        # Parse the JSON response to extract the feedback
+        response_json_hypothesis = json.loads(response_hypothesis)
+        return HypothesisFeedback(
+            observations=response_json_hypothesis.get("Observations", "No observations provided"),
+            hypothesis_evaluation=response_json_hypothesis.get("Feedback for Hypothesis", "No feedback provided"),
+            new_hypothesis=response_json_hypothesis.get("New Hypothesis", "No new hypothesis provided"),
+            reason=response_json_hypothesis.get("Reasoning", "No reasoning provided"),
+            decision=convert2bool(response_json_hypothesis.get("Decision", "false")),
+        )

rdagent/scenarios/data_mining/developer/model_coder.py

@@ -0,0 +1,3 @@
+from rdagent.components.coder.model_coder.CoSTEER import ModelCoSTEER
+
+DMModelCoSTEER = ModelCoSTEER

rdagent/scenarios/data_mining/developer/model_runner.py

@@ -0,0 +1,39 @@
+import shutil
+import uuid
+from pathlib import Path
+
+import pandas as pd
+
+from rdagent.components.coder.model_coder.model import ModelExperiment, ModelFBWorkspace
+from rdagent.components.runner import CachedRunner
+from rdagent.components.runner.conf import RUNNER_SETTINGS
+from rdagent.core.developer import Developer
+from rdagent.core.exception import ModelEmptyError
+from rdagent.log import rdagent_logger as logger
+from rdagent.scenarios.data_mining.experiment.model_experiment import DMModelExperiment
+from rdagent.utils.env import DMDockerEnv
+
+
+class DMModelRunner(CachedRunner[DMModelExperiment]):
+
+    def develop(self, exp: DMModelExperiment) -> DMModelExperiment:
+        if RUNNER_SETTINGS.cache_result:
+            cache_hit, result = self.get_cache_result(exp)
+            if cache_hit:
+                exp.result = result
+                return exp
+
+        if exp.sub_workspace_list[0].code_dict.get("model.py") is None:
+            raise ModelEmptyError("model.py is empty")
+        # to replace & inject code
+        exp.experiment_workspace.inject_code(**{"model.py": exp.sub_workspace_list[0].code_dict["model.py"]})
+
+        env_to_use = {"PYTHONPATH": "./"}
+
+        result = exp.experiment_workspace.execute(run_env=env_to_use)
+
+        exp.result = result
+        if RUNNER_SETTINGS.cache_result:
+            self.dump_cache_result(exp, result)
+
+        return exp

rdagent/scenarios/data_mining/docker/Dockerfile

@@ -0,0 +1,25 @@
+FROM pytorch/pytorch:latest
+# For GPU support, please choose the proper tag from https://hub.docker.com/r/pytorch/pytorch/tags
+
+RUN apt-get clean && apt-get update && apt-get install -y \  
+    curl \  
+    vim \  
+    git \  
+    build-essential \
+    && rm -rf /var/lib/apt/lists/* 
+
+WORKDIR /workspace
+
+RUN python -m pip install numpy
+RUN python -m pip install --upgrade cython
+# RUN python -m pip install -e .
+
+RUN python -m pip install pandas
+# RUN pip install pyg_lib torch_scatter torch_sparse torch_cluster -f https://data.pyg.org/whl/torch-2.3.0%2Bcu121.html
+RUN pip install torch_geometric
+RUN pip install ogb
+RUN pip install networkx
+RUN pip install scikit-learn
+RUN pip install catboost
+RUN pip install xgboost
+RUN pip install sparse

rdagent/scenarios/data_mining/experiment/model_experiment.py

@@ -0,0 +1,51 @@
+from pathlib import Path
+
+from rdagent.components.coder.model_coder.model import (
+    ModelExperiment,
+    ModelFBWorkspace,
+    ModelTask,
+)
+from rdagent.core.prompts import Prompts
+from rdagent.core.scenario import Scenario
+from rdagent.scenarios.data_mining.experiment.workspace import DMFBWorkspace
+
+prompt_dict = Prompts(file_path=Path(__file__).parent / "prompts.yaml")
+
+
+class DMModelExperiment(ModelExperiment[ModelTask, DMFBWorkspace, ModelFBWorkspace]):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.experiment_workspace = DMFBWorkspace(template_folder_path=Path(__file__).parent / "model_template")
+
+
+class DMModelScenario(Scenario):
+    @property
+    def background(self) -> str:
+        return prompt_dict["dm_model_background"]
+    
+    @property
+    def source_data(self) -> str:
+        raise NotImplementedError("source_data is not implemented")
+    
+    @property
+    def output_format(self) -> str:
+        return prompt_dict["dm_model_output_format"]
+
+    @property
+    def interface(self) -> str:
+        return prompt_dict["dm_model_interface"]
+
+    @property
+    def simulator(self) -> str:
+        return prompt_dict["dm_model_simulator"]
+
+    def get_scenario_all_desc(self) -> str:
+        return f"""Background of the scenario:
+{self.background}
+The interface you should follow to write the runnable code:
+{self.interface}
+The output of your code should be in the format:
+{self.output_format}
+The simulator user can use to test your model:
+{self.simulator}
+"""

rdagent/scenarios/data_mining/experiment/model_template/README.md

@@ -0,0 +1,3 @@
+## This folder is a template to be copied from for each model implementation & running process. 
+
+Components: Dummy model.py, versatile conf.yaml, and a result reader. 

rdagent/scenarios/data_mining/experiment/model_template/train.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from torchvision import transforms, datasets
+from torch.utils.data import DataLoader, Dataset
+import torch.nn as nn
+import sparse
+import random
+import os
+from model import model_cls
+from sklearn.metrics import accuracy_score, roc_auc_score
+import numpy as np
+
+# Set device for training
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+# device = torch.device("cpu")
+
+class MyDataset(Dataset):
+    def __init__(self, x, label, device):
+        self.x1 = x
+        self.label = label
+        self.device = device
+
+    def __len__(self):
+        return len(self.label)
+
+    def __getitem__(self, idx):
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+        return torch.FloatTensor(self.x1[idx]).to(self.device), \
+               torch.tensor(self.label[idx], dtype=torch.float).to(self.device)
+    
+def collate_fn(batch):
+    x, label = [], []
+    for data in batch:
+        x.append(data[0])
+        label.append(data[1])
+    return torch.stack(x, 0), torch.stack(label, 0)
+
+
+datapath = '/root/.data'
+# datapath = '/home/v-suhancui/RD-Agent/physionet.org/files/mimic-eicu-fiddle-feature/1.0.0/FIDDLE_mimic3'
+
+
+X = sparse.load_npz(datapath+'/features/ARF_12h/X.npz').todense()
+df_pop = pd.read_csv(datapath+'/population/ARF_12h.csv')['ARF_LABEL']
+
+X = X.transpose(0, 2, 1)
+
+indices = [i for i in range(len(df_pop))]
+random.shuffle(indices)
+split_point = int(0.7 * len(df_pop))
+
+X_train, y_train = X[indices[:split_point]], np.array(df_pop[indices[:split_point]])
+X_test, y_test = X[indices[split_point:]], np.array(df_pop[indices[split_point:]])
+
+
+train_dataloader = DataLoader(MyDataset(X_train, y_train, device), collate_fn=collate_fn, shuffle=True, drop_last=True, batch_size=64)
+test_dataloader = DataLoader(MyDataset(X_test, y_test, device), collate_fn=collate_fn, shuffle=False, drop_last=False, batch_size=64)
+
+num_features = 4816
+num_timesteps = 12
+# Define the optimizer and loss function
+model = model_cls(num_features=num_features, num_timesteps=num_timesteps).to(device)
+
+optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
+criterion = nn.CrossEntropyLoss()
+
+# Train the model
+
+for i in range(10):
+    for data in train_dataloader:
+        x, y = data
+        out = model(x)
+        optimizer.zero_grad()
+        loss = criterion(out.squeeze(), y)
+        loss.backward()
+        optimizer.step()
+
+y_pred = []
+for data in test_dataloader:
+    x, y = data
+    out = model(x)
+    y_pred.append(out.cpu().detach().numpy())
+
+acc = roc_auc_score(y_test, np.concatenate(y_pred))
+
+print(acc)
+# Save the predictions to submission.csv
+with open('./submission.txt', 'w') as f:
+    f.write(str(acc))

rdagent/scenarios/data_mining/experiment/prompts.yaml

@@ -0,0 +1,50 @@
+dm_model_background: |-
+  The model is a machine learning or deep learning structure used in clinical settings to predict whether the patient will suffer from acute respiratory failure (ARF) based on their vital signs monitored in ICU.
+  The data is extracted from MIMIC-III using FIDDLE pipeline, and we focus on the ARF_12h prediction task, which means we use the first 12 hours data to predict the onset of ARF on discharge.
+  The model is defined in the following parts:
+  1. Name: The name of the model.
+  2. Description: The description of the model.
+  3. Architecture: The detailed architecture of the model, such as neural network layers or tree structures.
+  The model should provide clear and detailed documentation of its architecture and hyperparameters. 
+
+dm_model_interface: |-
+  Your python code should follow the interface to better interact with the user's system.
+  You code should contain several parts:
+  1. The import part: import the necessary libraries.
+  2. A class which is a sub-class of pytorch.nn.Module. This class should should have a init function and a forward function which inputs a tensor and outputs a tensor.
+  3. Set a variable called "model_cls" to the class you defined.
+
+  The user will save your code into a python file called "model.py". Then the user imports model_cls in file "model.py" after setting the cwd into the directory:
+  ```python
+  from model import model_cls
+  ```
+  So your python code should follow the pattern:
+  ```python
+  class XXXModel(torch.nn.Module):
+      ...
+  model_cls = XXXModel
+  ```
+
+  The model has one type, "TimeSeries". The input shape to a time series model is (batch_size, num_features, num_timesteps). The output shape of the model should be (batch_size, 1).
+  The "batch_size" is a dynamic value which is determined by the input of forward function.
+  The "num_features" and "num_timesteps" are static which will be provided to the model through init function.
+  User will initialize the time series model with the following code:
+  ```python
+  model = model_cls(num_features=num_features, num_timesteps=num_timesteps)
+  ```
+  No other parameters will be passed to the model so give other parameters a default value or just make them static.
+
+  Remember to permute the input tensor since the input tensor is in the shape of (batch_size, num_features, num_timesteps) and a normal time series model is expecting the input tensor in the shape of (batch_size, num_timesteps, num_features).
+
+  Don't write any try-except block in your python code. The user will catch the exception message and provide the feedback to you. Also, don't write main function in your python code. The user will call the forward method in the model_cls to get the output tensor.
+
+  Please notice that your model should only use current features as input. The user will provide the input tensor to the model's forward function.
+
+
+dm_model_output_format: |-
+  Your output should be a tensor with shape (batch_size, 1). 
+  The output tensor should be saved in a file named "output.pth" in the same directory as your python file.
+  The user will evaluate the shape of the output tensor so the tensor read from "output.pth" should be 8 numbers.
+
+dm_model_simulator: |-
+  The models will be put to train on MIMIC-III dataset and evaluate their performance in terms of roc score (Area Under Curve Receiver Operating Characteristics Curve). Hypothesis is improved upon checking the feedback on the results. 

rdagent/scenarios/data_mining/experiment/workspace.py

@@ -0,0 +1,31 @@
+from pathlib import Path
+
+import pandas as pd
+
+from rdagent.core.experiment import FBWorkspace
+from rdagent.log import rdagent_logger as logger
+from rdagent.utils.env import DMDockerEnv
+from rdagent.app.data_mining.conf import PROP_SETTING
+
+class DMFBWorkspace(FBWorkspace):
+    def __init__(self, template_folder_path: Path, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.inject_code_from_folder(template_folder_path)
+
+    def execute(self, run_env: dict = {}, *args, **kwargs) -> str:
+        qtde = DMDockerEnv()
+        qtde.prepare(PROP_SETTING.username, PROP_SETTING.password)
+
+        execute_log = qtde.run(
+            local_path=str(self.workspace_path),
+            entry=f"python train.py",
+            env=run_env,
+        )
+
+        csv_path = self.workspace_path / "submission.txt"
+
+        if not csv_path.exists():
+            logger.error(f"File {csv_path} does not exist.")
+            return None
+        with open(self.workspace_path / "submission.txt", 'r') as f:
+            return f.read()

rdagent/scenarios/data_mining/proposal/model_proposal.py

@@ -0,0 +1,93 @@
+import json
+from pathlib import Path
+from typing import List, Tuple
+
+from jinja2 import Environment, StrictUndefined
+
+from rdagent.components.coder.model_coder.model import ModelExperiment, ModelTask
+from rdagent.components.proposal.model_proposal import (
+    ModelHypothesis,
+    ModelHypothesis2Experiment,
+    ModelHypothesisGen,
+)
+from rdagent.core.prompts import Prompts
+from rdagent.core.proposal import Hypothesis, Scenario, Trace
+from rdagent.scenarios.data_mining.experiment.model_experiment import DMModelExperiment
+
+prompt_dict = Prompts(file_path=Path(__file__).parent.parent.parent / "qlib" / "prompts.yaml")
+
+DMModelHypothesis = ModelHypothesis
+
+
+class DMModelHypothesisGen(ModelHypothesisGen):
+    """
+    # NOTE: we can share this class across different data mining scenarios
+    # It may better to move the class into components folder like `rdagent/components/proposal/model_proposal.py`
+    # Here is the use case:
+
+    .. code-block:: python
+
+        class XXXDMModelHypothesisGen(DMModelHypothesisGen):
+            prompts: Prompts = a_specifc_prompt_dict
+    """
+    def __init__(self, scen: Scenario) -> Tuple[dict, bool]:
+        super().__init__(scen)
+
+    def prepare_context(self, trace: Trace) -> Tuple[dict, bool]:
+        hypothesis_feedback = (
+            Environment(undefined=StrictUndefined)
+            .from_string(prompt_dict["hypothesis_and_feedback"])
+            .render(trace=trace)
+        )
+        context_dict = {
+            "hypothesis_and_feedback": hypothesis_feedback,
+            "RAG": "",
+            "hypothesis_output_format": prompt_dict["hypothesis_output_format"],
+            "hypothesis_specification": prompt_dict["model_hypothesis_specification"]
+        }
+        return context_dict, True
+
+    def convert_response(self, response: str) -> ModelHypothesis:
+        response_dict = json.loads(response)
+        hypothesis = DMModelHypothesis(hypothesis=response_dict["hypothesis"], reason=response_dict["reason"])
+        return hypothesis
+
+
+class DMModelHypothesis2Experiment(ModelHypothesis2Experiment):
+    def prepare_context(self, hypothesis: Hypothesis, trace: Trace) -> Tuple[dict, bool]:
+        scenario = trace.scen.get_scenario_all_desc()
+        experiment_output_format = prompt_dict["model_experiment_output_format"]
+
+        hypothesis_and_feedback = (
+            Environment(undefined=StrictUndefined)
+            .from_string(prompt_dict["hypothesis_and_feedback"])
+            .render(trace=trace)
+        )
+
+        experiment_list: List[ModelExperiment] = [t[1] for t in trace.hist]
+
+        model_list = []
+        for experiment in experiment_list:
+            model_list.extend(experiment.sub_tasks)
+
+        return {
+            "target_hypothesis": str(hypothesis),
+            "scenario": scenario,
+            "hypothesis_and_feedback": hypothesis_and_feedback,
+            "experiment_output_format": experiment_output_format,
+            "target_list": model_list,
+            "RAG": ...,
+        }, True
+
+    def convert_response(self, response: str, trace: Trace) -> ModelExperiment:
+        response_dict = json.loads(response)
+        tasks = []
+        for model_name in response_dict:
+            description = response_dict[model_name]["description"]
+            architecture = response_dict[model_name]["architecture"]
+            hyperparameters = response_dict[model_name]["hyperparameters"]
+            model_type = response_dict[model_name]["model_type"]
+            tasks.append(ModelTask(model_name, description, architecture, hyperparameters, model_type))
+        exp = DMModelExperiment(tasks)
+        exp.based_experiments = [t[1] for t in trace.hist if t[2]]
+        return exp

rdagent/scenarios/qlib/prompts.yaml

@@ -138,7 +138,7 @@ model_experiment_output_format: |-
   So far please only design one model to test the hypothesis! 
   The output should follow JSON format. The schema is as follows: 
   {
-    "model_name (The name of the model)": {
+    "model_name 1 (The name of the model)": {
         "description": "A detailed description of the model",
         "architecture": "A detailed description of the model's architecture, e.g., neural network layers or tree structures",
         "hyperparameters": {
@@ -147,6 +147,9 @@ model_experiment_output_format: |-
             "hyperparameter_name_3": "value of hyperparameter 3"
         },
         "model_type": "Tabular or TimeSeries"  # Should be one of "Tabular" or "TimeSeries"
+    },
+    "model_name 2 (The name of the model)": {
+        ...
     }
   }
   Usually a larger model works better than a smaller one. Hence, the parameters should be larger.

rdagent/utils/env.py

@@ -3,8 +3,8 @@ The motiviation of the utils is for environment management
 
 Tries to create uniform environment for the agent to run;
 - All the code and data is expected included in one folder
-
 """
+# TODO: move the scenario specific docker env into other folders.
 
 import os
 import subprocess
@@ -125,6 +125,7 @@ class DockerConf(BaseSettings):
     # So we just want to download it once.
     network: str | None = "bridge"  # the network mode for the docker
     shm_size: str | None = None
+    enable_gpu: bool = True  # because we will automatically disable GPU if not available. So we enable it by default.
 
 
 class QlibDockerConf(DockerConf):
@@ -141,6 +142,19 @@ class QlibDockerConf(DockerConf):
     enable_gpu: bool = True
 
 
+class DMDockerConf(DockerConf):
+    class Config:
+        env_prefix = "DM_DOCKER_"  
+
+    build_from_dockerfile: bool = True
+    dockerfile_folder_path: Path = Path(__file__).parent.parent / "scenarios" / "data_mining" / "docker"
+    image: str = "local_dm:latest"
+    mount_path: str = "/workspace/dm_workspace/"
+    default_entry: str = "python train.py"
+    extra_volumes: dict = {Path("~/.rdagent/.data/physionet.org/files/mimic-eicu-fiddle-feature/1.0.0/FIDDLE_mimic3/").expanduser().resolve(): "/root/.data/"}
+    shm_size: str | None = "16g"
+
+# physionet.org/files/mimic-eicu-fiddle-feature/1.0.0/FIDDLE_mimic3
 class DockerEnv(Env[DockerConf]):
     # TODO: Save the output into a specific file
 
@@ -243,3 +257,23 @@ class QTDockerEnv(DockerEnv):
             self.run(entry=cmd)
         else:
             logger.info("Data already exists. Download skipped.")
+
+
+class DMDockerEnv(DockerEnv):
+    """Qlib Torch Docker"""
+
+    def __init__(self, conf: DockerConf = DMDockerConf()):
+        super().__init__(conf)
+
+    def prepare(self, username: str, password: str):
+        """
+        Download image & data if it doesn't exist
+        """
+        super().prepare()
+        data_path = next(iter(self.conf.extra_volumes.keys()))
+        if not (Path(data_path)).exists():
+            logger.info("We are downloading!")
+            cmd = 'wget -r -N -c -np --user={} --password={} -P ~/.rdagent/.data/ https://physionet.org/files/mimic-eicu-fiddle-feature/1.0.0/'.format(username, password)
+            os.system(cmd)
+        else:
+            logger.info("Data already exists. Download skipped.")

rdagent/utils/workflow.py

@@ -13,7 +13,7 @@ from tqdm.auto import tqdm
 
 
 from collections import defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 import datetime
 from typing import Callable
 from rdagent.log import rdagent_logger as logger
@@ -21,15 +21,23 @@ from rdagent.log import rdagent_logger as logger
 
 class LoopMeta(type):
 
-    def __new__(cls, clsname, bases, attrs):
-
-        # move custommized steps into steps
+    @staticmethod
+    def _get_steps(bases):
+        """
+        get all the `steps` of base classes and combine them to a single one.
+        """
         steps = []
-        for name in attrs.keys():
-            if not name.startswith("__"):
+        for base in bases:
+            steps.extend(LoopMeta._get_steps(base.__bases__) + getattr(base,"steps", []))
+        return steps
+
+    def __new__(cls, clsname, bases, attrs):
+        # move custommized steps into steps
+        steps = LoopMeta._get_steps(bases)  # all the base classes of parents
+        for name, attr in attrs.items():
+            if not name.startswith("__") and isinstance(attr, Callable):
                 steps.append(name)
         attrs["steps"] = steps
-
         return super().__new__(cls, clsname, bases, attrs)
 
 
@@ -44,6 +52,8 @@ class LoopBase:
     steps: list[Callable]  # a list of steps to work on
     loop_trace: dict[int, list[LoopTrace]]
 
+    skip_loop_error: tuple[Exception]  = field(default_factory=tuple)  # you can define a list of error that will skip current loop
+
     def __init__(self):
         self.loop_idx = 0 # current loop index
         self.step_idx = 0 # the index of next step to be run
@@ -73,7 +83,14 @@ class LoopBase:
 
                 name = self.steps[si]
                 func = getattr(self, name)
-                self.loop_prev_out[name] = func(self.loop_prev_out)
+                try:
+                    self.loop_prev_out[name] = func(self.loop_prev_out)
+                    # TODO: Fix the error logger.exception(f"Skip loop {li} due to {e}")
+                except self.skip_loop_error as e:
+                    logger.warning(f"Skip loop {li} due to {e}")
+                    self.loop_idx += 1
+                    self.step_index = 0
+                    continue
 
                 end = datetime.datetime.now(datetime.timezone.utc)