chore: import upstream snapshot with attribution

rdagent/components/coder/model_coder/__init__.py

@@ -0,0 +1,21 @@
+from rdagent.components.coder.CoSTEER import CoSTEER
+from rdagent.components.coder.CoSTEER.config import CoSTEER_SETTINGS
+from rdagent.components.coder.CoSTEER.evaluators import CoSTEERMultiEvaluator
+from rdagent.components.coder.model_coder.evaluators import ModelCoSTEEREvaluator
+from rdagent.components.coder.model_coder.evolving_strategy import (
+    ModelMultiProcessEvolvingStrategy,
+)
+from rdagent.core.scenario import Scenario
+
+
+class ModelCoSTEER(CoSTEER):
+    def __init__(
+        self,
+        scen: Scenario,
+        *args,
+        **kwargs,
+    ) -> None:
+        eva = CoSTEERMultiEvaluator(ModelCoSTEEREvaluator(scen=scen), scen=scen)
+        es = ModelMultiProcessEvolvingStrategy(scen=scen, settings=CoSTEER_SETTINGS)
+
+        super().__init__(*args, settings=CoSTEER_SETTINGS, eva=eva, es=es, evolving_version=2, scen=scen, **kwargs)

rdagent/components/coder/model_coder/benchmark/eval.py

