Merge branch 'main' into Continous-Update-Branch

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/app/qlib_rd_loop/factor.py

@@ -36,7 +36,7 @@ qlib_factor_summarizer: HypothesisExperiment2Feedback = import_class(PROP_SETTIN
 trace = Trace(scen=scen)
 for _ in range(PROP_SETTING.evolving_n):
     try:
-        with logger.tag("r"):  # research
+        with logger.tag("r"):
             hypothesis = hypothesis_gen.gen(trace)
             logger.log_object(hypothesis, tag="hypothesis generation")
 
@@ -49,6 +49,9 @@ for _ in range(PROP_SETTING.evolving_n):
 
         with logger.tag("ef"):
             exp = qlib_factor_runner.develop(exp)
+            if exp is None:
+                logger.error(f"Factor extraction failed.")
+                continue
             logger.log_object(exp, tag="factor runner result")
             feedback = qlib_factor_summarizer.generate_feedback(exp, hypothesis, trace)
             logger.log_object(feedback, tag="feedback")

rdagent/app/qlib_rd_loop/factor_from_report_sh.py

@@ -7,7 +7,7 @@ from jinja2 import Environment, StrictUndefined
 import pandas as pd
 
 from rdagent.app.qlib_rd_loop.conf import PROP_SETTING
-from rdagent.components.document_reader.document_reader import load_and_process_pdfs_by_langchain
+from rdagent.components.document_reader.document_reader import extract_first_page_screenshot_from_pdf, load_and_process_pdfs_by_langchain
 from rdagent.core.prompts import Prompts
 from rdagent.core.scenario import Scenario
 from rdagent.core.utils import import_class
@@ -88,7 +88,11 @@ def extract_factors_and_implement(report_file_path: str) -> tuple:
             exp = FactorExperimentLoaderFromPDFfiles().load(report_file_path)
             if exp is None or exp.sub_tasks == []:
                 return None, None
-            
+
+        with logger.tag("load_pdf_screenshot"):
+            pdf_screenshot = extract_first_page_screenshot_from_pdf(report_file_path)
+            logger.log_object(pdf_screenshot)
+
     docs_dict = load_and_process_pdfs_by_langchain(Path(report_file_path))
 
     factor_result = {
@@ -118,19 +122,30 @@ try:
             report_file_path = Path(file_path.replace(PROP_SETTING.origin_report_path, PROP_SETTING.local_report_path))
             if report_file_path.exists():
                 logger.info(f"Processing {report_file_path}")
-                exp, hypothesis = extract_factors_and_implement(str(report_file_path))
-                if exp is None:
-                    continue
-                exp.based_experiments = [t[1] for t in trace.hist if t[2]]
-                if len(exp.based_experiments) == 0:
-                    exp.based_experiments.append(QlibFactorExperiment(sub_tasks=[]))
-                exp = qlib_factor_coder.develop(exp)
-                exp = qlib_factor_runner.develop(exp)
-                if exp is None:
-                    logger.error(f"Factor extraction failed for {report_file_path}. Skipping to the next report.")
-                    continue
-                feedback = qlib_factor_summarizer.generate_feedback(exp, hypothesis, trace)
+                
+                with logger.tag("r"):
+                    exp, hypothesis = extract_factors_and_implement(str(report_file_path))
+                    if exp is None:
+                        continue
+                    exp.based_experiments = [t[1] for t in trace.hist if t[2]]
+                    if len(exp.based_experiments) == 0:
+                        exp.based_experiments.append(QlibFactorExperiment(sub_tasks=[]))
+                    logger.log_object(hypothesis, tag="hypothesis generation")
+                    logger.log_object(exp.sub_tasks, tag="experiment generation")
+                
+                with logger.tag("d"):
+                    exp = qlib_factor_coder.develop(exp)
+                    logger.log_object(exp.sub_workspace_list)
 
+                with logger.tag("ef"):
+                    exp = qlib_factor_runner.develop(exp)
+                    if exp is None:
+                        logger.error(f"Factor extraction failed for {report_file_path}. Skipping to the next report.")
+                        continue
+                    logger.log_object(exp, tag="factor runner result")
+                    feedback = qlib_factor_summarizer.generate_feedback(exp, hypothesis, trace)
+                    logger.log_object(feedback, tag="feedback")
+                
                 trace.hist.append((hypothesis, exp, feedback))
                 logger.info(f"Processed {report_file_path}: Result: {exp}")
                 

rdagent/app/qlib_rd_loop/model_w_sc.py

@@ -67,20 +67,20 @@ class ModelLoop(LoopBase, metaclass=LoopMeta):
         self.trace.hist.append((prev_out["propose"],prev_out["running"] , feedback))
 
 
-def main(path=None):
+def main(path=None, step_n=None):
     """
     You can continue running session by
 
     .. code-block:: python
 
-        dotenv run -- python rdagent/app/qlib_rd_loop/model_w_sc.py $LOG_PATH/__session__/1/0_propose
+        dotenv run -- python rdagent/app/qlib_rd_loop/model_w_sc.py $LOG_PATH/__session__/1/0_propose  --step_n 1   # `step_n` is a optional paramter
 
     """
     if path is None:
         model_loop = ModelLoop()
     else:
         model_loop = ModelLoop.load(path)
-    model_loop.run()
+    model_loop.run(step_n=step_n)
 
 
 if __name__ == "__main__":

rdagent/components/coder/factor_coder/CoSTEER/evaluators.py

@@ -437,21 +437,36 @@ class FactorFinalDecisionEvaluator(Evaluator):
             else:
                 break
 
-        final_evaluation_dict = json.loads(
-            APIBackend().build_messages_and_create_chat_completion(
-                user_prompt=user_prompt,
-                system_prompt=system_prompt,
-                json_mode=True,
-            ),
-        )
-        if isinstance(final_evaluation_dict["final_decision"], str) and final_evaluation_dict[
-            "final_decision"
-        ].lower() in ("true", "false"):
-            final_evaluation_dict["final_decision"] = bool(final_evaluation_dict["final_decision"])
-        return (
-            final_evaluation_dict["final_decision"],
-            final_evaluation_dict["final_feedback"],
-        )
+        # TODO:  with retry_context(retry_n=3, except_list=[KeyError]):
+        final_evaluation_dict = None
+        attempts = 0
+        max_attempts = 3
+
+        while attempts < max_attempts:
+            try:
+                final_evaluation_dict = json.loads(
+                    APIBackend().build_messages_and_create_chat_completion(
+                        user_prompt=user_prompt,
+                        system_prompt=system_prompt,
+                        json_mode=True,
+                    ),
+                )
+                final_decision = final_evaluation_dict["final_decision"]
+                final_feedback = final_evaluation_dict["final_feedback"]
+
+                if isinstance(final_decision, str) and final_decision.lower() in ("true", "false"):
+                    final_decision = bool(final_decision)
+
+                return final_decision, final_feedback
+
+            except json.JSONDecodeError as e:
+                raise ValueError("Failed to decode JSON response from API.") from e
+            except KeyError as e:
+                attempts += 1
+                if attempts >= max_attempts:
+                    raise KeyError("Response from API is missing 'final_decision' or 'final_feedback' key after multiple attempts.") from e
+        
+        return None, None
 
 
 class FactorSingleFeedback:

rdagent/components/document_reader/document_reader.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 from pathlib import Path
+import fitz
+from PIL import Image
 from typing import TYPE_CHECKING
 
 from azure.ai.formrecognizer import DocumentAnalysisClient
@@ -96,3 +98,11 @@ def load_and_process_pdfs_by_azure_document_intelligence(path: Path) -> dict[str
                     RD_AGENT_SETTINGS.azure_document_intelligence_endpoint,
                 )
     return content_dict
+
+def extract_first_page_screenshot_from_pdf(pdf_path: Path) -> Image:
+    doc = fitz.open(pdf_path)
+    page = doc.load_page(0)
+    pix = page.get_pixmap()
+    image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+    
+    return image

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
@@ -82,6 +83,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/experiment/factor_template/conf_combined.yaml

@@ -21,7 +21,6 @@ data_handler_config: &data_handler_config
                             - ["LABEL0"]
                 - class: qlib.data.dataset.loader.StaticDataLoader
                   kwargs:
-                    # config: "/home/finco/v-yuanteli/RD-Agent/rdagent/scenarios/qlib/task_generator/factor_template/combined_factors_df.pkl"
                     config: "combined_factors_df.pkl"
 
     learn_processors:
@@ -90,4 +89,4 @@ task:
         - class: PortAnaRecord
           module_path: qlib.workflow.record_temp
           kwargs: 
-            config: *port_analysis_config
+            config: *port_analysis_config

rdagent/scenarios/qlib/experiment/workspace.py

@@ -17,12 +17,6 @@ class QlibFBWorkspace(FBWorkspace):
         qtde = QTDockerEnv()
         qtde.prepare()
 
-        # Run the Docker command
-        execute_log = qtde.run(
-            local_path=str(self.workspace_path),
-            entry="rm -r mlruns",
-            env=run_env,
-        )
         # Run the Qlib backtest
         execute_log = qtde.run(
             local_path=str(self.workspace_path),

rdagent/scenarios/qlib/factor_experiment_loader/pdf_loader.py

@@ -70,7 +70,7 @@ def classify_report_from_dict(
         if isinstance(value, str):
             content = value
         else:
-            logger.warning(f"输入格式不符合要求: {file_name}")
+            logger.warning(f"Input format does not meet the requirements: {file_name}")
             res_dict[file_name] = {"class": 0}
             continue
 
@@ -102,7 +102,7 @@ def classify_report_from_dict(
                 res = json.loads(res)
                 vote_list.append(int(res["class"]))
             except json.JSONDecodeError:
-                logger.warning(f"返回值无法解析: {file_name}")
+                logger.warning(f"Return value could not be parsed: {file_name}")
                 res_dict[file_name] = {"class": 0}
             count_0 = vote_list.count(0)
             count_1 = vote_list.count(1)
@@ -243,7 +243,7 @@ def extract_factors_from_report_dict(
     )
     for index, file_name in enumerate(file_name_list):
         final_report_factor_dict[file_name] = factor_dict_list[index]
-    logger.info(f"已经完成{len(final_report_factor_dict)}个报告的因子提取")
+    logger.info(f"Factor extraction completed for {len(final_report_factor_dict)} reports")
 
     return final_report_factor_dict
 
@@ -507,13 +507,24 @@ def deduplicate_factors_by_llm(  # noqa: C901, PLR0912
 
 class FactorExperimentLoaderFromPDFfiles(FactorExperimentLoader):
     def load(self, file_or_folder_path: Path) -> dict:
-        docs_dict = load_and_process_pdfs_by_langchain(Path(file_or_folder_path))
+        with logger.tag("docs"):
+            docs_dict = load_and_process_pdfs_by_langchain(Path(file_or_folder_path))
+            logger.log_object(docs_dict)
 
         selected_report_dict = classify_report_from_dict(report_dict=docs_dict, vote_time=1)
-        file_to_factor_result = extract_factors_from_report_dict(docs_dict, selected_report_dict)
-        factor_dict = merge_file_to_factor_dict_to_factor_dict(file_to_factor_result)
+        
+        with logger.tag("file_to_factor_result"):
+            file_to_factor_result = extract_factors_from_report_dict(docs_dict, selected_report_dict)
+            logger.log_object(file_to_factor_result)
+        
+        with logger.tag("factor_dict"):
+            factor_dict = merge_file_to_factor_dict_to_factor_dict(file_to_factor_result)
+            logger.log_object(factor_dict)
+
+        with logger.tag("filtered_factor_dict"):
+            factor_viability, filtered_factor_dict = check_factor_viability(factor_dict)
+            logger.log_object(filtered_factor_dict)
 
-        factor_viability, filtered_factor_dict = check_factor_viability(factor_dict)
         # factor_dict, duplication_names_list = deduplicate_factors_by_llm(factor_dict, factor_viability)
         
         return FactorExperimentLoaderFromDict().load(filtered_factor_dict)

rdagent/scenarios/qlib/prompts.yaml

@@ -58,6 +58,61 @@ factor_hypothesis_specification: |-
     - "Combine value and momentum factors using a weighted average approach."
     - "Filter stocks by market capitalization before calculating the factors."
 
+factor_hypothesis_specification: |-
+  Additional Specifications:
+    - Hypotheses should grow and evolve based on the previous hypothesis. If there is no previous hypothesis, start with something simple.
+    - Gradually build upon previous hypotheses and feedback.
+    - Ensure that the hypothesis focuses on the creation and selection of factors in quantitative finance.
+    - Each hypothesis should address specific factor characteristics such as type (momentum, value, quality), calculation methods, or inclusion criteria.
+    - Avoid hypotheses related to model architecture or optimization processes.
+    - If a hypothesis can be improved further, refine it. If it achieves the desired results, explore a new direction. Previous factors exceeding SOTA (State of the Art) are preserved and combined with new factors for subsequent evaluations.
+
+  Guiding Principles:
+  1. Diversity and Depth:
+    - Ensure a wide range of factor types, incorporating various financial dimensions (e.g., momentum, value, quality, volatility, sentiment).
+    - Explore different calculation methods and inclusion criteria to understand their impact.
+    - Consider combining multiple factors or filtering criteria for more sophisticated hypotheses.
+
+  2. Iterative Improvement:
+    - Build upon previous hypotheses, incorporating feedback and observed results.
+    - Aim for continuous refinement and complexity over iterations, starting from basic factors to more advanced combinations and techniques.
+
+  3. Contextual Relevance:
+    - Tailor hypotheses to the specific financial context and current market conditions.
+    - Leverage domain knowledge and recent financial research to inform hypothesis creation.
+
+  Sample Hypotheses (Use the format for guidance, not the specific content):
+  - "Include a momentum factor based on the last 12 months' returns."
+  - "Add a value factor calculated as the book-to-market ratio."
+  - "Incorporate a quality factor derived from return on equity (ROE)."
+  - "Use a volatility factor based on the standard deviation of returns over the past 6 months."
+  - "Include a sentiment factor derived from news sentiment scores."
+  - "The momentum factor should be calculated using a 6-month look-back period."
+  - "Combine value and momentum factors using a weighted average approach."
+  - "Filter stocks by market capitalization before calculating the factors."
+  - "Explore a liquidity factor based on the trading volume and bid-ask spread."
+  - "Investigate the impact of an earnings surprise factor calculated from recent earnings announcements."
+  - "Develop a composite factor integrating ESG (Environmental, Social, Governance) scores with traditional financial metrics."
+
+  Detailed Workflow:
+  1. Initial Hypothesis:
+    - Begin with a simple factor, such as "Include a momentum factor based on the last 12 months' returns."
+
+  2. Refine Hypothesis:
+    - If the initial hypothesis is promising, refine it further, e.g., "The momentum factor should be calculated using a 6-month look-back period."
+
+  3. Combine Factors:
+    - As individual factors show potential, combine them, e.g., "Combine value and momentum factors using a weighted average approach."
+
+  4. Contextual Adjustments:
+    - Adjust factors based on market conditions or new financial insights, e.g., "Incorporate a quality factor derived from return on equity (ROE)."
+
+  5. Advanced Hypotheses:
+    - Explore sophisticated combinations or new types of factors, e.g., "Develop a composite factor integrating ESG scores with traditional financial metrics."
+
+  Remember: If a hypothesis achieves the desired results, start a new direction while preserving the effective factors from previous hypotheses. New evaluations should combine the newly proposed factors with previously successful factors that surpassed SOTA.
+
+
 factor_experiment_output_format: |-
   The output should follow JSON format. The schema is as follows:
   {
@@ -84,7 +139,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",
         "formulation": "A LaTeX formula representing the model's formulation",
         "architecture": "A detailed description of the model's architecture, e.g., neural network layers or tree structures",
@@ -98,14 +153,17 @@ model_experiment_output_format: |-
             "hyperparameter_name_2": "value of hyperparameter 2",
             "hyperparameter_name_3": "value of hyperparameter 3"
         },
-        "model_type": "Tabular or TimeSeries  # Should be one of "Tabular" or "TimeSeries"
+        "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.
 
 factor_feedback_generation:
   system: |-
-    You are a professional result analysis assistant on data driven R&D.
+    You are a professional result analysis assistant in data-driven R&D.
     The task is described in the following scenario:
     {{ scenario }}
     You will receive a hypothesis, multiple tasks with their factors, and some results.
@@ -128,8 +186,7 @@ factor_feedback_generation:
     {{ combined_result }}
     Analyze the combined result in the context of its ability to:
     1. Support or refute the hypothesis.
-    2. Show improvement or deterioration compared to the last experiment.
-    3. Demonstrate positive or negative effects when compared to Alpha158.
+    2. Show improvement or deterioration compared to the SOTA experiment.
 
     Evaluation Metrics Explanations:
     Below are the financial meanings of each metric, which should be used to judge the results:
@@ -143,8 +200,17 @@ factor_feedback_generation:
     - IC: Measures the correlation between predicted returns (\hat{y}) and actual returns (y), using Pearson correlation.
     - 1day.excess_return_with_cost.information_ratio: Evaluates the excess return per unit of risk considering transaction costs.
 
-    When judging the results, prioritize metrics that consider transaction costs (with cost), as they provide a more accurate representation of real-world performance. Among these, the annualized return considering transaction costs is particularly important as it gives a clear picture of long-term profitability.
-    Provide detailed feedback and recommend whether to replace the best result if the new factor proves superior.
+    When judging the results:
+    1. Prioritize metrics that consider transaction costs (with cost):
+    - These metrics provide a more accurate representation of real-world performance.
+    2. Evaluate all metrics:
+        - Compare the combined results against the current best results across all metrics to get a comprehensive view of performance.
+    3. Focus on the annualized return considering transaction costs:
+        - This metric is particularly important as it gives a clear picture of long-term profitability.
+    4. Recommendation for replacement:
+        - If the new factor demonstrates a significant improvement in the annualized return considering transaction costs, it should be recommended to replace the current best result, even if other metrics show minor variations.
+
+    Please provide detailed feedback and recommend whether to replace the best result if the new factor proves superior.
 
 model_feedback_generation:
   system: |-

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
@@ -51,16 +61,36 @@ class LoopBase:
         self.loop_trace = defaultdict(list[LoopTrace])  # the key is the number of loop
         self.session_folder = logger.log_trace_path / "__session__"
 
-    def run(self):
+    def run(self, step_n: int | None = None):
+        """
+
+        Parameters
+        ----------
+        step_n : int | None
+            How many steps to run;
+            `None` indicates to run forever until error or KeyboardInterrupt
+        """
         with tqdm(total=len(self.steps), desc="Workflow Progress", unit="step") as pbar:
             while True:
+                if step_n is not None:
+                    if step_n <= 0:
+                        break
+                    step_n -= 1
+
                 li, si = self.loop_idx, self.step_idx
 
                 start = datetime.datetime.now(datetime.timezone.utc)
 
                 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)
 

requirements/package.txt

@@ -19,6 +19,7 @@ langchain
 tiktoken
 scikit-learn
 docker
+fitz  # Extract shotsreens from pdf
 
 # azure identity related
 azure.identity