【MIIT program】support matterchat

ppmat/datasets/__init__.py

@@ -48,6 +48,7 @@
 from ppmat.datasets.oc20_s2ef_dataset import OC20S2EFDataset  # noqa
 from ppmat.datasets.qm9_dataset import QM9Dataset # noqa
 from ppmat.datasets.omol25_dataset import OMol25Dataset
+from ppmat.models.matterchat.trainer import MTDataset  # noqa
 from ppmat.datasets.split_mptrj_data import none_to_zero
 from ppmat.datasets.transform import build_transforms
 from ppmat.utils import logger

ppmat/datasets/collate_fn.py

@@ -31,6 +31,8 @@
 from ppmat.datasets.geometric_data_type.batch import Batch
 from ppmat.datasets.geometric_data_type.data import Data
 
+from ppmat.models.matterchat.trainer import MTCollator # noqa
+
 
 class DefaultCollator(object):
     def __call__(self, batch: List[Any]) -> Any:

ppmat/models/__init__.py

@@ -42,6 +42,8 @@
 from ppmat.models.megnet.megnet import MEGNetPlus
 from ppmat.models.infgcn.infgcn import InfGCN
 from ppmat.models.mateno.mateno import MatENO
+from ppmat.models.matterchat.q_former.q_former_complete import Blip2MistralInstruct
+from ppmat.models.matterchat.trainer import MatterChatModule
 from ppmat.models.sfin.sfin import SFIN
 from ppmat.utils import download
 from ppmat.utils import logger
@@ -68,6 +70,8 @@
     "DiffNMR",
     "InfGCN",
     "MatENO",
+    "Blip2MistralInstruct",
+    "MatterChatModule",
     "SFIN",
 ]
 
@@ -113,6 +117,7 @@
     "mattergen_ml2ddb": "https://paddle-org.bj.bcebos.com/paddlematerial/workflow/ml2ddb/mattergen_ml2ddb.zip",
     "mattergen_ml2ddb_chemical_system": "https://paddle-org.bj.bcebos.com/paddlematerial/workflow/ml2ddb/mattergen_ml2ddb_chemical_system.zip",
     "mattergen_ml2ddb_space_group": "https://paddle-org.bj.bcebos.com/paddlematerial/workflow/ml2ddb/mattergen_ml2ddb_space_group.zip",
+    "matterchat_full": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/llm/MatterChat/matterchat_full.zip",
     "sfin_haadf_enhance": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_enhancement/sfin/sfin_haadf_enhance.zip",
     "sfin_haadf_detect": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_enhancement/sfin/sfin_haadf_detect.zip",
     "sfin_bf_enhance": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_enhancement/sfin/sfin_bf_enhance.zip",

