From d88ddc0c1cf39e538f0adfef6bf82d8ef6206bdf Mon Sep 17 00:00:00 2001 From: Xu Yang Date: Wed, 3 Jul 2024 17:42:07 +0800 Subject: [PATCH] Build model class inheritance (#44) * update all code * fix a typo --------- Co-authored-by: xuyang1 --- .../model_extraction_and_implementation.py | 17 +- rdagent/app/model_implementation/eval.py | 10 +- rdagent/components/loader/task_loader.py | 3 +- .../model_implementation/benchmark/eval.py | 9 +- .../{task.py => model.py} | 165 ++++++++---------- .../model_implementation/one_shot/__init__.py | 20 ++- .../{task_extraction.py => task_loader.py} | 24 ++- .../model_task_implementation/__init__.py | 6 + 8 files changed, 138 insertions(+), 116 deletions(-) rename rdagent/components/task_implementation/model_implementation/{task.py => model.py} (88%) rename rdagent/components/task_implementation/model_implementation/{task_extraction.py => task_loader.py} (83%) diff --git a/rdagent/app/model_extraction_and_implementation/model_extraction_and_implementation.py b/rdagent/app/model_extraction_and_implementation/model_extraction_and_implementation.py index d02f4854..830bc7bc 100644 --- a/rdagent/app/model_extraction_and_implementation/model_extraction_and_implementation.py +++ b/rdagent/app/model_extraction_and_implementation/model_extraction_and_implementation.py @@ -1,16 +1,21 @@ # %% from dotenv import load_dotenv -from rdagent.components.task_implementation.model_implementation.one_shot import ModelTaskGen -from rdagent.components.task_implementation.model_implementation.task_extraction import ModelImplementationTaskLoaderFromPDFfiles + +from rdagent.components.task_implementation.model_implementation.one_shot import ( + ModelCodeWriter, +) +from rdagent.components.task_implementation.model_implementation.task_loader import ( + ModelImplementationExperimentLoaderFromPDFfiles, +) - -def extract_models_and_implement(report_file_path: str="../test_doc") -> None: - factor_tasks = ModelImplementationTaskLoaderFromPDFfiles().load(report_file_path) - implementation_result = ModelTaskGen().generate(factor_tasks) +def extract_models_and_implement(report_file_path: str = "../test_doc") -> None: + factor_tasks = ModelImplementationExperimentLoaderFromPDFfiles().load(report_file_path) + implementation_result = ModelCodeWriter().generate(factor_tasks) return implementation_result import fire + if __name__ == "__main__": fire.Fire(extract_models_and_implement) diff --git a/rdagent/app/model_implementation/eval.py b/rdagent/app/model_implementation/eval.py index e28aed5b..67efee67 100644 --- a/rdagent/app/model_implementation/eval.py +++ b/rdagent/app/model_implementation/eval.py @@ -5,13 +5,13 @@ DIRNAME = Path(__file__).absolute().resolve().parent from rdagent.components.task_implementation.model_implementation.benchmark.eval import ( ModelImpValEval, ) -from rdagent.components.task_implementation.model_implementation.one_shot import ( - ModelTaskGen, -) -from rdagent.components.task_implementation.model_implementation.task import ( +from rdagent.components.task_implementation.model_implementation.model import ( ModelImpLoader, ModelTaskLoaderJson, ) +from rdagent.components.task_implementation.model_implementation.one_shot import ( + ModelCodeWriter, +) bench_folder = DIRNAME.parent.parent / "components" / "task_implementation" / "model_implementation" / "benchmark" mtl = ModelTaskLoaderJson(str(bench_folder / "model_dict.json")) @@ -20,7 +20,7 @@ task_l = mtl.load() task_l = [t for t in task_l if t.key == "A-DGN"] # FIXME: other models does not work well -mtg = ModelTaskGen() +mtg = ModelCodeWriter() impl_l = mtg.generate(task_l) diff --git a/rdagent/components/loader/task_loader.py b/rdagent/components/loader/task_loader.py index 862b557c..c651a66c 100644 --- a/rdagent/components/loader/task_loader.py +++ b/rdagent/components/loader/task_loader.py @@ -1,6 +1,7 @@ from rdagent.components.task_implementation.factor_implementation.factor import ( FactorTask, ) +from rdagent.components.task_implementation.model_implementation.model import ModelTask from rdagent.core.experiment import Loader @@ -8,5 +9,5 @@ class FactorTaskLoader(Loader[FactorTask]): pass -class ModelTaskLoader(Loader[FactorTask]): +class ModelTaskLoader(Loader[ModelTask]): pass diff --git a/rdagent/components/task_implementation/model_implementation/benchmark/eval.py b/rdagent/components/task_implementation/model_implementation/benchmark/eval.py index 4c1fb1b0..01519be9 100644 --- a/rdagent/components/task_implementation/model_implementation/benchmark/eval.py +++ b/rdagent/components/task_implementation/model_implementation/benchmark/eval.py @@ -1,6 +1,9 @@ # TODO: inherent from the benchmark base class import torch -from rdagent.components.task_implementation.model_implementation.task import ModelTaskImpl + +from rdagent.components.task_implementation.model_implementation.model import ( + ModelImplementation, +) def get_data_conf(init_val): @@ -15,7 +18,7 @@ def get_data_conf(init_val): class ModelImpValEval: """ - Evaluate the similarity of the model structure by changing the input and observate the output. + Evaluate the similarity of the model structure by changing the input and observe the output. Assumption: - If the model structure is similar, the output will change in similar way when we change the input. @@ -31,7 +34,7 @@ class ModelImpValEval: For each hidden output, we can calculate a correlation. The average correlation will be the metrics. """ - def evaluate(self, gt: ModelTaskImpl, gen: ModelTaskImpl): + def evaluate(self, gt: ModelImplementation, gen: ModelImplementation): round_n = 10 eval_pairs: list[tuple] = [] diff --git a/rdagent/components/task_implementation/model_implementation/task.py b/rdagent/components/task_implementation/model_implementation/model.py similarity index 88% rename from rdagent/components/task_implementation/model_implementation/task.py rename to rdagent/components/task_implementation/model_implementation/model.py index 85cef1f5..d7ed5e8b 100644 --- a/rdagent/components/task_implementation/model_implementation/task.py +++ b/rdagent/components/task_implementation/model_implementation/model.py @@ -10,11 +10,11 @@ from rdagent.components.task_implementation.model_implementation.conf import ( MODEL_IMPL_SETTINGS, ) from rdagent.core.exception import CodeFormatException -from rdagent.core.experiment import FBImplementation, ImpLoader, Task +from rdagent.core.experiment import Experiment, FBImplementation, ImpLoader, Task from rdagent.utils import get_module_by_module_path -class ModelImplTask(Task): +class ModelTask(Task): # TODO: it should change when the Task changes. name: str description: str @@ -49,95 +49,13 @@ key: {self.key} @staticmethod def from_dict(dict): - return ModelImplTask(**dict) + return ModelTask(**dict) def __repr__(self) -> str: return f"<{self.__class__.__name__} {self.name}>" -class ModelTaskLoaderJson(ModelTaskLoader): - # def __init__(self, json_uri: str, select_model: Optional[str] = None) -> None: - # super().__init__() - # self.json_uri = json_uri - # self.select_model = 'A-DGN' - - # def load(self, *argT, **kwargs) -> Sequence[ModelImplTask]: - # # json is supposed to be in the format of {model_name: dict{model_data}} - # model_dict = json.load(open(self.json_uri, "r")) - # if self.select_model is not None: - # assert self.select_model in model_dict - # model_name = self.select_model - # model_data = model_dict[self.select_model] - # else: - # model_name, model_data = list(model_dict.items())[0] - - # model_impl_task = ModelImplTask( - # name=model_name, - # description=model_data["description"], - # formulation=model_data["formulation"], - # variables=model_data["variables"], - # key=model_name - # ) - - # return [model_impl_task] - - def __init__(self, json_uri: str) -> None: - super().__init__() - self.json_uri = json_uri - - def load(self, *argT, **kwargs) -> Sequence[ModelImplTask]: - # json is supposed to be in the format of {model_name: dict{model_data}} - model_dict = json.load(open(self.json_uri, "r")) - - # FIXME: the model in the json file is not right due to extraction error - # We should fix them case by case in the future - # - # formula_info = { - # "name": "Anti-Symmetric Deep Graph Network (A-DGN)", - # "description": "A framework for stable and non-dissipative DGN design. It ensures long-range information preservation between nodes and prevents gradient vanishing or explosion during training.", - # "formulation": r"\mathbf{x}^{\prime}_i = \mathbf{x}_i + \epsilon \cdot \sigma \left( (\mathbf{W}-\mathbf{W}^T-\gamma \mathbf{I}) \mathbf{x}_i + \Phi(\mathbf{X}, \mathcal{N}_i) + \mathbf{b}\right),", - # "variables": { - # r"\mathbf{x}_i": "The state of node i at previous layer", - # r"\epsilon": "The step size in the Euler discretization", - # r"\sigma": "A monotonically non-decreasing activation function", - # r"\Phi": "A graph convolutional operator", - # r"W": "An anti-symmetric weight matrix", - # r"\mathbf{x}^{\prime}_i": "The node feature matrix at layer l-1", - # r"\mathcal{N}_i": "The set of neighbors of node u", - # r"\mathbf{b}": "A bias vector", - # }, - # "key": "A-DGN", - # } - model_impl_task_list = [] - for model_name, model_data in model_dict.items(): - model_impl_task = ModelImplTask( - name=model_name, - description=model_data["description"], - formulation=model_data["formulation"], - variables=model_data["variables"], - key=model_data["key"], - ) - model_impl_task_list.append(model_impl_task) - return model_impl_task_list - - -class ModelImplementationTaskLoaderFromDict(ModelTaskLoader): - def load(self, model_dict: dict) -> list: - """Load data from a dict.""" - task_l = [] - for model_name, model_data in model_dict.items(): - task = ModelImplTask( - name=model_name, - description=model_data["description"], - formulation=model_data["formulation"], - variables=model_data["variables"], - key=model_name, - ) - task_l.append(task) - return task_l - - -class ModelTaskImpl(FBImplementation): +class ModelImplementation(FBImplementation): """ It is a Pytorch model implementation task; All the things are placed in a folder. @@ -210,13 +128,82 @@ We'll import the model in the implementation in file `model.py` after setting th """ -class ModelImpLoader(ImpLoader[ModelImplTask, ModelTaskImpl]): +class ModelExperiment(Experiment[ModelTask, ModelImplementation]): ... + + +class ModelTaskLoaderJson(ModelTaskLoader): + # def __init__(self, json_uri: str, select_model: Optional[str] = None) -> None: + # super().__init__() + # self.json_uri = json_uri + # self.select_model = 'A-DGN' + + # def load(self, *argT, **kwargs) -> Sequence[ModelImplTask]: + # # json is supposed to be in the format of {model_name: dict{model_data}} + # model_dict = json.load(open(self.json_uri, "r")) + # if self.select_model is not None: + # assert self.select_model in model_dict + # model_name = self.select_model + # model_data = model_dict[self.select_model] + # else: + # model_name, model_data = list(model_dict.items())[0] + + # model_impl_task = ModelImplTask( + # name=model_name, + # description=model_data["description"], + # formulation=model_data["formulation"], + # variables=model_data["variables"], + # key=model_name + # ) + + # return [model_impl_task] + + def __init__(self, json_uri: str) -> None: + super().__init__() + self.json_uri = json_uri + + def load(self, *argT, **kwargs) -> Sequence[ModelTask]: + # json is supposed to be in the format of {model_name: dict{model_data}} + model_dict = json.load(open(self.json_uri, "r")) + + # FIXME: the model in the json file is not right due to extraction error + # We should fix them case by case in the future + # + # formula_info = { + # "name": "Anti-Symmetric Deep Graph Network (A-DGN)", + # "description": "A framework for stable and non-dissipative DGN design. It ensures long-range information preservation between nodes and prevents gradient vanishing or explosion during training.", + # "formulation": r"\mathbf{x}^{\prime}_i = \mathbf{x}_i + \epsilon \cdot \sigma \left( (\mathbf{W}-\mathbf{W}^T-\gamma \mathbf{I}) \mathbf{x}_i + \Phi(\mathbf{X}, \mathcal{N}_i) + \mathbf{b}\right),", + # "variables": { + # r"\mathbf{x}_i": "The state of node i at previous layer", + # r"\epsilon": "The step size in the Euler discretization", + # r"\sigma": "A monotonically non-decreasing activation function", + # r"\Phi": "A graph convolutional operator", + # r"W": "An anti-symmetric weight matrix", + # r"\mathbf{x}^{\prime}_i": "The node feature matrix at layer l-1", + # r"\mathcal{N}_i": "The set of neighbors of node u", + # r"\mathbf{b}": "A bias vector", + # }, + # "key": "A-DGN", + # } + model_impl_task_list = [] + for model_name, model_data in model_dict.items(): + model_impl_task = ModelTask( + name=model_name, + description=model_data["description"], + formulation=model_data["formulation"], + variables=model_data["variables"], + key=model_data["key"], + ) + model_impl_task_list.append(model_impl_task) + return model_impl_task_list + + +class ModelImpLoader(ImpLoader[ModelTask, ModelImplementation]): def __init__(self, path: Path) -> None: self.path = Path(path) - def load(self, task: ModelImplTask) -> ModelTaskImpl: + def load(self, task: ModelTask) -> ModelImplementation: assert task.key is not None - mti = ModelTaskImpl(task) + mti = ModelImplementation(task) mti.prepare() with open(self.path / f"{task.key}.py", "r") as f: code = f.read() diff --git a/rdagent/components/task_implementation/model_implementation/one_shot/__init__.py b/rdagent/components/task_implementation/model_implementation/one_shot/__init__.py index afad43db..4851bb21 100644 --- a/rdagent/components/task_implementation/model_implementation/one_shot/__init__.py +++ b/rdagent/components/task_implementation/model_implementation/one_shot/__init__.py @@ -4,22 +4,23 @@ from typing import Sequence from jinja2 import Environment, StrictUndefined -from rdagent.components.task_implementation.model_implementation.task import ( - ModelImplTask, - ModelTaskImpl, +from rdagent.components.task_implementation.model_implementation.model import ( + ModelExperiment, + ModelImplementation, + ModelTask, ) -from rdagent.core.task_generator import TaskGenerator from rdagent.core.prompts import Prompts +from rdagent.core.task_generator import TaskGenerator from rdagent.oai.llm_utils import APIBackend DIRNAME = Path(__file__).absolute().resolve().parent -class ModelTaskGen(TaskGenerator): - def generate(self, task_l: Sequence[ModelImplTask]) -> Sequence[ModelTaskImpl]: +class ModelCodeWriter(TaskGenerator[ModelExperiment]): + def generate(self, exp: ModelExperiment) -> ModelExperiment: mti_l = [] - for t in task_l: - mti = ModelTaskImpl(t) + for t in exp.sub_tasks: + mti = ModelImplementation(t) mti.prepare() pr = Prompts(file_path=DIRNAME / "prompt.yaml") @@ -42,4 +43,5 @@ class ModelTaskGen(TaskGenerator): code = match.group(1) mti.inject_code(**{"model.py": code}) mti_l.append(mti) - return mti_l + exp.sub_implementations = mti_l + return exp diff --git a/rdagent/components/task_implementation/model_implementation/task_extraction.py b/rdagent/components/task_implementation/model_implementation/task_loader.py similarity index 83% rename from rdagent/components/task_implementation/model_implementation/task_extraction.py rename to rdagent/components/task_implementation/model_implementation/task_loader.py index c39afadf..f11a1f33 100644 --- a/rdagent/components/task_implementation/model_implementation/task_extraction.py +++ b/rdagent/components/task_implementation/model_implementation/task_loader.py @@ -8,8 +8,10 @@ from rdagent.components.document_reader.document_reader import ( load_and_process_pdfs_by_langchain, ) from rdagent.components.loader.task_loader import ModelTaskLoader -from rdagent.components.task_implementation.model_implementation.task import ( +from rdagent.components.task_implementation.model_implementation.model import ( + ModelExperiment, ModelImplementationTaskLoaderFromDict, + ModelTask, ) from rdagent.core.log import RDAgentLog from rdagent.core.prompts import Prompts @@ -97,7 +99,23 @@ def extract_model_from_docs(docs_dict): return model_dict -class ModelImplementationTaskLoaderFromPDFfiles(ModelTaskLoader): +class ModelImplementationExperimentLoaderFromDict(ModelTaskLoader): + def load(self, model_dict: dict) -> list: + """Load data from a dict.""" + task_l = [] + for model_name, model_data in model_dict.items(): + task = ModelTask( + name=model_name, + description=model_data["description"], + formulation=model_data["formulation"], + variables=model_data["variables"], + key=model_name, + ) + task_l.append(task) + return ModelExperiment(sub_tasks=task_l) + + +class ModelImplementationExperimentLoaderFromPDFfiles(ModelTaskLoader): def load(self, file_or_folder_path: Path) -> dict: docs_dict = load_and_process_pdfs_by_langchain(Path(file_or_folder_path)) # dict{file_path:content} model_dict = extract_model_from_docs( @@ -106,7 +124,7 @@ class ModelImplementationTaskLoaderFromPDFfiles(ModelTaskLoader): model_dict = merge_file_to_model_dict_to_model_dict( model_dict ) # dict {model_name: dict{description, formulation, variables}} - return ModelImplementationTaskLoaderFromDict().load(model_dict) + return ModelImplementationExperimentLoaderFromDict().load(model_dict) def main(path="../test_doc"): diff --git a/rdagent/scenarios/qlib/model_task_implementation/__init__.py b/rdagent/scenarios/qlib/model_task_implementation/__init__.py index e69de29b..fcc823d4 100644 --- a/rdagent/scenarios/qlib/model_task_implementation/__init__.py +++ b/rdagent/scenarios/qlib/model_task_implementation/__init__.py @@ -0,0 +1,6 @@ +from rdagent.components.task_implementation.model_implementation.one_shot import ( + ModelCodeWriter, +) + + +class QlibModelCodeWriter(ModelCodeWriter): ...