调整mattergen-wrapper位置

mars_toolkit/compute/material_gen.py

@@ -28,14 +28,7 @@ from mars_toolkit.core.config import config
 import sys
 import os
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))
-from mattergen_wrapper import generator
+from ..core.mattergen_wrapper import *
-CrystalGenerator = generator.CrystalGenerator
-from mattergen.common.data.types import TargetProperty
-from mattergen.common.utils.eval_utils import MatterGenCheckpointInfo
-from mattergen.common.utils.data_classes import (
-    PRETRAINED_MODEL_NAME,
-    MatterGenCheckpointInfo,
-)
 
 logger = logging.getLogger(__name__)
 

mars_toolkit/core/config.py

@@ -27,6 +27,7 @@ class Config:
 
     # MatterGen
     MATTERGENMODEL_ROOT = '/home/ubuntu/50T/lzy/mars-mcp/pretrained_models/mattergen_ckpt'
+    MATTERGEN_ROOT='/home/ubuntu/50T/lzy/mars-mcp/mattergen'
     MATTERGENMODEL_RESULT_PATH = 'results/'
 
     # Dify

mars_toolkit/core/mattergen_wrapper.py

@@ -5,9 +5,9 @@ by modifying the Python path at runtime.
 import sys
 import os
 from pathlib import Path
-
+from .config import config
 # Add the mattergen directory to the Python path
-mattergen_dir = os.path.join(os.path.dirname(__file__), 'mattergen')
+mattergen_dir = config.MATTERGEN_ROOT
 sys.path.insert(0, mattergen_dir)
 
 # Import the necessary modules from the mattergen package
@@ -21,6 +21,6 @@ except ImportError as e:
     print(f"Error importing mattergen modules: {e}")
     print(f"Python path: {sys.path}")
     raise
-
+CrystalGenerator = generator.CrystalGenerator
 # Re-export the modules
-__all__ = ['generator', 'chemgraph', 'TargetProperty', 'MatterGenCheckpointInfo', 'PRETRAINED_MODEL_NAME']
+__all__ = ['generator', 'chemgraph', 'TargetProperty', 'MatterGenCheckpointInfo', 'PRETRAINED_MODEL_NAME','CrystalGenerator']

mars_toolkit/services/mattergen_service.py

@@ -14,14 +14,14 @@ from typing import Dict, Any, Optional, Union, List
 import threading
 
 # 导入mattergen相关模块
-import sys
+# import sys
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../')))
+# sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../')))
-from mattergen_wrapper import generator
+# from mars_toolkit.core.mattergen_wrapper import generator
-CrystalGenerator = generator.CrystalGenerator
+# CrystalGenerator = generator.CrystalGenerator
-from mattergen.common.data.types import TargetProperty
+# from mattergen.common.data.types import TargetProperty
-from mattergen.common.utils.eval_utils import MatterGenCheckpointInfo
+# from mattergen.common.utils.eval_utils import MatterGenCheckpointInfo
-from mattergen.common.utils.data_classes import PRETRAINED_MODEL_NAME
+# from mattergen.common.utils.data_classes import PRETRAINED_MODEL_NAME
-
+from ..core.mattergen_wrapper import *
 # 导入mars_toolkit配置
 from mars_toolkit.core.config import config
 

test_mars_toolkit.py

@@ -158,16 +158,16 @@ if __name__ == "__main__":
 
     # 测试工具函数列表
     tools_to_test = [
-        "get_current_time",                           # 基础工具
+        "get_current_time",                           # 0基础工具
-        "search_online",                              # 网络搜索工具
+        "search_online",                              # 1网络搜索工具
-        "search_material_property_from_material_project",  # 材料项目查询工具
+        "search_material_property_from_material_project",  # 2材料项目查询工具
-        "get_crystal_structures_from_materials_project",   # 晶体结构查询工具
+        "get_crystal_structures_from_materials_project",   # 3晶体结构查询工具
-        "get_mpid_from_formula",                      # 材料ID查询工具
+        "get_mpid_from_formula",                      # 4材料ID查询工具
-        "optimize_crystal_structure",                 # 晶体结构优化工具
+        "optimize_crystal_structure",                 # 5晶体结构优化工具
-        "generate_material",                          # 材料生成工具
+        "generate_material",                          # 6材料生成工具
-        "fetch_chemical_composition_from_OQMD",       # OQMD查询工具
+        "fetch_chemical_composition_from_OQMD",       # 7OQMD查询工具
-        "retrieval_from_knowledge_base",              # 知识库检索工具
+        "retrieval_from_knowledge_base",              # 8知识库检索工具
-        "predict_properties"                          # 属性预测工具
+        "predict_properties"                          # 9属性预测工具
     ]
 
     # 选择要测试的工具