ppmat/models/matterchat/__init__.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MatterChat: Crystal material generation based on Mistral-7B + Q-Former + CHGNet."""
+
+from ppmat.models.matterchat.mistral import MistralConfig, MistralForCausalLM, MistralModel
+from ppmat.models.matterchat.q_former import Blip2MistralInstruct, Blip2Base, BertConfig
+from ppmat.models.matterchat.trainer import (
+    LoRALinear,
+    MISTRAL_LORA_TARGETS,
+    MTCollator,
+    MTDataset,
+    MatterChatModule,
+    MatterChatTrainer,
+    apply_lora_to_mistral,
+)
+
+__all__ = [
+    "BertConfig",
+    "Blip2Base",
+    "Blip2MistralInstruct",
+    "LoRALinear",
+    "MISTRAL_LORA_TARGETS",
+    "MTCollator",
+    "MTDataset",
+    "MatterChatModule",
+    "MatterChatTrainer",
+    "MistralConfig",
+    "MistralForCausalLM",
+    "MistralModel",
+    "apply_lora_to_mistral",
+]

ppmat/models/matterchat/chgnet/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

ppmat/models/matterchat/chgnet/graph/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from ppmat.models.matterchat.chgnet.graph.converter import CrystalGraphConverter
+from ppmat.models.matterchat.chgnet.graph.crystalgraph import CrystalGraph
+
+__all__ = ["CrystalGraph", "CrystalGraphConverter"]

ppmat/models/matterchat/chgnet/graph/converter.py

@@ -0,0 +1,136 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import sys
+from typing import TYPE_CHECKING, Literal
+
+import paddle
+from paddle import Tensor
+import paddle.nn as nn
+
+from ppmat.models.matterchat.chgnet.graph.crystalgraph import CrystalGraph
+from ppmat.models.matterchat.chgnet.graph.graph import Graph, Node
+
+if TYPE_CHECKING:
+    from pymatgen.core import Structure
+
+datatype = paddle.float32
+
+
+class CrystalGraphConverter(nn.Layer):
+    """Convert a pymatgen.core.Structure to a CrystalGraph."""
+
+    def __init__(
+        self, atom_graph_cutoff: float = 5, bond_graph_cutoff: float = 3
+    ) -> None:
+        """Initialize the Crystal Graph Converter."""
+        super().__init__()
+        self.atom_graph_cutoff = atom_graph_cutoff
+        if bond_graph_cutoff is None:
+            self.bond_graph_cutoff = atom_graph_cutoff
+        else:
+            self.bond_graph_cutoff = bond_graph_cutoff
+
+    def forward(
+        self,
+        structure: Structure,
+        graph_id=None,
+        mp_id=None,
+        on_isolated_atoms: Literal["ignore", "warn", "error"] = "error",
+    ) -> CrystalGraph:
+        """Convert a structure, return a CrystalGraph."""
+        n_atoms = len(structure)
+        atomic_number = paddle.to_tensor(
+            [i.specie.Z for i in structure], dtype="int64"
+        )
+        atomic_number.stop_gradient = True
+        atom_frac_coord = paddle.to_tensor(
+            structure.frac_coords, dtype=datatype
+        )
+        atom_frac_coord.stop_gradient = False
+        lattice = paddle.to_tensor(
+            structure.lattice.matrix, dtype=datatype
+        )
+        lattice.stop_gradient = False
+        center_index, neighbor_index, image, distance = self.get_neighbors(structure)
+
+        graph = Graph([Node(index=i) for i in range(n_atoms)])
+        for ii, jj, img, dist in zip(center_index, neighbor_index, image, distance):
+            graph.add_edge(center_index=ii, neighbor_index=jj, image=img, distance=dist)
+
+        atom_graph, directed2undirected = graph.adjacency_list()
+        atom_graph = paddle.to_tensor(atom_graph, dtype="int64")
+        directed2undirected = paddle.to_tensor(directed2undirected, dtype="int64")
+
+        try:
+            bond_graph, undirected2directed = graph.line_graph_adjacency_list(
+                cutoff=self.bond_graph_cutoff
+            )
+        except Exception as exc:
+            structure.to(filename="bond_graph_error.cif")
+            raise SystemExit(
+                f"Failed creating bond graph for {graph_id}, check bond_graph_error.cif"
+            ) from exc
+        bond_graph = paddle.to_tensor(bond_graph, dtype="int64")
+        undirected2directed = paddle.to_tensor(undirected2directed, dtype="int64")
+
+        has_isolated_atom = not set(range(n_atoms)).issubset(center_index)
+        if has_isolated_atom:
+            r_cutoff = self.atom_graph_cutoff
+            msg = f"{graph_id=} has isolated atom with {r_cutoff=}, should be skipped"
+            if on_isolated_atoms == "ignore":
+                return None
+            if on_isolated_atoms == "warn":
+                print(msg, file=sys.stderr)
+                return None
+            raise ValueError(msg)
+
+        return CrystalGraph(
+            atomic_number=atomic_number,
+            atom_frac_coord=atom_frac_coord,
+            atom_graph=atom_graph,
+            neighbor_image=paddle.to_tensor(image, dtype=datatype),
+            directed2undirected=directed2undirected,
+            undirected2directed=undirected2directed,
+            bond_graph=bond_graph,
+            lattice=lattice,
+            graph_id=graph_id,
+            mp_id=mp_id,
+            composition=structure.composition.formula,
+            atom_graph_cutoff=self.atom_graph_cutoff,
+            bond_graph_cutoff=self.bond_graph_cutoff,
+        )
+
+    def get_neighbors(
+        self, structure: Structure
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+        """Get neighbor information from pymatgen utility function."""
+        center_index, neighbor_index, image, distance = structure.get_neighbor_list(
+            r=self.atom_graph_cutoff, sites=structure.sites, numerical_tol=1e-8
+        )
+        return center_index, neighbor_index, image, distance
+
+    def as_dict(self) -> dict[str, float]:
+        """Save the args of the graph converter."""
+        return {
+            "atom_graph_cutoff": self.atom_graph_cutoff,
+            "bond_graph_cutoff": self.bond_graph_cutoff,
+        }
+
+    @classmethod
+    def from_dict(cls, dict) -> CrystalGraphConverter:
+        """Create converter from dictionary."""
+        return CrystalGraphConverter(**dict)

ppmat/models/matterchat/chgnet/graph/crystalgraph.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import paddle
+from paddle import Tensor
+
+datatype = paddle.float32
+
+
+class CrystalGraph:
+    """A data class for crystal graph."""
+
+    def __init__(
+        self,
+        atomic_number: Tensor,
+        atom_frac_coord: Tensor,
+        atom_graph: Tensor,
+        atom_graph_cutoff: float,
+        neighbor_image: Tensor,
+        directed2undirected: Tensor,
+        undirected2directed: Tensor,
+        bond_graph: Tensor,
+        bond_graph_cutoff: float,
+        lattice: Tensor,
+        graph_id: str | None = None,
+        mp_id: str | None = None,
+        composition: str | None = None,
+    ) -> None:
+        """Initialize the crystal graph."""
+        super().__init__()
+        self.atomic_number = atomic_number
+        self.atom_frac_coord = atom_frac_coord
+        self.atom_graph = atom_graph
+        self.atom_graph_cutoff = atom_graph_cutoff
+        self.neighbor_image = neighbor_image
+        self.directed2undirected = directed2undirected
+        self.undirected2directed = undirected2directed
+        self.bond_graph = bond_graph
+        self.bond_graph_cutoff = bond_graph_cutoff
+        self.lattice = lattice
+        self.graph_id = graph_id
+        self.mp_id = mp_id
+        self.composition = composition
+        if len(directed2undirected) != 2 * len(undirected2directed):
+            raise ValueError(
+                f"{graph_id} number of directed indices != 2 * number of undirected indices!"
+            )
+
+    def to(self, device: str = "cpu") -> CrystalGraph:
+        """Move the graph to a device. Default = 'cpu'."""
+        return CrystalGraph(
+            atomic_number=self.atomic_number,
+            atom_frac_coord=self.atom_frac_coord,
+            atom_graph=self.atom_graph,
+            atom_graph_cutoff=self.atom_graph_cutoff,
+            neighbor_image=self.neighbor_image,
+            directed2undirected=self.directed2undirected,
+            undirected2directed=self.undirected2directed,
+            bond_graph=self.bond_graph,
+            bond_graph_cutoff=self.bond_graph_cutoff,
+            lattice=self.lattice,
+            graph_id=self.graph_id,
+            mp_id=self.mp_id,
+            composition=self.composition,
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert the graph to a dictionary."""
+        return {
+            "atomic_number": self.atomic_number,
+            "atom_frac_coord": self.atom_frac_coord,
+            "atom_graph": self.atom_graph,
+            "atom_graph_cutoff": self.atom_graph_cutoff,
+            "neighbor_image": self.neighbor_image,
+            "directed2undirected": self.directed2undirected,
+            "undirected2directed": self.undirected2directed,
+            "bond_graph": self.bond_graph,
+            "bond_graph_cutoff": self.bond_graph_cutoff,
+            "lattice": self.lattice,
+            "graph_id": self.graph_id,
+            "mp_id": self.mp_id,
+            "composition": self.composition,
+        }
+
+    def save(self, fname: str | None = None, save_dir: str = ".") -> str:
+        """Save the graph to a file."""
+        if fname is not None:
+            save_name = os.path.join(save_dir, fname)
+        elif self.graph_id is not None:
+            save_name = os.path.join(save_dir, f"{self.graph_id}.pdparams")
+        else:
+            save_name = os.path.join(save_dir, f"{self.composition}.pdparams")
+        paddle.save(self.to_dict(), path=save_name)
+        return save_name
+
+    @classmethod
+    def from_file(cls, file_name: str) -> CrystalGraph:
+        """Load a crystal graph from a file."""
+        return paddle.load(file_name)
+
+    @classmethod
+    def from_dict(cls, dic: dict[str, Any]) -> CrystalGraph:
+        """Load a CrystalGraph from a dictionary."""
+        return CrystalGraph(**dic)
+
+    def __repr__(self) -> str:
+        """Details of the graph."""
+        return (
+            f"Crystal Graph {self.composition} \n"
+            f"constructed using atom_graph_cutoff={self.atom_graph_cutoff}, "
+            f"bond_graph_cutoff={self.bond_graph_cutoff} \n"
+            f"(n_atoms={len(self.atomic_number)}, "
+            f"atom_graph={len(self.atom_graph)}, "
+            f"bond_graph={len(self.bond_graph)})"
+        )