Yn unimol2 (#288)

* add unimol2 into unimol_tools * update diff modelsize config * update get repr for unimol v2, and fix bug when input is a dict of atoms and coordinates * update verion number * remove unicore dependency * update args for diff model size
deepmodeling · Nov 21, 2024 · 69d7de5 · 69d7de5
1 parent 6fa4af5
commit 69d7de5
Show file tree

Hide file tree

Showing 15 changed files with 1,804 additions and 44 deletions.
diff --git a/unimol_tools/setup.py b/unimol_tools/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name="unimol_tools",
-    version="0.1.0.post4",
+    version="0.1.1",
     description=("unimol_tools is a Python package for property prediciton with Uni-Mol in molecule, materials and protein."),
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',

diff --git a/unimol_tools/unimol_tools/config/__init__.py b/unimol_tools/unimol_tools/config/__init__.py
@@ -1 +1 @@
-from .model_config import MODEL_CONFIG
+from .model_config import MODEL_CONFIG, MODEL_CONFIG_V2
diff --git a/unimol_tools/unimol_tools/config/model_config.py b/unimol_tools/unimol_tools/config/model_config.py
@@ -13,4 +13,14 @@
         "crystal": "mp.dict.txt",
         "oled": "oled.dict.txt",
     },
+}
+
+MODEL_CONFIG_V2 = {
+    'weight': {
+        '84m': 'modelzoo/84M/checkpoint.pt',
+        '164m': 'modelzoo/164M/checkpoint.pt',
+        '310m': 'modelzoo/310M/checkpoint.pt',
+        '570m': 'modelzoo/570M/checkpoint.pt',
+        '1.1B': 'modelzoo/1.1B/checkpoint.pt',
+    },
 }
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		from .model_config import MODEL_CONFIG
		from .model_config import MODEL_CONFIG, MODEL_CONFIG_V2