Feature request: transformer-based model for atomistic graph data #806

Transfer Learning/Atomistic graph data/config.json

@@ -0,0 +1,26 @@
+{
+  "_name_or_path": "/mnt/data2/s2ef_all_10epochs_weights/checkpoint-292950",
+  "architectures": [
+    "AtomformerModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_atomformer.AtomformerConfig",
+    "AutoModel": "modeling_atomformer.AtomformerModel"
+  },
+  "bos_token_id": 120,
+  "cls_token_id": 122,
+  "depth": 12,
+  "dim": 768,
+  "dropout": 0.0,
+  "eos_token_id": 121,
+  "gradient_checkpointing": false,
+  "k": 128,
+  "mask_token_id": 0,
+  "mlp_ratio": 4,
+  "model_type": "atomformer",
+  "num_heads": 32,
+  "pad_token_id": 119,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.0",
+  "vocab_size": 123
+}

Transfer Learning/Atomistic graph data/configuration_atomformer.py

@@ -0,0 +1,42 @@
+from transformers.configuration_utils import PretrainedConfig
+from typing import Any
+
+class AtomformerConfig(PretrainedConfig):  # type: ignore
+    r"""
+    Configuration of a :class:`~transform:class:`~transformers.AtomformerModel`.
+
+    It is used to instantiate an Atomformer model according to the specified arguments.
+    """
+
+    model_type = "atomformer"
+
+    def __init__(
+        self,
+        vocab_size: int = 123,
+        dim: int = 768,
+        num_heads: int = 32,
+        depth: int = 12,
+        mlp_ratio: int = 1,
+        k: int = 128,
+        dropout: float = 0.0,
+        mask_token_id: int = 0,
+        pad_token_id: int = 119,
+        bos_token_id: int = 120,
+        eos_token_id: int = 121,
+        cls_token_id: int = 122,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.dim = dim
+        self.num_heads = num_heads
+        self.depth = depth
+        self.mlp_ratio = mlp_ratio
+        self.k = k
+
+        self.dropout = dropout
+        self.mask_token_id = mask_token_id
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.cls_token_id = cls_token_id

Transfer Learning/Atomistic graph data/model.py

@@ -0,0 +1,6 @@
+import kagglehub
+
+# Download latest version
+path = kagglehub.model_download("tedlord/atomformer/pyTorch/default")
+
+print("Path to model files:", path)