Files
NexQuant/rdagent/scenarios/data_science/experiment/experiment.py
T
Xu Yang e82d38f88e fix: move mlebench check into runner (#556)
* abandon mlebench check in workflow

* move mlebench check to runner and put it into the exp
2025-02-06 15:37:41 +08:00

31 lines
1.0 KiB
Python

import re
from typing import Literal
import pandas as pd
from rdagent.core.experiment import Experiment, FBWorkspace, Task
COMPONENT = Literal["DataLoadSpec", "FeatureEng", "Model", "Ensemble", "Workflow"]
class DSExperiment(Experiment[Task, FBWorkspace, FBWorkspace]):
def __init__(self, pending_tasks_list: list, *args, **kwargs) -> None:
super().__init__(sub_tasks=[], *args, **kwargs)
self.experiment_workspace = FBWorkspace()
self.pending_tasks_list = pending_tasks_list
self.format_check_result = None
def next_component_required(self) -> COMPONENT | None:
files = list(self.experiment_workspace.file_dict.keys())
if "load_data.py" not in files:
return "DataLoadSpec"
if "feature.py" not in files:
return "FeatureEng"
if not any(re.match(r"model.*\.py", file) for file in files):
return "Model"
if "ensemble.py" not in files:
return "Ensemble"
if "main.py" not in files:
return "Workflow"
return None