refine class design and inheritance first version code (#41)

* refine class design and inheritance first version code

* fix all typos

---------

Co-authored-by: xuyang1 <xuyang1@microsoft.com>

rdagent/components/task_implementation/model_implementation/one_shot/__init__.py

@@ -2,13 +2,13 @@ import re
 from pathlib import Path
 from typing import Sequence
 
-from jinja2 import Template
+from jinja2 import Environment, StrictUndefined
 
 from rdagent.components.task_implementation.model_implementation.task import (
     ModelImplTask,
     ModelTaskImpl,
 )
-from rdagent.core.implementation import TaskGenerator
+from rdagent.core.task_generator import TaskGenerator
 from rdagent.core.prompts import Prompts
 from rdagent.oai.llm_utils import APIBackend
 
@@ -23,8 +23,8 @@ class ModelTaskGen(TaskGenerator):
             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_tpl = Environment(undefined=StrictUndefined).from_string(pr["code_implement_user"])
+            sys_prompt_tpl = Environment(undefined=StrictUndefined).from_string(pr["code_implement_sys"])
 
             user_prompt = user_prompt_tpl.render(
                 name=t.name,

rdagent/components/task_implementation/model_implementation/task.py

@@ -1,26 +1,21 @@
+import json
 import uuid
 from pathlib import Path
 from typing import Dict, Optional, Sequence
 
 import torch
 
+from rdagent.components.loader.task_loader import ModelTaskLoader
 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.core.experiment import FBImplementation, ImpLoader, Task
 from rdagent.utils import get_module_by_module_path
-import json
 
 
-class ModelImplTask(BaseTask):
-    # TODO: it should change when the BaseTask changes.
+class ModelImplTask(Task):
+    # TODO: it should change when the Task changes.
     name: str
     description: str
     formulation: str
@@ -43,6 +38,7 @@ class ModelImplTask(BaseTask):
         self.formulation = formulation
         self.variables = variables
         self.key = key
+
     def get_information(self):
         return f"""name: {self.name}
 description: {self.description}
@@ -50,16 +46,16 @@ 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):
+class ModelTaskLoaderJson(ModelTaskLoader):
     # def __init__(self, json_uri: str, select_model: Optional[str] = None) -> None:
     #     super().__init__()
     #     self.json_uri = json_uri
@@ -82,7 +78,7 @@ class ModelTaskLoderJson(TaskLoader):
     #         variables=model_data["variables"],
     #         key=model_name
     #     )
-        
+
     #     return [model_impl_task]
 
     def __init__(self, json_uri: str) -> None:
@@ -95,7 +91,7 @@ class ModelTaskLoderJson(TaskLoader):
 
         # 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.",
@@ -119,12 +115,13 @@ class ModelTaskLoderJson(TaskLoader):
                 description=model_data["description"],
                 formulation=model_data["formulation"],
                 variables=model_data["variables"],
-                key=model_data["key"]
+                key=model_data["key"],
             )
             model_impl_task_list.append(model_impl_task)
         return model_impl_task_list
-    
-class ModelImplementationTaskLoaderFromDict(TaskLoader):
+
+
+class ModelImplementationTaskLoaderFromDict(ModelTaskLoader):
     def load(self, model_dict: dict) -> list:
         """Load data from a dict."""
         task_l = []
@@ -134,13 +131,13 @@ class ModelImplementationTaskLoaderFromDict(TaskLoader):
                 description=model_data["description"],
                 formulation=model_data["formulation"],
                 variables=model_data["variables"],
-                key=model_name
+                key=model_name,
             )
             task_l.append(task)
         return task_l
 
 
-class ModelTaskImpl(FBTaskImplementation):
+class ModelTaskImpl(FBImplementation):
     """
     It is a Pytorch model implementation task;
     All the things are placed in a folder.
@@ -156,11 +153,11 @@ class ModelTaskImpl(FBTaskImplementation):
     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.
+    - And then verify the model.
 
     """
 
-    def __init__(self, target_task: BaseTask) -> None:
+    def __init__(self, target_task: Task) -> None:
         super().__init__(target_task)
         self.path = None
 
@@ -202,7 +199,7 @@ 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
+  - So your implemented code could follow the following pattern
     ```Python
     class XXXLayer(torch.nn.Module):
         ...

rdagent/components/task_implementation/model_implementation/task_extraction.py

@@ -1,26 +1,23 @@
 from __future__ import annotations
 
 import json
-import multiprocessing as mp
 import re
 from pathlib import Path
-from typing import Mapping
 
-import numpy as np
-import pandas as pd
-import tiktoken
-from jinja2 import Template
-from rdagent.core.conf import RD_AGENT_SETTINGS
+from rdagent.components.document_reader.document_reader import (
+    load_and_process_pdfs_by_langchain,
+)
+from rdagent.components.loader.task_loader import ModelTaskLoader
+from rdagent.components.task_implementation.model_implementation.task import (
+    ModelImplementationTaskLoaderFromDict,
+)
 from rdagent.core.log import RDAgentLog
 from rdagent.core.prompts import Prompts
-from rdagent.core.task import TaskLoader
-from rdagent.components.document_reader.document_reader import load_and_process_pdfs_by_langchain
-from rdagent.components.task_implementation.model_implementation.task import ModelImplementationTaskLoaderFromDict
-from rdagent.oai.llm_utils import APIBackend, create_embedding_with_multiprocessing
-
+from rdagent.oai.llm_utils import APIBackend
 
 document_process_prompts = Prompts(file_path=Path(__file__).parent / "prompts.yaml")
 
+
 def extract_model_from_doc(doc_content: str) -> dict:
     """
     Extract model information from document content.
@@ -67,11 +64,11 @@ def extract_model_from_doc(doc_content: str) -> dict:
             else:
                 break
 
-
     RDAgentLog().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:
@@ -80,7 +77,7 @@ def merge_file_to_model_dict_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:
@@ -100,18 +97,24 @@ def extract_model_from_docs(docs_dict):
     return model_dict
 
 
-class ModelImplementationTaskLoaderFromPDFfiles(TaskLoader):
+class ModelImplementationTaskLoaderFromPDFfiles(ModelTaskLoader):
     def load(self, file_or_folder_path: Path) -> dict:
-        docs_dict = load_and_process_pdfs_by_langchain(Path(file_or_folder_path)) # dict{file_path:content}
-        model_dict = extract_model_from_docs(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}}
+        docs_dict = load_and_process_pdfs_by_langchain(Path(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 ModelImplementationTaskLoaderFromDict().load(model_dict)
 
+
 def main(path="../test_doc"):
     doc_dict = load_and_process_pdfs_by_langchain(Path(path))
-    print(doc_dict.keys()) # if you run code like "python -u", the print content will be truncated
+    print(doc_dict.keys())  # if you run code like "python -u", the print content will be truncated
 
 
 import fire
+
 if __name__ == "__main__":
     fire.Fire(main)