example workflow code for model implementation (#17)

* upload example code for model implementation

* Refactor Model Implement

* add export

---------

Co-authored-by: Young <afe.young@gmail.com>

README.md

@@ -10,8 +10,39 @@ As the maintainer of this project, please make a few updates:
 - Understanding the security reporting process in SECURITY.MD
 - Remove this section from the README
 
+## Configuration:
+
+You can manually source the `.env` file in your shell before running the Python script:
+Most of the workflow are controlled by the environment variables.
+```sh
+# Export each variable in the .env file; Please note that it is different from `source .env` without export
+export $(grep -v '^#' .env | xargs)
+# Run the Python script
+python your_script.py
+```
+
+## Naming convention
+
+### File naming convention
+
+| Name      | Description       |
+| --        | --                |
+| `conf.py` | The configuration for the module & app & project  | 
+
+<!-- TODO: renaming files -->
+ 
+
 ## Contributing
 
+### Guidance
+This project welcomes contributions and suggestions.
+You can find issues in the issues list or simply running `grep -r "TODO:"`.
+
+Making contributions is not a hard thing. Solving an issue(maybe just answering a question raised in issues list ), fixing/issuing a bug, improving the documents and even fixing a typo are important contributions to RDAgent.
+
+
+### Policy
+
 This project welcomes contributions and suggestions.  Most contributions require you to agree to a
 Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us
 the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com.

rdagent/app/model_implementation/README.md

@@ -0,0 +1,39 @@
+
+# Preparation
+
+## Install Pytorch
+CPU CUDA will be enough for verify the implementation
+
+Please install pytorch based on your system.
+Here is an example on my system
+```bash
+pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+pip3 install torch_geometric
+
+```
+
+# Tasks
+
+## Task Extraction
+From paper to task.
+```bash
+python rdagent/app/model_implementation/task_extraction.py
+# It may based on rdagent/document_reader/document_reader.py
+```
+
+## Complete workflow
+From paper to implementation
+``` bash
+# Similar to
+# rdagent/app/factor_extraction_and_implementation/factor_extract_and_implement.py
+```
+
+## Paper benchmark
+```bash
+python rdagent/app/model_implementation/eval.py
+
+TODO:
+- Is evaluation reasonable
+```
+
+## Evolving

rdagent/app/model_implementation/eval.py

@@ -0,0 +1,29 @@
+from pathlib import Path
+
+DIRNAME = Path(__file__).absolute().resolve().parent
+
+from rdagent.model_implementation.benchmark.eval import ModelImpValEval
+from rdagent.model_implementation.one_shot import ModelTaskGen
+from rdagent.model_implementation.task import ModelImpLoader, ModelTaskLoderJson
+
+mtl = ModelTaskLoderJson("TODO: A Path to json")
+
+task_l = mtl.load()
+
+mtg = ModelTaskGen()
+
+impl_l = mtg.generate(task_l)
+
+# TODO: Align it with the benchmark framework after @wenjun's refine the evaluation part.
+# Currently, we just handcraft a workflow for fast evaluation.
+
+mil = ModelImpLoader(DIRNAME.parent.parent / "model_implementation" / "benchmark" /  "gt_code")
+
+mie = ModelImpValEval()
+# Evaluation:
+eval_l = []
+for impl in impl_l:
+    gt_impl = mil.load(impl.target_task)
+    eval_l.append(mie.evaluate(gt_impl, impl))
+
+print(eval_l)

rdagent/core/implementation.py

@@ -1,13 +1,21 @@
 from abc import ABC, abstractmethod
-from typing import List
+from typing import List, Sequence
 
 from rdagent.core.task import (
+    BaseTask,
     TaskImplementation,
 )
 
 class TaskGenerator(ABC):
     @abstractmethod
-    def generate(self, *args, **kwargs) -> List[TaskImplementation]:
+    def generate(self, task_l: Sequence[BaseTask]) -> Sequence[TaskImplementation]:
+        """
+        Task Generator should take in a sequence of tasks.
+
+        Because the schedule of different tasks is crucial for the final performance
+        due to it affects the learning process.
+
+        """
         raise NotImplementedError("generate method is not implemented.")
 
     def collect_feedback(self, feedback_obj_l: List[object]):

rdagent/core/task.py

@@ -1,27 +1,121 @@
 from abc import ABC, abstractmethod
-from typing import Tuple
+from pathlib import Path
+from typing import Generic, Optional, Sequence, Tuple, TypeVar
 import pandas as pd
-
 """
 This file contains the all the data class for rdagent task.
 """
 
 
 class BaseTask(ABC):
-    # 把name放在这里作为主键
+    # TODO: 把name放在这里作为主键
+    # Please refer to rdagent/model_implementation/task.py for the implementation
+    # I think the task version applies to the base class.
     pass
 
+ASpecificTask = TypeVar("ASpecificTask", bound=BaseTask)
 
-class TaskImplementation(ABC):
-    def __init__(self, target_task: BaseTask) -> None:
+
+class TaskImplementation(ABC, Generic[ASpecificTask]):
+
+    def __init__(self, target_task: ASpecificTask) -> None:
         self.target_task = target_task
 
     @abstractmethod
-    def execute(self, *args, **kwargs) -> Tuple[str, pd.DataFrame]:
-        raise NotImplementedError("__call__ method is not implemented.")
+    def execute(self, data=None, config: dict = {}) -> object:
+        """
+        The execution of the implementation can be dynamic.
+
+        So we may passin the data and config dynamically.
+        """
+        raise NotImplementedError("execute method is not implemented.")
+
+    @abstractmethod
+    def execute_desc(self):
+        """
+        return the description how we will execute the code in the folder.
+        """
+        raise NotImplementedError(f"This type of input is not supported")
+
+    # TODO:
+    # After execution, it should return some results.
+    # Some evaluators will input the results and output
+
+
+ASpecificTaskImp = TypeVar("ASpecificTaskImp", bound=TaskImplementation)
+
+
+class ImpLoader(ABC, Generic[ASpecificTask, ASpecificTaskImp]):
+
+    @abstractmethod
+    def load(self, task: ASpecificTask) -> ASpecificTaskImp:
+        raise NotImplementedError("load method is not implemented.")
+
+
+class FBTaskImplementation(TaskImplementation):
+    """
+    File-based task implementation
+    
+    The implemented task will be a folder which contains related elements.
+    - Data
+    - Code Implementation
+    - Output
+        - After execution, it will generate the final output as file.
+
+    A typical way to run the pipeline of FBTaskImplementation will be
+    (We didn't add it as a method due to that we may pass arguments into `prepare` or `execute` based on our requirements.)
+
+    .. code-block:: python
+
+        def run_pipline(self, **files: str):
+            self.prepare()
+            self.inject_code(**files)
+            self.execute()
+
+    """
+    # TODO:
+    # FileBasedFactorImplementation should inherient from it.
+    # Why not directly reuse FileBasedFactorImplementation.
+    #   Because it has too much concerete dependencies.
+    #   e.g.  dataframe, factors
+
+    path: Optional[Path]
+
+    @abstractmethod
+    def prepare(self, *args, **kwargs):
+        """
+        Prepare all the files except the injected code
+        - Data
+        - Documentation
+        - TODO: env?  Env is implicitly defined by the document?
+
+            typical usage of `*args, **kwargs`:
+                Different methods shares the same data. The data are passed by the arguments.
+        """
+
+    def inject_code(self, **files: str):
+        """
+        Inject the code into the folder.
+        {
+            "model.py": "<model code>"
+        }
+        """
+        for k, v in files.items():
+            with open(self.path / k, "w") as f:
+                f.write(v)
+
+    def get_files(self) -> list[Path]:
+        """
+        Get the environment description.
+
+        To be general, we only return a list of filenames.
+        How to summarize the environment is the responsibility of the TaskGenerator.
+        """
+        return list(self.path.iterdir())
 
 
 class TestCase:
+
     def __init__(
         self,
         target_task: BaseTask,
@@ -32,6 +126,7 @@ class TestCase:
 
 
 class TaskLoader:
+
     @abstractmethod
-    def load(self, *args, **kwargs) -> BaseTask | list[BaseTask]:
+    def load(self, *args, **kwargs) -> Sequence[BaseTask]:
         raise NotImplementedError("load method is not implemented.")

rdagent/factor_implementation/evolving/factor.py

@@ -30,6 +30,8 @@ from filelock import FileLock
 
 
 class FactorImplementTask(BaseTask):
+    # TODO:  generalized the attributes into the BaseTask
+    # - factor_* -> *
     def __init__(
         self,
         factor_name,
@@ -39,6 +41,7 @@ class FactorImplementTask(BaseTask):
         variables: dict = {},
         resource: str = None,
     ) -> None:
+        # TODO: remove the useless factor_formulation_description
         self.factor_name = factor_name
         self.factor_description = factor_description
         self.factor_formulation = factor_formulation

rdagent/model_implementation/benchmark/eval.py

@@ -0,0 +1,61 @@
+# TODO: inherent from the benchmark base class
+import torch
+from rdagent.model_implementation.task import ModelTaskImpl
+
+
+def get_data_conf(init_val):
+    # TODO: design this step in the workflow
+    in_dim = 1000
+    in_channels = 128
+    exec_config = {"model_eval_param_init": init_val}
+    node_feature = torch.randn(in_dim, in_channels)
+    edge_index = torch.randint(0, in_dim, (2, 2000))
+    return (node_feature, edge_index), exec_config
+
+
+class ModelImpValEval:
+    """
+    Evaluate the similarity of the model structure by changing the input and observate the output.
+
+    Assumption:
+    - If the model structure is similar, the output will change in similar way when we change the input.
+        - we try to initialize the model param in similar value. So only the model structure is different.
+    """
+
+    def evaluate(self, gt: ModelTaskImpl, gen: ModelTaskImpl):
+        round_n = 10
+
+        eval_pairs: list[tuple] = []
+
+        # run different input value
+        for _ in range(round_n):
+            # run different model initial parameters.
+            for init_val in [-0.2, -0.1, 0.1, 0.2]:
+                data, exec_config = get_data_conf(init_val)
+                gt_res = gt.execute(data=data, config=exec_config)
+                res = gen.execute(data=data, config=exec_config)
+                eval_pairs.append((res, gt_res))
+
+        # flat and concat the output
+        res_batch, gt_res_batch = [], []
+        for res, gt_res in eval_pairs:
+            res_batch.append(res.reshape(-1))
+            gt_res_batch.append(gt_res.reshape(-1))
+        res_batch = torch.stack(res_batch)
+        gt_res_batch = torch.stack(gt_res_batch)
+
+        res_batch = res_batch.detach().numpy()
+        gt_res_batch = gt_res_batch.detach().numpy()
+
+        # pearson correlation of each hidden output
+        def norm(x):
+            return (x - x.mean(axis=0)) / x.std(axis=0)
+        dim_corr = (norm(res_batch) * norm(gt_res_batch)).mean(axis=0)  # the correlation of each hidden output
+
+        # aggregate all the correlation
+        avr_corr = dim_corr.mean()
+        # FIXME: 
+        # It is too high(e.g. 0.944) .
+        # Check if it is not a good evaluation!! 
+        # Maybe all the same initial params will results in extreamly high correlation without regard to the model structure.
+        return avr_corr

rdagent/model_implementation/benchmark/gt_code/A-DGN.py

@@ -0,0 +1,135 @@
+import math
+from typing import Any, Callable, Dict, Optional, Union
+
+import torch
+from torch import Tensor
+from torch.nn import Parameter
+
+from torch_geometric.nn.conv import GCNConv, MessagePassing
+from torch_geometric.nn.inits import zeros
+from torch_geometric.nn.resolver import activation_resolver
+from torch_geometric.typing import Adj
+
+
+class AntiSymmetricConv(torch.nn.Module):
+    r"""The anti-symmetric graph convolutional operator from the
+    `"Anti-Symmetric DGN: a stable architecture for Deep Graph Networks"
+    <https://openreview.net/forum?id=J3Y7cgZOOS>`_ paper.
+
+    .. math::
+        \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),
+
+    where :math:`\Phi(\mathbf{X}, \mathcal{N}_i)` denotes a
+    :class:`~torch.nn.conv.MessagePassing` layer.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        phi (MessagePassing, optional): The message passing module
+            :math:`\Phi`. If set to :obj:`None`, will use a
+            :class:`~torch_geometric.nn.conv.GCNConv` layer as default.
+            (default: :obj:`None`)
+        num_iters (int, optional): The number of times the anti-symmetric deep
+            graph network operator is called. (default: :obj:`1`)
+        epsilon (float, optional): The discretization step size
+            :math:`\epsilon`. (default: :obj:`0.1`)
+        gamma (float, optional): The strength of the diffusion :math:`\gamma`.
+            It regulates the stability of the method. (default: :obj:`0.1`)
+        act (str, optional): The non-linear activation function :math:`\sigma`,
+            *e.g.*, :obj:`"tanh"` or :obj:`"relu"`. (default: :class:`"tanh"`)
+        act_kwargs (Dict[str, Any], optional): Arguments passed to the
+            respective activation function defined by :obj:`act`.
+            (default: :obj:`None`)
+        bias (bool, optional): If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+
+    Shapes:
+        - **input:**
+          node features :math:`(|\mathcal{V}|, F_{in})`,
+          edge indices :math:`(2, |\mathcal{E}|)`,
+          edge weights :math:`(|\mathcal{E}|)` *(optional)*
+        - **output:** node features :math:`(|\mathcal{V}|, F_{in})`
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        phi: Optional[MessagePassing] = None,
+        num_iters: int = 1,
+        epsilon: float = 0.1,
+        gamma: float = 0.1,
+        act: Union[str, Callable, None] = "tanh",
+        act_kwargs: Optional[Dict[str, Any]] = None,
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.num_iters = num_iters
+        self.gamma = gamma
+        self.epsilon = epsilon
+        self.act = activation_resolver(act, **(act_kwargs or {}))
+
+        if phi is None:
+            phi = GCNConv(in_channels, in_channels, bias=False)
+
+        self.W = Parameter(torch.empty(in_channels, in_channels))
+        self.register_buffer("eye", torch.eye(in_channels))
+        self.phi = phi
+
+        if bias:
+            self.bias = Parameter(torch.empty(in_channels))
+        else:
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        torch.nn.init.kaiming_uniform_(self.W, a=math.sqrt(5))
+        self.phi.reset_parameters()
+        zeros(self.bias)
+
+    def forward(self, x: Tensor, edge_index: Adj, *args, **kwargs) -> Tensor:
+        r"""Runs the forward pass of the module."""
+        antisymmetric_W = self.W - self.W.t() - self.gamma * self.eye
+
+        for _ in range(self.num_iters):
+            h = self.phi(x, edge_index, *args, **kwargs)
+            h = x @ antisymmetric_W.t() + h
+
+            if self.bias is not None:
+                h += self.bias
+
+            if self.act is not None:
+                h = self.act(h)
+
+            x = x + self.epsilon * h
+
+        return x
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"{self.in_channels}, "
+            f"phi={self.phi}, "
+            f"num_iters={self.num_iters}, "
+            f"epsilon={self.epsilon}, "
+            f"gamma={self.gamma})"
+        )
+
+
+model_cls = AntiSymmetricConv
+
+
+if __name__ == "__main__":
+    node_features = torch.load("node_features.pt")
+    edge_index = torch.load("edge_index.pt")
+
+    # Model instantiation and forward pass
+    model = AntiSymmetricConv(in_channels=node_features.size(-1))
+    output = model(node_features, edge_index)
+
+    # Save output to a file
+    torch.save(output, "gt_output.pt")

rdagent/model_implementation/benchmark/gt_code/dirgnn.py

@@ -0,0 +1,85 @@
+import copy
+
+import torch
+from torch import Tensor
+
+from torch_geometric.nn.conv import MessagePassing
+
+
+class DirGNNConv(torch.nn.Module):
+    r"""A generic wrapper for computing graph convolution on directed
+    graphs as described in the `"Edge Directionality Improves Learning on
+    Heterophilic Graphs" <https://arxiv.org/abs/2305.10498>`_ paper.
+    :class:`DirGNNConv` will pass messages both from source nodes to target
+    nodes and from target nodes to source nodes.
+
+    Args:
+        conv (MessagePassing): The underlying
+            :class:`~torch_geometric.nn.conv.MessagePassing` layer to use.
+        alpha (float, optional): The alpha coefficient used to weight the
+            aggregations of in- and out-edges as part of a convex combination.
+            (default: :obj:`0.5`)
+        root_weight (bool, optional): If set to :obj:`True`, the layer will add
+            transformed root node features to the output.
+            (default: :obj:`True`)
+    """
+    def __init__(
+        self,
+        conv: MessagePassing,
+        alpha: float = 0.5,
+        root_weight: bool = True,
+    ):
+        super().__init__()
+
+        self.alpha = alpha
+        self.root_weight = root_weight
+
+        self.conv_in = copy.deepcopy(conv)
+        self.conv_out = copy.deepcopy(conv)
+
+        if hasattr(conv, 'add_self_loops'):
+            self.conv_in.add_self_loops = False
+            self.conv_out.add_self_loops = False
+        if hasattr(conv, 'root_weight'):
+            self.conv_in.root_weight = False
+            self.conv_out.root_weight = False
+
+        if root_weight:
+            self.lin = torch.nn.Linear(conv.in_channels, conv.out_channels)
+        else:
+            self.lin = None
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        self.conv_in.reset_parameters()
+        self.conv_out.reset_parameters()
+        if self.lin is not None:
+            self.lin.reset_parameters()
+
+    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
+        """"""  # noqa: D419
+        x_in = self.conv_in(x, edge_index)
+        x_out = self.conv_out(x, edge_index.flip([0]))
+
+        out = self.alpha * x_out + (1 - self.alpha) * x_in
+
+        if self.root_weight:
+            out = out + self.lin(x)
+
+        return out
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.conv_in}, alpha={self.alpha})'
+    
+if __name__ == "__main__":
+    node_features = torch.load("node_features.pt")
+    edge_index = torch.load("edge_index.pt")
+
+    # Model instantiation and forward pass
+    model = DirGNNConv(MessagePassing())
+    output = model(node_features, edge_index)
+
+    # Save output to a file
+    torch.save(output, "gt_output.pt")

rdagent/model_implementation/benchmark/gt_code/gpsconv.py

@@ -0,0 +1,196 @@
+import inspect
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import Dropout, Linear, Sequential
+
+from torch_geometric.nn.attention import PerformerAttention
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.nn.inits import reset
+from torch_geometric.nn.resolver import (
+    activation_resolver,
+    normalization_resolver,
+)
+from torch_geometric.typing import Adj
+from torch_geometric.utils import to_dense_batch
+
+
+class GPSConv(torch.nn.Module):
+    r"""The general, powerful, scalable (GPS) graph transformer layer from the
+    `"Recipe for a General, Powerful, Scalable Graph Transformer"
+    <https://arxiv.org/abs/2205.12454>`_ paper.
+
+    The GPS layer is based on a 3-part recipe:
+
+    1. Inclusion of positional (PE) and structural encodings (SE) to the input
+       features (done in a pre-processing step via
+       :class:`torch_geometric.transforms`).
+    2. A local message passing layer (MPNN) that operates on the input graph.
+    3. A global attention layer that operates on the entire graph.
+
+    .. note::
+
+        For an example of using :class:`GPSConv`, see
+        `examples/graph_gps.py
+        <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
+        graph_gps.py>`_.
+
+    Args:
+        channels (int): Size of each input sample.
+        conv (MessagePassing, optional): The local message passing layer.
+        heads (int, optional): Number of multi-head-attentions.
+            (default: :obj:`1`)
+        dropout (float, optional): Dropout probability of intermediate
+            embeddings. (default: :obj:`0.`)
+        act (str or Callable, optional): The non-linear activation function to
+            use. (default: :obj:`"relu"`)
+        act_kwargs (Dict[str, Any], optional): Arguments passed to the
+            respective activation function defined by :obj:`act`.
+            (default: :obj:`None`)
+        norm (str or Callable, optional): The normalization function to
+            use. (default: :obj:`"batch_norm"`)
+        norm_kwargs (Dict[str, Any], optional): Arguments passed to the
+            respective normalization function defined by :obj:`norm`.
+            (default: :obj:`None`)
+        attn_type (str): Global attention type, :obj:`multihead` or
+            :obj:`performer`. (default: :obj:`multihead`)
+        attn_kwargs (Dict[str, Any], optional): Arguments passed to the
+            attention layer. (default: :obj:`None`)
+    """
+    def __init__(
+        self,
+        channels: int,
+        conv: Optional[MessagePassing],
+        heads: int = 1,
+        dropout: float = 0.0,
+        act: str = 'relu',
+        act_kwargs: Optional[Dict[str, Any]] = None,
+        norm: Optional[str] = 'batch_norm',
+        norm_kwargs: Optional[Dict[str, Any]] = None,
+        attn_type: str = 'multihead',
+        attn_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__()
+
+        self.channels = channels
+        self.conv = conv
+        self.heads = heads
+        self.dropout = dropout
+        self.attn_type = attn_type
+
+        attn_kwargs = attn_kwargs or {}
+        if attn_type == 'multihead':
+            self.attn = torch.nn.MultiheadAttention(
+                channels,
+                heads,
+                batch_first=True,
+                **attn_kwargs,
+            )
+        elif attn_type == 'performer':
+            self.attn = PerformerAttention(
+                channels=channels,
+                heads=heads,
+                **attn_kwargs,
+            )
+        else:
+            # TODO: Support BigBird
+            raise ValueError(f'{attn_type} is not supported')
+
+        self.mlp = Sequential(
+            Linear(channels, channels * 2),
+            activation_resolver(act, **(act_kwargs or {})),
+            Dropout(dropout),
+            Linear(channels * 2, channels),
+            Dropout(dropout),
+        )
+
+        norm_kwargs = norm_kwargs or {}
+        self.norm1 = normalization_resolver(norm, channels, **norm_kwargs)
+        self.norm2 = normalization_resolver(norm, channels, **norm_kwargs)
+        self.norm3 = normalization_resolver(norm, channels, **norm_kwargs)
+
+        self.norm_with_batch = False
+        if self.norm1 is not None:
+            signature = inspect.signature(self.norm1.forward)
+            self.norm_with_batch = 'batch' in signature.parameters
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        if self.conv is not None:
+            self.conv.reset_parameters()
+        self.attn._reset_parameters()
+        reset(self.mlp)
+        if self.norm1 is not None:
+            self.norm1.reset_parameters()
+        if self.norm2 is not None:
+            self.norm2.reset_parameters()
+        if self.norm3 is not None:
+            self.norm3.reset_parameters()
+
+    def forward(
+        self,
+        x: Tensor,
+        edge_index: Adj,
+        batch: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Tensor:
+        r"""Runs the forward pass of the module."""
+        hs = []
+        if self.conv is not None:  # Local MPNN.
+            h = self.conv(x, edge_index, **kwargs)
+            h = F.dropout(h, p=self.dropout, training=self.training)
+            h = h + x
+            if self.norm1 is not None:
+                if self.norm_with_batch:
+                    h = self.norm1(h, batch=batch)
+                else:
+                    h = self.norm1(h)
+            hs.append(h)
+
+        # Global attention transformer-style model.
+        h, mask = to_dense_batch(x, batch)
+
+        if isinstance(self.attn, torch.nn.MultiheadAttention):
+            h, _ = self.attn(h, h, h, key_padding_mask=~mask,
+                             need_weights=False)
+        elif isinstance(self.attn, PerformerAttention):
+            h = self.attn(h, mask=mask)
+
+        h = h[mask]
+        h = F.dropout(h, p=self.dropout, training=self.training)
+        h = h + x  # Residual connection.
+        if self.norm2 is not None:
+            if self.norm_with_batch:
+                h = self.norm2(h, batch=batch)
+            else:
+                h = self.norm2(h)
+        hs.append(h)
+
+        out = sum(hs)  # Combine local and global outputs.
+
+        out = out + self.mlp(out)
+        if self.norm3 is not None:
+            if self.norm_with_batch:
+                out = self.norm3(out, batch=batch)
+            else:
+                out = self.norm3(out)
+
+        return out
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.channels}, '
+                f'conv={self.conv}, heads={self.heads}, '
+                f'attn_type={self.attn_type})')
+
+if __name__ == "__main__":
+    node_features = torch.load("node_features.pt")
+    edge_index = torch.load("edge_index.pt")
+
+    # Model instantiation and forward pass
+    model = GPSConv(channels=node_features.size(-1),conv=MessagePassing())
+    output = model(node_features, edge_index)
+
+    # Save output to a file
+    torch.save(output, "gt_output.pt")

rdagent/model_implementation/benchmark/gt_code/linkx.py

@@ -0,0 +1,178 @@
+import math
+
+import torch
+from torch import Tensor
+from torch.nn import BatchNorm1d, Parameter
+
+from torch_geometric.nn import inits
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.nn.models import MLP
+from torch_geometric.typing import Adj, OptTensor
+from torch_geometric.utils import spmm
+
+
+class SparseLinear(MessagePassing):
+    def __init__(self, in_channels: int, out_channels: int, bias: bool = True):
+        super().__init__(aggr='add')
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.weight = Parameter(torch.empty(in_channels, out_channels))
+        if bias:
+            self.bias = Parameter(torch.empty(out_channels))
+        else:
+            self.register_parameter('bias', None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        inits.kaiming_uniform(self.weight, fan=self.in_channels,
+                              a=math.sqrt(5))
+        inits.uniform(self.in_channels, self.bias)
+
+    def forward(
+        self,
+        edge_index: Adj,
+        edge_weight: OptTensor = None,
+    ) -> Tensor:
+        # propagate_type: (weight: Tensor, edge_weight: OptTensor)
+        out = self.propagate(edge_index, weight=self.weight,
+                             edge_weight=edge_weight)
+
+        if self.bias is not None:
+            out = out + self.bias
+
+        return out
+
+    def message(self, weight_j: Tensor, edge_weight: OptTensor) -> Tensor:
+        if edge_weight is None:
+            return weight_j
+        else:
+            return edge_weight.view(-1, 1) * weight_j
+
+    def message_and_aggregate(self, adj_t: Adj, weight: Tensor) -> Tensor:
+        return spmm(adj_t, weight, reduce=self.aggr)
+
+
+class LINKX(torch.nn.Module):
+    r"""The LINKX model from the `"Large Scale Learning on Non-Homophilous
+    Graphs: New Benchmarks and Strong Simple Methods"
+    <https://arxiv.org/abs/2110.14446>`_ paper.
+
+    .. math::
+        \mathbf{H}_{\mathbf{A}} &= \textrm{MLP}_{\mathbf{A}}(\mathbf{A})
+
+        \mathbf{H}_{\mathbf{X}} &= \textrm{MLP}_{\mathbf{X}}(\mathbf{X})
+
+        \mathbf{Y} &= \textrm{MLP}_{f} \left( \sigma \left( \mathbf{W}
+        [\mathbf{H}_{\mathbf{A}}, \mathbf{H}_{\mathbf{X}}] +
+        \mathbf{H}_{\mathbf{A}} + \mathbf{H}_{\mathbf{X}} \right) \right)
+
+    .. note::
+
+        For an example of using LINKX, see `examples/linkx.py <https://
+        github.com/pyg-team/pytorch_geometric/blob/master/examples/linkx.py>`_.
+
+    Args:
+        num_nodes (int): The number of nodes in the graph.
+        in_channels (int): Size of each input sample, or :obj:`-1` to derive
+            the size from the first input(s) to the forward method.
+        hidden_channels (int): Size of each hidden sample.
+        out_channels (int): Size of each output sample.
+        num_layers (int): Number of layers of :math:`\textrm{MLP}_{f}`.
+        num_edge_layers (int, optional): Number of layers of
+            :math:`\textrm{MLP}_{\mathbf{A}}`. (default: :obj:`1`)
+        num_node_layers (int, optional): Number of layers of
+            :math:`\textrm{MLP}_{\mathbf{X}}`. (default: :obj:`1`)
+        dropout (float, optional): Dropout probability of each hidden
+            embedding. (default: :obj:`0.0`)
+    """
+    def __init__(
+        self,
+        num_nodes: int,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        num_layers: int,
+        num_edge_layers: int = 1,
+        num_node_layers: int = 1,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+
+        self.num_nodes = num_nodes
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_edge_layers = num_edge_layers
+
+        self.edge_lin = SparseLinear(num_nodes, hidden_channels)
+
+        if self.num_edge_layers > 1:
+            self.edge_norm = BatchNorm1d(hidden_channels)
+            channels = [hidden_channels] * num_edge_layers
+            self.edge_mlp = MLP(channels, dropout=0., act_first=True)
+        else:
+            self.edge_norm = None
+            self.edge_mlp = None
+
+        channels = [in_channels] + [hidden_channels] * num_node_layers
+        self.node_mlp = MLP(channels, dropout=0., act_first=True)
+
+        self.cat_lin1 = torch.nn.Linear(hidden_channels, hidden_channels)
+        self.cat_lin2 = torch.nn.Linear(hidden_channels, hidden_channels)
+
+        channels = [hidden_channels] * num_layers + [out_channels]
+        self.final_mlp = MLP(channels, dropout=dropout, act_first=True)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        self.edge_lin.reset_parameters()
+        if self.edge_norm is not None:
+            self.edge_norm.reset_parameters()
+        if self.edge_mlp is not None:
+            self.edge_mlp.reset_parameters()
+        self.node_mlp.reset_parameters()
+        self.cat_lin1.reset_parameters()
+        self.cat_lin2.reset_parameters()
+        self.final_mlp.reset_parameters()
+
+    def forward(
+        self,
+        x: OptTensor,
+        edge_index: Adj,
+        edge_weight: OptTensor = None,
+    ) -> Tensor:
+        """"""  # noqa: D419
+        out = self.edge_lin(edge_index, edge_weight)
+
+        if self.edge_norm is not None and self.edge_mlp is not None:
+            out = out.relu_()
+            out = self.edge_norm(out)
+            out = self.edge_mlp(out)
+
+        out = out + self.cat_lin1(out)
+
+        if x is not None:
+            x = self.node_mlp(x)
+            out = out + x
+            out = out + self.cat_lin2(x)
+
+        return self.final_mlp(out.relu_())
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}(num_nodes={self.num_nodes}, '
+                f'in_channels={self.in_channels}, '
+                f'out_channels={self.out_channels})')
+
+if __name__ == "__main__":
+    node_features = torch.load("node_features.pt")
+    edge_index = torch.load("edge_index.pt")
+
+    # Model instantiation and forward pass
+    model = LINKX(num_nodes=node_features.size(0), in_channels=node_features.size(1), hidden_channels=node_features.size(1), out_channels=node_features.size(1), num_layers=1)
+    output = model(node_features, edge_index)
+
+    # Save output to a file
+    torch.save(output, "gt_output.pt")

rdagent/model_implementation/benchmark/gt_code/pmlp.py

@@ -0,0 +1,112 @@
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from torch_geometric.nn import SimpleConv
+from torch_geometric.nn.dense.linear import Linear
+
+
+class PMLP(torch.nn.Module):
+    r"""The P(ropagational)MLP model from the `"Graph Neural Networks are
+    Inherently Good Generalizers: Insights by Bridging GNNs and MLPs"
+    <https://arxiv.org/abs/2212.09034>`_ paper.
+    :class:`PMLP` is identical to a standard MLP during training, but then
+    adopts a GNN architecture during testing.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        hidden_channels (int): Size of each hidden sample.
+        out_channels (int): Size of each output sample.
+        num_layers (int): The number of layers.
+        dropout (float, optional): Dropout probability of each hidden
+            embedding. (default: :obj:`0.`)
+        norm (bool, optional): If set to :obj:`False`, will not apply batch
+            normalization. (default: :obj:`True`)
+        bias (bool, optional): If set to :obj:`False`, the module
+            will not learn additive biases. (default: :obj:`True`)
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        num_layers: int,
+        dropout: float = 0.,
+        norm: bool = True,
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.bias = bias
+
+        self.lins = torch.nn.ModuleList()
+        self.lins.append(Linear(in_channels, hidden_channels, self.bias))
+        for _ in range(self.num_layers - 2):
+            lin = Linear(hidden_channels, hidden_channels, self.bias)
+            self.lins.append(lin)
+        self.lins.append(Linear(hidden_channels, out_channels, self.bias))
+
+        self.norm = None
+        if norm:
+            self.norm = torch.nn.BatchNorm1d(
+                hidden_channels,
+                affine=False,
+                track_running_stats=False,
+            )
+
+        self.conv = SimpleConv(aggr='mean', combine_root='self_loop')
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        for lin in self.lins:
+            torch.nn.init.xavier_uniform_(lin.weight, gain=1.414)
+            if self.bias:
+                torch.nn.init.zeros_(lin.bias)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        edge_index: Optional[Tensor] = None,
+    ) -> torch.Tensor:
+        """"""  # noqa: D419
+        if not self.training and edge_index is None:
+            raise ValueError(f"'edge_index' needs to be present during "
+                             f"inference in '{self.__class__.__name__}'")
+
+        for i in range(self.num_layers):
+            x = x @ self.lins[i].weight.t()
+            if not self.training:
+                x = self.conv(x, edge_index)
+            if self.bias:
+                x = x + self.lins[i].bias
+            if i != self.num_layers - 1:
+                if self.norm is not None:
+                    x = self.norm(x)
+                x = x.relu()
+                x = F.dropout(x, p=self.dropout, training=self.training)
+
+        return x
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, num_layers={self.num_layers})')
+
+if __name__ == "__main__":
+    node_features = torch.load("node_features.pt")
+    edge_index = torch.load("edge_index.pt")
+
+    # Model instantiation and forward pass
+    model = PMLP(in_channels=node_features.size(-1), hidden_channels=node_features.size(-1), node_features.size(-1), num_layers=1)
+    output = model(node_features, edge_index)
+
+    # Save output to a file
+    torch.save(output, "gt_output.pt")