@@ -0,0 +1,71 @@
+# TODO: inherent from the benchmark base class
+import torch
+
+from rdagent.components.coder.model_coder.model import ModelFBWorkspace
+
+
+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 observe the output.
+
+    Assumption:
+    - If the model structure is similar, the output will change in similar way when we change the input.
+
+    Challenge:
+    - The key difference between it and implementing models is that we have parameters in the layers (Model operators often have no parameters or are given parameters).
+    - we try to initialize the model param in similar value. So only the model structure is different.
+
+    Comparing the correlation of following sequences
+    - modelA[init1](input1).hidden_out1, modelA[init1](input2).hidden_out1, ...
+    - modelB[init1](input1).hidden_out1, modelB[init1](input2).hidden_out1, ...
+
+    For each hidden output, we can calculate a correlation. The average correlation will be the metrics.
+    """
+
+    def evaluate(self, gt: ModelFBWorkspace, gen: ModelFBWorkspace):
+        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]:
+                _, gt_res = gt.execute(input_value=init_val, param_init_value=init_val)
+                _, res = gen.execute(input_value=init_val, param_init_value=init_val)
+                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/components/coder/model_coder/benchmark/gt_code/A-DGN.py

@@ -0,0 +1,134 @@
+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/components/coder/model_coder/benchmark/gt_code/dirgnn.py

@@ -0,0 +1,89 @@
+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})"
+
+
+model_cls = DirGNNConv
+
+
+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/components/coder/model_coder/benchmark/gt_code/gpsconv.py

@@ -0,0 +1,198 @@
+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})"
+        )
+
+
+model_cls = GPSConv
+
+
+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/components/coder/model_coder/benchmark/gt_code/linkx.py

@@ -0,0 +1,187 @@
+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.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.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})"
+        )
+
+
+model_cls = LINKX
+
+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/components/coder/model_coder/benchmark/gt_code/pmlp.py

@@ -0,0 +1,118 @@
+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.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})"
+
+
+model_cls = PMLP
+
+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),
+        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/components/coder/model_coder/benchmark/model_dict.json

@@ -0,0 +1,80 @@
+{
+    "PMLP": {
+        "description": "`PMLP` is identical to a standard MLP during training, but then adopts a GNN architecture (add message passing) during testing.",
+        "formulation": "\\hat{y}_u = \\psi(\\text{MP}(\\{h^{(l-1)}_v\\}_{v \\in N_u \\cup \\{u\\}}))",
+        "variables": {
+            "\\hat{y}_u": "The predicted output for node u",
+            "\\psi": "A function representing the feed-forward process, consisting of a linear feature transformation followed by a non-linear activation",
+            "\\text{MP}": "Message Passing operation that aggregates neighbored information",
+            "h^{(l-1)}_v": "The feature representation of node v at layer (l-1)",
+            "N_u": "The set of neighbored nodes centered at node u"
+        },
+        "key": "pmlp",
+        "model_type": "TimeSeries"
+    },
+    "LINKX": {
+        "description": "A scalable model for node classification that separately embeds adjacency and node features, combines them with MLPs, and applies simple transformations.",
+        "formulation": "Y = MLP_f(\\sigma(W[h_A; h_X] + h_A + h_X))",
+        "variables": {
+            "Y": "The output predictions",
+            "\\sigma": "Non-linear activation function",
+            "W": "Learned weight matrix",
+            "h_A": "Embedding of the adjacency matrix",
+            "h_X": "Embedding of the node features",
+            "MLP_f": "Final multilayer perceptron for prediction"
+        },
+        "key": "linkx",
+        "model_type": "TimeSeries"
+    },
+    "GPSConv": {
+        "description": "A scalable and powerful graph transformer with linear complexity, capable of handling large graphs with state-of-the-art results across diverse benchmarks.",
+        "formulation": "X^{(l+1)} = \\text{MPNN}^{(l)}(X^{(l)}, A) + \\text{GlobalAttn}^{(l)}(X^{(l)})",
+        "variables": {
+            "X^{(l)}": "The node features at layer l",
+            "A": "The adjacency matrix of the graph",
+            "X^{(l+1)}": "The updated node features at layer l+1",
+            "MPNN^{(l)}": "The message-passing neural network function at layer l",
+            "GlobalAttn^{(l)}": "The global attention function at layer l"
+        },
+        "key": "gpsconv",
+        "model_type": "TimeSeries"
+    },
+    "ViSNet": {
+        "description": "ViSNet is an equivariant geometry-enhanced graph neural network designed for efficient molecular modeling[^1^][1][^2^][2]. It utilizes a Vector-Scalar interactive message passing mechanism to extract and utilize geometric features with low computational costs, achieving state-of-the-art performance on multiple molecular dynamics benchmarks.",
+        "formulation": "\\text{ViSNet}(G) = \\sum_{u \\in G} f(\\mathbf{h}_u, \\mathbf{e}_u, \\mathbf{v}_u)",
+        "variables": {
+            "\\mathbf{h}_u": "Node embedding for atom u",
+            "\\mathbf{e}_u": "Edge embedding associated with atom u",
+            "\\mathbf{v}_u": "Direction unit vector for atom u"
+        },
+        "key": "visnet",
+        "model_type": "TimeSeries"
+    },
+    "Dir-GNN": {
+        "description": "A framework for deep learning on directed graphs that extends MPNNs to incorporate edge directionality.",
+        "formulation": "x^{(k)}_i = COM^{(k)}\\left(x^{(k-1)}_i, m^{(k)}_{i,\\leftarrow}, m^{(k)}_{i,\\rightarrow}\\right)",
+        "variables": {
+            "x^{(k)}_i": "The feature representation of node i at layer k",
+            "m^{(k)}_{i,\\leftarrow}": "The aggregated incoming messages to node i at layer k",
+            "m^{(k)}_{i,\\rightarrow}": "The aggregated outgoing messages from node i at layer k"
+        },
+        "key": "dirgnn",
+        "model_type": "TimeSeries"
+    },
+    "A-DGN": {
+        "description": "A framework for stable and non-dissipative DGN design, conceived through the lens of ordinary differential equations (ODEs). It ensures long-range information preservation between nodes and prevents gradient vanishing or explosion during training.",
+        "formulation": "\\frac{\\partial x_u(t)}{\\partial t} = \\sigma(W^T x_u(t) + \\Phi(X(t), N_u) + b)",
+        "variables": {
+            "x_u(t)": "The state of node u at time t",
+            "\\frac{\\partial x_u(t)}{\\partial t}": "The rate of change of the state of node u at time t",
+            "\\sigma": "A monotonically non-decreasing activation function",
+            "W": "A weight matrix",
+            "b": "A bias vector",
+            "\\Phi(X(t), N_u)": "The aggregation function for the states of the nodes in the neighborhood of u",
+            "X(t)": "The node feature matrix of the whole graph at time t",
+            "N_u": "The set of neighboring nodes of u"
+        },
+        "key": "A-DGN",
+        "model_type": "TimeSeries"
+    }
+}

rdagent/components/coder/model_coder/conf.py

@@ -0,0 +1,38 @@
+from typing import Optional
+
+from pydantic_settings import SettingsConfigDict
+
+from rdagent.components.coder.CoSTEER.config import CoSTEERSettings
+from rdagent.utils.env import Env, QlibCondaConf, QlibCondaEnv, QTDockerEnv
+
+
+class ModelCoSTEERSettings(CoSTEERSettings):
+    model_config = SettingsConfigDict(env_prefix="MODEL_CoSTEER_")
+
+    env_type: str = "conda"  # or "docker"
+    """Environment to run model code in coder and runner: 'conda' for local conda env, 'docker' for Docker container"""
+
+
+def get_model_env(
+    conf_type: Optional[str] = None,
+    extra_volumes: dict = {},
+    running_timeout_period: int = 600,
+    enable_cache: Optional[bool] = None,
+) -> Env:
+    conf = ModelCoSTEERSettings()
+    if conf.env_type == "docker":
+        env = QTDockerEnv()
+    elif conf.env_type == "conda":
+        env = QlibCondaEnv(conf=QlibCondaConf())
+    else:
+        raise ValueError(f"Unknown env type: {conf.env_type}")
+
+    env.conf.extra_volumes = extra_volumes.copy()
+    env.conf.running_timeout_period = running_timeout_period
+    if enable_cache is not None:
+        env.conf.enable_cache = enable_cache
+    env.prepare()
+    return env
+
+
+MODEL_COSTEER_SETTINGS = ModelCoSTEERSettings()

rdagent/components/coder/model_coder/eva_utils.py

@@ -0,0 +1,166 @@
+import json
+from typing import Dict, Tuple
+
+import numpy as np
+
+from rdagent.components.coder.CoSTEER.evaluators import CoSTEEREvaluator
+from rdagent.components.coder.model_coder.model import ModelFBWorkspace, ModelTask
+from rdagent.core.experiment import Task, Workspace
+from rdagent.oai.llm_conf import LLM_SETTINGS
+from rdagent.oai.llm_utils import APIBackend
+from rdagent.utils.agent.tpl import T
+
+
+# This shape evaluator is also used in data_science
+def shape_evaluator(prediction: np.ndarray, target_shape: Tuple = None) -> Tuple[str, bool]:
+    if target_shape is None or prediction is None:
+        return (
+            "No output generated from the model. No shape evaluation conducted.",
+            False,
+        )
+    pre_shape = prediction.shape
+
+    if pre_shape == target_shape:
+        return "The shape of the output is correct.", True
+    else:
+        return (
+            f"The shape of the output is incorrect. Expected {target_shape}, but got {pre_shape}.",
+            False,
+        )
+
+
+def value_evaluator(
+    prediction: np.ndarray,
+    target: np.ndarray,
+) -> Tuple[np.ndarray, bool]:
+    if prediction is None:
+        return "No output generated from the model. Skip value evaluation", False
+    elif target is None:
+        return (
+            "No ground truth output provided. Value evaluation not impractical",
+            False,
+        )
+    else:
+        # Calculate the mean absolute difference
+        diff = np.mean(np.abs(target - prediction))
+        return (
+            f"The value of the output is correct. The mean absolute difference is {diff}.",
+            diff < 0.1,
+        )
+
+
+class ModelCodeEvaluator(CoSTEEREvaluator):
+    def evaluate(
+        self,
+        target_task: Task,
+        implementation: Workspace,
+        gt_implementation: Workspace,
+        model_execution_feedback: str = "",
+        model_value_feedback: str = "",
+    ):
+        assert isinstance(target_task, ModelTask)
+        assert isinstance(implementation, ModelFBWorkspace)
+        if gt_implementation is not None:
+            assert isinstance(gt_implementation, ModelFBWorkspace)
+
+        model_task_information = target_task.get_task_information()
+        code = implementation.all_codes
+
+        system_prompt = T(".prompts:evaluator_code_feedback.system").r(
+            scenario=(
+                self.scen.get_scenario_all_desc(target_task, filtered_tag=target_task.model_type)
+                if self.scen is not None
+                else "No scenario description."
+            )
+        )
+        execution_feedback_to_render = model_execution_feedback
+        for _ in range(10):  # 10 times to split the content is enough
+            user_prompt = T(".prompts:evaluator_code_feedback.user").r(
+                model_information=model_task_information,
+                code=code,
+                model_execution_feedback=execution_feedback_to_render,
+                model_value_feedback=model_value_feedback,
+                gt_code=gt_implementation.all_codes if gt_implementation else None,
+            )
+            if (
+                APIBackend().build_messages_and_calculate_token(
+                    user_prompt=user_prompt,
+                    system_prompt=system_prompt,
+                )
+                > APIBackend().chat_token_limit
+            ):
+                execution_feedback_to_render = execution_feedback_to_render[len(execution_feedback_to_render) // 2 :]
+            else:
+                break
+
+        critic_response = APIBackend().build_messages_and_create_chat_completion(
+            user_prompt=user_prompt,
+            system_prompt=system_prompt,
+            json_mode=False,
+        )
+
+        return critic_response, None
+
+
+class ModelFinalEvaluator(CoSTEEREvaluator):
+    def evaluate(
+        self,
+        target_task: Task,
+        implementation: Workspace,
+        gt_implementation: Workspace,
+        model_execution_feedback: str,
+        model_shape_feedback: str,
+        model_value_feedback: str,
+        model_code_feedback: str,
+    ):
+        assert isinstance(target_task, ModelTask)
+        assert isinstance(implementation, ModelFBWorkspace)
+        if gt_implementation is not None:
+            assert isinstance(gt_implementation, ModelFBWorkspace)
+
+        system_prompt = T(".prompts:evaluator_final_feedback.system").r(
+            scenario=(
+                self.scen.get_scenario_all_desc(target_task, filtered_tag=target_task.model_type)
+                if self.scen is not None
+                else "No scenario description."
+            )
+        )
+
+        execution_feedback_to_render = model_execution_feedback
+
+        for _ in range(10):  # 10 times to split the content is enough
+            user_prompt = T(".prompts:evaluator_final_feedback.user").r(
+                model_information=target_task.get_task_information(),
+                model_execution_feedback=execution_feedback_to_render,
+                model_shape_feedback=model_shape_feedback,
+                model_code_feedback=model_code_feedback,
+                model_value_feedback=model_value_feedback,
+            )
+
+            if (
+                APIBackend().build_messages_and_calculate_token(
+                    user_prompt=user_prompt,
+                    system_prompt=system_prompt,
+                )
+                > APIBackend().chat_token_limit
+            ):
+                execution_feedback_to_render = execution_feedback_to_render[len(execution_feedback_to_render) // 2 :]
+            else:
+                break
+
+        final_evaluation_dict = json.loads(
+            APIBackend().build_messages_and_create_chat_completion(
+                user_prompt=user_prompt,
+                system_prompt=system_prompt,
+                json_mode=True,
+                json_target_type=Dict[str, str | bool | int],
+            ),
+        )
+        if isinstance(final_evaluation_dict["final_decision"], str) and final_evaluation_dict[
+            "final_decision"
+        ].lower() in ("true", "false"):
+            final_evaluation_dict["final_decision"] = bool(final_evaluation_dict["final_decision"])
+        return (
+            final_evaluation_dict["final_feedback"],
+            final_evaluation_dict["final_decision"],
+        )

rdagent/components/coder/model_coder/evaluators.py

@@ -0,0 +1,104 @@
+from rdagent.components.coder.CoSTEER.evaluators import (
+    CoSTEEREvaluator,
+    CoSTEERMultiFeedback,
+    CoSTEERSingleFeedbackDeprecated,
+)
+from rdagent.components.coder.model_coder.eva_utils import (
+    ModelCodeEvaluator,
+    ModelFinalEvaluator,
+    shape_evaluator,
+    value_evaluator,
+)
+from rdagent.components.coder.model_coder.model import ModelFBWorkspace, ModelTask
+from rdagent.core.evolving_framework import QueriedKnowledge
+from rdagent.core.experiment import Task, Workspace
+
+ModelSingleFeedback = CoSTEERSingleFeedbackDeprecated
+ModelMultiFeedback = CoSTEERMultiFeedback
+
+
+class ModelCoSTEEREvaluator(CoSTEEREvaluator):
+    def evaluate(
+        self,
+        target_task: Task,
+        implementation: Workspace,
+        gt_implementation: Workspace,
+        queried_knowledge: QueriedKnowledge = None,
+        **kwargs,
+    ) -> ModelSingleFeedback:
+        target_task_information = target_task.get_task_information()
+        if (
+            queried_knowledge is not None
+            and target_task_information in queried_knowledge.success_task_to_knowledge_dict
+        ):
+            return queried_knowledge.success_task_to_knowledge_dict[target_task_information].feedback
+        elif queried_knowledge is not None and target_task_information in queried_knowledge.failed_task_info_set:
+            return ModelSingleFeedback(
+                execution_feedback="This task has failed too many times, skip implementation.",
+                shape_feedback="This task has failed too many times, skip implementation.",
+                value_feedback="This task has failed too many times, skip implementation.",
+                code_feedback="This task has failed too many times, skip implementation.",
+                final_feedback="This task has failed too many times, skip implementation.",
+                final_decision=False,
+            )
+        assert isinstance(target_task, ModelTask)
+
+        # NOTE: Use fixed input to test the model to avoid randomness
+        batch_size = 8
+        num_features = 30
+        num_timesteps = 40
+        input_value = 0.4
+        param_init_value = 0.6
+
+        assert isinstance(implementation, ModelFBWorkspace)
+        model_execution_feedback, gen_np_array = implementation.execute(
+            batch_size=batch_size,
+            num_features=num_features,
+            num_timesteps=num_timesteps,
+            input_value=input_value,
+            param_init_value=param_init_value,
+        )
+        if gt_implementation is not None:
+            assert isinstance(gt_implementation, ModelFBWorkspace)
+            _, gt_np_array = gt_implementation.execute(
+                batch_size=batch_size,
+                num_features=num_features,
+                num_timesteps=num_timesteps,
+                input_value=input_value,
+                param_init_value=param_init_value,
+            )
+        else:
+            gt_np_array = None
+
+        shape_feedback, shape_decision = shape_evaluator(
+            gen_np_array,
+            (batch_size, self.scen.model_output_channel if hasattr(self.scen, "model_output_channel") else 1),
+        )
+        value_feedback, value_decision = value_evaluator(gen_np_array, gt_np_array)
+        code_feedback, _ = ModelCodeEvaluator(scen=self.scen).evaluate(
+            target_task=target_task,
+            implementation=implementation,
+            gt_implementation=gt_implementation,
+            model_execution_feedback=model_execution_feedback,
+            model_value_feedback="\n".join([shape_feedback, value_feedback]),
+        )
+        final_feedback, final_decision = ModelFinalEvaluator(scen=self.scen).evaluate(
+            target_task=target_task,
+            implementation=implementation,
+            gt_implementation=gt_implementation,
+            model_execution_feedback=model_execution_feedback,
+            model_shape_feedback=shape_feedback,
+            model_value_feedback=value_feedback,
+            model_code_feedback=code_feedback,
+        )
+
+        return ModelSingleFeedback(
+            execution_feedback=model_execution_feedback,
+            shape_feedback=shape_feedback,
+            value_feedback=value_feedback,
+            code_feedback=code_feedback,
+            final_feedback=final_feedback,
+            final_decision=final_decision,
+            value_generated_flag=(gen_np_array is not None),
+            final_decision_based_on_gt=(gt_implementation is not None),
+        )

rdagent/components/coder/model_coder/evolving_strategy.py

@@ -0,0 +1,89 @@
+import json
+from typing import Dict
+
+from rdagent.components.coder.CoSTEER.config import CoSTEER_SETTINGS
+from rdagent.components.coder.CoSTEER.evaluators import CoSTEERSingleFeedback
+from rdagent.components.coder.CoSTEER.evolving_strategy import (
+    MultiProcessEvolvingStrategy,
+)
+from rdagent.components.coder.CoSTEER.knowledge_management import (
+    CoSTEERQueriedKnowledge,
+    CoSTEERQueriedKnowledgeV2,
+)
+from rdagent.components.coder.model_coder.model import ModelFBWorkspace, ModelTask
+from rdagent.core.experiment import FBWorkspace
+from rdagent.oai.llm_conf import LLM_SETTINGS
+from rdagent.oai.llm_utils import APIBackend
+from rdagent.utils.agent.tpl import T
+
+
+class ModelMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
+    def implement_one_task(
+        self,
+        target_task: ModelTask,
+        queried_knowledge: CoSTEERQueriedKnowledge = None,
+        workspace: FBWorkspace | None = None,
+        prev_task_feedback: CoSTEERSingleFeedback | None = None,
+    ) -> str:
+        model_information_str = target_task.get_task_information()
+
+        queried_similar_successful_knowledge = (
+            queried_knowledge.task_to_similar_task_successful_knowledge[model_information_str]
+            if queried_knowledge is not None
+            else []
+        )
+        queried_former_failed_knowledge = (
+            queried_knowledge.task_to_former_failed_traces[model_information_str]
+            if queried_knowledge is not None
+            else []
+        )
+
+        queried_former_failed_knowledge_to_render = (
+            queried_former_failed_knowledge[0]
+            if isinstance(queried_knowledge, CoSTEERQueriedKnowledgeV2)
+            else queried_former_failed_knowledge
+        )
+        system_prompt = T(".prompts:evolving_strategy_model_coder.system").r(
+            scenario=self.scen.get_scenario_all_desc(filtered_tag="model"),
+            queried_former_failed_knowledge=queried_former_failed_knowledge_to_render,
+            current_code=workspace.file_dict.get("model.py"),
+        )
+
+        queried_similar_successful_knowledge_to_render = queried_similar_successful_knowledge
+        for _ in range(10):  # max attempt to reduce the length of user_prompt
+            user_prompt = T(".prompts:evolving_strategy_model_coder.user").r(
+                model_information_str=model_information_str,
+                queried_similar_successful_knowledge=queried_similar_successful_knowledge_to_render,
+                queried_former_failed_knowledge=queried_former_failed_knowledge_to_render,
+            )
+            if (
+                APIBackend().build_messages_and_calculate_token(
+                    user_prompt=user_prompt,
+                    system_prompt=system_prompt,
+                )
+                < APIBackend().chat_token_limit
+            ):
+                break
+            elif len(queried_former_failed_knowledge_to_render) > 1:
+                queried_former_failed_knowledge_to_render = queried_former_failed_knowledge_to_render[1:]
+            elif len(queried_similar_successful_knowledge_to_render) > 1:
+                queried_similar_successful_knowledge_to_render = queried_similar_successful_knowledge_to_render[1:]
+
+        code = json.loads(
+            APIBackend(use_chat_cache=CoSTEER_SETTINGS.coder_use_cache).build_messages_and_create_chat_completion(
+                user_prompt=user_prompt,
+                system_prompt=system_prompt,
+                json_mode=True,
+                json_target_type=Dict[str, str],
+            ),
+        )["code"]
+        return code
+
+    def assign_code_list_to_evo(self, code_list, evo):
+        for index in range(len(evo.sub_tasks)):
+            if code_list[index] is None:
+                continue
+            if evo.sub_workspace_list[index] is None:
+                evo.sub_workspace_list[index] = ModelFBWorkspace(target_task=evo.sub_tasks[index])
+            evo.sub_workspace_list[index].inject_files(**{"model.py": code_list[index]})
+        return evo

rdagent/components/coder/model_coder/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/components/coder/model_coder/model.py

@@ -0,0 +1,163 @@
+import pickle
+import site
+import traceback
+from pathlib import Path
+from typing import Dict, Optional
+
+from rdagent.components.coder.CoSTEER.task import CoSTEERTask
+from rdagent.components.coder.model_coder.conf import MODEL_COSTEER_SETTINGS
+from rdagent.core.experiment import Experiment, FBWorkspace
+from rdagent.core.utils import cache_with_pickle
+from rdagent.oai.llm_utils import md5_hash
+from rdagent.utils.env import KGDockerEnv, QlibCondaConf, QlibCondaEnv, QTDockerEnv
+
+
+class ModelTask(CoSTEERTask):
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        architecture: str,
+        *args,
+        hyperparameters: Dict[str, str],
+        training_hyperparameters: Dict[str, str],
+        formulation: str = None,
+        variables: Dict[str, str] = None,
+        model_type: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        self.formulation: str = formulation
+        self.architecture: str = architecture
+        self.variables: str = variables
+        self.hyperparameters: str = hyperparameters
+        self.training_hyperparameters: str = training_hyperparameters
+        self.model_type: str = (
+            model_type  # Tabular for tabular model, TimesSeries for time series model, Graph for graph model, XGBoost for XGBoost model
+        )
+        super().__init__(name=name, description=description, *args, **kwargs)
+
+    def get_task_information(self):
+        task_desc = f"""name: {self.name}
+description: {self.description}
+"""
+        task_desc += f"formulation: {self.formulation}\n" if self.formulation else ""
+        task_desc += f"architecture: {self.architecture}\n"
+        task_desc += f"variables: {self.variables}\n" if self.variables else ""
+        task_desc += f"hyperparameters: {self.hyperparameters}\n"
+        task_desc += f"training_hyperparameters: {self.training_hyperparameters}\n"
+        task_desc += f"model_type: {self.model_type}\n"
+        return task_desc
+
+    def get_task_brief_information(self):
+        task_desc = f"""name: {self.name}
+description: {self.description}
+"""
+        task_desc += f"architecture: {self.architecture}\n"
+        task_desc += f"hyperparameters: {self.hyperparameters}\n"
+        task_desc += f"training_hyperparameters: {self.training_hyperparameters}\n"
+        task_desc += f"model_type: {self.model_type}\n"
+        return task_desc
+
+    @staticmethod
+    def from_dict(dict):
+        return ModelTask(**dict)
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__} {self.name}>"
+
+
+class ModelFBWorkspace(FBWorkspace):
+    """
+    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 support two ways of interface:
+        (version 1) for qlib we'll make a script to 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 model.
+
+        (version 2) for kaggle we'll make a script to call the fit and predict function in the implementation in file `model.py` after setting the cwd into the directory
+    """
+
+    def hash_func(
+        self,
+        batch_size: int = 8,
+        num_features: int = 10,
+        num_timesteps: int = 4,
+        num_edges: int = 20,
+        input_value: float = 1.0,
+        param_init_value: float = 1.0,
+    ) -> str:
+        target_file_name = f"{batch_size}_{num_features}_{num_timesteps}_{input_value}_{param_init_value}"
+        for code_file_name in sorted(list(self.file_dict.keys())):
+            target_file_name = f"{target_file_name}_{self.file_dict[code_file_name]}"
+        return md5_hash(target_file_name)
+
+    @cache_with_pickle(hash_func)
+    def execute(
+        self,
+        batch_size: int = 8,
+        num_features: int = 10,
+        num_timesteps: int = 4,
+        num_edges: int = 20,
+        input_value: float = 1.0,
+        param_init_value: float = 1.0,
+    ):
+        self.before_execute()
+        try:
+            if self.target_task.version == 1:
+                if MODEL_COSTEER_SETTINGS.env_type == "docker":
+                    qtde = QTDockerEnv()
+                elif MODEL_COSTEER_SETTINGS.env_type == "conda":
+                    qtde = QlibCondaEnv(conf=QlibCondaConf())
+                else:
+                    raise ValueError(f"Unknown env_type: {MODEL_COSTEER_SETTINGS.env_type}")
+            else:
+                qtde = KGDockerEnv()
+            qtde.prepare()
+
+            if self.target_task.version == 1:
+                dump_code = f"""
+MODEL_TYPE = "{self.target_task.model_type}"
+BATCH_SIZE = {batch_size}
+NUM_FEATURES = {num_features}
+NUM_TIMESTEPS = {num_timesteps}
+NUM_EDGES = {num_edges}
+INPUT_VALUE = {input_value}
+PARAM_INIT_VALUE = {param_init_value}
+{(Path(__file__).parent / 'model_execute_template_v1.txt').read_text()}
+"""
+            elif self.target_task.version == 2:
+                dump_code = (Path(__file__).parent / "model_execute_template_v2.txt").read_text()
+
+            log, results = qtde.dump_python_code_run_and_get_results(
+                code=dump_code,
+                dump_file_names=["execution_feedback_str.pkl", "execution_model_output.pkl"],
+                local_path=str(self.workspace_path),
+                env={},
+                code_dump_file_py_name="model_test",
+            )
+            if len(results) == 0:
+                raise RuntimeError(f"Error in running the model code: {log}")
+            [execution_feedback_str, execution_model_output] = results
+
+        except Exception as e:
+            execution_feedback_str = f"Execution error: {e}\nTraceback: {traceback.format_exc()}"
+            execution_model_output = None
+
+        if len(execution_feedback_str) > 2000:
+            execution_feedback_str = (
+                execution_feedback_str[:1000] + "....hidden long error message...." + execution_feedback_str[-1000:]
+            )
+        return execution_feedback_str, execution_model_output
+
+
+ModelExperiment = Experiment

rdagent/components/coder/model_coder/model_execute_template_v1.txt

@@ -0,0 +1,44 @@
+# MODEL_TYPE = "Tabular"
+# BATCH_SIZE = 32
+# NUM_FEATURES = 10
+# NUM_TIMESTEPS = 4
+# NUM_EDGES = 20
+# INPUT_VALUE = 1.0
+# PARAM_INIT_VALUE = 1.0
+
+import pickle
+
+import torch
+from model import model_cls
+
+if 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 MODEL_TYPE == "TimeSeries":
+    input_shape = (BATCH_SIZE, NUM_TIMESTEPS, NUM_FEATURES)
+    m = model_cls(num_features=input_shape[2], num_timesteps=input_shape[1])
+    data = torch.full(input_shape, INPUT_VALUE)
+elif MODEL_TYPE == "Graph":
+    node_feature = torch.randn(BATCH_SIZE, NUM_FEATURES)
+    edge_index = torch.randint(0, BATCH_SIZE, (2, NUM_EDGES))
+    m = model_cls(num_features=NUM_FEATURES)
+    data = (node_feature, edge_index)
+else:
+    raise ValueError(f"Unsupported model type: {MODEL_TYPE}")
+
+# Initialize all parameters of `m` to `param_init_value`
+for _, param in m.named_parameters():
+    param.data.fill_(PARAM_INIT_VALUE)
+
+# Execute the model
+if MODEL_TYPE == "Graph":
+    out = m(*data)
+else:
+    out = m(data)
+
+execution_model_output = out.cpu().detach().numpy()
+execution_feedback_str = f"Execution successful, output tensor shape: {execution_model_output.shape}"
+
+pickle.dump(execution_model_output, open("execution_model_output.pkl", "wb"))
+pickle.dump(execution_feedback_str, open("execution_feedback_str.pkl", "wb"))

rdagent/components/coder/model_coder/model_execute_template_v2.txt

@@ -0,0 +1,24 @@
+import os
+import pickle
+
+import numpy as np
+import pandas as pd
+import torch
+from model import fit, predict
+
+train_X = pd.DataFrame(np.random.randn(8, 30), columns=[f"{i}" for i in range(30)])
+train_y = pd.Series(np.random.randint(0, 2, 8))
+valid_X = pd.DataFrame(np.random.randn(8, 30), columns=[f"{i}" for i in range(30)])
+valid_y = pd.Series(np.random.randint(0, 2, 8))
+
+model = fit(train_X, train_y, valid_X, valid_y)
+execution_model_output = predict(model, valid_X)
+
+if isinstance(execution_model_output, torch.Tensor):
+    execution_model_output = execution_model_output.cpu().detach().numpy()
+
+
+execution_feedback_str = f"Execution successful, output numpy ndarray shape: {execution_model_output.shape}"
+
+pickle.dump(execution_model_output, open("execution_model_output.pkl", "wb"))
+pickle.dump(execution_feedback_str, open("execution_feedback_str.pkl", "wb"))

rdagent/components/coder/model_coder/one_shot/__init__.py

@@ -0,0 +1,35 @@
+import re
+from pathlib import Path
+
+from rdagent.components.coder.model_coder.model import ModelExperiment, ModelFBWorkspace
+from rdagent.core.developer import Developer
+from rdagent.oai.llm_utils import APIBackend
+from rdagent.utils.agent.tpl import T
+
+DIRNAME = Path(__file__).absolute().resolve().parent
+
+
+class ModelCodeWriter(Developer[ModelExperiment]):
+    def develop(self, exp: ModelExperiment) -> ModelExperiment:
+        mti_l = []
+        for t in exp.sub_tasks:
+            mti = ModelFBWorkspace(t)
+            mti.prepare()
+
+            user_prompt = T(".prompts:code_implement_user").r(
+                name=t.name,
+                description=t.description,
+                formulation=t.formulation,
+                variables=t.variables,
+            )
+            system_prompt = T(".prompts:code_implement_sys").r()
+
+            resp = APIBackend().build_messages_and_create_chat_completion(user_prompt, system_prompt)
+
+            # Extract the code part from the response
+            match = re.search(r".*```[Pp]ython\n(.*)\n```.*", resp, re.DOTALL)
+            code = match.group(1)
+            mti.inject_files(**{"model.py": code})
+            mti_l.append(mti)
+        exp.sub_workspace_list = mti_l
+        return exp

rdagent/components/coder/model_coder/one_shot/prompt.yaml

@@ -0,0 +1,27 @@
+
+
+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 executed in the follow way:
+  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 implemented 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.

rdagent/components/coder/model_coder/prompts.yaml

@@ -0,0 +1,156 @@
+extract_model_formulation_system: |-
+    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 (The name of the model)": {
+            "description": "A detailed description of the model",
+            "formulation": "A LaTeX formula representing the model's formulation",
+            "architecture": "A detailed description of the model's architecture, e.g., neural network layers or tree structures",
+            "variables": {
+                "\\hat{y}_u": "The predicted output for node u",
+                "variable_name_2": "Description of variable 2",
+                "variable_name_3": "Description of variable 3"
+            },
+            "hyperparameters": {
+                "hyperparameter_name_1": "value of hyperparameter 1",
+                "hyperparameter_name_2": "value of hyperparameter 2",
+                "hyperparameter_name_3": "value of hyperparameter 3"
+            },
+            "training_hyperparameters" {  # All values are for reference; you can set them yourself
+                "n_epochs": "100",
+                "lr": "1e-3",
+                "early_stop": 10,
+                "batch_size": 256,
+                "weight_decay": 1e-4,
+            }
+            "model_type": "Tabular or TimeSeries or Graph or XGBoost"  # Should be one of "Tabular", "TimeSeries", "Graph", or "XGBoost"
+        }
+    }
+    such format content should be begin with ```json and end with ``` and the content should be in json format.
+
+evolving_strategy_model_coder:
+    system: |-
+        User is trying to implement some pytorch models in the following scenario:
+        {{ scenario }}
+        Your code is expected to align the scenario in any form which means The user needs to get the prediction of the model based on the input data.
+
+        To help you write the correct code, the user might provide multiple information that helps you write the correct code:
+        1. The user might provide you the correct code to similar models. Your should learn from these code to write the correct code.
+        2. The user might provide you the failed former code and the corresponding feedback to the code. The feedback contains to the execution, the code and the model output value. You should analyze the feedback and try to correct the latest code.
+        3. The user might provide you the suggestion to the latest fail code and some similar fail to correct pairs. Each pair contains the fail code with similar error and the corresponding corrected version code. You should learn from these suggestion to write the correct code.
+
+        Your must write your code based on your former latest attempt below which consists of your former code and code feedback, you should read the former attempt carefully and must not modify the right part of your former code.
+
+        {% if current_code is not none %}
+        User has write some code before. You should write the new code based on this code. Here is the latest code:
+        ```python
+        {{ current_code }}
+        ```
+        Your code should be very similar to the former code which means your code should be ninety more percent same as the former code! You should not modify the right part of the code.
+        {% else %}
+        User has not write any code before. You should write the new code from scratch.
+        {% endif %}
+
+        {% if queried_former_failed_knowledge|length != 0 %}
+        --------------Your former latest attempt:---------------
+        =====Code to the former implementation=====
+        {{ queried_former_failed_knowledge[-1].implementation.all_codes }}
+        =====Feedback to the former implementation=====
+        {{ queried_former_failed_knowledge[-1].feedback }}
+        {% endif %}
+        
+        Please response the code in the following json format. Here is an example structure for the JSON output:
+        {
+            "code": "The Python code as a string."
+        }
+
+    user: |-
+        --------------Target model information:---------------
+        {{ model_information_str }}
+
+        {% if queried_similar_successful_knowledge|length != 0 %}
+        --------------Correct code to similar models:---------------
+        {% for similar_successful_knowledge in queried_similar_successful_knowledge %}
+        =====Model {{loop.index}}:=====
+        {{ similar_successful_knowledge.target_task.get_task_information() }}
+        =====Code:=====
+        {{ similar_successful_knowledge.implementation.all_codes }}
+        {% endfor %}
+        {% endif %}
+
+        {% if queried_former_failed_knowledge|length != 0 %}
+        --------------Former failed code:---------------
+        {% for former_failed_knowledge in queried_former_failed_knowledge %}
+        =====Code to implementation {{ loop.index }}=====
+        {{ former_failed_knowledge.implementation.all_codes }}
+        =====Feedback to implementation {{ loop.index }}=====
+        {{ former_failed_knowledge.feedback }}
+        {% endfor %}
+        {% endif %}
+
+evaluator_code_feedback:
+    system: |-
+        User is trying to implement some models in the following scenario:
+        {{ scenario }}
+        User will provide you the information of the model.
+
+        Your job is to check whether user's code is align with the model information and the scenario.
+        The user will provide the source python code and the execution error message if execution failed.
+        The user might provide you the ground truth code for you to provide the critic. You should not leak the ground truth code to the user in any form but you can use it to provide the critic.
+
+        User has also compared the output generated by the user's code and the ground truth code. The user will provide you some analysis results comparing two output. You may find some error in the code which caused the difference between the two output.
+
+        If the ground truth code is provided, your critic should only consider checking whether the user's code is align with the ground truth code since the ground truth is definitely correct.
+        If the ground truth code is not provided, your critic should consider checking whether the user's code is reasonable and correct to the description and to the scenario.
+
+        Notice that your critics are not for user to debug the code. They are sent to the coding agent to correct the code. So don't give any following items for the user to check like "Please check the code line XXX".
+
+        You suggestion should not include any code, just some clear and short suggestions. Please point out very critical issues in your response, ignore non-important issues to avoid confusion. If no big issue found in the code, you can response "No critics found".
+
+        You should provide the suggestion to each of your critic to help the user improve the code. Please response the critic in the following format. Here is an example structure for the output:
+        critic 1: The critic message to critic 1
+        critic 2: The critic message to critic 2
+    
+    user: |-
+        --------------Model information:---------------
+        {{ model_information }}
+        --------------Python code:---------------
+        {{ code }}
+        --------------Execution feedback:---------------
+        {{ model_execution_feedback }}
+        {% if model_value_feedback is not none %}
+        --------------Model value feedback:---------------
+        {{ model_value_feedback }}
+        {% endif %}
+        {% if gt_code is not none %}
+        --------------Ground truth Python code:---------------
+        {{ gt_code }}
+        {% endif %}
+
+
+evaluator_final_feedback:
+    system: |-
+        User is trying to implement a model in the following scenario:
+        {{ scenario }}
+        User has finished evaluation and got some feedback from the evaluator.
+        The evaluator run the code and get the output and provide several feedback regarding user's code and code output. You should analyze the feedback and considering the scenario and model description to give a final decision about the evaluation result. The final decision concludes whether the model is implemented correctly and if not, detail feedback containing reason and suggestion if the final decision is False.
+
+        The implementation final decision is considered in the following logic:
+        1. If the value and the ground truth value are exactly the same under a small tolerance, the implementation is considered correct.
+        2. If no ground truth value is not provided, the implementation is considered correct if the code execution is successful and the code feedback is align with the scenario and model description.
+
+        Please response the critic in the json format. Here is an example structure for the JSON output, please strictly follow the format:
+        {
+            "final_decision": True,
+            "final_feedback": "The final feedback message",
+        }
+    user: |-
+        --------------Model information:---------------
+        {{ model_information }}
+        --------------Model Execution feedback:---------------
+        {{ model_execution_feedback }}
+        --------------Model shape feedback:---------------
+        {{ model_shape_feedback }}
+        --------------Model Code feedback:---------------
+        {{ model_code_feedback }}
+        --------------Model value feedback:---------------
+        {{ model_value_feedback }}

rdagent/components/coder/model_coder/task_loader.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import json
+import re
+
+from pydantic import BaseModel, Field
+
+from rdagent.components.coder.model_coder.model import ModelTask
+from rdagent.components.document_reader.document_reader import (
+    load_and_process_pdfs_by_langchain,
+)
+from rdagent.components.loader.task_loader import ModelTaskLoader
+from rdagent.log import rdagent_logger as logger
+from rdagent.oai.llm_utils import APIBackend
+from rdagent.scenarios.qlib.experiment.model_experiment import QlibModelExperiment
+from rdagent.utils.agent.tpl import T
+from rdagent.utils.workflow import wait_retry
+
+
+def extract_model_from_doc(doc_content: str) -> dict:
+    """
+    Extract model information from document content.
+
+    Parameters
+    ----------
+    doc_content : str
+        Document content.
+
+    Returns
+    -------
+    dict
+        {model_name: dict{description, formulation, variables}}
+    """
+    session = APIBackend().build_chat_session(
+        session_system_prompt=T(".prompts:extract_model_formulation_system").r(),
+    )
+    current_user_prompt = doc_content
+
+    # Extract model information from document content.
+    model_dict = {}
+
+    for _ in range(10):
+        # try to extract model information from the document content, retry at most 10 times.
+        extract_result_resp = session.build_chat_completion(
+            user_prompt=current_user_prompt,
+            json_mode=False,
+        )
+        re_search_res = re.search(r"```json(.*)```", extract_result_resp, re.S)
+        ret_json_str = re_search_res.group(1) if re_search_res is not None else ""
+        try:
+            ret_dict = json.loads(ret_json_str)
+            parse_success = bool(isinstance(ret_dict, dict))
+        except json.JSONDecodeError:
+            parse_success = False
+        if ret_json_str is None or not parse_success:
+            current_user_prompt = "Your response didn't follow the instruction might be wrong json format. Try again."
+        else:
+            for name, formulation_and_description in ret_dict.items():
+                if name not in model_dict:
+                    model_dict[name] = formulation_and_description
+            if len(model_dict) == 0:
+                current_user_prompt = "No model extracted. Please try again."
+            else:
+                break
+
+    logger.info(f"已经完成{len(model_dict)}个模型的提取")
+
+    return model_dict
+
+
+def merge_file_to_model_dict_to_model_dict(
+    file_to_model_dict: dict[str, dict],
+) -> dict:
+    model_dict = {}
+    for file_name in file_to_model_dict:
+        for model_name in file_to_model_dict[file_name]:
+            model_dict.setdefault(model_name, [])
+            model_dict[model_name].append(file_to_model_dict[file_name][model_name])
+
+    model_dict_simple_deduplication = {}
+    for model_name in model_dict:
+        if len(model_dict[model_name]) > 1:
+            model_dict_simple_deduplication[model_name] = max(
+                model_dict[model_name],
+                key=lambda x: len(x["formulation"]),
+            )
+        else:
+            model_dict_simple_deduplication[model_name] = model_dict[model_name][0]
+    return model_dict_simple_deduplication
+
+
+def extract_model_from_docs(docs_dict):
+    model_dict = {}
+    for doc_name, doc_content in docs_dict.items():
+        model_dict[doc_name] = extract_model_from_doc(doc_content)
+    return model_dict
+
+
+class ModelExperimentLoaderFromDict(ModelTaskLoader):
+    def load(self, model_dict: dict) -> QlibModelExperiment:
+        """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"],
+                architecture=model_data["architecture"],
+                variables=model_data["variables"],
+                hyperparameters=model_data["hyperparameters"],
+                training_hyperparameters=model_data["training_hyperparameters"],
+                model_type=model_data["model_type"],
+            )
+            task_l.append(task)
+        return QlibModelExperiment(sub_tasks=task_l)
+
+
+class ModelExperimentLoaderFromPDFfiles(ModelTaskLoader):
+    @wait_retry(retry_n=5)
+    def load(self, file_or_folder_path: str) -> QlibModelExperiment:
+        docs_dict = load_and_process_pdfs_by_langchain(file_or_folder_path)  # dict{file_path:content}
+        model_dict = extract_model_from_docs(
+            docs_dict
+        )  # dict{file_name: dict{model_name: dict{description, formulation, variables}}}
+        model_dict = merge_file_to_model_dict_to_model_dict(
+            model_dict
+        )  # dict {model_name: dict{description, formulation, variables}}
+        return ModelExperimentLoaderFromDict().load(model_dict)