From 509b4d7745ddbae3f94dfa18eeb82222fe9787e9 Mon Sep 17 00:00:00 2001 From: Xisen-Wang Date: Wed, 24 Jul 2024 09:14:20 +0000 Subject: [PATCH] Revised relevant code to enable graph input --- .../model_extraction_and_implementation.py | 2 +- rdagent/components/coder/model_coder/model.py | 32 +++++++++++++++---- .../components/coder/model_coder/prompts.yaml | 17 ++++++++-- .../coder/model_coder/task_loader.py | 3 +- 4 files changed, 44 insertions(+), 10 deletions(-) diff --git a/rdagent/app/model_extraction_and_code/model_extraction_and_implementation.py b/rdagent/app/model_extraction_and_code/model_extraction_and_implementation.py index 1114cf13..ec22b6ae 100644 --- a/rdagent/app/model_extraction_and_code/model_extraction_and_implementation.py +++ b/rdagent/app/model_extraction_and_code/model_extraction_and_implementation.py @@ -8,7 +8,7 @@ from rdagent.scenarios.qlib.developer.model_coder import QlibModelCoSTEER from rdagent.scenarios.qlib.experiment.model_experiment import QlibModelScenario -def extract_models_and_implement(report_file_path: str = "../test_doc") -> None: +def extract_models_and_implement(report_file_path: str = "/home/xiaoyang/PaperImpBench/raw_paper/4096.pdf") -> None: scenario = QlibModelScenario() exp = ModelExperimentLoaderFromPDFfiles().load(report_file_path) exp = QlibModelCoSTEER(scenario).develop(exp) diff --git a/rdagent/components/coder/model_coder/model.py b/rdagent/components/coder/model_coder/model.py index 67fe2609..75a7c66e 100644 --- a/rdagent/components/coder/model_coder/model.py +++ b/rdagent/components/coder/model_coder/model.py @@ -16,18 +16,20 @@ from rdagent.utils import get_module_by_module_path class ModelTask(Task): def __init__( - self, name: str, description: str, formulation: str, variables: Dict[str, str], model_type: Optional[str] = None + self, name: str, description: str, formulation: str, architecture: str, variables: Dict[str, str], model_type: Optional[str] = None ) -> None: self.name: str = name self.description: str = description self.formulation: str = formulation + self.architecture: str = architecture self.variables: str = variables - self.model_type: str = model_type # Tabular for tabular model, TimesSeries for time series model + self.model_type: str = model_type # Tabular for tabular model, TimesSeries for time series model, Graph for graph model def get_task_information(self): return f"""name: {self.name} description: {self.description} formulation: {self.formulation} +architecture: {self.architecture} variables: {self.variables} model_type: {self.model_type} """ @@ -65,6 +67,8 @@ class ModelFBWorkspace(FBWorkspace): batch_size: int = 8, num_features: int = 10, num_timesteps: int = 4, + num_nodes: int = 50, + num_edges: int = 100, input_value: float = 1.0, param_init_value: float = 1.0, ): @@ -85,21 +89,37 @@ class ModelFBWorkspace(FBWorkspace): if self.target_task.model_type == "Tabular": input_shape = (batch_size, num_features) m = model_cls(num_features=input_shape[1]) + data = torch.full(input_shape, input_value) elif self.target_task.model_type == "TimeSeries": input_shape = (batch_size, num_features, num_timesteps) m = model_cls(num_features=input_shape[1], num_timesteps=input_shape[2]) - data = torch.full(input_shape, input_value) + data = torch.full(input_shape, input_value) + elif self.target_task.model_type == "Graph": + node_feature = torch.randn(batch_size, num_nodes, num_features) + edge_index = torch.randint(0, num_nodes, (2, num_edges)) + m = model_cls(num_nodes=num_nodes, num_features=num_features) + data = (node_feature, edge_index) + else: + raise ValueError(f"Unsupported model type: {self.target_task.model_type}") - # initialize all parameters of `m` to `param_init_value` + # Initialize all parameters of `m` to `param_init_value` for _, param in m.named_parameters(): param.data.fill_(param_init_value) - out = m(data) + + # Execute the model + if self.target_task.model_type == "Graph": + out = m(*data) + else: + out = m(data) + execution_model_output = out.cpu().detach() execution_feedback_str = f"Execution successful, output tensor shape: {execution_model_output.shape}" + if MODEL_IMPL_SETTINGS.enable_execution_cache: pickle.dump((execution_feedback_str, execution_model_output), open(cache_file_path, "wb")) + return execution_feedback_str, execution_model_output - + except Exception as e: return f"Execution error: {e}", None diff --git a/rdagent/components/coder/model_coder/prompts.yaml b/rdagent/components/coder/model_coder/prompts.yaml index b5679941..107c4189 100644 --- a/rdagent/components/coder/model_coder/prompts.yaml +++ b/rdagent/components/coder/model_coder/prompts.yaml @@ -1,11 +1,24 @@ extract_model_formulation_system: |- - offer description of the proposed model in this paper, write a latex formula with variable of the model. the format should be like - "Model Name": { + offer description of the proposed model in this paper, write a latex formula with variable as well as the architecture of the model. the format should be like + "Model Name (replace with a reasonable model type)": { "description": "", "formulation": "", + "architecture": "", "variables": { "\\hat{y}_u": "The predicted output for node u", }" + "model_type":"Tabular", "TimeSeries", or "Graph", (depending on which type of the model it is) + } + Eg. + "ABC Model":{ + "description": "", + "formulation": "", + "architecture": "", + "variables": { + "\\hat{y}_u": "The predicted output for node u", + }" + "model_type":"Tabular", + } such format content should be begin with ```json and end with ``` and the content should be in json format. evolving_strategy_model_coder: diff --git a/rdagent/components/coder/model_coder/task_loader.py b/rdagent/components/coder/model_coder/task_loader.py index 5929f10e..ad047c1a 100644 --- a/rdagent/components/coder/model_coder/task_loader.py +++ b/rdagent/components/coder/model_coder/task_loader.py @@ -105,8 +105,9 @@ class ModelExperimentLoaderFromDict(ModelTaskLoader): name=model_name, description=model_data["description"], formulation=model_data["formulation"], + architecture=model_data["architecture"], variables=model_data["variables"], - key=model_name, + model_type=model_data["model_type"], ) task_l.append(task) return QlibModelExperiment(sub_tasks=task_l)