NexQuant/rdagent/components/task_implementation/model_implementation/task.py

import uuid
from pathlib import Path
from typing import Dict, Optional, Sequence

import torch

from rdagent.components.task_implementation.model_implementation.conf import (
    MODEL_IMPL_SETTINGS,
)
from rdagent.core.exception import CodeFormatException
from rdagent.core.task import (
    BaseTask,
    FBTaskImplementation,
    ImpLoader,
    TaskImplementation,
    TaskLoader,
)
from rdagent.utils import get_module_by_module_path
import json


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
    def get_information(self):
        return f"""name: {self.name}
description: {self.description}
formulation: {self.formulation}
variables: {self.variables}
key: {self.key}
"""
    
    @staticmethod
    def from_dict(dict):
        return ModelImplTask(**dict)
    
    def __repr__(self) -> str:
        return f"<{self.__class__.__name__} {self.name}>"


class ModelTaskLoderJson(TaskLoader):
    # def __init__(self, json_uri: str, select_model: Optional[str] = None) -> None:
    #     super().__init__()
    #     self.json_uri = json_uri
    #     self.select_model = 'A-DGN'

    # def load(self, *argT, **kwargs) -> Sequence[ModelImplTask]:
    #     # json is supposed to be in the format of {model_name: dict{model_data}}
    #     model_dict = json.load(open(self.json_uri, "r"))
    #     if self.select_model is not None:
    #         assert self.select_model in model_dict
    #         model_name = self.select_model
    #         model_data = model_dict[self.select_model]
    #     else:
    #         model_name, model_data = list(model_dict.items())[0]

    #     model_impl_task = ModelImplTask(
    #         name=model_name,
    #         description=model_data["description"],
    #         formulation=model_data["formulation"],
    #         variables=model_data["variables"],
    #         key=model_name
    #     )
        
    #     return [model_impl_task]

    def __init__(self, json_uri: str) -> None:
        super().__init__()
        self.json_uri = json_uri

    def load(self, *argT, **kwargs) -> Sequence[ModelImplTask]:
        # json is supposed to be in the format of {model_name: dict{model_data}}
        model_dict = json.load(open(self.json_uri, "r"))

        # FIXME: the model in the json file is not right due to extraction error
        #       We should fix them case by case in the future
        # 
        # formula_info = {
        #     "name": "Anti-Symmetric Deep Graph Network (A-DGN)",
        #     "description": "A framework for stable and non-dissipative DGN design. It ensures long-range information preservation between nodes and prevents gradient vanishing or explosion during training.",
        #     "formulation": r"\mathbf{x}^{\prime}_i = \mathbf{x}_i + \epsilon \cdot \sigma \left( (\mathbf{W}-\mathbf{W}^T-\gamma \mathbf{I}) \mathbf{x}_i + \Phi(\mathbf{X}, \mathcal{N}_i) + \mathbf{b}\right),",
        #     "variables": {
        #         r"\mathbf{x}_i": "The state of node i at previous layer",
        #         r"\epsilon": "The step size in the Euler discretization",
        #         r"\sigma": "A monotonically non-decreasing activation function",
        #         r"\Phi": "A graph convolutional operator",
        #         r"W": "An anti-symmetric weight matrix",
        #         r"\mathbf{x}^{\prime}_i": "The node feature matrix at layer l-1",
        #         r"\mathcal{N}_i": "The set of neighbors of node u",
        #         r"\mathbf{b}": "A bias vector",
        #     },
        #     "key": "A-DGN",
        # }
        model_impl_task_list = []
        for model_name, model_data in model_dict.items():
            model_impl_task = ModelImplTask(
                name=model_name,
                description=model_data["description"],
                formulation=model_data["formulation"],
                variables=model_data["variables"],
                key=model_data["key"]
            )
            model_impl_task_list.append(model_impl_task)
        return model_impl_task_list
    
class ModelImplementationTaskLoaderFromDict(TaskLoader):
    def load(self, model_dict: dict) -> list:
        """Load data from a dict."""
        task_l = []
        for model_name, model_data in model_dict.items():
            task = ModelImplTask(
                name=model_name,
                description=model_data["description"],
                formulation=model_data["formulation"],
                variables=model_data["variables"],
                key=model_name
            )
            task_l.append(task)
        return task_l


class ModelTaskImpl(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