rdagent/model_implementation/conf.py

@@ -0,0 +1,10 @@
+from pathlib import Path
+from pydantic_settings import BaseSettings
+
+class ModelImplSettings(BaseSettings):
+    workspace_path: Path = Path("./git_ignore_folder/model_imp_workspace/")  # Added type annotation for work_space
+    
+    class Config:
+        env_prefix = 'MODEL_IMPL_'  # Use MODEL_IMPL_ as prefix for environment variables
+
+MODEL_IMPL_SETTINGS = ModelImplSettings()

rdagent/model_implementation/evaluator.py

@@ -0,0 +1,63 @@
+import torch
+import numpy as np
+
+
+def shape_evaluator(target, prediction):
+    if target is None or prediction is None:
+        return None, 0
+    tar_shape = target.shape
+    pre_shape = prediction.shape
+
+    diff = []
+    for i in range(max(len(tar_shape), len(pre_shape))):
+        dim_tar = tar_shape[i] if i < len(tar_shape) else 0
+        dim_pre = pre_shape[i] if i < len(pre_shape) else 0
+        diff.append(abs(dim_tar - dim_pre))
+
+    metric = 1 / (np.exp(np.mean(diff)) + 1)
+    return diff, metric
+
+
+def reshape_tensor(original_tensor, target_shape):
+    new_tensor = torch.zeros(target_shape)
+    for i, dim in enumerate(original_tensor.shape):
+        new_tensor = new_tensor.narrow(i, 0, dim).copy_(original_tensor)
+
+    return new_tensor
+
+
+def value_evaluator(target, prediction):
+    if target is None or prediction is None:
+        return None, 0
+    tar_shape = target.shape
+    pre_shape = prediction.shape
+
+    # Determine the shape of the padded tensors
+    dims = [
+        max(s1, s2)
+        for s1, s2 in zip(
+            tar_shape + (1,) * (len(pre_shape) - len(tar_shape)),
+            pre_shape + (1,) * (len(tar_shape) - len(pre_shape)),
+        )
+    ]
+    # Reshape both tensors to the determined shape
+    target = target.reshape(
+        *tar_shape, *(1,) * (max(len(tar_shape), len(pre_shape)) - len(tar_shape))
+    )
+    prediction = prediction.reshape(
+        *pre_shape, *(1,) * (max(len(tar_shape), len(pre_shape)) - len(pre_shape))
+    )
+    target_padded = reshape_tensor(target, dims)
+    prediction_padded = reshape_tensor(prediction, dims)
+
+    # Calculate the mean absolute difference
+    diff = torch.abs(target_padded - prediction_padded)
+    metric = 1 / (1 + np.exp(torch.mean(diff).item()))
+    return diff, metric
+
+
+if __name__ == "__main__":
+    tar = torch.rand(4, 5, 5)
+    pre = torch.rand(4, 1)
+    print(shape_evaluator(tar, pre))
+    print(value_evaluator(tar, pre)[1])

