gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-07-28 16:07:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Build model class inheritance (#44)

Browse Source 
* update all code

* fix a typo

---------

Co-authored-by: xuyang1 <xuyang1@microsoft.com>
This commit is contained in:
Xu Yang
2024-07-03 17:42:07 +08:00
committed by GitHub
parent ab2f61fe9d
commit d88ddc0c1c
8 changed files with 138 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rdagent/app/model_extraction_and_implementation/model_extraction_and_implementation.py
+11 -6
View File
@@ -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)
rdagent/app/model_implementation/eval.py
+5 -5
View File
@@ -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)
rdagent/components/loader/task_loader.py
+2 -1
View File
@@ -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
rdagent/components/task_implementation/model_implementation/benchmark/eval.py
+6 -3
View File
@@ -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] = []
rdagent/components/task_implementation/model_implementation/task.py → rdagent/components/task_implementation/model_implementation/model.py
+76 -89
View File
@@ -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()
rdagent/components/task_implementation/model_implementation/one_shot/__init__.py
+11 -9
View File
@@ -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
rdagent/components/task_implementation/model_implementation/task_extraction.py → rdagent/components/task_implementation/model_implementation/task_loader.py
+21 -3
View File
@@ -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"):
rdagent/scenarios/qlib/model_task_implementation/__init__.py
+6
View File
@@ -0,0 +1,6 @@
from rdagent.components.task_implementation.model_implementation.one_shot import (
ModelCodeWriter,
)
class QlibModelCodeWriter(ModelCodeWriter): ...