diff --git a/mars_toolkit/compute/__pycache__/material_gen.cpython-310.pyc b/mars_toolkit/compute/__pycache__/material_gen.cpython-310.pyc index 319fd14..e929dd8 100644 Binary files a/mars_toolkit/compute/__pycache__/material_gen.cpython-310.pyc and b/mars_toolkit/compute/__pycache__/material_gen.cpython-310.pyc differ diff --git a/mars_toolkit/compute/material_gen.py b/mars_toolkit/compute/material_gen.py index ac60b4a..b1fa653 100644 --- a/mars_toolkit/compute/material_gen.py +++ b/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 -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, -) +from ..core.mattergen_wrapper import * logger = logging.getLogger(__name__) diff --git a/mars_toolkit/core/__pycache__/config.cpython-310.pyc b/mars_toolkit/core/__pycache__/config.cpython-310.pyc index 7b77453..b7698f1 100644 Binary files a/mars_toolkit/core/__pycache__/config.cpython-310.pyc and b/mars_toolkit/core/__pycache__/config.cpython-310.pyc differ diff --git a/mars_toolkit/core/__pycache__/mattergen_wrapper.cpython-310.pyc b/mars_toolkit/core/__pycache__/mattergen_wrapper.cpython-310.pyc new file mode 100644 index 0000000..a0a2b88 Binary files /dev/null and b/mars_toolkit/core/__pycache__/mattergen_wrapper.cpython-310.pyc differ diff --git a/mars_toolkit/core/config.py b/mars_toolkit/core/config.py index 29cd4ef..4babf2c 100644 --- a/mars_toolkit/core/config.py +++ b/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 diff --git a/mattergen_wrapper.py b/mars_toolkit/core/mattergen_wrapper.py similarity index 82% rename from mattergen_wrapper.py rename to mars_toolkit/core/mattergen_wrapper.py index 6bbe6de..8642ff8 100644 --- a/mattergen_wrapper.py +++ b/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'] diff --git a/mars_toolkit/services/__pycache__/mattergen_service.cpython-310.pyc b/mars_toolkit/services/__pycache__/mattergen_service.cpython-310.pyc index 77575bd..f5cdb5d 100644 Binary files a/mars_toolkit/services/__pycache__/mattergen_service.cpython-310.pyc and b/mars_toolkit/services/__pycache__/mattergen_service.cpython-310.pyc differ diff --git a/mars_toolkit/services/mattergen_service.py b/mars_toolkit/services/mattergen_service.py index 1578268..2693811 100644 --- a/mars_toolkit/services/mattergen_service.py +++ b/mars_toolkit/services/mattergen_service.py @@ -14,14 +14,14 @@ from typing import Dict, Any, Optional, Union, List import threading # 导入mattergen相关模块 -import sys -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../'))) -from mattergen_wrapper import generator -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 - +# import sys +# sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../'))) +# from mars_toolkit.core.mattergen_wrapper import generator +# 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 +from ..core.mattergen_wrapper import * # 导入mars_toolkit配置 from mars_toolkit.core.config import config diff --git a/test_mars_toolkit.py b/test_mars_toolkit.py index f3b997d..ec87396 100644 --- a/test_mars_toolkit.py +++ b/test_mars_toolkit.py @@ -158,16 +158,16 @@ if __name__ == "__main__": # 测试工具函数列表 tools_to_test = [ - "get_current_time", # 基础工具 - "search_online", # 网络搜索工具 - "search_material_property_from_material_project", # 材料项目查询工具 - "get_crystal_structures_from_materials_project", # 晶体结构查询工具 - "get_mpid_from_formula", # 材料ID查询工具 - "optimize_crystal_structure", # 晶体结构优化工具 - "generate_material", # 材料生成工具 - "fetch_chemical_composition_from_OQMD", # OQMD查询工具 - "retrieval_from_knowledge_base", # 知识库检索工具 - "predict_properties" # 属性预测工具 + "get_current_time", # 0基础工具 + "search_online", # 1网络搜索工具 + "search_material_property_from_material_project", # 2材料项目查询工具 + "get_crystal_structures_from_materials_project", # 3晶体结构查询工具 + "get_mpid_from_formula", # 4材料ID查询工具 + "optimize_crystal_structure", # 5晶体结构优化工具 + "generate_material", # 6材料生成工具 + "fetch_chemical_composition_from_OQMD", # 7OQMD查询工具 + "retrieval_from_knowledge_base", # 8知识库检索工具 + "predict_properties" # 9属性预测工具 ] # 选择要测试的工具