rdagent/model_implementation/gt_code.py

@@ -0,0 +1,136 @@
+"""
+This is just an exmaple.
+It will be replaced wtih a list of ground truth tasks.
+"""
+import math
+from typing import Any, Callable, Dict, Optional, Union
+
+import torch
+from torch import Tensor
+from torch.nn import Parameter
+
+from torch_geometric.nn.conv import GCNConv, MessagePassing
+from torch_geometric.nn.inits import zeros
+from torch_geometric.nn.resolver import activation_resolver
+from torch_geometric.typing import Adj
+
+
+class AntiSymmetricConv(torch.nn.Module):
+    r"""The anti-symmetric graph convolutional operator from the
+    `"Anti-Symmetric DGN: a stable architecture for Deep Graph Networks"
+    <https://openreview.net/forum?id=J3Y7cgZOOS>`_ paper.
+
+    .. math::
+        \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),
+
+    where :math:`\Phi(\mathbf{X}, \mathcal{N}_i)` denotes a
+    :class:`~torch.nn.conv.MessagePassing` layer.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        phi (MessagePassing, optional): The message passing module
+            :math:`\Phi`. If set to :obj:`None`, will use a
+            :class:`~torch_geometric.nn.conv.GCNConv` layer as default.
+            (default: :obj:`None`)
+        num_iters (int, optional): The number of times the anti-symmetric deep
+            graph network operator is called. (default: :obj:`1`)
+        epsilon (float, optional): The discretization step size
+            :math:`\epsilon`. (default: :obj:`0.1`)
+        gamma (float, optional): The strength of the diffusion :math:`\gamma`.
+            It regulates the stability of the method. (default: :obj:`0.1`)
+        act (str, optional): The non-linear activation function :math:`\sigma`,
+            *e.g.*, :obj:`"tanh"` or :obj:`"relu"`. (default: :class:`"tanh"`)
+        act_kwargs (Dict[str, Any], optional): Arguments passed to the
+            respective activation function defined by :obj:`act`.
+            (default: :obj:`None`)
+        bias (bool, optional): If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+
+    Shapes:
+        - **input:**
+          node features :math:`(|\mathcal{V}|, F_{in})`,
+          edge indices :math:`(2, |\mathcal{E}|)`,
+          edge weights :math:`(|\mathcal{E}|)` *(optional)*
+        - **output:** node features :math:`(|\mathcal{V}|, F_{in})`
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        phi: Optional[MessagePassing] = None,
+        num_iters: int = 1,
+        epsilon: float = 0.1,
+        gamma: float = 0.1,
+        act: Union[str, Callable, None] = "tanh",
+        act_kwargs: Optional[Dict[str, Any]] = None,
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.num_iters = num_iters
+        self.gamma = gamma
+        self.epsilon = epsilon
+        self.act = activation_resolver(act, **(act_kwargs or {}))
+
+        if phi is None:
+            phi = GCNConv(in_channels, in_channels, bias=False)
+
+        self.W = Parameter(torch.empty(in_channels, in_channels))
+        self.register_buffer("eye", torch.eye(in_channels))
+        self.phi = phi
+
+        if bias:
+            self.bias = Parameter(torch.empty(in_channels))
+        else:
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        torch.nn.init.kaiming_uniform_(self.W, a=math.sqrt(5))
+        self.phi.reset_parameters()
+        zeros(self.bias)
+
+    def forward(self, x: Tensor, edge_index: Adj, *args, **kwargs) -> Tensor:
+        r"""Runs the forward pass of the module."""
+        antisymmetric_W = self.W - self.W.t() - self.gamma * self.eye
+
+        for _ in range(self.num_iters):
+            h = self.phi(x, edge_index, *args, **kwargs)
+            h = x @ antisymmetric_W.t() + h
+
+            if self.bias is not None:
+                h += self.bias
+
+            if self.act is not None:
+                h = self.act(h)
+
+            x = x + self.epsilon * h
+
+        return x
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"{self.in_channels}, "
+            f"phi={self.phi}, "
+            f"num_iters={self.num_iters}, "
+            f"epsilon={self.epsilon}, "
+            f"gamma={self.gamma})"
+        )
+
+
+if __name__ == "__main__":
+    node_features = torch.load("node_features.pt")
+    edge_index = torch.load("edge_index.pt")
+
+    # Model instantiation and forward pass
+    model = AntiSymmetricConv(in_channels=node_features.size(-1))
+    output = model(node_features, edge_index)
+
+    # Save output to a file
+    torch.save(output, "gt_output.pt")

rdagent/model_implementation/main.py

@@ -0,0 +1,91 @@
+"""
+This file will be removed in the future and replaced by
+- rdagent/app/model_implementation/eval.py
+"""
+from dotenv import load_dotenv
+from rdagent.oai.llm_utils import APIBackend
+
+# randomly generate a input graph, node_feature and edge_index
+# 1000 nodes, 128 dim node feature, 2000 edges
+import torch
+import os
+
+assert load_dotenv()
+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": "x_u^{(l)} = x_u^{(l-1)} + \\epsilon \\sigma \\left( W^T x_u^{(l-1)} + \\Phi(X^{(l-1)}, N_u) + b \\right)",
+    "variables": {
+        "x_u^{(l)}": "The state of node u at layer l",
+        "\\epsilon": "The step size in the Euler discretization",
+        "\\sigma": "A monotonically non-decreasing activation function",
+        "W": "An anti-symmetric weight matrix",
+        "X^{(l-1)}": "The node feature matrix at layer l-1",
+        "N_u": "The set of neighbors of node u",
+        "b": "A bias vector",
+    },
+}
+
+system_prompt = "You are an assistant whose job is to answer user's question."
+user_prompt = "With the following given information, write a python code using pytorch and torch_geometric to implement the model. This model is in the graph learning field, only have one layer. The input will be node_feature [num_nodes, dim_feature] and edge_index [2, num_edges], and they should be loaded from the files 'node_features.pt' and 'edge_index.pt'. There is not edge attribute or edge weight as input. The model should detect the node_feature and edge_index shape, if there is Linear transformation layer in the model, the input and output shape should be consistent. The in_channels is the dimension of the node features. You code should contain additional 'if __name__ == '__main__', where you should load the node_feature and edge_index from the files and run the model, and save the output to a file 'llm_output.pt'. Implement the model forward function based on the following information: model formula information. 1. model name: {}, 2. model description: {}, 3. model formulation: {}, 4. model variables: {}. You must complete the forward function as far as you can do.".format(
+    formula_info["name"],
+    formula_info["description"],
+    formula_info["formulation"],
+    formula_info["variables"],
+)
+
+resp = APIBackend(use_chat_cache=False).build_messages_and_create_chat_completion(
+    user_prompt, system_prompt
+)
+
+print(resp)
+
+# take the code part from the response and save it to a file, the code is covered in the ```python``` block
+code = resp.split("```python")[1].split("```")[0]
+with open("llm_code.py", "w") as f:
+    f.write(code)
+
+average_shape_eval = []
+average_value_eval = []
+for test_mode in ["zeros", "ones", "randn"]:
+
+    if test_mode == "zeros":
+        node_feature = torch.zeros(1000, 128)
+    elif test_mode == "ones":
+        node_feature = torch.ones(1000, 128)
+    elif test_mode == "randn":
+        node_feature = torch.randn(1000, 128)
+    edge_index = torch.randint(0, 1000, (2, 2000))
+
+    torch.save(node_feature, "node_features.pt")
+    torch.save(edge_index, "edge_index.pt")
+
+    try:
+        os.system("python llm_code.py")
+    except:
+        print("Error in running the LLM code")
+    os.system("python gt_code.py")
+    os.system("rm edge_index.pt")
+    os.system("rm node_features.pt")
+    # load the output and print the shape
+
+    from evaluator import shape_evaluator, value_evaluator
+
+    try:
+        llm_output = torch.load("llm_output.pt")
+    except:
+        llm_output = None
+    gt_output = torch.load("gt_output.pt")
+
+    average_shape_eval.append(shape_evaluator(llm_output, gt_output)[1])
+    average_value_eval.append(value_evaluator(llm_output, gt_output)[1])
+
+    print("Shape evaluation: ", average_shape_eval[-1])
+    print("Value evaluation:super().generate(task_l) ", average_value_eval[-1])
+
+    os.system("rm llm_output.pt")
+    os.system("rm gt_output.pt")
+os.system("rm llm_code.py")
+
+print("Average shape evaluation: ", sum(average_shape_eval) / len(average_shape_eval))
+print("Average value evaluation: ", sum(average_value_eval) / len(average_value_eval))

rdagent/model_implementation/one_shot/__init__.py

@@ -0,0 +1,42 @@
+from typing import Sequence
+from rdagent.oai.llm_utils import APIBackend
+
+from jinja2 import Template
+from rdagent.core.implementation import TaskGenerator
+from rdagent.core.prompts import Prompts
+from rdagent.model_implementation.task import ModelImplTask, ModelTaskImpl
+
+from pathlib import Path
+DIRNAME = Path(__file__).absolute().resolve().parent
+
+
+class ModelTaskGen(TaskGenerator):
+
+    def generate(self, task_l: Sequence[ModelImplTask]) -> Sequence[ModelTaskImpl]:
+        mti_l = []
+        for t in task_l:
+            mti = ModelTaskImpl(t)
+            mti.prepare()
+            pr = Prompts(file_path=DIRNAME / "prompt.yaml")
+
+            user_prompt_tpl = Template(pr["code_implement_user"])
+            sys_prompt_tpl = Template(pr["code_implement_sys"])
+
+            user_prompt = user_prompt_tpl.render(
+                name=t.name,
+                description=t.description,
+                formulation=t.formulation,
+                variables=t.variables,
+                execute_desc=mti.execute_desc()
+            )
+            system_prompt = sys_prompt_tpl.render()
+
+            resp = APIBackend().build_messages_and_create_chat_completion(
+                user_prompt, system_prompt
+            )
+
+            # Extract the code part from the response
+            code = resp.split("```python")[1].split("```")[0]
+            mti.inject_code(**{"model.py": code})
+            mti_l.append(mti)
+        return mti_l

rdagent/model_implementation/one_shot/prompt.yaml

@@ -0,0 +1,18 @@
+
+
+code_implement_sys: -|
+  You are an assistant whose job is to answer user's question."
+code_implement_user: -|
+  With the following given information, write a python code using pytorch and torch_geometric to implement the model.
+  This model is in the graph learning field, only have one layer.
+  The input will be node_feature [num_nodes, dim_feature] and edge_index [2, num_edges]  (It would be the input of the forward model)
+  There is not edge attribute or edge weight as input. The model should detect the node_feature and edge_index shape, if there is Linear transformation layer in the model, the input and output shape should be consistent. The in_channels is the dimension of the node features.
+  Implement the model forward function based on the following information:model formula information.
+  1. model name:{{name}}
+  2. model description:{{description}}
+  3. model formulation:{{formulation}}
+  4. model variables:{{variables}}.
+  You must complete the forward function as far as you can do.
+  \# Execution
+  Your implemented code will be exectued in the follow way
+  {{execute_desc}}

rdagent/model_implementation/task.py

@@ -0,0 +1,144 @@
+import torch
+from pathlib import Path
+import uuid
+from typing import Dict, Optional, Sequence
+from rdagent.core.exception import CodeFormatException
+from rdagent.core.task import BaseTask, FBTaskImplementation, ImpLoader, TaskImplementation, TaskLoader
+from rdagent.model_implementation.conf import MODEL_IMPL_SETTINGS
+from rdagent.utils import get_module_by_module_path
+
+
+class ModelImplTask(BaseTask):
+    # TODO: it should change when the BaseTask changes.
+    name: str
+    description: str
+    formulation: str
+    variables: Dict[str, str]  # map the variable name to the variable description
+
+    def __init__(self, name: str, description: str, formulation: str, variables: Dict[str, str], key: Optional[str] = None) -> None:
+        """
+
+        Parameters
+        ----------
+
+        key : Optional[str]
+            Key is a string to identify the task.
+            It will be used to connect to other information(e.g. ground truth).
+        """
+        self.name = name
+        self.description = description
+        self.formulation = formulation
+        self.variables = variables
+        self.key = key
+
+
+class ModelTaskLoderJson(TaskLoader):
+    def __init__(self, json_uri: str) -> None:
+        super().__init__()
+        # TODO: the json should be loaded from URI.
+        self.json_uri = json_uri
+
+    def load(self, *argT, **kwargs) -> Sequence[ModelImplTask]:
+        # TODO: we should load the tasks from json;
+
+        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": "x_u^{(l)} = x_u^{(l-1)} + \\epsilon \\sigma \\left( W^T x_u^{(l-1)} + \\Phi(X^{(l-1)}, N_u) + b \\right)",
+            "variables": {
+                "x_u^{(l)}": "The state of node u at layer l",
+                "\\epsilon": "The step size in the Euler discretization",
+                "\\sigma": "A monotonically non-decreasing activation function",
+                "W": "An anti-symmetric weight matrix",
+                "X^{(l-1)}": "The node feature matrix at layer l-1",
+                "N_u": "The set of neighbors of node u",
+                "b": "A bias vector",
+            },
+            "key": "A-DGN",
+        }
+        return [ModelImplTask(**formula_info)]
+
+
+class ModelTaskImpl(FBTaskImplementation):
+    """
+    It is a Pytorch model implementation task;
+    All the things are placed in a folder.
+
+    Folder
+    - data source and documents prepared by `prepare`
+        - Please note that new data may be passed in dynamically in `execute`
+    - code (file `model.py` ) injected by `inject_code`
+        - the `model.py` that contains a variable named `model_cls` which indicates the implemented model structure
+            - `model_cls` is a instance of `torch.nn.Module`;
+
+    
+    We'll import the model in the implementation in file `model.py` after setting the cwd into the directory
+    - from model import model_cls
+    - initialize the model by initializing it `model_cls(input_dim=INPUT_DIM)`
+    - And then verify the modle.
+
+    """
+    def __init__(self, target_task: BaseTask) -> None:
+        super().__init__(target_task)
+        self.path = None
+
+    def prepare(self) -> None:
+        """
+        Prepare for the workspace;
+        """
+        unique_id = uuid.uuid4()
+        self.path = MODEL_IMPL_SETTINGS.workspace_path / f"M{unique_id}"
+        # start with `M` so that it can be imported via python
+        self.path.mkdir(parents=True, exist_ok=True)
+
+    def execute(self, data=None, config: dict = {}):
+        mod = get_module_by_module_path(str(self.path / "model.py"))
+        try:
+            model_cls = mod.model_cls
+        except AttributeError:
+            raise CodeFormatException("The model_cls is not implemented in the model.py")
+        # model_init = 
+
+        assert isinstance(data, tuple)
+        node_feature, _ = data
+        in_channels = node_feature.size(-1)
+        m = model_cls(in_channels)
+
+        # TODO: initialize all the parameters of `m` to `model_eval_param_init`
+        model_eval_param_init: float = config["model_eval_param_init"]
+
+        # initialize all parameters of `m` to `model_eval_param_init`
+        for _, param in m.named_parameters():
+            param.data.fill_(model_eval_param_init)
+
+        assert isinstance(data, tuple)
+        return m(*data)
+
+    def execute_desc(self) -> str:
+        return """
+The the implemented code will be placed in a file like <uuid>/model.py
+
+We'll import the model in the implementation in file `model.py` after setting the cwd into the directory
+- from model import model_cls (So you must have a variable named `model_cls` in the file)
+  - So your implelemented code could follow the following pattern
+    ```Python
+    class XXXLayer(torch.nn.Module):
+        ...
+    model_cls = XXXLayer
+    ```
+- initialize the model by initializing it `model_cls(input_dim=INPUT_DIM)`
+- And then verify the model by comparing the output tensors by feeding specific input tensor.
+"""
+
+class ModelImpLoader(ImpLoader[ModelImplTask, ModelTaskImpl]):
+    def __init__(self, path: Path) -> None:
+        self.path = Path(path)
+
+    def load(self, task: ModelImplTask) -> ModelTaskImpl:
+        assert task.key is not None
+        mti = ModelTaskImpl(task)
+        mti.prepare()
+        with open(self.path / f"{task.key}.py", "r") as f:
+            code = f.read()
+        mti.inject_code(**{"model.py": code})
+        return mti

rdagent/utils/__init__.py

@@ -0,0 +1,36 @@
+"""
+This is some common utils functions.
+it is not binding to the scenarios or framework (So it is not placed in rdagent.core.utils)
+"""
+# TODO: merge the common utils in `rdagent.core.utils` into this folder
+# TODO: split the utils in this module into different modules in the future.
+
+import importlib
+import re
+import sys
+from types import ModuleType
+from typing import Union
+
+
+def get_module_by_module_path(module_path: Union[str, ModuleType]):
+    """Load module from path like a/b/c/d.py or a.b.c.d
+
+    :param module_path:
+    :return:
+    :raises: ModuleNotFoundError
+    """
+    if module_path is None:
+        raise ModuleNotFoundError("None is passed in as parameters as module_path")
+
+    if isinstance(module_path, ModuleType):
+        module = module_path
+    else:
+        if module_path.endswith(".py"):
+            module_name = re.sub("^[^a-zA-Z_]+", "", re.sub("[^0-9a-zA-Z_]", "", module_path[:-3].replace("/", "_")))
+            module_spec = importlib.util.spec_from_file_location(module_name, module_path)
+            module = importlib.util.module_from_spec(module_spec)
+            sys.modules[module_name] = module
+            module_spec.loader.exec_module(module)
+        else:
+            module = importlib.import_module(module_path)
+    return module