diff --git a/ppmat/datasets/__init__.py b/ppmat/datasets/__init__.py
index 05d3de49..a70a3b24 100644
--- a/ppmat/datasets/__init__.py
+++ b/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
diff --git a/ppmat/datasets/collate_fn.py b/ppmat/datasets/collate_fn.py
index 9073af4c..26e09b1c 100644
--- a/ppmat/datasets/collate_fn.py
+++ b/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:
diff --git a/ppmat/models/__init__.py b/ppmat/models/__init__.py
index d258bd2a..1489575b 100644
--- a/ppmat/models/__init__.py
+++ b/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",
diff --git a/ppmat/models/matterchat/__init__.py b/ppmat/models/matterchat/__init__.py
new file mode 100644
index 00000000..38ff2c5e
--- /dev/null
+++ b/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",
+]
diff --git a/ppmat/models/matterchat/chgnet/__init__.py b/ppmat/models/matterchat/chgnet/__init__.py
new file mode 100644
index 00000000..290f972c
--- /dev/null
+++ b/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.
diff --git a/ppmat/models/matterchat/chgnet/graph/__init__.py b/ppmat/models/matterchat/chgnet/graph/__init__.py
new file mode 100644
index 00000000..82ba909a
--- /dev/null
+++ b/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"]
diff --git a/ppmat/models/matterchat/chgnet/graph/converter.py b/ppmat/models/matterchat/chgnet/graph/converter.py
new file mode 100644
index 00000000..8441ce0b
--- /dev/null
+++ b/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)
diff --git a/ppmat/models/matterchat/chgnet/graph/crystalgraph.py b/ppmat/models/matterchat/chgnet/graph/crystalgraph.py
new file mode 100644
index 00000000..0aca3c01
--- /dev/null
+++ b/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)})"
+ )
diff --git a/ppmat/models/matterchat/chgnet/graph/graph.py b/ppmat/models/matterchat/chgnet/graph/graph.py
new file mode 100644
index 00000000..c2b1811e
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/graph/graph.py
@@ -0,0 +1,273 @@
+# 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.utils import write_json
+
+
+class Node:
+ """A node in a graph."""
+
+ def __init__(self, index: int, info: dict = None) -> None:
+ self.index = index
+ self.info = info
+ self.neighbors: dict[int, list] = {}
+
+ def add_neighbor(self, index, edge):
+ """Draw a directed edge between self and the node specified by index."""
+ if index not in self.neighbors:
+ self.neighbors[index] = [edge]
+ else:
+ self.neighbors[index].append(edge)
+
+
+class UndirectedEdge:
+ """An undirected/bi-directed edge in a graph."""
+
+ def __init__(
+ self, nodes: list, index: int | None = None, info: dict | None = None
+ ) -> None:
+ self.nodes = nodes
+ self.index = index
+ self.info = info
+
+ def __repr__(self):
+ return (
+ f"UndirectedEdge between Nodes{self.nodes}, "
+ f"info={self.info}, index={self.index}"
+ )
+
+ def __eq__(self, other):
+ return set(self.nodes) == set(other.nodes) and self.info == other.info
+
+
+class DirectedEdge:
+ """A directed edge in a graph."""
+
+ def __init__(
+ self, nodes: list, index: int | None = None, info: dict | None = None
+ ) -> None:
+ self.nodes = nodes
+ self.index = index
+ self.info = info
+
+ def make_undirected(self, index, info=None):
+ """Make a directed edge undirected."""
+ if info is None:
+ info = {}
+ info["distance"] = self.info["distance"]
+ return UndirectedEdge(self.nodes, index, info)
+
+ def __eq__(self, other) -> bool:
+ """Check if two directed edges are equal or reverse of each other."""
+ import numpy as np
+
+ self_image = self.info["image"]
+ other_image = other.info["image"]
+
+ if isinstance(self_image, np.ndarray):
+ eq_check = np.array_equal(self_image, other_image)
+ else:
+ eq_check = (self_image == other_image)
+
+ if (
+ self.nodes == other.nodes
+ and eq_check
+ ):
+ print(
+ "!!!!!! the two directed edges are equal but this operation is "
+ "not supposed to happen"
+ )
+ return True
+
+ neg_other_image = -1 * other_image
+ if isinstance(neg_other_image, np.ndarray):
+ neg_eq_check = np.array_equal(self_image, neg_other_image)
+ else:
+ neg_eq_check = (self_image == neg_other_image)
+
+ if (
+ self.nodes == other.nodes[::-1]
+ and neg_eq_check
+ ):
+ return True
+ return False
+
+ def __repr__(self):
+ return (
+ f"DirectedEdge between Nodes{self.nodes}, "
+ f"info={self.info}, index={self.index}"
+ )
+
+
+class Graph:
+ """A graph for storing the neighbor information of atoms."""
+
+ def __init__(self, nodes: list[Node]) -> None:
+ self.nodes = nodes
+ self.directed_edges: dict[frozenset[int], list[DirectedEdge]] = {}
+ self.directed_edges_list: list[DirectedEdge] = []
+ self.undirected_edges: dict[frozenset[int], list[UndirectedEdge]] = {}
+ self.undirected_edges_list: list[UndirectedEdge] = []
+
+ def add_edge(self, center_index, neighbor_index, image, distance) -> None:
+ """Add a directed edge to the graph."""
+ directed_edge_index = len(self.directed_edges_list)
+ this_directed_edge = DirectedEdge(
+ [center_index, neighbor_index],
+ index=directed_edge_index,
+ info={"image": image, "distance": distance},
+ )
+
+ tmp = frozenset([center_index, neighbor_index])
+ if tmp not in self.undirected_edges:
+ this_directed_edge.info["undirected_edge_index"] = len(
+ self.undirected_edges_list
+ )
+ this_undirected_edge = this_directed_edge.make_undirected(
+ index=len(self.undirected_edges_list),
+ info={"directed_edge_index": [directed_edge_index]},
+ )
+ self.undirected_edges[tmp] = [this_undirected_edge]
+ self.undirected_edges_list.append(this_undirected_edge)
+ self.nodes[center_index].add_neighbor(neighbor_index, this_directed_edge)
+ self.directed_edges_list.append(this_directed_edge)
+ else:
+ # Check if this is the reverse directed edge of an existing undirected edge
+ for undirected_edge in self.undirected_edges[tmp]:
+ if (
+ abs(undirected_edge.info["distance"] - distance) < 1e-6
+ and len(undirected_edge.info["directed_edge_index"]) == 1
+ ):
+ added_DE = self.directed_edges_list[
+ undirected_edge.info["directed_edge_index"][0]
+ ]
+ if added_DE == this_directed_edge:
+ this_directed_edge.info[
+ "undirected_edge_index"
+ ] = added_DE.info["undirected_edge_index"]
+ self.nodes[center_index].add_neighbor(
+ neighbor_index, this_directed_edge
+ )
+ self.directed_edges_list.append(this_directed_edge)
+ undirected_edge.info["directed_edge_index"].append(
+ directed_edge_index
+ )
+ return
+
+ # No matching undirected edge; create a new one
+ this_directed_edge.info["undirected_edge_index"] = len(
+ self.undirected_edges_list
+ )
+ this_undirected_edge = this_directed_edge.make_undirected(
+ index=len(self.undirected_edges_list),
+ info={"directed_edge_index": [directed_edge_index]},
+ )
+ self.undirected_edges[tmp].append(this_undirected_edge)
+ self.undirected_edges_list.append(this_undirected_edge)
+ self.nodes[center_index].add_neighbor(neighbor_index, this_directed_edge)
+ self.directed_edges_list.append(this_directed_edge)
+
+ def adjacency_list(self):
+ """Get the adjacency list."""
+ graph = [edge.nodes for edge in self.directed_edges_list]
+ directed2undirected = [
+ edge.info["undirected_edge_index"] for edge in self.directed_edges_list
+ ]
+ return graph, directed2undirected
+
+ def line_graph_adjacency_list(self, cutoff):
+ """Get the line graph adjacency list."""
+ assert len(self.directed_edges_list) == 2 * len(self.undirected_edges_list), (
+ f"Error: number of directed edges={len(self.directed_edges_list)} != 2 * "
+ f"number of undirected edges={len(self.undirected_edges_list)}!"
+ f"This indicates directed edges are not complete"
+ )
+ line_graph = []
+ undirected2directed = []
+ for u_edge in self.undirected_edges_list:
+ undirected2directed.append(u_edge.info["directed_edge_index"][0])
+ if u_edge.info["distance"] > cutoff:
+ continue
+ center1, center2 = list(u_edge.nodes)
+ try:
+ directed_edge1, directed_edge2 = u_edge.info["directed_edge_index"]
+ except ValueError:
+ print("Did not find 2 Directed_edges !!!")
+ print(u_edge)
+ print(
+ "edge.info['directed_edge_index'] = ",
+ u_edge.info["directed_edge_index"],
+ )
+ print()
+ print("len directed_edges_list = ", len(self.directed_edges_list))
+ print("len undirected_edges_list = ", len(self.undirected_edges_list))
+ for directed_edges in self.nodes[center1].neighbors.values():
+ for directed_edge in directed_edges:
+ if directed_edge.index == directed_edge1:
+ continue
+ if directed_edge.info["distance"] < cutoff:
+ line_graph.append(
+ [
+ center1,
+ u_edge.index,
+ directed_edge1,
+ directed_edge.info["undirected_edge_index"],
+ directed_edge.index,
+ ]
+ )
+ for directed_edges in self.nodes[center2].neighbors.values():
+ for directed_edge in directed_edges:
+ if directed_edge.index == directed_edge2:
+ continue
+ if directed_edge.info["distance"] < cutoff:
+ line_graph.append(
+ [
+ center2,
+ u_edge.index,
+ directed_edge2,
+ directed_edge.info["undirected_edge_index"],
+ directed_edge.index,
+ ]
+ )
+ return line_graph, undirected2directed
+
+ def undirected2directed(self):
+ """Map undirected_edge index to one of its directed_edge indices."""
+ out = []
+ for undirected_edge in self.undirected_edges_list:
+ out.append(undirected_edge.info["directed_edge_index"][0])
+ return out
+
+ def as_dict(self):
+ """Return dictionary serialization of a Graph."""
+ return {
+ "nodes": self.nodes,
+ "directed_edges": self.directed_edges,
+ "directed_edges_list": self.directed_edges_list,
+ "undirected_edges": self.undirected_edges,
+ "undirected_edges_list": self.undirected_edges_list,
+ }
+
+ def to(self, filename="graph.json"):
+ """Save graph dictionary to file."""
+ write_json(self.as_dict(), filename)
+
+ def __repr__(self) -> str:
+ """Return string representation of the Graph."""
+ num_nodes = len(self.nodes)
+ num_directed_edges = len(self.directed_edges_list)
+ num_undirected_edges = len(self.undirected_edges_list)
+ return f"Graph({num_nodes=}, {num_directed_edges=}, {num_undirected_edges=})"
diff --git a/ppmat/models/matterchat/chgnet/model/__init__.py b/ppmat/models/matterchat/chgnet/model/__init__.py
new file mode 100644
index 00000000..311e388a
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/model/__init__.py
@@ -0,0 +1,24 @@
+# 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.model.dynamics import (
+ CHGNetCalculator,
+ MolecularDynamics,
+ StructOptimizer,
+)
+from ppmat.models.matterchat.chgnet.model.model import CHGNet
+
+__all__ = ["CHGNet", "StructOptimizer", "MolecularDynamics", "CHGNetCalculator"]
diff --git a/ppmat/models/matterchat/chgnet/model/_local_basis.py b/ppmat/models/matterchat/chgnet/model/_local_basis.py
new file mode 100644
index 00000000..f1f26f88
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/model/_local_basis.py
@@ -0,0 +1,42 @@
+# 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.
+
+import paddle
+from paddle import Tensor, nn
+
+
+class GaussianExpansion(nn.Layer):
+ """Expand distance by Gaussian basis functions."""
+
+ def __init__(
+ self,
+ min: float = 0,
+ max: float = 5,
+ step: float = 0.5,
+ var: float | None = None,
+ ) -> None:
+ super().__init__()
+ assert min < max
+ assert max - min > step
+ self.register_buffer(
+ "gaussian_centers", paddle.arange(min, max + step, step)
+ )
+ if var is None:
+ var = step
+ self.var = var
+
+ def expand(self, features: Tensor) -> Tensor:
+ return paddle.exp(
+ -((features.reshape([-1, 1]) - self.gaussian_centers) ** 2) / self.var**2
+ )
diff --git a/ppmat/models/matterchat/chgnet/model/basis.py b/ppmat/models/matterchat/chgnet/model/basis.py
new file mode 100644
index 00000000..fe3a853d
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/model/basis.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.
+
+# GaussianExpansion is kept local because framework variants have a different
+# interface (initial/final/num_centers/width) than CHGNet's (min/max/step/var).
+from ppmat.models.chgnet.chgnet import CutoffPolynomial, Fourier, RadialBessel
+from ppmat.models.matterchat.chgnet.model._local_basis import GaussianExpansion
+
+__all__ = ["CutoffPolynomial", "Fourier", "GaussianExpansion", "RadialBessel"]
diff --git a/ppmat/models/matterchat/chgnet/model/composition_model.py b/ppmat/models/matterchat/chgnet/model/composition_model.py
new file mode 100644
index 00000000..24c0faf0
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/model/composition_model.py
@@ -0,0 +1,380 @@
+# 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 collections
+from typing import TYPE_CHECKING, Sequence
+
+import numpy as np
+import paddle
+from pymatgen.core import Structure
+from paddle import Tensor, nn
+
+from ppmat.models.matterchat.chgnet.model.functions import GatedMLP, find_activation
+
+if TYPE_CHECKING:
+ from ppmat.models.matterchat.chgnet.graph.crystalgraph import CrystalGraph
+
+
+class CompositionModel(nn.Layer):
+ """A simple FC model that takes in a chemical composition (no structure info)
+ and outputs energy.
+ """
+
+ def __init__(
+ self,
+ atom_fea_dim: int = 64,
+ activation: str = "silu",
+ is_intensive: bool = True,
+ max_num_elements: int = 94,
+ ) -> None:
+ """Initialize a CompositionModel."""
+ super().__init__()
+ self.is_intensive = is_intensive
+ self.max_num_elements = max_num_elements
+ self.fc1 = nn.Linear(max_num_elements, atom_fea_dim)
+ self.activation = find_activation(activation)
+ self.gated_mlp = GatedMLP(
+ input_dim=atom_fea_dim,
+ output_dim=atom_fea_dim,
+ hidden_dim=atom_fea_dim,
+ activation=activation,
+ )
+ self.fc2 = nn.Linear(atom_fea_dim, 1)
+
+ def _get_energy(self, composition_feas: Tensor) -> Tensor:
+ """Predict the energy given composition encoding."""
+ composition_feas = self.activation(self.fc1(composition_feas))
+ composition_feas = composition_feas + self.gated_mlp(composition_feas)
+ return self.fc2(composition_feas).reshape([-1])
+
+ def forward(self, graphs: list[CrystalGraph]) -> Tensor:
+ """Get the energy of a list of CrystalGraphs as Tensor."""
+ composition_feas = self._assemble_graphs(graphs)
+ return self._get_energy(composition_feas)
+
+ def _assemble_graphs(self, graphs: list[CrystalGraph]):
+ """Assemble a list of graphs into one-hot composition encodings."""
+ composition_feas = []
+ for _graph_idx, graph in enumerate(graphs):
+ composition_fea = paddle.bincount(
+ graph.atomic_number - 1, minlength=self.max_num_elements
+ )
+ if self.is_intensive:
+ n_atom = graph.atomic_number.shape[0]
+ composition_fea = composition_fea.cast(paddle.float32) / n_atom
+ composition_feas.append(composition_fea)
+ return paddle.stack(composition_feas, axis=0)
+
+
+class AtomRef(nn.Layer):
+ """A linear regression for elemental energy."""
+
+ def __init__(self, is_intensive: bool = True, max_num_elements: int = 94) -> None:
+ """Initialize an AtomRef model."""
+ super().__init__()
+ self.is_intensive = is_intensive
+ self.max_num_elements = max_num_elements
+ self.fc = nn.Linear(max_num_elements, 1, bias_attr=False)
+ self.fitted = False
+
+ def forward(self, graphs: list[CrystalGraph]):
+ """Get the energy of a list of CrystalGraphs."""
+ assert self.fitted is True, "composition model need to be fitted first!"
+ composition_feas = self._assemble_graphs(graphs)
+ return self._get_energy(composition_feas)
+
+ def _get_energy(self, composition_feas: Tensor) -> Tensor:
+ """Predict the energy given composition encoding."""
+ return self.fc(composition_feas).reshape([-1])
+
+ def fit(
+ self,
+ structures_or_graphs: Sequence[Structure | CrystalGraph],
+ energies: Sequence[float],
+ ) -> None:
+ """Fit the model to a list of crystals and energies."""
+ num_data = len(energies)
+ composition_feas = paddle.zeros([num_data, self.max_num_elements])
+ e = paddle.zeros([num_data])
+ for index, (structure, energy) in enumerate(
+ zip(structures_or_graphs, energies)
+ ):
+ if isinstance(structure, Structure):
+ atomic_number = paddle.to_tensor(
+ [i.specie.Z for i in structure], dtype="int32"
+ )
+ else:
+ atomic_number = structure.atomic_number
+ composition_fea = paddle.bincount(
+ atomic_number - 1, minlength=self.max_num_elements
+ )
+ if self.is_intensive:
+ composition_fea = composition_fea.cast(paddle.float32) / atomic_number.shape[0]
+ composition_feas[index, :] = composition_fea
+ e[index] = energy
+
+ self.feature_matrix = composition_feas.detach().numpy()
+ self.energies = e.detach().numpy()
+ state_dict = collections.OrderedDict()
+ weight = (
+ np.linalg.pinv(self.feature_matrix.T @ self.feature_matrix)
+ @ self.feature_matrix.T
+ @ self.energies
+ )
+ state_dict["weight"] = paddle.to_tensor(weight).reshape([94, 1])
+ self.fc.load_state_dict(state_dict)
+ self.fitted = True
+
+ def _assemble_graphs(self, graphs: list[CrystalGraph]):
+ """Assemble a list of graphs into one-hot composition encodings."""
+ composition_feas = []
+ for _graph_idx, graph in enumerate(graphs):
+ composition_fea = paddle.bincount(
+ graph.atomic_number - 1, minlength=self.max_num_elements
+ )
+ if self.is_intensive:
+ n_atom = graph.atomic_number.shape[0]
+ composition_fea = composition_fea.cast(paddle.float32) / n_atom
+ composition_feas.append(composition_fea)
+ return paddle.stack(composition_feas, axis=0).cast("float32")
+
+ def initialize_from(self, dataset: str):
+ """Initialize pre-fitted weights from a dataset."""
+ if dataset in ["MPtrj", "MPtrj_e"]:
+ self.initialize_from_MPtrj()
+ elif dataset in ["MPF"]:
+ self.initialize_from_MPF()
+ else:
+ raise NotImplementedError(f"{dataset=} not supported yet")
+
+ def initialize_from_MPtrj(self):
+ """Initialize pre-fitted weights from MPtrj dataset."""
+ state_dict = collections.OrderedDict()
+ state_dict["weight"] = paddle.to_tensor(
+ [
+ -3.4431,
+ -0.1279,
+ -2.8300,
+ -3.4737,
+ -7.4946,
+ -8.2354,
+ -8.1611,
+ -8.3861,
+ -5.7498,
+ -0.0236,
+ -1.7406,
+ -1.6788,
+ -4.2833,
+ -6.2002,
+ -6.1315,
+ -5.8405,
+ -3.8795,
+ -0.0703,
+ -1.5668,
+ -3.4451,
+ -7.0549,
+ -9.1465,
+ -9.2594,
+ -9.3514,
+ -8.9843,
+ -8.0228,
+ -6.4955,
+ -5.6057,
+ -3.4002,
+ -0.9217,
+ -3.2499,
+ -4.9164,
+ -4.7810,
+ -5.0191,
+ -3.3316,
+ 0.5130,
+ -1.4043,
+ -3.2175,
+ -7.4994,
+ -9.3816,
+ -10.4386,
+ -9.9539,
+ -7.9555,
+ -8.5440,
+ -7.3245,
+ -5.2771,
+ -1.9014,
+ -0.4034,
+ -2.6002,
+ -4.0054,
+ -4.1156,
+ -3.9928,
+ -2.7003,
+ 2.2170,
+ -1.9671,
+ -3.7180,
+ -6.8133,
+ -7.3502,
+ -6.0712,
+ -6.1699,
+ -5.1471,
+ -6.1925,
+ -11.5829,
+ -15.8841,
+ -5.9994,
+ -6.0798,
+ -5.9513,
+ -6.0400,
+ -5.9773,
+ -2.5091,
+ -6.0767,
+ -10.6666,
+ -11.8761,
+ -11.8491,
+ -10.7397,
+ -9.6100,
+ -8.4755,
+ -6.2070,
+ -3.0337,
+ 0.4726,
+ -1.6425,
+ -3.1295,
+ -3.3328,
+ -0.1221,
+ -0.3448,
+ -0.4364,
+ -0.1661,
+ -0.3680,
+ -4.1869,
+ -8.4233,
+ -10.0467,
+ -12.0953,
+ -12.5228,
+ -14.2530,
+ ]
+ ).reshape([94, 1])
+ self.fc.load_state_dict(state_dict)
+ self.is_intensive = True
+ self.fitted = True
+
+ def initialize_from_MPF(self):
+ """Initialize pre-fitted weights from MPF dataset."""
+ state_dict = collections.OrderedDict()
+ state_dict["weight"] = paddle.to_tensor(
+ [
+ -3.4654e00,
+ -6.2617e-01,
+ -3.4622e00,
+ -4.7758e00,
+ -8.0362e00,
+ -8.4038e00,
+ -7.7681e00,
+ -7.3892e00,
+ -4.9472e00,
+ -5.4833e00,
+ -2.4783e00,
+ -2.0202e00,
+ -5.1548e00,
+ -7.9121e00,
+ -6.9135e00,
+ -4.6228e00,
+ -3.0155e00,
+ -2.1285e00,
+ -2.3174e00,
+ -4.7595e00,
+ -8.1742e00,
+ -1.1421e01,
+ -8.9229e00,
+ -8.4901e00,
+ -8.1664e00,
+ -6.5826e00,
+ -5.2614e00,
+ -4.4844e00,
+ -3.2737e00,
+ -1.3498e00,
+ -3.6264e00,
+ -4.6727e00,
+ -4.1316e00,
+ -3.6755e00,
+ -2.8030e00,
+ 6.4728e00,
+ -2.2469e00,
+ -4.2510e00,
+ -1.0245e01,
+ -1.1666e01,
+ -1.1802e01,
+ -8.6551e00,
+ -9.3641e00,
+ -7.5716e00,
+ -5.6990e00,
+ -4.9716e00,
+ -1.8871e00,
+ -6.7951e-01,
+ -2.7488e00,
+ -3.7945e00,
+ -3.3883e00,
+ -2.5588e00,
+ -1.9621e00,
+ 9.9793e00,
+ -2.5566e00,
+ -4.8803e00,
+ -8.8604e00,
+ -9.0537e00,
+ -7.9431e00,
+ -8.1259e00,
+ -6.3212e00,
+ -8.3025e00,
+ -1.2289e01,
+ -1.7310e01,
+ -7.5512e00,
+ -8.1959e00,
+ -8.3493e00,
+ -7.2591e00,
+ -8.4170e00,
+ -3.3873e00,
+ -7.6823e00,
+ -1.2630e01,
+ -1.3626e01,
+ -9.5299e00,
+ -1.1840e01,
+ -9.7990e00,
+ -7.5561e00,
+ -5.4690e00,
+ -2.6508e00,
+ 4.1746e-01,
+ -2.3255e00,
+ -3.4830e00,
+ -3.1808e00,
+ -1.6934e-02,
+ -3.6191e-02,
+ -1.0842e-02,
+ 1.3170e-02,
+ -6.5371e-02,
+ -5.4892e00,
+ -1.0335e01,
+ -1.1130e01,
+ -1.4312e01,
+ -1.4700e01,
+ -1.5473e01,
+ ]
+ ).reshape([94, 1])
+ self.fc.load_state_dict(state_dict)
+ self.is_intensive = False
+ self.fitted = True
+
+ def initialize_from_numpy(self, file_name):
+ """Initialize pre-fitted weights from numpy file."""
+ atom_ref_np = np.load(file_name)
+ state_dict = collections.OrderedDict()
+ state_dict["weight"] = paddle.to_tensor(atom_ref_np).reshape([94, 1])
+ self.fc.load_state_dict(state_dict)
+ self.is_intensive = False
+ self.fitted = True
diff --git a/ppmat/models/matterchat/chgnet/model/dynamics.py b/ppmat/models/matterchat/chgnet/model/dynamics.py
new file mode 100644
index 00000000..58df4b94
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/model/dynamics.py
@@ -0,0 +1,413 @@
+# 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 contextlib
+import io
+import pickle
+import sys
+from typing import TYPE_CHECKING
+
+import numpy as np
+import paddle
+from ase import Atoms, units
+from ase.calculators.calculator import Calculator, all_changes, all_properties
+from ase.constraints import ExpCellFilter
+from ase.md.nptberendsen import Inhomogeneous_NPTBerendsen, NPTBerendsen
+from ase.md.nvtberendsen import NVTBerendsen
+from ase.optimize.bfgs import BFGS
+from ase.optimize.bfgslinesearch import BFGSLineSearch
+from ase.optimize.fire import FIRE
+from ase.optimize.lbfgs import LBFGS, LBFGSLineSearch
+from ase.optimize.mdmin import MDMin
+from ase.optimize.sciopt import SciPyFminBFGS, SciPyFminCG
+from pymatgen.analysis.eos import BirchMurnaghan
+from pymatgen.core.structure import Molecule, Structure
+from pymatgen.io.ase import AseAtomsAdaptor
+
+from ppmat.models.matterchat.chgnet.model.model import CHGNet
+from ppmat.utils import logger
+
+if TYPE_CHECKING:
+ from ase.io import Trajectory
+ from ase.optimize.optimize import Optimizer
+
+OPTIMIZERS = {
+ "FIRE": FIRE,
+ "BFGS": BFGS,
+ "LBFGS": LBFGS,
+ "LBFGSLineSearch": LBFGSLineSearch,
+ "MDMin": MDMin,
+ "SciPyFminCG": SciPyFminCG,
+ "SciPyFminBFGS": SciPyFminBFGS,
+ "BFGSLineSearch": BFGSLineSearch,
+}
+
+
+class CHGNetCalculator(Calculator):
+ """CHGNet Calculator for ASE applications."""
+
+ implemented_properties = ["energy", "forces", "stress", "magmoms"]
+
+ def __init__(
+ self,
+ model: CHGNet | None = None,
+ use_device: str | None = None,
+ stress_weight: float | None = 1 / 160.21766208,
+ **kwargs,
+ ) -> None:
+ """Provide a CHGNet instance to calculate various atomic properties using ASE."""
+ super().__init__(**kwargs)
+
+ if use_device == "mps":
+ raise NotImplementedError("mps is not supported yet")
+ self.device = use_device or (
+ "gpu" if paddle.is_compiled_with_cuda() else "cpu"
+ )
+
+ if model is None:
+ model = CHGNet.load()
+ paddle.set_device(self.device)
+ self.model = model
+ self.stress_weight = stress_weight
+ logger.message(f"CHGNet will run on {self.device}")
+
+ def calculate(
+ self,
+ atoms: Atoms | None = None,
+ properties: list | None = None,
+ system_changes: list | None = None,
+ ) -> None:
+ """Calculate various properties of the atoms using CHGNet."""
+ properties = properties or all_properties
+ system_changes = system_changes or all_changes
+ super().calculate(
+ atoms=atoms,
+ properties=properties,
+ system_changes=system_changes,
+ )
+
+ structure = AseAtomsAdaptor.get_structure(atoms)
+ graph = self.model.graph_converter(structure)
+ model_prediction = self.model.predict_graph(
+ graph.to(self.device), task="efsm"
+ )
+
+ factor = (
+ 1
+ if not self.model.is_intensive
+ else structure.composition.num_atoms
+ )
+ self.results.update(
+ energy=model_prediction["e"] * factor,
+ forces=model_prediction["f"],
+ free_energy=model_prediction["e"] * factor,
+ magmoms=model_prediction["m"],
+ stress=model_prediction["s"] * self.stress_weight,
+ )
+
+
+class StructOptimizer:
+ """Wrapper class for structural relaxation."""
+
+ def __init__(
+ self,
+ model: CHGNet | None = None,
+ optimizer_class: Optimizer | str | None = "FIRE",
+ use_device: str | None = None,
+ stress_weight: float = 1 / 160.21766208,
+ ) -> None:
+ """Provide a trained CHGNet model and an optimizer to relax crystal structures."""
+ if isinstance(optimizer_class, str):
+ if optimizer_class in OPTIMIZERS:
+ optimizer_class = OPTIMIZERS[optimizer_class]
+ else:
+ raise ValueError(
+ f"Optimizer instance not found. Select one from {list(OPTIMIZERS)}"
+ )
+
+ self.optimizer_class: Optimizer = optimizer_class
+ self.calculator = CHGNetCalculator(
+ model=model, stress_weight=stress_weight, use_device=use_device
+ )
+
+ def relax(
+ self,
+ atoms: Structure | Atoms,
+ fmax: float | None = 0.1,
+ steps: int | None = 500,
+ relax_cell: bool | None = True,
+ save_path: str | None = None,
+ trajectory_save_interval: int | None = 1,
+ verbose: bool = True,
+ **kwargs,
+ ) -> dict[str, Structure | TrajectoryObserver]:
+ """Relax the Structure/Atoms until maximum force is smaller than fmax."""
+ if isinstance(atoms, Structure):
+ atoms = AseAtomsAdaptor.get_atoms(atoms)
+
+ atoms.calc = self.calculator
+
+ stream = sys.stdout if verbose else io.StringIO()
+ with contextlib.redirect_stdout(stream):
+ obs = TrajectoryObserver(atoms)
+ if relax_cell:
+ atoms = ExpCellFilter(atoms)
+ optimizer = self.optimizer_class(atoms, **kwargs)
+ optimizer.attach(obs, interval=trajectory_save_interval)
+ optimizer.run(fmax=fmax, steps=steps)
+ obs()
+
+ if save_path is not None:
+ obs.save(save_path)
+
+ if isinstance(atoms, ExpCellFilter):
+ atoms = atoms.atoms
+ struct = AseAtomsAdaptor.get_structure(atoms)
+ for k in struct.site_properties:
+ struct.remove_site_property(property_name=k)
+ struct.add_site_property(
+ "magmom", [float(i) for i in atoms.get_magnetic_moments()]
+ )
+ return {"final_structure": struct, "trajectory": obs}
+
+
+class TrajectoryObserver:
+ """Trajectory observer is a hook in the relaxation process that saves the
+ intermediate structures.
+ """
+
+ def __init__(self, atoms: Atoms) -> None:
+ """Create a TrajectoryObserver from an Atoms object."""
+ self.atoms = atoms
+ self.energies: list[float] = []
+ self.forces: list[np.ndarray] = []
+ self.stresses: list[np.ndarray] = []
+ self.magmoms: list[np.ndarray] = []
+ self.atom_positions: list[np.ndarray] = []
+ self.cells: list[np.ndarray] = []
+
+ def __call__(self):
+ """The logic for saving the properties of an Atoms during the relaxation."""
+ self.energies.append(self.compute_energy())
+ self.forces.append(self.atoms.get_forces())
+ self.stresses.append(self.atoms.get_stress())
+ self.magmoms.append(self.atoms.get_magnetic_moments())
+ self.atom_positions.append(self.atoms.get_positions())
+ self.cells.append(self.atoms.get_cell()[:])
+
+ def __len__(self) -> int:
+ """The number of steps in the trajectory."""
+ return len(self.energies)
+
+ def compute_energy(self) -> float:
+ """Calculate the potential energy."""
+ return self.atoms.get_potential_energy()
+
+ def save(self, filename: str) -> None:
+ """Save the trajectory to file."""
+ out_pkl = {
+ "energy": self.energies,
+ "forces": self.forces,
+ "stresses": self.stresses,
+ "magmoms": self.magmoms,
+ "atom_positions": self.atom_positions,
+ "cell": self.cells,
+ "atomic_number": self.atoms.get_atomic_numbers(),
+ }
+ with open(filename, "wb") as file:
+ pickle.dump(out_pkl, file)
+
+
+class MolecularDynamics:
+ """Molecular dynamics class."""
+
+ def __init__(
+ self,
+ atoms: Atoms | Structure,
+ model: CHGNet | None = None,
+ ensemble: str = "nvt",
+ temperature: int = 300,
+ timestep: float = 2.0,
+ pressure: float = 1.01325 * units.bar,
+ taut: float | None = None,
+ taup: float | None = None,
+ compressibility_au: float | None = None,
+ trajectory: str | Trajectory | None = None,
+ logfile: str | None = None,
+ loginterval: int = 1,
+ append_trajectory: bool = False,
+ use_device: str | None = None,
+ ) -> None:
+ """Initialize the MD class."""
+ if isinstance(atoms, (Structure, Molecule)):
+ atoms = AseAtomsAdaptor.get_atoms(atoms)
+
+ self.atoms = atoms
+ self.atoms.calc = CHGNetCalculator(model, use_device=use_device)
+
+ if taut is None:
+ taut = 100 * timestep * units.fs
+ if taup is None:
+ taup = 1000 * timestep * units.fs
+
+ if ensemble.lower() == "nvt":
+ self.dyn = NVTBerendsen(
+ atoms=self.atoms,
+ timestep=timestep * units.fs,
+ temperature_K=temperature,
+ taut=taut,
+ trajectory=trajectory,
+ logfile=logfile,
+ loginterval=loginterval,
+ append_trajectory=append_trajectory,
+ )
+ else:
+ if compressibility_au is None:
+ eos = EquationOfState(
+ model=model,
+ use_device=use_device,
+ )
+ eos.fit(atoms=atoms, steps=500, fmax=0.1)
+ compressibility_au = eos.get_compressibility(unit="A^3/eV")
+ logger.message(
+ f"Done compressibility calculation: "
+ f"b = {round(compressibility_au, 3)} A^3/eV"
+ )
+
+ if ensemble.lower() == "npt":
+ self.dyn = Inhomogeneous_NPTBerendsen(
+ atoms=self.atoms,
+ timestep=timestep * units.fs,
+ temperature_K=temperature,
+ pressure_au=pressure,
+ taut=taut,
+ taup=taup,
+ compressibility_au=compressibility_au,
+ trajectory=trajectory,
+ logfile=logfile,
+ loginterval=loginterval,
+ )
+
+ elif ensemble.lower() == "npt_berendsen":
+ self.dyn = NPTBerendsen(
+ atoms=self.atoms,
+ timestep=timestep * units.fs,
+ temperature_K=temperature,
+ pressure_au=pressure,
+ taut=taut,
+ taup=taup,
+ compressibility_au=compressibility_au,
+ trajectory=trajectory,
+ logfile=logfile,
+ loginterval=loginterval,
+ append_trajectory=append_trajectory,
+ )
+
+ else:
+ raise ValueError("Ensemble not supported")
+
+ self.trajectory = trajectory
+ self.logfile = logfile
+ self.loginterval = loginterval
+ self.timestep = timestep
+
+ def run(self, steps: int):
+ """Thin wrapper of ase MD run."""
+ self.dyn.run(steps)
+
+ def set_atoms(self, atoms: Atoms):
+ """Set new atoms to run MD."""
+ calculator = self.atoms.calc
+ self.atoms = atoms
+ self.dyn.atoms = atoms
+ self.dyn.atoms.calc = calculator
+
+
+class EquationOfState:
+ """Class to calculate equation of state."""
+
+ def __init__(
+ self,
+ model: CHGNet | None = None,
+ optimizer_class: Optimizer | str | None = "FIRE",
+ use_device: str | None = None,
+ stress_weight: float = 1 / 160.21766208,
+ ) -> None:
+ """Initialize a structure optimizer object for calculation of bulk modulus."""
+ self.relaxer = StructOptimizer(
+ model=model,
+ optimizer_class=optimizer_class,
+ use_device=use_device,
+ stress_weight=stress_weight,
+ )
+ self.fitted = False
+
+ def fit(
+ self,
+ atoms: Structure | Atoms,
+ n_points: int = 11,
+ fmax: float | None = 0.1,
+ steps: int | None = 500,
+ **kwargs,
+ ):
+ """Relax the Structure/Atoms and fit the Birch-Murnaghan equation of state."""
+ if isinstance(atoms, Atoms):
+ atoms = AseAtomsAdaptor.get_structure(atoms)
+ volumes, energies = [], []
+ for i in np.linspace(-0.1, 0.1, n_points):
+ structure_strained = atoms.copy()
+ structure_strained.apply_strain([i, i, i])
+ result = self.relaxer.relax(
+ structure_strained,
+ relax_cell=False,
+ fmax=fmax,
+ steps=steps,
+ verbose=False,
+ **kwargs,
+ )
+ volumes.append(result["final_structure"].volume)
+ energies.append(result["trajectory"].energies[-1])
+ self.bm = BirchMurnaghan(volumes=volumes, energies=energies)
+ self.bm.fit()
+ self.fitted = True
+
+ def get_bulk_mudulus(self, unit: str = "eV/A^3"):
+ """Get the bulk modulus from the fitted Birch-Murnaghan equation of state."""
+ if self.fitted is False:
+ raise ValueError(
+ "Equation of state needs to be fitted first through self.fit()"
+ )
+ if unit == "eV/A^3":
+ return self.bm.b0
+ if unit == "GPa":
+ return self.bm.b0_GPa
+ raise NotImplementedError("unit has to be eV/A^3 or GPa")
+
+ def get_compressibility(self, unit: str = "A^3/eV"):
+ """Get the compressibility from the fitted Birch-Murnaghan equation of state."""
+ if self.fitted is False:
+ raise ValueError(
+ "Equation of state needs to be fitted first through self.fit()"
+ )
+ if unit == "A^3/eV":
+ return 1 / self.bm.b0
+ if unit == "GPa^-1":
+ return 1 / self.bm.b0_GPa
+ if unit in ["Pa^-1", "m^2/N"]:
+ return 1 / (self.bm.b0_GPa * 1e9)
+ raise NotImplementedError(
+ "unit has to be one of A^3/eV, GPa^-1 Pa^-1 or m^2/N"
+ )
diff --git a/ppmat/models/matterchat/chgnet/model/encoders.py b/ppmat/models/matterchat/chgnet/model/encoders.py
new file mode 100644
index 00000000..df0e1149
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/model/encoders.py
@@ -0,0 +1,96 @@
+# 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 paddle
+from paddle import Tensor, nn
+
+from ppmat.models.matterchat.chgnet.model.basis import Fourier, RadialBessel
+
+
+class AtomEmbedding(nn.Layer):
+ """Encode an atom by its atomic number."""
+
+ def __init__(self, atom_feature_dim: int, max_num_elements: int = 94) -> None:
+ super().__init__()
+ self.embedding = nn.Embedding(max_num_elements, atom_feature_dim)
+
+ def forward(self, atomic_numbers: Tensor) -> Tensor:
+ return self.embedding(atomic_numbers)
+
+
+class BondEncoder(nn.Layer):
+ """Encode bonds using Gaussian distance expansion."""
+
+ def __init__(
+ self,
+ atom_graph_cutoff: float = 5,
+ bond_graph_cutoff: float = 3,
+ num_radial: int = 9,
+ cutoff_coeff: int = 5,
+ learnable: bool = False,
+ ) -> None:
+ super().__init__()
+ self.rbf_expansion_ag = RadialBessel(
+ num_radial=num_radial,
+ cutoff=atom_graph_cutoff,
+ smooth_cutoff=cutoff_coeff,
+ learnable=learnable,
+ )
+ self.rbf_expansion_bg = RadialBessel(
+ num_radial=num_radial,
+ cutoff=bond_graph_cutoff,
+ smooth_cutoff=cutoff_coeff,
+ learnable=learnable,
+ )
+
+ def forward(
+ self,
+ center: Tensor,
+ neighbor: Tensor,
+ undirected2directed: Tensor,
+ image: Tensor,
+ lattice: Tensor,
+ ) -> tuple[Tensor, Tensor, Tensor]:
+ neighbor = neighbor + image @ lattice
+ bond_vectors = center - neighbor
+ bond_lengths = paddle.norm(bond_vectors, axis=1)
+ bond_vectors = bond_vectors / bond_lengths[:, None]
+
+ undirected_bond_lengths = paddle.index_select(
+ bond_lengths, 0, undirected2directed
+ )
+ bond_basis_ag = self.rbf_expansion_ag(undirected_bond_lengths)
+ bond_basis_bg = self.rbf_expansion_bg(undirected_bond_lengths)
+ return bond_basis_ag, bond_basis_bg, bond_vectors
+
+
+class AngleEncoder(nn.Layer):
+ """Encode angles using Fourier expansion."""
+
+ def __init__(self, num_angular: int = 9, learnable: bool = True) -> None:
+ super().__init__()
+ if num_angular % 2 != 1:
+ raise ValueError(f"{num_angular=} must be an odd integer")
+ circular_harmonics_order = (num_angular - 1) // 2
+ self.fourier_expansion = Fourier(
+ order=circular_harmonics_order, learnable=learnable
+ )
+
+ def forward(self, bond_i: Tensor, bond_j: Tensor) -> Tensor:
+ # 1 - 1e-6 for acos stability
+ cosine_ij = paddle.sum(bond_i * bond_j, axis=1) * (1 - 1e-6)
+ angle = paddle.acos(cosine_ij)
+ return self.fourier_expansion(angle)
diff --git a/ppmat/models/matterchat/chgnet/model/functions.py b/ppmat/models/matterchat/chgnet/model/functions.py
new file mode 100644
index 00000000..61db31a3
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/model/functions.py
@@ -0,0 +1,127 @@
+# 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 typing import Sequence
+
+from paddle import Tensor, nn
+
+from ppmat.models.chgnet.chgnet import aggregate
+from ppmat.models.matterchat.utils.activations import ScaledSiLU
+from ppmat.models.matterchat.utils.activations import get_activation as find_activation
+
+
+class MLP(nn.Layer):
+ """Multi-Layer Perceptron with configurable activation.
+
+ Kept local because the framework CHGNet MLP hardcodes paddle.nn.Silu().
+ """
+
+ def __init__(
+ self,
+ input_dim: int,
+ output_dim: int = 1,
+ hidden_dim: int | Sequence[int] | None = (64, 64),
+ dropout: float = 0,
+ activation: str = "silu",
+ ) -> None:
+ super().__init__()
+ if hidden_dim in (None, 0):
+ layers = [nn.Dropout(dropout), nn.Linear(input_dim, output_dim)]
+ elif type(hidden_dim) == int:
+ layers = [
+ nn.Linear(input_dim, hidden_dim),
+ find_activation(activation),
+ nn.Dropout(dropout),
+ nn.Linear(hidden_dim, output_dim),
+ ]
+ elif isinstance(hidden_dim, Sequence):
+ layers = [nn.Linear(input_dim, hidden_dim[0]), find_activation(activation)]
+ if len(hidden_dim) != 1:
+ for h_in, h_out in zip(hidden_dim[0:-1], hidden_dim[1:]):
+ layers.append(nn.Linear(h_in, h_out))
+ layers.append(find_activation(activation))
+ layers.append(nn.Dropout(dropout))
+ layers.append(nn.Linear(hidden_dim[-1], output_dim))
+ else:
+ raise TypeError(
+ f"{hidden_dim=} must be an integer, a list of integers, or None."
+ )
+ self.layers = nn.Sequential(*layers)
+
+ def forward(self, X: Tensor) -> Tensor:
+ return self.layers(X)
+
+
+class GatedMLP(nn.Layer):
+ """Gated MLP (core * sigmoid(gate)). Used in CGCNN and M3GNet."""
+
+ def __init__(
+ self,
+ input_dim: int,
+ output_dim: int,
+ hidden_dim: int | list[int] | None = None,
+ dropout=0,
+ activation="silu",
+ norm="batch",
+ ) -> None:
+ super().__init__()
+ self.mlp_core = MLP(
+ input_dim=input_dim,
+ output_dim=output_dim,
+ hidden_dim=hidden_dim,
+ dropout=dropout,
+ activation=activation,
+ )
+ self.mlp_gate = MLP(
+ input_dim=input_dim,
+ output_dim=output_dim,
+ hidden_dim=hidden_dim,
+ dropout=dropout,
+ activation=activation,
+ )
+ self.activation = find_activation(activation)
+ self.sigmoid = nn.Sigmoid()
+ self.norm = norm
+ self.bn1 = find_normalization(name=norm, dim=output_dim)
+ self.bn2 = find_normalization(name=norm, dim=output_dim)
+
+ def forward(self, X: Tensor) -> Tensor:
+ if self.norm is None:
+ core = self.activation(self.mlp_core(X))
+ gate = self.sigmoid(self.mlp_gate(X))
+ else:
+ core = self.activation(self.bn1(self.mlp_core(X)))
+ gate = self.sigmoid(self.bn2(self.mlp_gate(X)))
+ return core * gate
+
+
+def find_normalization(name: str, dim: int = None) -> nn.Layer | None:
+ if name is None:
+ return None
+ return {
+ "batch": nn.BatchNorm1D(dim),
+ "layer": nn.LayerNorm(dim),
+ }.get(name.lower(), None)
+
+
+__all__ = [
+ "MLP",
+ "GatedMLP",
+ "aggregate",
+ "find_normalization",
+ "find_activation",
+ "ScaledSiLU",
+]
diff --git a/ppmat/models/matterchat/chgnet/model/layers.py b/ppmat/models/matterchat/chgnet/model/layers.py
new file mode 100644
index 00000000..aad2f104
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/model/layers.py
@@ -0,0 +1,251 @@
+# 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.
+
+# Reference: raw-code/Model/chgnet_lib/model/layers.py
+from __future__ import annotations
+
+import paddle
+from paddle import Tensor, nn
+
+from ppmat.models.matterchat.chgnet.model.functions import (
+ MLP,
+ GatedMLP,
+ aggregate,
+ find_activation,
+ find_normalization,
+)
+
+
+class AtomConv(nn.Layer):
+ """Convolution layer to update atom features."""
+
+ def __init__(
+ self,
+ atom_fea_dim: int,
+ bond_fea_dim: int,
+ hidden_dim: int = 64,
+ dropout: float = 0,
+ activation: str = "silu",
+ norm: str | None = None,
+ use_mlp_out: bool = True,
+ resnet: bool = True,
+ gMLP_norm: str = "batch",
+ ) -> None:
+ super().__init__()
+ self.use_mlp_out = use_mlp_out
+ self.resnet = resnet
+ self.activation = find_activation(activation)
+ self.twoBody_atom = GatedMLP(
+ input_dim=2 * atom_fea_dim + bond_fea_dim,
+ output_dim=atom_fea_dim,
+ hidden_dim=hidden_dim,
+ dropout=dropout,
+ norm=gMLP_norm,
+ activation=activation,
+ )
+ if self.use_mlp_out:
+ self.mlp_out = MLP(
+ input_dim=atom_fea_dim, output_dim=atom_fea_dim, hidden_dim=0
+ )
+ self.atom_norm = find_normalization(name=norm, dim=atom_fea_dim)
+
+ def forward(
+ self,
+ atom_feas: Tensor,
+ bond_feas: Tensor,
+ bond_weights: Tensor,
+ atom_graph: Tensor,
+ directed2undirected: Tensor,
+ ) -> Tensor:
+ center_atoms = paddle.index_select(atom_feas, 0, atom_graph[:, 0])
+ nbr_atoms = paddle.index_select(atom_feas, 0, atom_graph[:, 1])
+ bonds = paddle.index_select(bond_feas, 0, directed2undirected)
+ messages = paddle.concat([center_atoms, bonds, nbr_atoms], axis=1)
+ messages = self.twoBody_atom(messages)
+
+ bond_weight = paddle.index_select(bond_weights, 0, directed2undirected)
+ messages = messages * bond_weight
+
+ new_atom_feas = aggregate(messages, atom_graph[:, 0], average=False)
+
+ if self.use_mlp_out:
+ new_atom_feas = self.mlp_out(new_atom_feas)
+ if self.resnet:
+ new_atom_feas += atom_feas
+ if self.atom_norm is not None:
+ new_atom_feas = self.atom_norm(new_atom_feas)
+ return new_atom_feas
+
+
+class BondConv(nn.Layer):
+ """Convolution layer to update bond features."""
+
+ def __init__(
+ self,
+ atom_fea_dim: int,
+ bond_fea_dim: int,
+ angle_fea_dim: int,
+ hidden_dim: int = 64,
+ dropout: float = 0,
+ activation: str = "silu",
+ norm: str | None = None,
+ use_mlp_out: bool = True,
+ resnet=True,
+ **kwargs,
+ ) -> None:
+ super().__init__()
+ self.use_mlp_out = use_mlp_out
+ self.resnet = resnet
+ self.activation = find_activation(activation)
+ self.twoBody_bond = GatedMLP(
+ input_dim=atom_fea_dim + 2 * bond_fea_dim + angle_fea_dim,
+ output_dim=bond_fea_dim,
+ hidden_dim=hidden_dim,
+ dropout=dropout,
+ norm=kwargs.pop("gMLP_norm", "batch"),
+ activation=activation,
+ )
+ if self.use_mlp_out:
+ self.mlp_out = MLP(
+ input_dim=bond_fea_dim, output_dim=bond_fea_dim, hidden_dim=0
+ )
+ self.bond_norm = find_normalization(name=norm, dim=bond_fea_dim)
+
+ def forward(
+ self,
+ atom_feas: Tensor,
+ bond_feas: Tensor,
+ bond_weights: Tensor,
+ angle_feas: Tensor,
+ bond_graph: Tensor,
+ ) -> Tensor:
+ center_atoms = paddle.index_select(atom_feas, 0, bond_graph[:, 0])
+ bond_feas_i = paddle.index_select(bond_feas, 0, bond_graph[:, 1])
+ bond_feas_j = paddle.index_select(bond_feas, 0, bond_graph[:, 2])
+ total_fea = paddle.concat(
+ [bond_feas_i, bond_feas_j, angle_feas, center_atoms], axis=1
+ )
+ bond_update = self.twoBody_bond(total_fea)
+
+ bond_weights_i = paddle.index_select(bond_weights, 0, bond_graph[:, 1])
+ bond_weights_j = paddle.index_select(bond_weights, 0, bond_graph[:, 2])
+ bond_update = bond_update * bond_weights_i * bond_weights_j
+
+ new_bond_feas = aggregate(
+ bond_update, bond_graph[:, 1], average=False, num_owner=len(bond_feas)
+ )
+
+ if self.use_mlp_out:
+ new_bond_feas = self.mlp_out(new_bond_feas)
+ if self.resnet:
+ new_bond_feas += bond_feas
+ if self.bond_norm is not None:
+ new_bond_feas = self.bond_norm(new_bond_feas)
+ return new_bond_feas
+
+
+class AngleUpdate(nn.Layer):
+ """Update angle features."""
+
+ def __init__(
+ self,
+ atom_fea_dim: int,
+ bond_fea_dim: int,
+ angle_fea_dim: int,
+ hidden_dim: int = 0,
+ dropout: float = 0,
+ activation: str = "silu",
+ norm: str | None = None,
+ resnet: bool = True,
+ **kwargs,
+ ) -> None:
+ super().__init__()
+ self.resnet = resnet
+ self.activation = find_activation(activation)
+ self.twoBody_bond = GatedMLP(
+ input_dim=atom_fea_dim + 2 * bond_fea_dim + angle_fea_dim,
+ output_dim=angle_fea_dim,
+ hidden_dim=hidden_dim,
+ dropout=dropout,
+ norm=kwargs.pop("gMLP_norm", "batch"),
+ activation=activation,
+ )
+ self.angle_norm = find_normalization(norm, dim=angle_fea_dim)
+
+ def forward(
+ self,
+ atom_feas: Tensor,
+ bond_feas: Tensor,
+ angle_feas: Tensor,
+ bond_graph: Tensor,
+ ) -> Tensor:
+ center_atoms = paddle.index_select(atom_feas, 0, bond_graph[:, 0])
+ bond_feas_i = paddle.index_select(bond_feas, 0, bond_graph[:, 1])
+ bond_feas_j = paddle.index_select(bond_feas, 0, bond_graph[:, 2])
+ total_fea = paddle.concat(
+ [bond_feas_i, bond_feas_j, angle_feas, center_atoms], axis=1
+ )
+
+ new_angle_feas = self.twoBody_bond(total_fea)
+
+ if self.resnet:
+ new_angle_feas += angle_feas
+ if self.angle_norm is not None:
+ new_angle_feas = self.angle_norm(new_angle_feas)
+ return new_angle_feas
+
+
+class GraphPooling(nn.Layer):
+ """Pool sub-graphs in a batched graph."""
+
+ def __init__(self, average: bool = False) -> None:
+ super().__init__()
+ self.average = average
+
+ def forward(self, atom_feas: Tensor, atom_owner: Tensor) -> Tensor:
+ output = aggregate(atom_feas, atom_owner, average=False)
+ if self.average:
+ bc = paddle.bincount(atom_owner.cast("int32"))
+ bc = paddle.where(bc != 0, bc, paddle.ones([1], dtype=bc.dtype))
+ output = (output.T / bc.cast(output.dtype)).T
+ return output
+
+
+class GraphAttentionReadOut(nn.Layer):
+ """Multi-head attention read-out: atom_feas -> crystal_fea."""
+
+ def __init__(
+ self, atom_fea_dim: int, num_head: int = 3, hidden_dim: int = 32, average=False
+ ) -> None:
+ super().__init__()
+ self.key = MLP(
+ input_dim=atom_fea_dim, output_dim=num_head, hidden_dim=hidden_dim
+ )
+ self.softmax = nn.Softmax(axis=0)
+ self.average = average
+
+ def forward(self, atom_feas: Tensor, atom_owner: Tensor) -> Tensor:
+ crystal_feas = []
+ weights = self.key(atom_feas)
+ bin_count = paddle.bincount(atom_owner)
+ start_index = 0
+ for n_atom in bin_count:
+ atom_fea = atom_feas[start_index : start_index + n_atom, :]
+ weight = self.softmax(weights[start_index : start_index + n_atom, :])
+ crystal_fea = (atom_fea.T @ weight).reshape([-1])
+ if self.average:
+ crystal_fea /= n_atom.cast(crystal_fea.dtype)
+ crystal_feas.append(crystal_fea)
+ start_index += n_atom
+ return paddle.stack(crystal_feas, axis=0)
diff --git a/ppmat/models/matterchat/chgnet/model/model.py b/ppmat/models/matterchat/chgnet/model/model.py
new file mode 100644
index 00000000..192c5519
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/model/model.py
@@ -0,0 +1,638 @@
+# 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 math
+import os
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal, Sequence
+
+import paddle
+from pymatgen.core import Structure
+from paddle import Tensor, nn
+
+from ppmat.models.matterchat.chgnet.graph import CrystalGraph, CrystalGraphConverter
+from ppmat.models.matterchat.chgnet.graph.crystalgraph import datatype
+from ppmat.models.matterchat.chgnet.model.composition_model import AtomRef
+from ppmat.models.matterchat.chgnet.model.encoders import (
+ AngleEncoder,
+ AtomEmbedding,
+ BondEncoder,
+)
+from ppmat.models.matterchat.chgnet.model.functions import (
+ MLP,
+ GatedMLP,
+ find_normalization,
+)
+from ppmat.models.matterchat.chgnet.model.layers import (
+ AngleUpdate,
+ AtomConv,
+ BondConv,
+ GraphAttentionReadOut,
+ GraphPooling,
+)
+from ppmat.utils import logger
+
+if TYPE_CHECKING:
+ PredTask = str
+
+
+class CHGNet(nn.Layer):
+ """Crystal Hamiltonian Graph neural Network
+ A model that takes in a crystal graph and output energy, force, magmom, stress.
+ """
+
+ def __init__(
+ self,
+ atom_fea_dim: int = 64,
+ bond_fea_dim: int = 64,
+ angle_fea_dim: int = 64,
+ composition_model: str | nn.Layer = "MPtrj",
+ num_radial: int = 9,
+ num_angular: int = 9,
+ n_conv: int = 4,
+ atom_conv_hidden_dim: Sequence[int] | int = 64,
+ update_bond: bool = True,
+ bond_conv_hidden_dim: Sequence[int] | int = 64,
+ update_angle: bool = True,
+ angle_layer_hidden_dim: Sequence[int] | int = 0,
+ conv_dropout: float = 0,
+ read_out: str = "ave",
+ mlp_hidden_dims: Sequence[int] | int = (64, 64),
+ mlp_dropout: float = 0,
+ mlp_first: bool = True,
+ is_intensive: bool = True,
+ non_linearity: Literal["silu", "relu", "tanh", "gelu"] = "relu",
+ atom_graph_cutoff: int = 5,
+ bond_graph_cutoff: int = 3,
+ cutoff_coeff: int = 5,
+ learnable_rbf: bool = True,
+ **kwargs,
+ ) -> None:
+ """Initialize the CHGNet."""
+ self.model_args = {
+ k: v
+ for k, v in locals().items()
+ if k not in ["self", "__class__", "kwargs"]
+ }
+ self.model_args.update(kwargs)
+
+ super().__init__()
+ self.atom_fea_dim = atom_fea_dim
+ self.bond_fea_dim = bond_fea_dim
+ self.is_intensive = is_intensive
+ self.n_conv = n_conv
+
+ if isinstance(composition_model, nn.Layer):
+ self.composition_model = composition_model
+ elif isinstance(composition_model, str):
+ self.composition_model = AtomRef(is_intensive=is_intensive)
+ self.composition_model.initialize_from(composition_model)
+ else:
+ self.composition_model = None
+
+ if self.composition_model is not None:
+ for param in self.composition_model.parameters():
+ param.stop_gradient = True
+
+ self.graph_converter = CrystalGraphConverter(
+ atom_graph_cutoff=atom_graph_cutoff, bond_graph_cutoff=bond_graph_cutoff
+ )
+
+ self.atom_embedding = AtomEmbedding(atom_feature_dim=atom_fea_dim)
+ self.bond_basis_expansion = BondEncoder(
+ atom_graph_cutoff=atom_graph_cutoff,
+ bond_graph_cutoff=bond_graph_cutoff,
+ num_radial=num_radial,
+ cutoff_coeff=cutoff_coeff,
+ learnable=learnable_rbf,
+ )
+ self.bond_embedding = nn.Linear(
+ in_features=num_radial, out_features=bond_fea_dim, bias_attr=False
+ )
+ self.bond_weights_ag = nn.Linear(
+ in_features=num_radial, out_features=atom_fea_dim, bias_attr=False
+ )
+ self.bond_weights_bg = nn.Linear(
+ in_features=num_radial, out_features=bond_fea_dim, bias_attr=False
+ )
+ self.angle_basis_expansion = AngleEncoder(
+ num_angular=num_angular, learnable=learnable_rbf
+ )
+ self.angle_embedding = nn.Linear(
+ in_features=num_angular, out_features=angle_fea_dim, bias_attr=False
+ )
+
+ conv_norm = kwargs.pop("conv_norm", None)
+ gMLP_norm = kwargs.pop("gMLP_norm", None)
+ atom_graph_layers = []
+ for _i in range(n_conv):
+ atom_graph_layers.append(
+ AtomConv(
+ atom_fea_dim=atom_fea_dim,
+ bond_fea_dim=bond_fea_dim,
+ hidden_dim=atom_conv_hidden_dim,
+ dropout=conv_dropout,
+ activation=non_linearity,
+ norm=conv_norm,
+ gMLP_norm=gMLP_norm,
+ use_mlp_out=True,
+ resnet=True,
+ )
+ )
+ self.atom_conv_layers = nn.LayerList(atom_graph_layers)
+
+ if update_bond is True:
+ bond_graph_layers = [
+ BondConv(
+ atom_fea_dim=atom_fea_dim,
+ bond_fea_dim=bond_fea_dim,
+ angle_fea_dim=angle_fea_dim,
+ hidden_dim=bond_conv_hidden_dim,
+ dropout=conv_dropout,
+ activation=non_linearity,
+ norm=conv_norm,
+ gMLP_norm=gMLP_norm,
+ use_mlp_out=True,
+ resnet=True,
+ )
+ for _ in range(n_conv - 1)
+ ]
+ self.bond_conv_layers = nn.LayerList(bond_graph_layers)
+ else:
+ self.bond_conv_layers = [None for _ in range(n_conv - 1)]
+
+ if update_angle is True:
+ angle_layers = [
+ AngleUpdate(
+ atom_fea_dim=atom_fea_dim,
+ bond_fea_dim=bond_fea_dim,
+ angle_fea_dim=angle_fea_dim,
+ hidden_dim=angle_layer_hidden_dim,
+ dropout=conv_dropout,
+ activation=non_linearity,
+ norm=conv_norm,
+ gMLP_norm=gMLP_norm,
+ resnet=True,
+ )
+ for _ in range(n_conv - 1)
+ ]
+ self.angle_layers = nn.LayerList(angle_layers)
+ else:
+ self.angle_layers = [None for _ in range(n_conv - 1)]
+
+ self.site_wise = nn.Linear(atom_fea_dim, 1)
+ self.readout_norm = find_normalization(
+ name=kwargs.pop("readout_norm", None), dim=atom_fea_dim
+ )
+ self.mlp_first = mlp_first
+ if mlp_first:
+ self.read_out_type = "sum"
+ input_dim = atom_fea_dim
+ self.pooling = GraphPooling(average=False)
+ elif read_out in ["attn", "weighted"]:
+ self.read_out_type = "attn"
+ num_heads = kwargs.pop("num_heads", 3)
+ self.pooling = GraphAttentionReadOut(
+ atom_fea_dim, num_head=num_heads, average=True
+ )
+ input_dim = atom_fea_dim * num_heads
+ else:
+ self.read_out_type = "ave"
+ input_dim = atom_fea_dim
+ self.pooling = GraphPooling(average=True)
+ if kwargs.pop("final_mlp", "MLP") in ["normal", "MLP"]:
+ self.mlp = MLP(
+ input_dim=input_dim,
+ hidden_dim=mlp_hidden_dims,
+ output_dim=1,
+ dropout=mlp_dropout,
+ activation=non_linearity,
+ )
+ else:
+ self.mlp = nn.Sequential(
+ GatedMLP(
+ input_dim=input_dim,
+ hidden_dim=mlp_hidden_dims,
+ output_dim=mlp_hidden_dims[-1],
+ dropout=mlp_dropout,
+ activation=non_linearity,
+ ),
+ nn.Linear(
+ in_features=mlp_hidden_dims[-1], out_features=1
+ ),
+ )
+
+ logger.message(
+ f"CHGNet initialized with {sum(p.numel() for p in self.parameters()):,} "
+ f"parameters"
+ )
+
+ def forward(
+ self,
+ graphs: Sequence[CrystalGraph],
+ task: PredTask = "e",
+ return_atom_feas: bool = False,
+ return_crystal_feas: bool = False,
+ ) -> dict:
+ """Get prediction associated with input graphs."""
+ compute_force = "f" in task
+ compute_stress = "s" in task
+ site_wise = "m" in task
+
+ comp_energy = (
+ 0 if self.composition_model is None else self.composition_model(graphs)
+ )
+
+ batched_graph = BatchedGraph.from_graphs(
+ graphs,
+ bond_basis_expansion=self.bond_basis_expansion,
+ angle_basis_expansion=self.angle_basis_expansion,
+ compute_stress=compute_stress,
+ )
+
+ prediction = self._compute(
+ batched_graph,
+ site_wise=site_wise,
+ compute_force=compute_force,
+ compute_stress=compute_stress,
+ return_atom_feas=return_atom_feas,
+ return_crystal_feas=return_crystal_feas,
+ )
+ prediction["e"] += comp_energy
+ return prediction
+
+ def _compute(
+ self,
+ g,
+ site_wise: bool = False,
+ compute_force: bool = False,
+ compute_stress: bool = False,
+ return_atom_feas: bool = False,
+ return_crystal_feas: bool = False,
+ ) -> dict:
+ """Get Energy, Force, Stress, Magmom associated with input graphs."""
+ prediction = {}
+ atoms_per_graph = paddle.bincount(g.atom_owners)
+ prediction["atoms_per_graph"] = atoms_per_graph
+
+ # H is the first embedding column
+ atom_feas = self.atom_embedding(g.atomic_numbers - 1)
+ bond_feas = self.bond_embedding(g.bond_bases_ag)
+ bond_weights_ag = self.bond_weights_ag(g.bond_bases_ag)
+ bond_weights_bg = self.bond_weights_bg(g.bond_bases_bg)
+ if len(g.angle_bases) != 0:
+ angle_feas = self.angle_embedding(g.angle_bases)
+
+ for idx, (atom_layer, bond_layer, angle_layer) in enumerate(
+ zip(self.atom_conv_layers[:-1], self.bond_conv_layers, self.angle_layers)
+ ):
+ atom_feas = atom_layer(
+ atom_feas=atom_feas,
+ bond_feas=bond_feas,
+ bond_weights=bond_weights_ag,
+ atom_graph=g.batched_atom_graph,
+ directed2undirected=g.directed2undirected,
+ )
+
+ if len(g.angle_bases) != 0 and bond_layer is not None:
+ bond_feas = bond_layer(
+ atom_feas=atom_feas,
+ bond_feas=bond_feas,
+ bond_weights=bond_weights_bg,
+ angle_feas=angle_feas,
+ bond_graph=g.batched_bond_graph,
+ )
+
+ if angle_layer is not None:
+ angle_feas = angle_layer(
+ atom_feas=atom_feas,
+ bond_feas=bond_feas,
+ angle_feas=angle_feas,
+ bond_graph=g.batched_bond_graph,
+ )
+ if idx == self.n_conv - 2:
+ if return_atom_feas is True:
+ prediction["atom_fea"] = paddle.split(
+ atom_feas, atoms_per_graph.tolist()
+ )
+ if site_wise:
+ magmom = paddle.abs(self.site_wise(atom_feas))
+ prediction["m"] = list(
+ paddle.split(magmom.reshape([-1]), atoms_per_graph.tolist())
+ )
+
+ atom_feas = self.atom_conv_layers[-1](
+ atom_feas=atom_feas,
+ bond_feas=bond_feas,
+ bond_weights=bond_weights_ag,
+ atom_graph=g.batched_atom_graph,
+ directed2undirected=g.directed2undirected,
+ )
+ if self.readout_norm is not None:
+ atom_feas = self.readout_norm(atom_feas)
+
+ if self.mlp_first:
+ energies = self.mlp(atom_feas)
+ energy = self.pooling(energies, g.atom_owners).reshape([-1])
+ else:
+ crystal_feas = self.pooling(atom_feas, g.atom_owners)
+ energy = self.mlp(crystal_feas).reshape([-1]) * atoms_per_graph
+ if return_crystal_feas is True:
+ prediction["crystal_fea"] = crystal_feas
+
+ if compute_force:
+ force = paddle.grad(
+ energy.sum(), g.atom_positions, create_graph=True, retain_graph=True
+ )
+ force = [-1 * i for i in force]
+ prediction["f"] = force
+
+ if compute_stress:
+ stress = paddle.grad(
+ energy.sum(), g.strains, create_graph=True, retain_graph=True
+ )
+ # Convert eV/A^3 to GPa
+ scale = 1 / g.volumes * 160.21766208
+ stress = [i * j for i, j in zip(stress, scale)]
+ prediction["s"] = stress
+
+ if self.is_intensive:
+ energy = energy / atoms_per_graph.cast(energy.dtype)
+ prediction["e"] = energy
+
+ return prediction
+
+ def predict_structure(
+ self,
+ structure: Structure | Sequence[Structure],
+ task: PredTask = "efsm",
+ return_atom_feas: bool = False,
+ return_crystal_feas: bool = False,
+ batch_size: int = 100,
+ ) -> dict[str, Tensor]:
+ """Predict from pymatgen.core.Structure."""
+ assert (
+ self.graph_converter is not None
+ ), "self.graph_converter needs to be initialized first!"
+ if type(structure) == Structure:
+ graph = self.graph_converter(structure)
+ return self.predict_graph(
+ graph,
+ task=task,
+ return_atom_feas=return_atom_feas,
+ return_crystal_feas=return_crystal_feas,
+ batch_size=batch_size,
+ )
+ if type(structure) == list:
+ graphs = [self.graph_converter(i) for i in structure]
+ return self.predict_graph(
+ graphs,
+ task=task,
+ return_atom_feas=return_atom_feas,
+ return_crystal_feas=return_crystal_feas,
+ batch_size=batch_size,
+ )
+ raise Exception("input should either be a structure or list of structures!")
+
+ def predict_graph(
+ self,
+ graph: CrystalGraph | Sequence[CrystalGraph],
+ task: PredTask = "efsm",
+ return_atom_feas: bool = False,
+ return_crystal_feas: bool = False,
+ batch_size: int = 100,
+ ) -> dict[str, Tensor]:
+ """Predict from CrystalGraph."""
+ model_device = next(self.parameters()).place
+ if type(graph) == CrystalGraph:
+ self.eval()
+ prediction = self.forward(
+ [graph.to(model_device)],
+ task=task,
+ return_atom_feas=return_atom_feas,
+ return_crystal_feas=return_crystal_feas,
+ )
+ out = {}
+ for key, pred in prediction.items():
+ if key == "e":
+ out[key] = pred.item()
+ elif key in ["f", "s", "m", "atom_fea"]:
+ assert len(pred) == 1
+ out[key] = pred[0].detach().numpy()
+ elif key == "crystal_fea":
+ out[key] = pred.reshape([-1]).detach().numpy()
+ return out
+ if type(graph) == list:
+ self.eval()
+ predictions: list[dict[str, Tensor]] = [{} for _ in range(len(graph))]
+ n_steps = math.ceil(len(graph) / batch_size)
+ for n in range(n_steps):
+ prediction = self.forward(
+ [
+ g.to(model_device)
+ for g in graph[batch_size * n : batch_size * (n + 1)]
+ ],
+ task=task,
+ return_atom_feas=return_atom_feas,
+ return_crystal_feas=return_crystal_feas,
+ )
+ for key, pred in prediction.items():
+ if key in ["e"]:
+ for i, e in enumerate(pred.detach().numpy()):
+ predictions[n * batch_size + i][key] = e
+ elif key in ["f", "s", "m"]:
+ for i, tmp in enumerate(pred):
+ predictions[n * batch_size + i][key] = (
+ tmp.detach().numpy()
+ )
+ elif key == "atom_fea":
+ for i, atom_fea in enumerate(pred):
+ predictions[n * batch_size + i][key] = (
+ atom_fea.detach().numpy()
+ )
+ elif key == "crystal_fea":
+ for i, crystal_fea in enumerate(pred.detach().numpy()):
+ predictions[n * batch_size + i][key] = crystal_fea
+ return predictions
+ raise Exception("input should either be a graph or list of graphs!")
+
+ @staticmethod
+ def split(x: Tensor, n: Tensor) -> Sequence[Tensor]:
+ """Split a batched result Tensor into a list of Tensors."""
+ start = 0
+ result = []
+ for i in n:
+ result.append(x[start : start + i])
+ start += i
+ assert start == len(x), "Error: source tensor not correctly split!"
+ return result
+
+ def as_dict(self):
+ """Return the CHGNet weights and args in a dictionary."""
+ return {"state_dict": self.state_dict(), "model_args": self.model_args}
+
+ @classmethod
+ def from_dict(cls, dict, **kwargs):
+ """Build a CHGNet from a saved dictionary."""
+ chgnet = CHGNet(**dict["model_args"])
+ chgnet.set_state_dict(dict["state_dict"])
+ return chgnet
+
+ @classmethod
+ def from_file(cls, path, **kwargs):
+ """Build a CHGNet from a saved file."""
+ state = paddle.load(path)
+ return CHGNet.from_dict(state["model"], **kwargs)
+
+ @classmethod
+ def load(cls, model_name="MPtrj-efsm"):
+ """Load pretrained CHGNet."""
+ current_dir = os.path.dirname(os.path.abspath(__file__))
+ if model_name == "MPtrj-efsm":
+ return cls.from_file(
+ os.path.join(current_dir, "../pretrained/e30f77s348m32.pth.tar")
+ )
+ raise Exception("model_name not supported")
+
+
+@dataclass
+class BatchedGraph:
+ """Batched crystal graph for parallel computing."""
+
+ atomic_numbers: Tensor
+ bond_bases_ag: Tensor
+ bond_bases_bg: Tensor
+ angle_bases: Tensor
+ batched_atom_graph: Tensor
+ batched_bond_graph: Tensor
+ atom_owners: Tensor
+ directed2undirected: Tensor
+ atom_positions: Sequence[Tensor]
+ strains: Sequence[Tensor]
+ volumes: Sequence[Tensor]
+
+ @classmethod
+ def from_graphs(
+ cls,
+ graphs: Sequence[CrystalGraph],
+ bond_basis_expansion: nn.Layer,
+ angle_basis_expansion: nn.Layer,
+ compute_stress: bool = False,
+ ) -> BatchedGraph:
+ """Featurize and assemble a list of graphs."""
+ atomic_numbers, atom_positions = [], []
+ strains, volumes = [], []
+ bond_bases_ag, bond_bases_bg, angle_bases = [], [], []
+ batched_atom_graph, batched_bond_graph = [], []
+ directed2undirected = []
+ atom_owners = []
+ atom_offset_idx = 0
+ n_undirected = 0
+
+ for graph_idx, graph in enumerate(graphs):
+ n_atom = graph.atomic_number.shape[0]
+ atomic_numbers.append(graph.atomic_number)
+
+ if compute_stress:
+ strain = paddle.zeros([3, 3])
+ strain.stop_gradient = False
+ lattice = graph.lattice @ (
+ paddle.eye(3).cast(strain.dtype) + strain
+ )
+ else:
+ strain = None
+ lattice = graph.lattice
+ volumes.append(paddle.linalg.det(lattice))
+ strains.append(strain)
+
+ atom_cart_coords = graph.atom_frac_coord @ lattice
+ bond_basis_ag, bond_basis_bg, bond_vectors = bond_basis_expansion(
+ center=atom_cart_coords[graph.atom_graph[:, 0]],
+ neighbor=atom_cart_coords[graph.atom_graph[:, 1]],
+ undirected2directed=graph.undirected2directed,
+ image=graph.neighbor_image,
+ lattice=lattice,
+ )
+ atom_positions.append(atom_cart_coords)
+ bond_bases_ag.append(bond_basis_ag)
+ bond_bases_bg.append(bond_basis_bg)
+
+ batched_atom_graph.append(graph.atom_graph + atom_offset_idx)
+ directed2undirected.append(
+ graph.directed2undirected + n_undirected
+ )
+
+ if len(graph.bond_graph) != 0:
+ bond_vecs_i = paddle.index_select(
+ bond_vectors, 0, graph.bond_graph[:, 2]
+ )
+ bond_vecs_j = paddle.index_select(
+ bond_vectors, 0, graph.bond_graph[:, 4]
+ )
+ angle_basis = angle_basis_expansion(bond_vecs_i, bond_vecs_j)
+ angle_bases.append(angle_basis)
+
+ bond_graph = paddle.zeros(
+ [graph.bond_graph.shape[0], 3], dtype=graph.bond_graph.dtype
+ )
+ bond_graph[:, 0] = graph.bond_graph[:, 0] + atom_offset_idx
+ bond_graph[:, 1] = graph.bond_graph[:, 1] + n_undirected
+ bond_graph[:, 2] = graph.bond_graph[:, 3] + n_undirected
+ batched_bond_graph.append(bond_graph)
+
+ atom_owners.append(
+ paddle.ones([n_atom], dtype="float32") * graph_idx
+ )
+ atom_offset_idx += n_atom
+ n_undirected += len(bond_basis_ag)
+
+ atomic_numbers = paddle.concat(atomic_numbers, axis=0)
+ bond_bases_ag = paddle.concat(bond_bases_ag, axis=0)
+ bond_bases_bg = paddle.concat(bond_bases_bg, axis=0)
+ angle_bases = (
+ paddle.concat(angle_bases, axis=0)
+ if len(angle_bases) != 0
+ else paddle.to_tensor([])
+ )
+ batched_atom_graph = paddle.concat(batched_atom_graph, axis=0)
+ if batched_bond_graph != []:
+ batched_bond_graph = paddle.concat(batched_bond_graph, axis=0)
+ else:
+ batched_bond_graph = paddle.to_tensor([])
+ atom_owners = (
+ paddle.concat(atom_owners, axis=0)
+ .cast("int32")
+ )
+ if atomic_numbers.place.is_gpu_place():
+ atom_owners = atom_owners.to(atomic_numbers.place)
+ directed2undirected = paddle.concat(directed2undirected, axis=0)
+ volumes = paddle.to_tensor(
+ volumes, dtype=datatype
+ )
+
+ return cls(
+ atomic_numbers=atomic_numbers,
+ bond_bases_ag=bond_bases_ag,
+ bond_bases_bg=bond_bases_bg,
+ angle_bases=angle_bases,
+ batched_atom_graph=batched_atom_graph,
+ batched_bond_graph=batched_bond_graph,
+ atom_owners=atom_owners,
+ directed2undirected=directed2undirected,
+ atom_positions=atom_positions,
+ strains=strains,
+ volumes=volumes,
+ )
diff --git a/ppmat/models/matterchat/chgnet/model/model_embedding.py b/ppmat/models/matterchat/chgnet/model/model_embedding.py
new file mode 100644
index 00000000..9321acb9
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/model/model_embedding.py
@@ -0,0 +1,575 @@
+# 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.
+
+"""CHGNet crystal graph neural network for material structure embedding.
+
+Cannot reuse Paddle built-in layers: crystal graph convolution (message passing
+between atoms, bonds, and angles), RBF basis expansions, and crystal graph batching
+(BatchedGraph) are domain-specific to materials science. PaddlePaddle has no
+equivalent graph operations for periodic crystal structures.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal, Sequence
+
+import paddle
+from paddle import Tensor, nn
+from pymatgen.core import Structure
+
+from ppmat.models.matterchat.chgnet.graph import CrystalGraph, CrystalGraphConverter
+from ppmat.models.matterchat.chgnet.graph.crystalgraph import datatype
+from ppmat.models.matterchat.chgnet.model.composition_model import AtomRef
+from ppmat.models.matterchat.chgnet.model.encoders import (
+ AngleEncoder,
+ AtomEmbedding,
+ BondEncoder,
+)
+from ppmat.models.matterchat.chgnet.model.functions import (
+ MLP,
+ GatedMLP,
+ find_normalization,
+)
+from ppmat.models.matterchat.chgnet.model.layers import (
+ AngleUpdate,
+ AtomConv,
+ BondConv,
+ GraphAttentionReadOut,
+ GraphPooling,
+)
+from ppmat.utils import logger
+
+if TYPE_CHECKING:
+ PredTask = str
+
+
+class CHGNet(nn.Layer):
+ """Crystal Hamiltonian Graph neural Network
+ A model that takes in a crystal graph and outputs atom feature embeddings.
+ """
+
+ def __init__(
+ self,
+ atom_fea_dim: int = 64,
+ bond_fea_dim: int = 64,
+ angle_fea_dim: int = 64,
+ composition_model: str | nn.Layer = "MPtrj",
+ num_radial: int = 9,
+ num_angular: int = 9,
+ n_conv: int = 4,
+ atom_conv_hidden_dim: Sequence[int] | int = 64,
+ update_bond: bool = True,
+ bond_conv_hidden_dim: Sequence[int] | int = 64,
+ update_angle: bool = True,
+ angle_layer_hidden_dim: Sequence[int] | int = 0,
+ conv_dropout: float = 0,
+ read_out: str = "ave",
+ mlp_hidden_dims: Sequence[int] | int = (64, 64),
+ mlp_dropout: float = 0,
+ mlp_first: bool = True,
+ is_intensive: bool = True,
+ non_linearity: Literal["silu", "relu", "tanh", "gelu"] = "silu",
+ atom_graph_cutoff: int = 5,
+ bond_graph_cutoff: int = 3,
+ cutoff_coeff: int = 5,
+ learnable_rbf: bool = True,
+ **kwargs,
+ ) -> None:
+ """Initialize the CHGNet."""
+ self.model_args = {
+ k: v
+ for k, v in locals().items()
+ if k not in ["self", "__class__", "kwargs"]
+ }
+ self.model_args.update(kwargs)
+
+ super().__init__()
+ self.atom_fea_dim = atom_fea_dim
+ self.bond_fea_dim = bond_fea_dim
+ self.is_intensive = is_intensive
+ self.n_conv = n_conv
+
+ if isinstance(composition_model, nn.Layer):
+ self.composition_model = composition_model
+ elif isinstance(composition_model, str):
+ self.composition_model = AtomRef(is_intensive=is_intensive)
+ self.composition_model.initialize_from(composition_model)
+ else:
+ self.composition_model = None
+
+ if self.composition_model is not None:
+ for param in self.composition_model.parameters():
+ param.stop_gradient = True
+
+ self.graph_converter = CrystalGraphConverter(
+ atom_graph_cutoff=atom_graph_cutoff, bond_graph_cutoff=bond_graph_cutoff
+ )
+
+ self.atom_embedding = AtomEmbedding(atom_feature_dim=atom_fea_dim)
+ self.bond_basis_expansion = BondEncoder(
+ atom_graph_cutoff=atom_graph_cutoff,
+ bond_graph_cutoff=bond_graph_cutoff,
+ num_radial=num_radial,
+ cutoff_coeff=cutoff_coeff,
+ learnable=learnable_rbf,
+ )
+ self.bond_embedding = nn.Linear(
+ in_features=num_radial, out_features=bond_fea_dim, bias_attr=False
+ )
+ self.bond_weights_ag = nn.Linear(
+ in_features=num_radial, out_features=atom_fea_dim, bias_attr=False
+ )
+ self.bond_weights_bg = nn.Linear(
+ in_features=num_radial, out_features=bond_fea_dim, bias_attr=False
+ )
+ self.angle_basis_expansion = AngleEncoder(
+ num_angular=num_angular, learnable=learnable_rbf
+ )
+ self.angle_embedding = nn.Linear(
+ in_features=num_angular, out_features=angle_fea_dim, bias_attr=False
+ )
+
+ conv_norm = kwargs.pop("conv_norm", None)
+ gMLP_norm = kwargs.pop("gMLP_norm", None)
+ atom_graph_layers = []
+ for _i in range(n_conv):
+ atom_graph_layers.append(
+ AtomConv(
+ atom_fea_dim=atom_fea_dim,
+ bond_fea_dim=bond_fea_dim,
+ hidden_dim=atom_conv_hidden_dim,
+ dropout=conv_dropout,
+ activation=non_linearity,
+ norm=conv_norm,
+ gMLP_norm=gMLP_norm,
+ use_mlp_out=True,
+ resnet=True,
+ )
+ )
+ self.atom_conv_layers = nn.LayerList(atom_graph_layers)
+
+ if update_bond is True:
+ bond_graph_layers = [
+ BondConv(
+ atom_fea_dim=atom_fea_dim,
+ bond_fea_dim=bond_fea_dim,
+ angle_fea_dim=angle_fea_dim,
+ hidden_dim=bond_conv_hidden_dim,
+ dropout=conv_dropout,
+ activation=non_linearity,
+ norm=conv_norm,
+ gMLP_norm=gMLP_norm,
+ use_mlp_out=True,
+ resnet=True,
+ )
+ for _ in range(n_conv - 1)
+ ]
+ self.bond_conv_layers = nn.LayerList(bond_graph_layers)
+ else:
+ self.bond_conv_layers = [None for _ in range(n_conv - 1)]
+
+ if update_angle is True:
+ angle_layers = [
+ AngleUpdate(
+ atom_fea_dim=atom_fea_dim,
+ bond_fea_dim=bond_fea_dim,
+ angle_fea_dim=angle_fea_dim,
+ hidden_dim=angle_layer_hidden_dim,
+ dropout=conv_dropout,
+ activation=non_linearity,
+ norm=conv_norm,
+ gMLP_norm=gMLP_norm,
+ resnet=True,
+ )
+ for _ in range(n_conv - 1)
+ ]
+ self.angle_layers = nn.LayerList(angle_layers)
+ else:
+ self.angle_layers = [None for _ in range(n_conv - 1)]
+
+ self.site_wise = nn.Linear(atom_fea_dim, 1)
+ self.readout_norm = find_normalization(
+ name=kwargs.pop("readout_norm", None), dim=atom_fea_dim
+ )
+ self.mlp_first = mlp_first
+ if mlp_first:
+ self.read_out_type = "sum"
+ input_dim = atom_fea_dim
+ self.pooling = GraphPooling(average=False)
+ elif read_out in ["attn", "weighted"]:
+ self.read_out_type = "attn"
+ num_heads = kwargs.pop("num_heads", 3)
+ self.pooling = GraphAttentionReadOut(
+ atom_fea_dim, num_head=num_heads, average=True
+ )
+ input_dim = atom_fea_dim * num_heads
+ else:
+ self.read_out_type = "ave"
+ input_dim = atom_fea_dim
+ self.pooling = GraphPooling(average=True)
+ if kwargs.pop("final_mlp", "MLP") in ["normal", "MLP"]:
+ self.mlp = MLP(
+ input_dim=input_dim,
+ hidden_dim=mlp_hidden_dims,
+ output_dim=1,
+ dropout=mlp_dropout,
+ activation=non_linearity,
+ )
+ else:
+ self.mlp = nn.Sequential(
+ GatedMLP(
+ input_dim=input_dim,
+ hidden_dim=mlp_hidden_dims,
+ output_dim=mlp_hidden_dims[-1],
+ dropout=mlp_dropout,
+ activation=non_linearity,
+ ),
+ nn.Linear(
+ in_features=mlp_hidden_dims[-1], out_features=1
+ ),
+ )
+
+ logger.message(
+ f"CHGNet initialized with {sum(p.numel() for p in self.parameters()):,} "
+ f"parameters"
+ )
+
+ def forward(
+ self,
+ graphs: Sequence[CrystalGraph],
+ task: PredTask = "e",
+ return_atom_feas: bool = False,
+ return_crystal_feas: bool = False,
+ ) -> dict:
+ """Get prediction associated with input graphs."""
+ compute_force = "f" in task
+ compute_stress = "s" in task
+ site_wise = "m" in task
+
+ comp_energy = (
+ 0 if self.composition_model is None else self.composition_model(graphs)
+ )
+
+ batched_graph = BatchedGraph.from_graphs(
+ graphs,
+ bond_basis_expansion=self.bond_basis_expansion,
+ angle_basis_expansion=self.angle_basis_expansion,
+ compute_stress=compute_stress,
+ )
+
+ feas_embedding = self._compute_embedding(
+ batched_graph,
+ site_wise=site_wise,
+ compute_force=compute_force,
+ compute_stress=compute_stress,
+ return_atom_feas=return_atom_feas,
+ return_crystal_feas=return_crystal_feas,
+ )
+ return feas_embedding
+
+ def _compute_embedding(
+ self,
+ g,
+ site_wise: bool = False,
+ compute_force: bool = False,
+ compute_stress: bool = False,
+ return_atom_feas: bool = False,
+ return_crystal_feas: bool = False,
+ ) -> dict:
+ """Get atom feature embeddings from input graphs."""
+ prediction = {}
+ atoms_per_graph = paddle.bincount(g.atom_owners)
+ prediction["atoms_per_graph"] = atoms_per_graph
+
+ # H is the first embedding column
+ atom_feas = self.atom_embedding(g.atomic_numbers - 1)
+ bond_feas = self.bond_embedding(g.bond_bases_ag)
+ bond_weights_ag = self.bond_weights_ag(g.bond_bases_ag)
+ bond_weights_bg = self.bond_weights_bg(g.bond_bases_bg)
+ if len(g.angle_bases) != 0:
+ angle_feas = self.angle_embedding(g.angle_bases)
+
+ for idx, (atom_layer, bond_layer, angle_layer) in enumerate(
+ zip(self.atom_conv_layers[:-1], self.bond_conv_layers, self.angle_layers)
+ ):
+ atom_feas = atom_layer(
+ atom_feas=atom_feas,
+ bond_feas=bond_feas,
+ bond_weights=bond_weights_ag,
+ atom_graph=g.batched_atom_graph,
+ directed2undirected=g.directed2undirected,
+ )
+
+ if len(g.angle_bases) != 0 and bond_layer is not None:
+ bond_feas = bond_layer(
+ atom_feas=atom_feas,
+ bond_feas=bond_feas,
+ bond_weights=bond_weights_bg,
+ angle_feas=angle_feas,
+ bond_graph=g.batched_bond_graph,
+ )
+
+ if angle_layer is not None:
+ angle_feas = angle_layer(
+ atom_feas=atom_feas,
+ bond_feas=bond_feas,
+ angle_feas=angle_feas,
+ bond_graph=g.batched_bond_graph,
+ )
+ if idx == self.n_conv - 2:
+ if return_atom_feas is True:
+ prediction["atom_fea"] = paddle.split(
+ atom_feas, atoms_per_graph.tolist()
+ )
+ if site_wise:
+ magmom = paddle.abs(self.site_wise(atom_feas))
+ prediction["m"] = list(
+ paddle.split(
+ magmom.reshape([-1]), atoms_per_graph.tolist()
+ )
+ )
+
+ atom_feas = self.atom_conv_layers[-1](
+ atom_feas=atom_feas,
+ bond_feas=bond_feas,
+ bond_weights=bond_weights_ag,
+ atom_graph=g.batched_atom_graph,
+ directed2undirected=g.directed2undirected,
+ )
+ if self.readout_norm is not None:
+ atom_feas = self.readout_norm(atom_feas)
+
+ return atom_feas
+
+ def predict_structure_embedding(
+ self,
+ structure: Structure | Sequence[Structure],
+ task: PredTask = "efsm",
+ return_atom_feas: bool = False,
+ return_crystal_feas: bool = False,
+ batch_size: int = 100,
+ ) -> dict[str, Tensor]:
+ """Predict embeddings from pymatgen.core.Structure."""
+ assert (
+ self.graph_converter is not None
+ ), "self.graph_converter needs to be initialized first!"
+ if type(structure) == Structure:
+ graph = self.graph_converter(structure)
+ return self.predict_graph(
+ graph,
+ task=task,
+ return_atom_feas=return_atom_feas,
+ return_crystal_feas=return_crystal_feas,
+ batch_size=batch_size,
+ )
+ if type(structure) == list:
+ graphs = [self.graph_converter(i) for i in structure]
+ return self.predict_graph(
+ graphs,
+ task=task,
+ return_atom_feas=return_atom_feas,
+ return_crystal_feas=return_crystal_feas,
+ batch_size=batch_size,
+ )
+ raise Exception("input should either be a structure or list of structures!")
+
+ def predict_graph(
+ self,
+ graph: CrystalGraph | Sequence[CrystalGraph],
+ task: PredTask = "efsm",
+ return_atom_feas: bool = False,
+ return_crystal_feas: bool = False,
+ batch_size: int = 100,
+ ) -> dict[str, Tensor]:
+ """Predict from CrystalGraph."""
+ model_device = next(iter(self.parameters())).place
+ self.eval()
+ output = self.forward(
+ [graph.to(model_device)],
+ task=task,
+ return_atom_feas=return_atom_feas,
+ return_crystal_feas=return_crystal_feas,
+ )
+
+ return output
+
+ @staticmethod
+ def split(x: Tensor, n: Tensor) -> Sequence[Tensor]:
+ """Split a batched result Tensor into a list of Tensors."""
+ start = 0
+ result = []
+ for i in n:
+ result.append(x[start : start + i])
+ start += i
+ assert start == len(x), "Error: source tensor not correctly split!"
+ return result
+
+ def as_dict(self):
+ """Return the CHGNet weights and args in a dictionary."""
+ return {"state_dict": self.state_dict(), "model_args": self.model_args}
+
+ @classmethod
+ def from_dict(cls, dict, **kwargs):
+ """Build a CHGNet from a saved dictionary."""
+ chgnet = CHGNet(**dict["model_args"])
+ chgnet.set_state_dict(dict["state_dict"])
+ return chgnet
+
+ @classmethod
+ def from_file(cls, path, **kwargs):
+ """Build a CHGNet from a saved file."""
+ state = paddle.load(path)
+ return CHGNet.from_dict(state["model"], **kwargs)
+
+ @classmethod
+ def load(cls, model_name="MPtrj-efsm"):
+ """Load pretrained CHGNet."""
+ current_dir = os.path.dirname(os.path.abspath(__file__))
+ if model_name == "MPtrj-efsm":
+ return cls.from_file(
+ os.path.join(current_dir, "../pretrained/e30f77s348m32.pth.tar")
+ )
+ raise Exception("model_name not supported")
+
+
+@dataclass
+class BatchedGraph:
+ """Batched crystal graph for parallel computing."""
+
+ atomic_numbers: Tensor
+ bond_bases_ag: Tensor
+ bond_bases_bg: Tensor
+ angle_bases: Tensor
+ batched_atom_graph: Tensor
+ batched_bond_graph: Tensor
+ atom_owners: Tensor
+ directed2undirected: Tensor
+ atom_positions: Sequence[Tensor]
+ strains: Sequence[Tensor]
+ volumes: Sequence[Tensor]
+
+ @classmethod
+ def from_graphs(
+ cls,
+ graphs: Sequence[CrystalGraph],
+ bond_basis_expansion: nn.Layer,
+ angle_basis_expansion: nn.Layer,
+ compute_stress: bool = False,
+ ) -> BatchedGraph:
+ """Featurize and assemble a list of graphs."""
+ atomic_numbers, atom_positions = [], []
+ strains, volumes = [], []
+ bond_bases_ag, bond_bases_bg, angle_bases = [], [], []
+ batched_atom_graph, batched_bond_graph = [], []
+ directed2undirected = []
+ atom_owners = []
+ atom_offset_idx = 0
+ n_undirected = 0
+
+ for graph_idx, graph in enumerate(graphs):
+ n_atom = graph.atomic_number.shape[0]
+ atomic_numbers.append(graph.atomic_number)
+
+ if compute_stress:
+ strain = paddle.zeros([3, 3])
+ strain.stop_gradient = False
+ lattice = graph.lattice @ (
+ paddle.eye(3).cast(strain.dtype) + strain
+ )
+ else:
+ strain = None
+ lattice = graph.lattice
+ volumes.append(paddle.linalg.det(lattice))
+ strains.append(strain)
+
+ atom_cart_coords = graph.atom_frac_coord @ lattice
+ bond_basis_ag, bond_basis_bg, bond_vectors = bond_basis_expansion(
+ center=atom_cart_coords[graph.atom_graph[:, 0]],
+ neighbor=atom_cart_coords[graph.atom_graph[:, 1]],
+ undirected2directed=graph.undirected2directed,
+ image=graph.neighbor_image,
+ lattice=lattice,
+ )
+ atom_positions.append(atom_cart_coords)
+ bond_bases_ag.append(bond_basis_ag)
+ bond_bases_bg.append(bond_basis_bg)
+
+ batched_atom_graph.append(graph.atom_graph + atom_offset_idx)
+ directed2undirected.append(
+ graph.directed2undirected + n_undirected
+ )
+
+ if len(graph.bond_graph) != 0:
+ bond_vecs_i = paddle.index_select(
+ bond_vectors, 0, graph.bond_graph[:, 2]
+ )
+ bond_vecs_j = paddle.index_select(
+ bond_vectors, 0, graph.bond_graph[:, 4]
+ )
+ angle_basis = angle_basis_expansion(bond_vecs_i, bond_vecs_j)
+ angle_bases.append(angle_basis)
+
+ bond_graph = paddle.zeros(
+ [graph.bond_graph.shape[0], 3], dtype=graph.bond_graph.dtype
+ )
+ bond_graph[:, 0] = graph.bond_graph[:, 0] + atom_offset_idx
+ bond_graph[:, 1] = graph.bond_graph[:, 1] + n_undirected
+ bond_graph[:, 2] = graph.bond_graph[:, 3] + n_undirected
+ batched_bond_graph.append(bond_graph)
+
+ atom_owners.append(
+ paddle.ones([n_atom], dtype="float32") * graph_idx
+ )
+ atom_offset_idx += n_atom
+ n_undirected += len(bond_basis_ag)
+
+ atomic_numbers = paddle.concat(atomic_numbers, axis=0)
+ bond_bases_ag = paddle.concat(bond_bases_ag, axis=0)
+ bond_bases_bg = paddle.concat(bond_bases_bg, axis=0)
+ angle_bases = (
+ paddle.concat(angle_bases, axis=0)
+ if len(angle_bases) != 0
+ else paddle.to_tensor([])
+ )
+ batched_atom_graph = paddle.concat(batched_atom_graph, axis=0)
+ if batched_bond_graph != []:
+ batched_bond_graph = paddle.concat(batched_bond_graph, axis=0)
+ else:
+ batched_bond_graph = paddle.to_tensor([])
+ atom_owners = (
+ paddle.concat(atom_owners, axis=0)
+ .cast("int32")
+ )
+ if atomic_numbers.place.is_gpu_place():
+ atom_owners = atom_owners.to(atomic_numbers.place)
+ directed2undirected = paddle.concat(directed2undirected, axis=0)
+ volumes = paddle.to_tensor(
+ volumes, dtype=datatype
+ )
+
+ return cls(
+ atomic_numbers=atomic_numbers,
+ bond_bases_ag=bond_bases_ag,
+ bond_bases_bg=bond_bases_bg,
+ angle_bases=angle_bases,
+ batched_atom_graph=batched_atom_graph,
+ batched_bond_graph=batched_bond_graph,
+ atom_owners=atom_owners,
+ directed2undirected=directed2undirected,
+ atom_positions=atom_positions,
+ strains=strains,
+ volumes=volumes,
+ )
diff --git a/ppmat/models/matterchat/chgnet/utils/__init__.py b/ppmat/models/matterchat/chgnet/utils/__init__.py
new file mode 100644
index 00000000..ddd1e061
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/utils/__init__.py
@@ -0,0 +1,28 @@
+# 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.utils.common_utils import AverageMeter, mae, mkdir, read_json, write_json
+from ppmat.models.matterchat.chgnet.utils.vasp_utils import parse_vasp_dir, solve_charge_by_mag
+
+__all__ = [
+ "AverageMeter",
+ "mae",
+ "read_json",
+ "write_json",
+ "mkdir",
+ "parse_vasp_dir",
+ "solve_charge_by_mag",
+]
diff --git a/ppmat/models/matterchat/chgnet/utils/common_utils.py b/ppmat/models/matterchat/chgnet/utils/common_utils.py
new file mode 100644
index 00000000..fea8ee33
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/utils/common_utils.py
@@ -0,0 +1,33 @@
+# 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
+
+import paddle
+
+from ppmat.utils.io import read_json, write_json
+from ppmat.utils.misc import AverageMeter
+
+
+def mae(prediction: paddle.Tensor, target: paddle.Tensor) -> paddle.Tensor:
+ """Computes the mean absolute error between prediction and target."""
+ return paddle.mean(paddle.abs(target - prediction))
+
+
+def mkdir(path: str):
+ """Make directory."""
+ os.makedirs(path, exist_ok=True)
+ return path
diff --git a/ppmat/models/matterchat/chgnet/utils/vasp_utils.py b/ppmat/models/matterchat/chgnet/utils/vasp_utils.py
new file mode 100644
index 00000000..13b0ce07
--- /dev/null
+++ b/ppmat/models/matterchat/chgnet/utils/vasp_utils.py
@@ -0,0 +1,174 @@
+# 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 re
+from typing import TYPE_CHECKING
+
+from monty.io import reverse_readfile
+from pymatgen.io.vasp.outputs import Oszicar, Vasprun
+
+from ppmat.utils import logger
+
+if TYPE_CHECKING:
+ from pymatgen.core import Structure
+
+
+def parse_vasp_dir(file_root):
+ """Parse VASP output files into structures and labels."""
+ try:
+ oszicar = Oszicar(file_root + "/OSZICAR")
+ vasprun_orig = Vasprun(file_root + "/vasprun.xml", exception_on_bad_xml=False)
+ outcar_filename = file_root + "/OUTCAR"
+ except Exception:
+ oszicar = Oszicar(file_root + "/OSZICAR.gz")
+ vasprun_orig = Vasprun(
+ file_root + "/vasprun.xml.gz", exception_on_bad_xml=False
+ )
+ outcar_filename = file_root + "/OUTCAR.gz"
+
+ charge = []
+ mag_x = []
+ mag_y = []
+ mag_z = []
+ header = []
+ all_lines = []
+
+ for line in reverse_readfile(outcar_filename):
+ clean = line.strip()
+ all_lines.append(clean)
+
+ all_lines.reverse()
+ read_charge = False
+ read_mag_x = False
+ read_mag_y = False
+ read_mag_z = False
+ mag_x_all = []
+ ion_step_count = 0
+
+ for clean in all_lines:
+ if "magnetization (x)" in clean:
+ ion_step_count += 1
+ if read_charge or read_mag_x or read_mag_y or read_mag_z:
+ if clean.startswith("# of ion"):
+ header = re.split(r"\s{2,}", clean.strip())
+ header.pop(0)
+ else:
+ m = re.match(r"\s*(\d+)\s+(([\d\.\-]+)\s+)+", clean)
+ if m:
+ toks = [float(i) for i in re.findall(r"[\d\.\-]+", clean)]
+ toks.pop(0)
+ if read_charge:
+ charge.append(dict(zip(header, toks)))
+ elif read_mag_x:
+ mag_x.append(dict(zip(header, toks)))
+ elif read_mag_y:
+ mag_y.append(dict(zip(header, toks)))
+ elif read_mag_z:
+ mag_z.append(dict(zip(header, toks)))
+ elif clean.startswith("tot"):
+ if ion_step_count == (len(mag_x_all) + 1):
+ mag_x_all.append(mag_x)
+ read_charge = False
+ read_mag_x = False
+ read_mag_y = False
+ read_mag_z = False
+ if clean == "total charge":
+ read_charge = True
+ read_mag_x, read_mag_y, read_mag_z = False, False, False
+ elif clean == "magnetization (x)":
+ mag_x = []
+ read_mag_x = True
+ read_charge, read_mag_y, read_mag_z = False, False, False
+ elif clean == "magnetization (y)":
+ mag_y = []
+ read_mag_y = True
+ read_charge, read_mag_x, read_mag_z = False, False, False
+ elif clean == "magnetization (z)":
+ mag_z = []
+ read_mag_z = True
+ read_charge, read_mag_x, read_mag_y = False, False, False
+ elif re.search("electrostatic", clean):
+ read_charge, read_mag_x, read_mag_y, read_mag_z = (
+ False,
+ False,
+ False,
+ False,
+ )
+
+ if len(oszicar.ionic_steps) == len(mag_x_all):
+ logger.warning("Unfinished OUTCAR")
+ mag_x_all = mag_x_all
+ elif len(oszicar.ionic_steps) == (len(mag_x_all) - 1):
+ mag_x_all.pop(-1)
+
+ n_atoms = len(vasprun_orig.ionic_steps[0]["structure"])
+ dataset = {
+ "structures": [i["structure"] for i in vasprun_orig.ionic_steps],
+ "uncorrected_total_energies": [
+ i["e_0_energy"] for i in vasprun_orig.ionic_steps
+ ],
+ "energies_per_atom": [
+ i["e_0_energy"] / n_atoms for i in vasprun_orig.ionic_steps
+ ],
+ "forces": [i["forces"] for i in vasprun_orig.ionic_steps],
+ "magmoms": [[i["tot"] for i in j] for j in mag_x_all],
+ }
+ if "stress" in vasprun_orig.ionic_steps[0]:
+ dataset["stress"] = [i["stress"] for i in vasprun_orig.ionic_steps]
+ else:
+ dataset["stress"] = None
+
+ return dataset
+
+
+def solve_charge_by_mag(
+ structure: Structure,
+ default_ox: dict[str, float] | None = None,
+ ox_ranges: dict[str, dict[tuple[float, float], int]] | None = None,
+):
+ """Solve oxidation states by magmom."""
+ ox_list = []
+ solved_ox = True
+ default_ox = default_ox or {"Li": 1, "O": -2}
+ ox_ranges = ox_ranges or {
+ "Mn": {(0.5, 1.5): 2, (1.5, 2.5): 3, (2.5, 3.5): 4, (3.5, 4.2): 3, (4.2, 5): 2}
+ }
+
+ mag_key = (
+ "final_magmom" if "final_magmom" in structure.site_properties else "magmom"
+ )
+
+ mag = structure.site_properties[mag_key]
+
+ for site_i, site in enumerate(structure.sites):
+ assigned = False
+ if site.species_string in ox_ranges:
+ for (minmag, maxmag), magox in ox_ranges[site.species_string].items():
+ if mag[site_i] >= minmag and mag[site_i] < maxmag:
+ ox_list.append(magox)
+ assigned = True
+ break
+ elif site.species_string in default_ox:
+ ox_list.append(default_ox[site.species_string])
+ assigned = True
+ if not assigned:
+ solved_ox = False
+
+ if solved_ox:
+ logger.debug(ox_list)
+ structure.add_oxidation_state_by_site(ox_list)
+ return structure
+ return None
diff --git a/ppmat/models/matterchat/mistral/__init__.py b/ppmat/models/matterchat/mistral/__init__.py
new file mode 100644
index 00000000..1ad9550d
--- /dev/null
+++ b/ppmat/models/matterchat/mistral/__init__.py
@@ -0,0 +1,27 @@
+# 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 ppmat.models.matterchat.mistral.configuration_mistral import MistralConfig
+from ppmat.models.matterchat.mistral.modeling_mistral import (
+ MistralForCausalLM,
+ MistralModel,
+ DynamicCache,
+)
+
+__all__ = [
+ "DynamicCache",
+ "MistralConfig",
+ "MistralForCausalLM",
+ "MistralModel",
+]
diff --git a/ppmat/models/matterchat/mistral/configuration_mistral.py b/ppmat/models/matterchat/mistral/configuration_mistral.py
new file mode 100644
index 00000000..c6b94f64
--- /dev/null
+++ b/ppmat/models/matterchat/mistral/configuration_mistral.py
@@ -0,0 +1,79 @@
+# 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.
+
+"""Mistral model configuration for PaddlePaddle."""
+
+
+class MistralConfig:
+ """Configuration class for Mistral model."""
+
+ model_type = "mistral"
+
+ def __init__(
+ self,
+ vocab_size=32769,
+ hidden_size=4096,
+ intermediate_size=14336,
+ num_hidden_layers=32,
+ num_attention_heads=32,
+ num_key_value_heads=8,
+ head_dim=None,
+ hidden_act="silu",
+ max_position_embeddings=4096 * 32,
+ initializer_range=0.02,
+ rms_norm_eps=1e-5,
+ use_cache=True,
+ pad_token_id=None,
+ bos_token_id=1,
+ eos_token_id=2,
+ tie_word_embeddings=False,
+ rope_theta=1000000.0,
+ sliding_window=4096,
+ attention_dropout=0.0,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ self.max_position_embeddings = max_position_embeddings
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.sliding_window = sliding_window
+ self.head_dim = head_dim or hidden_size // num_attention_heads
+
+ if num_key_value_heads is None:
+ num_key_value_heads = num_attention_heads
+
+ self.num_key_value_heads = num_key_value_heads
+ self.hidden_act = hidden_act
+ self.initializer_range = initializer_range
+ self.rms_norm_eps = rms_norm_eps
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.attention_dropout = attention_dropout
+ self.pad_token_id = pad_token_id
+ self.bos_token_id = bos_token_id
+ self.eos_token_id = eos_token_id
+ self.tie_word_embeddings = tie_word_embeddings
+
+ # SDPA not supported for autoregressive generation
+ self._attn_implementation = kwargs.pop("attn_implementation", "eager")
+
+ self.output_attentions = kwargs.pop("output_attentions", False)
+ self.output_hidden_states = kwargs.pop("output_hidden_states", False)
+ self.use_return_dict = kwargs.pop("use_return_dict", True)
+ self.num_labels = kwargs.pop("num_labels", 2)
+
+ for key, value in kwargs.items():
+ setattr(self, key, value)
diff --git a/ppmat/models/matterchat/mistral/modeling_mistral.py b/ppmat/models/matterchat/mistral/modeling_mistral.py
new file mode 100644
index 00000000..c6ad1dee
--- /dev/null
+++ b/ppmat/models/matterchat/mistral/modeling_mistral.py
@@ -0,0 +1,1060 @@
+# 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.
+
+"""PaddlePaddle Mistral-7B model implementation.
+
+Cannot reuse Paddle built-in layers: GQA, RoPE, RMSNorm, Sliding Window Attention,
+and Dynamic KV Cache have no PaddlePaddle equivalent.
+"""
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+from ppmat.models.matterchat.mistral.configuration_mistral import MistralConfig
+from ppmat.models.matterchat.utils.activations import get_activation
+from ppmat.models.matterchat.utils.weight_init import default_init_weights
+from ppmat.utils import logger
+
+
+class _SubscriptableOutput:
+ """Base class for model outputs that support tuple-like indexing."""
+
+ def __getitem__(self, index):
+ fields = [f for f in self.__dataclass_fields__ if not f.startswith('_')]
+ return getattr(self, fields[index])
+
+ def __len__(self):
+ return len([f for f in self.__dataclass_fields__ if not f.startswith('_')])
+
+
+@dataclass
+class BaseModelOutputWithPast(_SubscriptableOutput):
+ last_hidden_state: paddle.Tensor = None
+ past_key_values: Optional[Tuple] = None
+ hidden_states: Optional[Tuple] = None
+ attentions: Optional[Tuple] = None
+
+
+@dataclass
+class CausalLMOutputWithPast(_SubscriptableOutput):
+ loss: Optional[paddle.Tensor] = None
+ logits: paddle.Tensor = None
+ past_key_values: Optional[Tuple] = None
+ hidden_states: Optional[Tuple] = None
+ attentions: Optional[Tuple] = None
+
+
+class DynamicCache:
+ """Simple dynamic KV cache for autoregressive generation."""
+
+ def __init__(self):
+ self._seen_tokens = 0
+ self.key_cache: List[paddle.Tensor] = []
+ self.value_cache: List[paddle.Tensor] = []
+
+ def __getitem__(self, layer_idx):
+ if layer_idx < len(self.key_cache):
+ return (self.key_cache[layer_idx], self.value_cache[layer_idx])
+ raise IndexError(f"layer_idx {layer_idx} out of range")
+
+ def update(
+ self,
+ key_states: paddle.Tensor,
+ value_states: paddle.Tensor,
+ layer_idx: int,
+ cache_kwargs: Optional[dict] = None,
+ ):
+ if layer_idx == 0:
+ self._seen_tokens += key_states.shape[-2]
+
+ if len(self.key_cache) <= layer_idx:
+ self.key_cache.append(key_states)
+ self.value_cache.append(value_states)
+ else:
+ self.key_cache[layer_idx] = paddle.concat(
+ [self.key_cache[layer_idx], key_states], axis=-2
+ )
+ self.value_cache[layer_idx] = paddle.concat(
+ [self.value_cache[layer_idx], value_states], axis=-2
+ )
+
+ return self.key_cache[layer_idx], self.value_cache[layer_idx]
+
+ def get_seq_length(self, layer_idx=0):
+ if len(self.key_cache) <= layer_idx:
+ return 0
+ return self.key_cache[layer_idx].shape[-2]
+
+ def get_max_length(self):
+ return None
+
+ @classmethod
+ def from_legacy_cache(cls, past_key_values):
+ cache = cls()
+ if past_key_values is not None:
+ for layer_past in past_key_values:
+ cache.key_cache.append(layer_past[0])
+ cache.value_cache.append(layer_past[1])
+ cache._seen_tokens = layer_past[0].shape[-2]
+ return cache
+
+ def to_legacy_cache(self):
+ legacy_cache = ()
+ for layer_idx in range(len(self.key_cache)):
+ legacy_cache += ((self.key_cache[layer_idx], self.value_cache[layer_idx]),)
+ return legacy_cache
+
+
+class MistralRMSNorm(nn.Layer):
+ def __init__(self, hidden_size, eps=1e-6):
+ super().__init__()
+ self.weight = self.create_parameter(
+ shape=[hidden_size],
+ default_initializer=nn.initializer.Constant(1.0),
+ )
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states):
+ input_dtype = hidden_states.dtype
+ hidden_states = hidden_states.cast(paddle.float32)
+ variance = paddle.pow(hidden_states, 2).mean(axis=-1, keepdim=True)
+ hidden_states = hidden_states * paddle.rsqrt(variance + self.variance_epsilon)
+ return self.weight * hidden_states.cast(input_dtype)
+
+
+class MistralRotaryEmbedding(nn.Layer):
+ def __init__(self, dim, max_position_embeddings=2048, base=10000):
+ super().__init__()
+ self.dim = dim
+ self.max_position_embeddings = max_position_embeddings
+ self.base = base
+ inv_freq = 1.0 / (
+ self.base
+ ** (
+ paddle.arange(0, self.dim, 2, dtype=paddle.int64).cast(paddle.float32)
+ / self.dim
+ )
+ )
+ self.register_buffer("inv_freq", inv_freq, persistable=False)
+
+ @paddle.no_grad()
+ def forward(self, x, position_ids):
+ inv_freq_expanded = (
+ self.inv_freq[None, :, None]
+ .cast(paddle.float32)
+ .expand([position_ids.shape[0], -1, 1])
+ )
+ position_ids_expanded = position_ids[:, None, :].cast(paddle.float32)
+ freqs = (inv_freq_expanded.cast(paddle.float32) @ position_ids_expanded.cast(paddle.float32)).transpose([0, 2, 1])
+ emb = paddle.concat((freqs, freqs), axis=-1)
+ cos = emb.cos()
+ sin = emb.sin()
+ return cos.cast(x.dtype), sin.cast(x.dtype)
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return paddle.concat((-x2, x1), axis=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+ cos = cos.unsqueeze(axis=unsqueeze_dim)
+ sin = sin.unsqueeze(axis=unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+class MistralMLP(nn.Layer):
+ def __init__(self, config):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias_attr=False)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias_attr=False)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias_attr=False)
+ self.act_fn = get_activation(config.hidden_act)
+
+ def forward(self, hidden_state):
+ return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
+
+
+def repeat_kv(hidden_states: paddle.Tensor, n_rep: int) -> paddle.Tensor:
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+ if n_rep == 1:
+ return hidden_states
+ hidden_states = hidden_states[:, :, None, :, :].expand(
+ [batch, num_key_value_heads, n_rep, slen, head_dim]
+ )
+ return hidden_states.reshape([batch, num_key_value_heads * n_rep, slen, head_dim])
+
+
+class MistralAttention(nn.Layer):
+ def __init__(self, config: MistralConfig, layer_idx: Optional[int] = None):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ if layer_idx is None:
+ logger.warning(
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended."
+ )
+
+ self.attention_dropout = config.attention_dropout
+ self.hidden_size = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = self.hidden_size // self.num_heads
+ self.num_key_value_heads = config.num_key_value_heads
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+ self.max_position_embeddings = config.max_position_embeddings
+ self.rope_theta = config.rope_theta
+ self.is_causal = True
+
+ if (self.head_dim * self.num_heads) != self.hidden_size:
+ raise ValueError(
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+ f" and `num_heads`: {self.num_heads})."
+ )
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias_attr=False)
+ self.k_proj = nn.Linear(
+ self.hidden_size, self.num_key_value_heads * self.head_dim, bias_attr=False
+ )
+ self.v_proj = nn.Linear(
+ self.hidden_size, self.num_key_value_heads * self.head_dim, bias_attr=False
+ )
+ self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias_attr=False)
+
+ self.rotary_emb = MistralRotaryEmbedding(
+ self.head_dim,
+ max_position_embeddings=self.max_position_embeddings,
+ base=self.rope_theta,
+ )
+
+ def forward(
+ self,
+ hidden_states: paddle.Tensor,
+ attention_mask: Optional[paddle.Tensor] = None,
+ position_ids: Optional[paddle.Tensor] = None,
+ past_key_value: Optional[DynamicCache] = None,
+ output_attentions: bool = False,
+ use_cache: bool = False,
+ cache_position: Optional[paddle.Tensor] = None,
+ ) -> Tuple[paddle.Tensor, Optional[paddle.Tensor], Optional[Tuple[paddle.Tensor]]]:
+ bsz, q_len, _ = hidden_states.shape
+
+ query_states = self.q_proj(hidden_states)
+ key_states = self.k_proj(hidden_states)
+ value_states = self.v_proj(hidden_states)
+
+ query_states = query_states.reshape([bsz, q_len, self.num_heads, self.head_dim]).transpose([0, 2, 1, 3])
+ key_states = key_states.reshape(
+ [bsz, q_len, self.num_key_value_heads, self.head_dim]
+ ).transpose([0, 2, 1, 3])
+ value_states = value_states.reshape(
+ [bsz, q_len, self.num_key_value_heads, self.head_dim]
+ ).transpose([0, 2, 1, 3])
+
+ cos, sin = self.rotary_emb(value_states, position_ids)
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_value.update(
+ key_states, value_states, self.layer_idx, cache_kwargs
+ )
+
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+ # softmax in fp32 for numerical stability, then cast back
+ attn_weights = paddle.matmul(query_states, key_states.transpose([0, 1, 3, 2])) / math.sqrt(
+ self.head_dim
+ )
+
+ if attention_mask is not None:
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+ attn_weights = attn_weights + causal_mask
+
+ attn_weights = F.softmax(attn_weights, axis=-1, dtype=paddle.float32).cast(query_states.dtype)
+ attn_output = paddle.matmul(attn_weights, value_states)
+
+ if attn_output.shape != [bsz, self.num_heads, q_len, self.head_dim]:
+ raise ValueError(
+ f"`attn_output` should be of size {[bsz, self.num_heads, q_len, self.head_dim]}, but is"
+ f" {attn_output.shape}"
+ )
+
+ attn_output = attn_output.transpose([0, 2, 1, 3])
+ attn_output = attn_output.reshape([bsz, q_len, -1])
+ attn_output = self.o_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value
+
+
+class MistralSdpaAttention(MistralAttention):
+ """Mistral attention using paddle.nn.functional.scaled_dot_product_attention."""
+
+ def forward(
+ self,
+ hidden_states: paddle.Tensor,
+ attention_mask: Optional[paddle.Tensor] = None,
+ position_ids: Optional[paddle.Tensor] = None,
+ past_key_value: Optional["DynamicCache"] = None,
+ output_attentions: bool = False,
+ use_cache: bool = False,
+ cache_position: Optional[paddle.Tensor] = None,
+ **kwargs,
+ ) -> Tuple[paddle.Tensor, Optional[paddle.Tensor], Optional[Tuple[paddle.Tensor]]]:
+ if output_attentions:
+ return super().forward(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_value,
+ output_attentions=output_attentions,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ )
+
+ bsz, q_len, _ = hidden_states.shape
+
+ query_states = self.q_proj(hidden_states)
+ key_states = self.k_proj(hidden_states)
+ value_states = self.v_proj(hidden_states)
+
+ query_states = query_states.reshape(
+ [bsz, q_len, self.num_heads, self.head_dim]
+ ).transpose([0, 2, 1, 3])
+ key_states = key_states.reshape(
+ [bsz, q_len, self.num_key_value_heads, self.head_dim]
+ ).transpose([0, 2, 1, 3])
+ value_states = value_states.reshape(
+ [bsz, q_len, self.num_key_value_heads, self.head_dim]
+ ).transpose([0, 2, 1, 3])
+
+ cos, sin = self.rotary_emb(value_states, position_ids)
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_value.update(
+ key_states, value_states, self.layer_idx, cache_kwargs
+ )
+
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+ # Use SDPA for prefill; fall back to manual attention for generation
+ # because Paddle SDPA does not support q_len != k_len.
+ if q_len == key_states.shape[2]:
+ attn_output = paddle.nn.functional.scaled_dot_product_attention(
+ query_states,
+ key_states,
+ value_states,
+ attn_mask=None,
+ dropout_p=self.attention_dropout if self.training else 0.0,
+ is_causal=True,
+ training=self.training,
+ enable_gqa=False,
+ )
+ else:
+ attn_weights = paddle.matmul(query_states, key_states.transpose([0, 1, 3, 2])) / math.sqrt(
+ self.head_dim
+ )
+ attn_weights = F.softmax(attn_weights, axis=-1, dtype=paddle.float32).cast(query_states.dtype)
+ attn_output = paddle.matmul(attn_weights, value_states)
+
+ attn_output = attn_output.transpose([0, 2, 1, 3])
+ attn_output = attn_output.reshape([bsz, q_len, -1])
+ attn_output = self.o_proj(attn_output)
+
+ return attn_output, None, past_key_value
+
+
+MISTRAL_ATTENTION_CLASSES = {
+ "eager": MistralAttention,
+ "sdpa": MistralSdpaAttention,
+}
+
+
+class MistralDecoderLayer(nn.Layer):
+ def __init__(self, config: MistralConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ attn_cls = MISTRAL_ATTENTION_CLASSES.get(
+ getattr(config, "_attn_implementation", "eager"), MistralAttention
+ )
+ self.self_attn = attn_cls(config=config, layer_idx=layer_idx)
+ self.mlp = MistralMLP(config)
+ self.input_layernorm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: paddle.Tensor,
+ attention_mask: Optional[paddle.Tensor] = None,
+ position_ids: Optional[paddle.Tensor] = None,
+ past_key_value: Optional[DynamicCache] = None,
+ output_attentions: Optional[bool] = False,
+ use_cache: Optional[bool] = False,
+ cache_position: Optional[paddle.Tensor] = None,
+ **kwargs,
+ ) -> Tuple[paddle.Tensor, Optional[Tuple[paddle.Tensor, paddle.Tensor]]]:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_value,
+ output_attentions=output_attentions,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ outputs = (hidden_states,)
+ if output_attentions:
+ outputs += (self_attn_weights,)
+ if use_cache:
+ outputs += (present_key_value,)
+
+ return outputs
+
+
+class MistralPreTrainedModel(nn.Layer):
+ config_class = MistralConfig
+ base_model_prefix = "model"
+
+ def __init__(self, config: MistralConfig):
+ super().__init__()
+ self.config = config
+
+ def _init_weights(self, module):
+ default_init_weights(module, self.config.initializer_range)
+
+ def init_weights(self):
+ self.apply(self._init_weights)
+
+
+def _apply_repetition_penalty(logits, gen_ids, penalty):
+ """Vectorized repetition penalty following HF GenerationMixin."""
+ if gen_ids is None or len(gen_ids) == 0:
+ return logits
+ unique_ids = paddle.unique(gen_ids)
+ gathered = logits[:, unique_ids]
+ penalized = paddle.where(
+ gathered > 0, gathered / penalty, gathered * penalty
+ )
+ logits = logits.clone()
+ logits[:, unique_ids] = penalized
+ return logits
+
+
+def _create_causal_mask(
+ attention_mask: Optional[paddle.Tensor],
+ input_tensor: paddle.Tensor,
+ cache_position: paddle.Tensor,
+ past_key_values: Optional[DynamicCache],
+):
+ dtype = input_tensor.dtype
+ min_dtype = paddle.finfo(dtype).min
+ sequence_length = input_tensor.shape[1]
+
+ past_seen_tokens = cache_position[0] if past_key_values is not None else 0
+ target_length = (
+ attention_mask.shape[-1]
+ if attention_mask is not None and isinstance(attention_mask, paddle.Tensor)
+ else past_seen_tokens + sequence_length + 1
+ )
+
+ if attention_mask is not None and len(attention_mask.shape) == 4:
+ causal_mask = attention_mask
+ else:
+ causal_mask = paddle.full(
+ [sequence_length, target_length], fill_value=min_dtype, dtype=dtype
+ )
+ exclude_mask = paddle.arange(target_length) > cache_position.reshape([-1, 1])
+ sliding_window = getattr(input_tensor, "_sliding_window", None)
+ if sliding_window is not None:
+ if sequence_length > sliding_window:
+ exclude_mask = paddle.bitwise_or(
+ exclude_mask,
+ paddle.arange(target_length)
+ <= (cache_position.reshape([-1, 1]) - sliding_window),
+ )
+ causal_mask = causal_mask * exclude_mask.cast(dtype)
+ causal_mask = causal_mask[None, None, :, :].expand(
+ [input_tensor.shape[0], 1, -1, -1]
+ )
+ if attention_mask is not None:
+ causal_mask = causal_mask.clone()
+ if len(attention_mask.shape) == 2:
+ mask_length = attention_mask.shape[-1]
+ attention_mask_float = attention_mask.cast(dtype)
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask_float[:, None, None, :]
+ padding_mask = padding_mask == 0
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+ padding_mask, min_dtype
+ )
+
+ return causal_mask
+
+
+class MistralModel(MistralPreTrainedModel):
+ def __init__(self, config: MistralConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=self.padding_idx)
+ self.layers = nn.LayerList(
+ [MistralDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.init_weights()
+
+ def get_input_embeddings(self):
+ return self.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.embed_tokens = value
+
+ def forward(
+ self,
+ input_ids: paddle.Tensor = None,
+ attention_mask: Optional[paddle.Tensor] = None,
+ position_ids: Optional[paddle.Tensor] = None,
+ past_key_values: Optional[Union[DynamicCache, List[paddle.Tensor]]] = None,
+ inputs_embeds: Optional[paddle.Tensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ cache_position: Optional[paddle.Tensor] = None,
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ if (input_ids is None) == (inputs_embeds is None):
+ raise ValueError(
+ "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+ )
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ return_legacy_cache = False
+ if use_cache and not isinstance(past_key_values, DynamicCache):
+ if past_key_values is not None:
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+ else:
+ past_key_values = DynamicCache()
+ return_legacy_cache = True
+
+ if cache_position is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ cache_position = paddle.arange(
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1]
+ )
+
+ if position_ids is None:
+ position_ids = cache_position.unsqueeze(0)
+
+ causal_mask = _create_causal_mask(
+ attention_mask, inputs_embeds, cache_position, past_key_values
+ )
+
+ hidden_states = inputs_embeds
+
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+ next_decoder_cache = None
+
+ for decoder_layer in self.layers:
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_values,
+ output_attentions=output_attentions,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ )
+
+ hidden_states = layer_outputs[0]
+ if use_cache:
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+ if output_attentions:
+ all_self_attns += (layer_outputs[1],)
+
+ hidden_states = self.norm(hidden_states)
+
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ next_cache = next_decoder_cache if use_cache else None
+ if return_legacy_cache and next_cache is not None:
+ if isinstance(next_cache, DynamicCache):
+ next_cache = next_cache.to_legacy_cache()
+
+ if not return_dict:
+ return tuple(
+ v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None
+ )
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ )
+
+
+class MistralForCausalLM(MistralPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = MistralModel(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias_attr=False)
+ self.init_weights()
+
+ def get_input_embeddings(self):
+ return self.model.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.embed_tokens = value
+
+ def get_output_embeddings(self):
+ return self.lm_head
+
+ def set_output_embeddings(self, new_embeddings):
+ self.lm_head = new_embeddings
+
+ def set_decoder(self, decoder):
+ self.model = decoder
+
+ def get_decoder(self):
+ return self.model
+
+ def forward(
+ self,
+ input_ids: paddle.Tensor = None,
+ attention_mask: Optional[paddle.Tensor] = None,
+ position_ids: Optional[paddle.Tensor] = None,
+ past_key_values: Optional[Union[DynamicCache, List[paddle.Tensor]]] = None,
+ inputs_embeds: Optional[paddle.Tensor] = None,
+ labels: Optional[paddle.Tensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ cache_position: Optional[paddle.Tensor] = None,
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ cache_position=cache_position,
+ )
+
+ hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
+ logits = self.lm_head(hidden_states)
+ logits = logits.cast(paddle.float32)
+
+ loss = None
+ if labels is not None:
+ shift_logits = logits[..., :-1, :]
+ shift_labels = labels[..., 1:]
+ shift_logits = shift_logits.reshape([-1, self.config.vocab_size])
+ shift_labels = shift_labels.reshape([-1])
+ loss_fct = nn.CrossEntropyLoss()
+ loss = loss_fct(shift_logits, shift_labels)
+
+ if not return_dict:
+ output = (logits,) + outputs[1:]
+ return (loss,) + output if loss is not None else output
+
+ past_kv = outputs.past_key_values if hasattr(outputs, "past_key_values") else None
+ hidden = outputs.hidden_states if hasattr(outputs, "hidden_states") else None
+ attns = outputs.attentions if hasattr(outputs, "attentions") else None
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=past_kv,
+ hidden_states=hidden,
+ attentions=attns,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ cache_position=None,
+ position_ids=None,
+ use_cache=True,
+ **kwargs,
+ ):
+ if past_key_values is not None:
+ if inputs_embeds is not None:
+ input_ids = input_ids[:, -cache_position.shape[0] :]
+ elif input_ids.shape[1] != cache_position.shape[0]:
+ input_ids = input_ids[:, cache_position]
+
+ if attention_mask is not None and position_ids is None:
+ position_ids = attention_mask.cast(paddle.int64).cumsum(axis=-1) - 1
+ position_ids = position_ids.masked_fill(attention_mask == 0, 1)
+ if past_key_values:
+ position_ids = position_ids[:, -input_ids.shape[1] :]
+
+ if inputs_embeds is not None and cache_position[0] == 0:
+ model_inputs = {"inputs_embeds": inputs_embeds}
+ else:
+ model_inputs = {"input_ids": input_ids}
+
+ model_inputs.update(
+ {
+ "position_ids": position_ids,
+ "cache_position": cache_position,
+ "past_key_values": past_key_values,
+ "use_cache": use_cache,
+ "attention_mask": attention_mask,
+ }
+ )
+ return model_inputs
+
+ def _generate_beam_search(
+ self,
+ inputs_embeds,
+ attention_mask,
+ max_length=256,
+ min_length=1,
+ num_beams=5,
+ repetition_penalty=1.5,
+ length_penalty=1.0,
+ temperature=1.0,
+ _suppress_token_id=None,
+ ):
+ """Beam search generation following HF GenerationMixin pattern."""
+ seq_len = inputs_embeds.shape[1]
+ eos_token_id = getattr(self.config, "eos_token_id", 2)
+ vocab_size = self.config.vocab_size
+
+ past_key_values = DynamicCache()
+ outputs = self.model(
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ use_cache=True,
+ cache_position=paddle.arange(seq_len),
+ )
+ logits = self.lm_head(outputs[0][:, -1:, :])[:, 0, :]
+ logits = logits / max(temperature, 1e-8)
+
+ if _suppress_token_id is not None:
+ logits[:, _suppress_token_id] = float("-inf")
+
+ log_probs = F.log_softmax(logits, axis=-1)
+ topk_scores, topk_ids = paddle.topk(log_probs[0], num_beams)
+
+ beam_tokens = topk_ids.unsqueeze(1)
+ beam_scores = topk_scores
+
+ # Expand KV cache from batch=1 to batch=num_beams
+ # repeat_interleave has no float16 GPU kernel; cast through float32.
+ for layer_idx in range(len(past_key_values.key_cache)):
+ k_cache = past_key_values.key_cache[layer_idx]
+ v_cache = past_key_values.value_cache[layer_idx]
+ past_key_values.key_cache[layer_idx] = paddle.repeat_interleave(
+ k_cache.cast(paddle.float32), num_beams, axis=0
+ ).cast(k_cache.dtype)
+ past_key_values.value_cache[layer_idx] = paddle.repeat_interleave(
+ v_cache.cast(paddle.float32), num_beams, axis=0
+ ).cast(v_cache.dtype)
+
+ attention_mask = paddle.repeat_interleave(attention_mask, num_beams, axis=0)
+
+ finished_hypos = []
+ max_new_tokens = max_length - seq_len
+
+ for step in range(1, max_new_tokens):
+ num_active = beam_tokens.shape[0]
+
+ attention_mask = paddle.concat(
+ [
+ attention_mask,
+ paddle.ones([num_active, 1], dtype=paddle.int64),
+ ],
+ axis=-1,
+ )
+
+ cache_position = paddle.arange(seq_len + step - 1, seq_len + step)
+
+ step_outputs = self.model(
+ input_ids=beam_tokens[:, -1:],
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ use_cache=True,
+ cache_position=cache_position,
+ )
+
+ logits = self.lm_head(step_outputs[0][:, -1:, :])[:, 0, :]
+ logits = logits / max(temperature, 1e-8)
+
+ if _suppress_token_id is not None:
+ logits[:, _suppress_token_id] = float("-inf")
+
+ # Per-beam repetition penalty
+ if repetition_penalty != 1.0:
+ for bi in range(num_active):
+ ids = paddle.unique(beam_tokens[bi])
+ gathered = logits[bi, ids]
+ penalized = paddle.where(
+ gathered > 0, gathered / repetition_penalty,
+ gathered * repetition_penalty,
+ )
+ logits[bi] = logits[bi].clone()
+ logits[bi, ids] = penalized
+
+ next_log_probs = F.log_softmax(logits, axis=-1)
+ candidate_scores = beam_scores.unsqueeze(-1) + next_log_probs
+
+ # Suppress EOS below min_length
+ if step < min_length:
+ candidate_scores[:, eos_token_id] = float("-inf")
+
+ n_select = num_beams + len(finished_hypos)
+ flat_scores = candidate_scores.reshape([-1])
+ topk_scores, topk_indices = paddle.topk(flat_scores, min(n_select, flat_scores.shape[0]))
+
+ beam_indices = topk_indices // vocab_size
+ token_indices = topk_indices % vocab_size
+
+ new_tokens_list = []
+ new_scores_list = []
+ new_beam_map = []
+
+ for i in range(topk_scores.shape[0]):
+ bi = beam_indices[i].item()
+ ti = token_indices[i].item()
+ sc = topk_scores[i].item()
+
+ if ti == eos_token_id:
+ seq = paddle.concat([beam_tokens[bi], token_indices[i:i + 1]])
+ length = seq.shape[0]
+ lp = ((5.0 + length) / (5.0 + 1.0)) ** length_penalty
+ finished_hypos.append((sc / lp, seq))
+ continue
+
+ new_seq = paddle.concat([beam_tokens[bi], token_indices[i:i + 1]])
+ new_tokens_list.append(new_seq)
+ new_scores_list.append(sc)
+ new_beam_map.append(bi)
+
+ if len(new_tokens_list) >= num_beams:
+ break
+
+ if not new_tokens_list:
+ break
+
+ beam_tokens = paddle.stack(new_tokens_list)
+ beam_scores = paddle.to_tensor(new_scores_list, dtype="float32")
+
+ # Reorder KV cache to match new beam assignment
+ reorder_idx = paddle.to_tensor(new_beam_map, dtype="int64")
+ for layer_idx in range(len(past_key_values.key_cache)):
+ past_key_values.key_cache[layer_idx] = paddle.index_select(
+ past_key_values.key_cache[layer_idx], reorder_idx, axis=0
+ )
+ past_key_values.value_cache[layer_idx] = paddle.index_select(
+ past_key_values.value_cache[layer_idx], reorder_idx, axis=0
+ )
+
+ attention_mask = paddle.index_select(attention_mask, reorder_idx, axis=0)
+
+ if finished_hypos:
+ finished_hypos.sort(key=lambda x: x[0], reverse=True)
+ best_seq = finished_hypos[0][1]
+ else:
+ best_idx = int(paddle.argmax(beam_scores).item())
+ best_seq = beam_tokens[best_idx]
+
+ return best_seq.unsqueeze(0)
+
+ def generate(
+ self,
+ input_ids=None,
+ inputs_embeds=None,
+ attention_mask=None,
+ max_length=256,
+ min_length=1,
+ num_beams=1,
+ do_sample=False,
+ top_p=0.9,
+ temperature=1.0,
+ repetition_penalty=1.5,
+ length_penalty=1.0,
+ num_return_sequences=1,
+ **kwargs,
+ ):
+ """Autoregressive generation loop."""
+ if inputs_embeds is not None:
+ hidden = inputs_embeds
+ else:
+ hidden = self.model.embed_tokens(input_ids)
+
+ batch_size = hidden.shape[0]
+ seq_len = hidden.shape[1]
+
+ if attention_mask is None:
+ attention_mask = paddle.ones(
+ [batch_size, seq_len], dtype=paddle.int64
+ )
+
+ _suppress_token_id = getattr(self.config, "pad_token_id", None)
+
+ if num_beams > 1:
+ return self._generate_beam_search(
+ inputs_embeds=hidden,
+ attention_mask=attention_mask,
+ max_length=max_length,
+ min_length=min_length,
+ num_beams=num_beams,
+ repetition_penalty=repetition_penalty,
+ length_penalty=length_penalty,
+ temperature=temperature,
+ _suppress_token_id=_suppress_token_id,
+ )
+
+ # Greedy / sampling path (num_beams == 1)
+ past_key_values = DynamicCache()
+ cache_position = paddle.arange(seq_len)
+
+ _suppress_token_id = getattr(self.config, "pad_token_id", None)
+
+ outputs = self.model(
+ inputs_embeds=hidden,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ use_cache=True,
+ cache_position=cache_position,
+ )
+ logits = self.lm_head(outputs[0][:, -1:, :])[:, 0, :]
+ logits = logits / max(temperature, 1e-8)
+
+ if _suppress_token_id is not None:
+ logits[:, _suppress_token_id] = float("-inf")
+
+ if repetition_penalty != 1.0:
+ logits = _apply_repetition_penalty(
+ logits, None, repetition_penalty
+ )
+
+ if do_sample:
+ next_tokens = paddle.multinomial(
+ F.softmax(logits, axis=-1), num_samples=1
+ )
+ else:
+ next_tokens = paddle.argmax(logits, axis=-1, keepdim=True)
+
+ generated = [next_tokens]
+ gen_ids_tensor = next_tokens.reshape([-1])
+
+ max_new_tokens = max_length - seq_len
+ eos_token_id = getattr(self.config, "eos_token_id", 2)
+
+ for step in range(1, max_new_tokens):
+ attention_mask = paddle.concat(
+ [
+ attention_mask,
+ paddle.ones([batch_size, 1], dtype=paddle.int64),
+ ],
+ axis=-1,
+ )
+
+ past_len = seq_len + step
+ cache_position = paddle.arange(past_len - 1, past_len)
+
+ step_outputs = self.model(
+ input_ids=next_tokens,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ use_cache=True,
+ cache_position=cache_position,
+ )
+
+ logits = self.lm_head(step_outputs[0][:, -1:, :])[:, 0, :]
+ logits = logits / max(temperature, 1e-8)
+
+ if _suppress_token_id is not None:
+ logits[:, _suppress_token_id] = float("-inf")
+
+ if repetition_penalty != 1.0:
+ logits = _apply_repetition_penalty(
+ logits, gen_ids_tensor, repetition_penalty
+ )
+
+ if do_sample:
+ next_tokens = paddle.multinomial(
+ F.softmax(logits, axis=-1), num_samples=1
+ )
+ else:
+ next_tokens = paddle.argmax(logits, axis=-1, keepdim=True)
+
+ generated.append(next_tokens)
+ gen_ids_tensor = paddle.concat(
+ [gen_ids_tensor, next_tokens.reshape([-1])]
+ )
+
+ if next_tokens.reshape([-1])[0].item() == eos_token_id:
+ break
+
+ output_ids = paddle.concat(generated, axis=-1)
+ return output_ids
diff --git a/ppmat/models/matterchat/q_former/__init__.py b/ppmat/models/matterchat/q_former/__init__.py
new file mode 100644
index 00000000..cfe95a27
--- /dev/null
+++ b/ppmat/models/matterchat/q_former/__init__.py
@@ -0,0 +1,31 @@
+# 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 ppmat.models.matterchat.q_former.q_former_base import (
+ BertConfig,
+ BertLMHeadModel,
+ BertModel,
+ BertForMaskedLM,
+ Blip2Base,
+)
+from ppmat.models.matterchat.q_former.q_former_complete import Blip2MistralInstruct
+
+__all__ = [
+ "BertConfig",
+ "BertForMaskedLM",
+ "BertLMHeadModel",
+ "BertModel",
+ "Blip2Base",
+ "Blip2MistralInstruct",
+]
diff --git a/ppmat/models/matterchat/q_former/q_former_base.py b/ppmat/models/matterchat/q_former/q_former_base.py
new file mode 100644
index 00000000..6fb3d6f9
--- /dev/null
+++ b/ppmat/models/matterchat/q_former/q_former_base.py
@@ -0,0 +1,1179 @@
+# 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.
+
+"""Q-Former (BERT variant with cross-attention) for MatterChat.
+
+Cannot reuse paddle.nn / PaddleNLP BertModel: Q-Former injects cross-attention layers,
+accepts query_embeds as input (bypassing word embeddings), and splits each layer's FFN
+into two parallel paths (query vs. text). These modifications are incompatible with the
+standard BERT implementation.
+"""
+
+import contextlib
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+from ppmat.models.matterchat.chgnet.model.model_embedding import CHGNet
+from ppmat.models.matterchat.utils.activations import get_activation as _get_activation
+from ppmat.models.matterchat.utils.weight_init import default_init_weights
+
+
+class BertConfig:
+ """Configuration for BERT-like Q-Former model."""
+
+ def __init__(
+ self,
+ vocab_size=30522,
+ hidden_size=768,
+ num_hidden_layers=12,
+ num_attention_heads=12,
+ intermediate_size=3072,
+ hidden_act="gelu",
+ hidden_dropout_prob=0.1,
+ attention_probs_dropout_prob=0.1,
+ max_position_embeddings=512,
+ type_vocab_size=2,
+ initializer_range=0.02,
+ layer_norm_eps=1e-12,
+ pad_token_id=0,
+ chunk_size_feed_forward=0,
+ add_cross_attention=False,
+ cross_attention_freq=2,
+ encoder_width=768,
+ query_length=32,
+ output_attentions=False,
+ output_hidden_states=False,
+ use_return_dict=True,
+ gradient_checkpointing=False,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ self.hidden_size = hidden_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.intermediate_size = intermediate_size
+ self.hidden_act = hidden_act
+ self.hidden_dropout_prob = hidden_dropout_prob
+ self.attention_probs_dropout_prob = attention_probs_dropout_prob
+ self.max_position_embeddings = max_position_embeddings
+ self.type_vocab_size = type_vocab_size
+ self.initializer_range = initializer_range
+ self.layer_norm_eps = layer_norm_eps
+ self.pad_token_id = pad_token_id
+ self.chunk_size_feed_forward = chunk_size_feed_forward
+ self.add_cross_attention = add_cross_attention
+ self.cross_attention_freq = cross_attention_freq
+ self.encoder_width = encoder_width
+ self.query_length = query_length
+ self.output_attentions = output_attentions
+ self.output_hidden_states = output_hidden_states
+ self.use_return_dict = use_return_dict
+ self.gradient_checkpointing = gradient_checkpointing
+
+ @classmethod
+ def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+ """Create config from pretrained model name. Returns default config."""
+ return cls(**kwargs)
+
+
+class _SubscriptableOutput:
+ """Base class for model outputs that support tuple-like indexing."""
+
+ def __getitem__(self, index):
+ fields = [f for f in self.__dataclass_fields__ if not f.startswith('_')]
+ return getattr(self, fields[index])
+
+ def __len__(self):
+ return len([f for f in self.__dataclass_fields__ if not f.startswith('_')])
+
+
+@dataclass
+class BaseModelOutputWithPastAndCrossAttentions(_SubscriptableOutput):
+ last_hidden_state: paddle.Tensor = None
+ past_key_values: Optional[Tuple] = None
+ hidden_states: Optional[Tuple] = None
+ attentions: Optional[Tuple] = None
+ cross_attentions: Optional[Tuple] = None
+
+
+@dataclass
+class BaseModelOutputWithPoolingAndCrossAttentions(_SubscriptableOutput):
+ last_hidden_state: paddle.Tensor = None
+ pooler_output: Optional[paddle.Tensor] = None
+ past_key_values: Optional[Tuple] = None
+ hidden_states: Optional[Tuple] = None
+ attentions: Optional[Tuple] = None
+ cross_attentions: Optional[Tuple] = None
+
+
+@dataclass
+class CausalLMOutputWithCrossAttentions(_SubscriptableOutput):
+ loss: Optional[paddle.Tensor] = None
+ logits: paddle.Tensor = None
+ past_key_values: Optional[Tuple] = None
+ hidden_states: Optional[Tuple] = None
+ attentions: Optional[Tuple] = None
+ cross_attentions: Optional[Tuple] = None
+
+
+@dataclass
+class MaskedLMOutput:
+ loss: Optional[paddle.Tensor] = None
+ logits: paddle.Tensor = None
+ hidden_states: Optional[Tuple] = None
+ attentions: Optional[Tuple] = None
+
+
+def apply_chunking_to_forward(forward_fn, chunk_size, chunk_dim, *input_tensors):
+ """Apply forward function. Simplified from transformers version."""
+ return forward_fn(*input_tensors)
+
+
+def find_pruneable_heads_and_indices(heads, n_heads, head_size, already_pruned_heads):
+ """Find pruneable heads and their indices."""
+ mask = paddle.ones([n_heads, head_size])
+ heads = set(heads) - already_pruned_heads
+ for head in heads:
+ mask[head] = 0
+ index = paddle.nonzero(mask.reshape([-1]) == 1).squeeze(-1)
+ index = index[: len(heads) * head_size]
+ return heads, index
+
+
+def prune_linear_layer(layer, index, dim=0):
+ """Prune linear layer (stub)."""
+ return layer
+
+
+class BertEmbeddings(nn.Layer):
+ """Construct the embeddings from word and position embeddings."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.word_embeddings = nn.Embedding(
+ config.vocab_size,
+ config.hidden_size,
+ padding_idx=config.pad_token_id,
+ )
+ self.position_embeddings = nn.Embedding(
+ config.max_position_embeddings,
+ config.hidden_size,
+ )
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ self.register_buffer(
+ "position_ids",
+ paddle.arange(config.max_position_embeddings).reshape([1, -1]),
+ )
+ self.position_embedding_type = getattr(
+ config, "position_embedding_type", "absolute"
+ )
+ self.config = config
+
+ def forward(
+ self,
+ input_ids=None,
+ position_ids=None,
+ query_embeds=None,
+ past_key_values_length=0,
+ ):
+ if input_ids is not None:
+ seq_length = input_ids.shape[1]
+ else:
+ seq_length = 0
+
+ if position_ids is None:
+ position_ids = self.position_ids[
+ :, past_key_values_length : seq_length + past_key_values_length
+ ].clone()
+
+ if input_ids is not None:
+ embeddings = self.word_embeddings(input_ids)
+ if self.position_embedding_type == "absolute":
+ position_embeddings = self.position_embeddings(position_ids)
+ embeddings = embeddings + position_embeddings
+ if query_embeds is not None:
+ embeddings = paddle.concat((query_embeds, embeddings), axis=1)
+ else:
+ embeddings = query_embeds
+
+ embeddings = self.LayerNorm(embeddings)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+
+class BertSelfAttention(nn.Layer):
+ def __init__(self, config, is_cross_attention):
+ super().__init__()
+ self.config = config
+
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
+ config, "embedding_size"
+ ):
+ raise ValueError(
+ "The hidden size (%d) is not a multiple of the number of attention "
+ "heads (%d)" % (config.hidden_size, config.num_attention_heads)
+ )
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+
+ if is_cross_attention:
+ self.key = nn.Linear(config.encoder_width, self.all_head_size)
+ self.value = nn.Linear(config.encoder_width, self.all_head_size)
+ else:
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+ self.position_embedding_type = getattr(
+ config, "position_embedding_type", "absolute"
+ )
+ if (
+ self.position_embedding_type == "relative_key"
+ or self.position_embedding_type == "relative_key_query"
+ ):
+ self.max_position_embeddings = config.max_position_embeddings
+ self.distance_embedding = nn.Embedding(
+ 2 * config.max_position_embeddings - 1, self.attention_head_size
+ )
+
+ self.save_attention = False
+
+ def save_attn_gradients(self, attn_gradients):
+ self.attn_gradients = attn_gradients
+
+ def get_attn_gradients(self):
+ return self.attn_gradients
+
+ def save_attention_map(self, attention_map):
+ self.attention_map = attention_map
+
+ def get_attention_map(self):
+ return self.attention_map
+
+ def transpose_for_scores(self, x):
+ new_x_shape = x.shape[:-1] + (self.num_attention_heads, self.attention_head_size)
+ x = x.reshape(new_x_shape)
+ return x.transpose([0, 2, 1, 3])
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ head_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_value=None,
+ output_attentions=False,
+ ):
+ is_cross_attention = encoder_hidden_states is not None
+
+ if is_cross_attention:
+ key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
+ value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
+ attention_mask = encoder_attention_mask
+ elif past_key_value is not None:
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
+ key_layer = paddle.concat([past_key_value[0], key_layer], axis=2)
+ value_layer = paddle.concat([past_key_value[1], value_layer], axis=2)
+ else:
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
+
+ mixed_query_layer = self.query(hidden_states)
+ query_layer = self.transpose_for_scores(mixed_query_layer)
+ past_key_value = (key_layer, value_layer)
+
+ attention_scores = paddle.matmul(query_layer, key_layer.transpose([0, 1, 3, 2]))
+
+ if (
+ self.position_embedding_type == "relative_key"
+ or self.position_embedding_type == "relative_key_query"
+ ):
+ seq_length = hidden_states.shape[1]
+ position_ids_l = paddle.arange(seq_length, dtype=paddle.int64).reshape([-1, 1])
+ position_ids_r = paddle.arange(seq_length, dtype=paddle.int64).reshape([1, -1])
+ distance = position_ids_l - position_ids_r
+ positional_embedding = self.distance_embedding(
+ distance + self.max_position_embeddings - 1
+ )
+ positional_embedding = positional_embedding.cast(query_layer.dtype)
+
+ if self.position_embedding_type == "relative_key":
+ relative_position_scores = paddle.einsum(
+ "bhld,lrd->bhlr", query_layer, positional_embedding
+ )
+ attention_scores = attention_scores + relative_position_scores
+ elif self.position_embedding_type == "relative_key_query":
+ relative_position_scores_query = paddle.einsum(
+ "bhld,lrd->bhlr", query_layer, positional_embedding
+ )
+ relative_position_scores_key = paddle.einsum(
+ "bhrd,lrd->bhlr", key_layer, positional_embedding
+ )
+ attention_scores = (
+ attention_scores
+ + relative_position_scores_query
+ + relative_position_scores_key
+ )
+
+ attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+ if attention_mask is not None:
+ attention_scores = attention_scores + attention_mask
+
+ attention_probs = F.softmax(attention_scores, axis=-1)
+
+ if is_cross_attention and self.save_attention:
+ self.save_attention_map(attention_probs)
+
+ attention_probs_dropped = self.dropout(attention_probs)
+
+ if head_mask is not None:
+ attention_probs_dropped = attention_probs_dropped * head_mask
+
+ context_layer = paddle.matmul(attention_probs_dropped, value_layer)
+ context_layer = context_layer.transpose([0, 2, 1, 3])
+ new_context_layer_shape = context_layer.shape[:-2] + (self.all_head_size,)
+ context_layer = context_layer.reshape(new_context_layer_shape)
+
+ outputs = (
+ (context_layer, attention_probs) if output_attentions else (context_layer,)
+ )
+ outputs = outputs + (past_key_value,)
+ return outputs
+
+
+class BertSelfOutput(nn.Layer):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states, input_tensor):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+class BertAttention(nn.Layer):
+ def __init__(self, config, is_cross_attention=False):
+ super().__init__()
+ self.self = BertSelfAttention(config, is_cross_attention)
+ self.output = BertSelfOutput(config)
+ self.pruned_heads = set()
+
+ def prune_heads(self, heads):
+ if len(heads) == 0:
+ return
+ heads, index = find_pruneable_heads_and_indices(
+ heads,
+ self.self.num_attention_heads,
+ self.self.attention_head_size,
+ self.pruned_heads,
+ )
+ self.self.query = prune_linear_layer(self.self.query, index)
+ self.self.key = prune_linear_layer(self.self.key, index)
+ self.self.value = prune_linear_layer(self.self.value, index)
+ self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+ self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
+ self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
+ self.pruned_heads = self.pruned_heads.union(heads)
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ head_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_value=None,
+ output_attentions=False,
+ ):
+ self_outputs = self.self(
+ hidden_states,
+ attention_mask,
+ head_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ past_key_value,
+ output_attentions,
+ )
+ attention_output = self.output(self_outputs[0], hidden_states)
+ outputs = (attention_output,) + self_outputs[1:]
+ return outputs
+
+
+class BertIntermediate(nn.Layer):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = _get_activation(config.hidden_act)
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+class BertOutput(nn.Layer):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states, input_tensor):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+class BertLayer(nn.Layer):
+ def __init__(self, config, layer_num):
+ super().__init__()
+ self.config = config
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = BertAttention(config)
+ self.layer_num = layer_num
+ if (
+ self.config.add_cross_attention
+ and layer_num % self.config.cross_attention_freq == 0
+ ):
+ self.crossattention = BertAttention(
+ config, is_cross_attention=self.config.add_cross_attention
+ )
+ self.has_cross_attention = True
+ else:
+ self.has_cross_attention = False
+
+ self.intermediate = BertIntermediate(config)
+ self.output = BertOutput(config)
+ self.intermediate_query = BertIntermediate(config)
+ self.output_query = BertOutput(config)
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ head_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_value=None,
+ output_attentions=False,
+ query_length=0,
+ ):
+ self_attn_past_key_value = (
+ past_key_value[:2] if past_key_value is not None else None
+ )
+ self_attention_outputs = self.attention(
+ hidden_states,
+ attention_mask,
+ head_mask,
+ output_attentions=output_attentions,
+ past_key_value=self_attn_past_key_value,
+ )
+ attention_output = self_attention_outputs[0]
+ outputs = self_attention_outputs[1:-1]
+ present_key_value = self_attention_outputs[-1]
+
+ if query_length > 0:
+ query_attention_output = attention_output[:, :query_length, :]
+
+ if self.has_cross_attention:
+ assert (
+ encoder_hidden_states is not None
+ ), "encoder_hidden_states must be given for cross-attention layers"
+ cross_attention_outputs = self.crossattention(
+ query_attention_output,
+ attention_mask,
+ head_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ output_attentions=output_attentions,
+ )
+ query_attention_output = cross_attention_outputs[0]
+ outputs = outputs + cross_attention_outputs[1:-1]
+
+ layer_output = apply_chunking_to_forward(
+ self.feed_forward_chunk_query,
+ self.chunk_size_feed_forward,
+ self.seq_len_dim,
+ query_attention_output,
+ )
+ if attention_output.shape[1] > query_length:
+ layer_output_text = apply_chunking_to_forward(
+ self.feed_forward_chunk,
+ self.chunk_size_feed_forward,
+ self.seq_len_dim,
+ attention_output[:, query_length:, :],
+ )
+ layer_output = paddle.concat([layer_output, layer_output_text], axis=1)
+ else:
+ layer_output = apply_chunking_to_forward(
+ self.feed_forward_chunk,
+ self.chunk_size_feed_forward,
+ self.seq_len_dim,
+ attention_output,
+ )
+ outputs = (layer_output,) + outputs
+ outputs = outputs + (present_key_value,)
+ return outputs
+
+ def feed_forward_chunk(self, attention_output):
+ intermediate_output = self.intermediate(attention_output)
+ layer_output = self.output(intermediate_output, attention_output)
+ return layer_output
+
+ def feed_forward_chunk_query(self, attention_output):
+ intermediate_output = self.intermediate_query(attention_output)
+ layer_output = self.output_query(intermediate_output, attention_output)
+ return layer_output
+
+
+class BertEncoder(nn.Layer):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.LayerList(
+ [BertLayer(config, i) for i in range(config.num_hidden_layers)]
+ )
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ head_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_values=None,
+ use_cache=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ query_length=0,
+ ):
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attentions = () if output_attentions else None
+ all_cross_attentions = (
+ () if output_attentions and self.config.add_cross_attention else None
+ )
+ next_decoder_cache = () if use_cache else None
+
+ for i in range(self.config.num_hidden_layers):
+ layer_module = self.layer[i]
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ layer_head_mask = head_mask[i] if head_mask is not None else None
+ past_key_value = past_key_values[i] if past_key_values is not None else None
+
+ layer_outputs = layer_module(
+ hidden_states,
+ attention_mask,
+ layer_head_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ past_key_value,
+ output_attentions,
+ query_length,
+ )
+
+ hidden_states = layer_outputs[0]
+ if use_cache:
+ next_decoder_cache += (layer_outputs[-1],)
+ if output_attentions:
+ all_self_attentions = all_self_attentions + (layer_outputs[1],)
+ all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
+
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [
+ hidden_states,
+ next_decoder_cache,
+ all_hidden_states,
+ all_self_attentions,
+ all_cross_attentions,
+ ]
+ if v is not None
+ )
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=next_decoder_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ cross_attentions=all_cross_attentions,
+ )
+
+
+class BertPooler(nn.Layer):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.activation = nn.Tanh()
+
+ def forward(self, hidden_states):
+ first_token_tensor = hidden_states[:, 0]
+ pooled_output = self.dense(first_token_tensor)
+ pooled_output = self.activation(pooled_output)
+ return pooled_output
+
+
+class BertPredictionHeadTransform(nn.Layer):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ if isinstance(config.hidden_act, str):
+ self.transform_act_fn = _get_activation(config.hidden_act)
+ else:
+ self.transform_act_fn = config.hidden_act
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps)
+
+ def forward(self, hidden_states):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.transform_act_fn(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states)
+ return hidden_states
+
+
+class BertLMPredictionHead(nn.Layer):
+ def __init__(self, config):
+ super().__init__()
+ self.transform = BertPredictionHeadTransform(config)
+ self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias_attr=False)
+ self.bias = self.create_parameter(
+ shape=[config.vocab_size],
+ is_bias=True,
+ default_initializer=nn.initializer.Constant(0.0),
+ )
+ self.decoder.bias = self.bias
+
+ def forward(self, hidden_states):
+ hidden_states = self.transform(hidden_states)
+ hidden_states = self.decoder(hidden_states)
+ return hidden_states
+
+
+class BertOnlyMLMHead(nn.Layer):
+ def __init__(self, config):
+ super().__init__()
+ self.predictions = BertLMPredictionHead(config)
+
+ def forward(self, sequence_output):
+ prediction_scores = self.predictions(sequence_output)
+ return prediction_scores
+
+
+class BertPreTrainedModel(nn.Layer):
+ config_class = BertConfig
+ base_model_prefix = "bert"
+
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+
+ def _init_weights(self, module):
+ default_init_weights(module, self.config.initializer_range)
+
+ def init_weights(self):
+ self.apply(self._init_weights)
+
+ def get_head_mask(self, head_mask, num_hidden_layers):
+ if head_mask is not None:
+ if isinstance(head_mask, paddle.Tensor):
+ if len(head_mask.shape) == 1:
+ head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+ head_mask = head_mask.expand([num_hidden_layers, -1, -1, -1, -1])
+ return head_mask
+ return [None] * num_hidden_layers
+
+ def invert_attention_mask(self, encoder_attention_mask):
+ if isinstance(encoder_attention_mask, paddle.Tensor):
+ if len(encoder_attention_mask.shape) == 2:
+ encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
+ elif len(encoder_attention_mask.shape) == 3:
+ encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
+ else:
+ encoder_extended_attention_mask = encoder_attention_mask
+ encoder_extended_attention_mask = encoder_extended_attention_mask.cast(paddle.float32)
+ encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * paddle.finfo(paddle.float32).min
+ return encoder_extended_attention_mask
+ return encoder_attention_mask
+
+
+class BertModel(BertPreTrainedModel):
+ def __init__(self, config, add_pooling_layer=False):
+ super().__init__(config)
+ self.config = config
+ self.embeddings = BertEmbeddings(config)
+ self.encoder = BertEncoder(config)
+ self.pooler = BertPooler(config) if add_pooling_layer else None
+ self.init_weights()
+
+ def get_input_embeddings(self):
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value):
+ self.embeddings.word_embeddings = value
+
+ def _prune_heads(self, heads_to_prune):
+ for layer, heads in heads_to_prune.items():
+ self.encoder.layer[layer].attention.prune_heads(heads)
+
+ def get_extended_attention_mask(
+ self,
+ attention_mask,
+ input_shape,
+ is_decoder,
+ has_query=False,
+ ):
+ if len(attention_mask.shape) == 3:
+ extended_attention_mask = attention_mask[:, None, :, :]
+ elif len(attention_mask.shape) == 2:
+ if is_decoder:
+ batch_size, seq_length = input_shape
+ seq_ids = paddle.arange(seq_length)
+ causal_mask = (
+ seq_ids[None, None, :].expand([batch_size, seq_length, seq_length])
+ <= seq_ids[None, :, None]
+ )
+ causal_mask = causal_mask.cast(attention_mask.dtype)
+
+ if causal_mask.shape[1] < attention_mask.shape[1]:
+ prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
+ if has_query:
+ causal_mask = paddle.concat(
+ [
+ paddle.zeros(
+ [batch_size, prefix_seq_len, seq_length],
+ dtype=causal_mask.dtype,
+ ),
+ causal_mask,
+ ],
+ axis=1,
+ )
+ causal_mask = paddle.concat(
+ [
+ paddle.ones(
+ [batch_size, causal_mask.shape[1], prefix_seq_len],
+ dtype=causal_mask.dtype,
+ ),
+ causal_mask,
+ ],
+ axis=-1,
+ )
+ extended_attention_mask = (
+ causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
+ )
+ else:
+ extended_attention_mask = attention_mask[:, None, None, :]
+ else:
+ raise ValueError(
+ "Wrong shape for attention_mask (shape {})".format(attention_mask.shape)
+ )
+
+ extended_attention_mask = extended_attention_mask.cast(paddle.float32)
+ extended_attention_mask = (1.0 - extended_attention_mask) * paddle.finfo(paddle.float32).min
+ return extended_attention_mask
+
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ position_ids=None,
+ head_mask=None,
+ query_embeds=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_values=None,
+ use_cache=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ is_decoder=False,
+ ):
+ output_attentions = (
+ output_attentions if output_attentions is not None else self.config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+
+ if input_ids is None:
+ assert (
+ query_embeds is not None
+ ), "You have to specify query_embeds when input_ids is None"
+
+ past_key_values_length = (
+ past_key_values[0][0].shape[2] - self.config.query_length
+ if past_key_values is not None
+ else 0
+ )
+ query_length = query_embeds.shape[1] if query_embeds is not None else 0
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ query_embeds=query_embeds,
+ past_key_values_length=past_key_values_length,
+ )
+
+ input_shape = embedding_output.shape[:-1]
+ batch_size, seq_length = input_shape
+
+ if attention_mask is None:
+ attention_mask = paddle.ones([batch_size, seq_length + past_key_values_length])
+
+ if is_decoder:
+ extended_attention_mask = self.get_extended_attention_mask(
+ attention_mask,
+ input_ids.shape if input_ids is not None else (batch_size, seq_length),
+ is_decoder,
+ has_query=(query_embeds is not None),
+ )
+ else:
+ extended_attention_mask = self.get_extended_attention_mask(
+ attention_mask, input_shape, is_decoder
+ )
+
+ if encoder_hidden_states is not None:
+ if isinstance(encoder_hidden_states, list):
+ encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].shape
+ else:
+ (
+ encoder_batch_size,
+ encoder_sequence_length,
+ _,
+ ) = encoder_hidden_states.shape
+ encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+
+ if isinstance(encoder_attention_mask, list):
+ encoder_extended_attention_mask = [
+ self.invert_attention_mask(mask) for mask in encoder_attention_mask
+ ]
+ elif encoder_attention_mask is None:
+ encoder_attention_mask = paddle.ones(encoder_hidden_shape)
+ encoder_extended_attention_mask = self.invert_attention_mask(
+ encoder_attention_mask
+ )
+ else:
+ encoder_extended_attention_mask = self.invert_attention_mask(
+ encoder_attention_mask
+ )
+ else:
+ encoder_extended_attention_mask = None
+
+ head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=extended_attention_mask,
+ head_mask=head_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_extended_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ query_length=query_length,
+ )
+ sequence_output = encoder_outputs[0]
+ pooled_output = (
+ self.pooler(sequence_output) if self.pooler is not None else None
+ )
+
+ if not return_dict:
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPoolingAndCrossAttentions(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ past_key_values=encoder_outputs.past_key_values if return_dict else None,
+ hidden_states=encoder_outputs.hidden_states if return_dict else None,
+ attentions=encoder_outputs.attentions if return_dict else None,
+ cross_attentions=encoder_outputs.cross_attentions if return_dict else None,
+ )
+
+
+class BertLMHeadModel(BertPreTrainedModel):
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.bert = BertModel(config, add_pooling_layer=False)
+ self.cls = BertOnlyMLMHead(config)
+ self.init_weights()
+
+ def get_output_embeddings(self):
+ return self.cls.predictions.decoder
+
+ def set_output_embeddings(self, new_embeddings):
+ self.cls.predictions.decoder = new_embeddings
+
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ position_ids=None,
+ head_mask=None,
+ query_embeds=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ labels=None,
+ past_key_values=None,
+ use_cache=True,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ return_logits=False,
+ is_decoder=True,
+ reduction="mean",
+ ):
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+ if labels is not None:
+ use_cache = False
+ if past_key_values is not None:
+ query_embeds = None
+
+ outputs = self.bert(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ head_mask=head_mask,
+ query_embeds=query_embeds,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ is_decoder=is_decoder,
+ )
+
+ sequence_output = outputs[0]
+ if query_embeds is not None:
+ sequence_output = outputs[0][:, query_embeds.shape[1] :, :]
+
+ prediction_scores = self.cls(sequence_output)
+
+ if return_logits:
+ return prediction_scores[:, :-1, :]
+
+ lm_loss = None
+ if labels is not None:
+ shifted_prediction_scores = prediction_scores[:, :-1, :]
+ labels = labels[:, 1:]
+ loss_fct = nn.CrossEntropyLoss(reduction=reduction)
+ lm_loss = loss_fct(
+ shifted_prediction_scores.reshape([-1, self.config.vocab_size]),
+ labels.reshape([-1]),
+ )
+ if reduction == "none":
+ lm_loss = lm_loss.reshape([prediction_scores.shape[0], -1]).sum(1)
+
+ if not return_dict:
+ output = (prediction_scores,) + outputs[2:]
+ return ((lm_loss,) + output) if lm_loss is not None else output
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=lm_loss,
+ logits=prediction_scores,
+ past_key_values=outputs.past_key_values if hasattr(outputs, "past_key_values") else None,
+ hidden_states=outputs.hidden_states if hasattr(outputs, "hidden_states") else None,
+ attentions=outputs.attentions if hasattr(outputs, "attentions") else None,
+ cross_attentions=outputs.cross_attentions if hasattr(outputs, "cross_attentions") else None,
+ )
+
+ def prepare_inputs_for_generation(
+ self, input_ids, query_embeds, past=None, attention_mask=None, **model_kwargs
+ ):
+ if attention_mask is None:
+ attention_mask = paddle.ones(input_ids.shape)
+ query_mask = paddle.ones(query_embeds.shape[:-1])
+ attention_mask = paddle.concat([query_mask, attention_mask], axis=-1)
+
+ if past is not None:
+ input_ids = input_ids[:, -1:]
+
+ return {
+ "input_ids": input_ids,
+ "query_embeds": query_embeds,
+ "attention_mask": attention_mask,
+ "past_key_values": past,
+ "encoder_hidden_states": model_kwargs.get("encoder_hidden_states", None),
+ "encoder_attention_mask": model_kwargs.get("encoder_attention_mask", None),
+ "is_decoder": True,
+ }
+
+ def _reorder_cache(self, past, beam_idx):
+ reordered_past = ()
+ for layer_past in past:
+ reordered_past += (
+ tuple(
+ paddle.index_sample(past_state, beam_idx) for past_state in layer_past
+ ),
+ )
+ return reordered_past
+
+
+class BertForMaskedLM(BertPreTrainedModel):
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.bert = BertModel(config, add_pooling_layer=False)
+ self.cls = BertOnlyMLMHead(config)
+ self.init_weights()
+
+ def get_output_embeddings(self):
+ return self.cls.predictions.decoder
+
+ def set_output_embeddings(self, new_embeddings):
+ self.cls.predictions.decoder = new_embeddings
+
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ position_ids=None,
+ head_mask=None,
+ query_embeds=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ labels=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ return_logits=False,
+ is_decoder=False,
+ ):
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+
+ outputs = self.bert(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ head_mask=head_mask,
+ query_embeds=query_embeds,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ is_decoder=is_decoder,
+ )
+
+ if query_embeds is not None:
+ sequence_output = outputs[0][:, query_embeds.shape[1] :, :]
+ else:
+ sequence_output = outputs[0]
+ prediction_scores = self.cls(sequence_output)
+
+ if return_logits:
+ return prediction_scores
+
+ masked_lm_loss = None
+ if labels is not None:
+ loss_fct = nn.CrossEntropyLoss()
+ masked_lm_loss = loss_fct(
+ prediction_scores.reshape([-1, self.config.vocab_size]),
+ labels.reshape([-1]),
+ )
+
+ if not return_dict:
+ output = (prediction_scores,) + outputs[2:]
+ return (
+ ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+ )
+
+ return MaskedLMOutput(
+ loss=masked_lm_loss,
+ logits=prediction_scores,
+ hidden_states=outputs.hidden_states if hasattr(outputs, "hidden_states") else None,
+ attentions=outputs.attentions if hasattr(outputs, "attentions") else None,
+ )
+
+
+class LayerNorm(nn.LayerNorm):
+ """Subclass of paddle LayerNorm to support mixed precision."""
+
+ def forward(self, x):
+ orig_dtype = x.dtype
+ x = super().forward(x.cast(paddle.float32))
+ return x.cast(orig_dtype)
+
+
+class Blip2Base(nn.Layer):
+ """Base class for Blip2 models."""
+
+ @classmethod
+ def init_tokenizer(cls, truncation_side="right"):
+ """Initialize tokenizer. Returns None - tokenizer is managed externally."""
+ return None
+
+ @classmethod
+ def init_Qformer(cls, num_query_token, vision_width, cross_attention_freq=2):
+ config = BertConfig()
+ config.encoder_width = vision_width
+ config.add_cross_attention = True
+ config.cross_attention_freq = cross_attention_freq
+ config.query_length = num_query_token
+
+ Qformer = BertLMHeadModel(config)
+
+ query_tokens = paddle.create_parameter(
+ shape=[1, num_query_token, config.hidden_size],
+ dtype=paddle.float32,
+ default_initializer=nn.initializer.Normal(
+ mean=0.0, std=config.initializer_range
+ ),
+ )
+ return Qformer, query_tokens
+
+ def init_material_encoder(self):
+ """Initialize CHGNet material encoder without pretrained weights."""
+ return CHGNet()
+
+ def maybe_autocast(self, dtype=paddle.float16):
+ """Context manager for mixed precision."""
+ if not paddle.is_compiled_with_cuda():
+ return contextlib.nullcontext()
+ return paddle.amp.auto_cast(dtype=dtype)
diff --git a/ppmat/models/matterchat/q_former/q_former_complete.py b/ppmat/models/matterchat/q_former/q_former_complete.py
new file mode 100644
index 00000000..14fa1713
--- /dev/null
+++ b/ppmat/models/matterchat/q_former/q_former_complete.py
@@ -0,0 +1,472 @@
+# 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.
+
+"""Complete MatterChat model: Blip2MistralInstruct."""
+
+import paddle
+import paddle.nn as nn
+
+from ppmat.models.matterchat.mistral.configuration_mistral import MistralConfig
+from ppmat.models.matterchat.mistral.modeling_mistral import MistralForCausalLM
+from ppmat.models.matterchat.q_former.q_former_base import Blip2Base
+
+
+class Blip2MistralInstruct(Blip2Base):
+ def __init__(
+ self,
+ num_query_token=32,
+ prompt="test",
+ max_txt_len=512,
+ max_output_txt_len=1024,
+ qformer_text_input=False,
+ llm_tokenizer=None,
+ llm_model=None,
+ tokenizer_path=None,
+ ):
+ super().__init__()
+
+ if llm_tokenizer is None and tokenizer_path is not None and tokenizer_path:
+ from ppmat.models.matterchat.utils.tokenizer import MistralTokenizerWrapper
+ import os as _os
+
+ _resolved = _os.path.expanduser(tokenizer_path)
+ llm_tokenizer = MistralTokenizerWrapper(_resolved)
+
+ self.tokenizer = self.init_tokenizer(truncation_side="left")
+ self.material_encoder = self.init_material_encoder()
+ for param in self.material_encoder.parameters():
+ param.stop_gradient = True
+
+ self.Qformer, self.query_tokens = self.init_Qformer(num_query_token, 64)
+
+ self.qformer_text_input = qformer_text_input
+ if not qformer_text_input:
+ self.Qformer.bert.embeddings.word_embeddings = None
+ self.Qformer.bert.embeddings.position_embeddings = None
+ for layer in self.Qformer.bert.encoder.layer:
+ layer.output = None
+ layer.intermediate = None
+ else:
+ if llm_tokenizer is not None:
+ self.Qformer.resize_token_embeddings(len(llm_tokenizer))
+ self.Qformer.cls = None
+
+ self.llm_tokenizer = llm_tokenizer
+ self.llm_model = (
+ llm_model
+ if llm_model is not None
+ else MistralForCausalLM(self._get_default_mistral_config())
+ )
+
+ for param in self.llm_model.parameters():
+ param.stop_gradient = True
+
+ self.llm_proj = nn.Linear(
+ self.Qformer.config.hidden_size, self.llm_model.config.hidden_size
+ )
+
+ self.max_txt_len = max_txt_len
+ self.max_output_txt_len = max_output_txt_len
+ self.prompt = prompt
+
+ if self.llm_tokenizer is not None:
+ prompt_tokens = self.llm_tokenizer(self.prompt, return_tensors="pd")
+ self.prompt_length = prompt_tokens.attention_mask.sum(1).item()
+ else:
+ self.prompt_length = 0
+
+ @staticmethod
+ def _get_default_mistral_config():
+ return MistralConfig()
+
+ @classmethod
+ def from_pretrained(cls, path):
+ """Load MatterChat from a MODEL_REGISTRY-style directory.
+
+ Directory structure:
+ path/
+ ├── *.yaml
+ ├── tokenizer.json
+ └── checkpoints/
+ ├── best.pdparams (or model.pdparams.index.json + shards)
+
+ Args:
+ path (str): Path to the extracted MODEL_REGISTRY directory.
+
+ Returns:
+ Blip2MistralInstruct: Loaded model with tokenizer and weights.
+ """
+ import os as _os, json as _json
+
+ from omegaconf import OmegaConf
+
+ yaml_files = [f for f in _os.listdir(path) if f.endswith(".yaml") or f.endswith(".yml")]
+ if not yaml_files:
+ raise FileNotFoundError(f"No yaml config found in {path}")
+ config_path = _os.path.join(path, yaml_files[0])
+ config = OmegaConf.to_container(OmegaConf.load(config_path), resolve=True)
+
+ init_params = config["Model"]["__init_params__"]
+
+ tokenizer_path = _os.path.join(path, "tokenizer.json")
+ if _os.path.exists(tokenizer_path):
+ init_params["tokenizer_path"] = tokenizer_path
+
+ model = cls(**init_params)
+
+ model._load_weights_from_registry(path)
+
+ return model
+
+ def _load_weights_from_registry(self, weight_dir):
+ import os as _os, json as _json
+
+ ckpt_dir = _os.path.join(weight_dir, "checkpoints")
+ if not _os.path.isdir(ckpt_dir):
+ ckpt_dir = weight_dir
+
+ index_path = _os.path.join(ckpt_dir, "model.pdparams.index.json")
+ if _os.path.exists(index_path):
+ self._load_sharded_weights(ckpt_dir, index_path)
+ return
+
+ best_path = _os.path.join(ckpt_dir, "best.pdparams")
+ if _os.path.exists(best_path):
+ self._load_single_weight(best_path)
+ return
+
+ raise FileNotFoundError(
+ f"No weights found in {ckpt_dir}: "
+ f"expected best.pdparams or model.pdparams.index.json"
+ )
+
+ def _load_sharded_weights(self, ckpt_dir, index_path):
+ import json as _json, os as _os
+
+ with open(index_path) as f:
+ idx = _json.load(f)
+
+ model_sd = self.state_dict()
+ shard_files = sorted(set(idx["weight_map"].values()))
+ for shard_file in shard_files:
+ shard_path = _os.path.join(ckpt_dir, shard_file)
+ shard = paddle.load(shard_path, return_numpy=True)
+ for key, val in shard.items():
+ if key in model_sd:
+ w = paddle.to_tensor(val)
+ target = model_sd[key]
+ if w.dtype != target.dtype:
+ w = w.cast(target.dtype)
+ target.set_value(w)
+ del shard
+
+ def _load_single_weight(self, weight_path):
+ state_dict = paddle.load(weight_path)
+ if "state_dict" in state_dict:
+ state_dict = state_dict["state_dict"]
+ self.set_state_dict(state_dict)
+
+ def _run_qformer(self, query_tokens, material_embeds, material_atts, prompt_text=None):
+ """Run Q-Former with optional text input, shared by forward/generate/generate_followup.
+
+ Args:
+ query_tokens: [B, num_query_tokens, hidden_size]
+ material_embeds: [B, N_atoms, encoder_width]
+ material_atts: [B, N_atoms]
+ prompt_text: list of strings or None
+
+ Returns:
+ Q-Former output (BaseModelOutputWithPoolingAndCrossAttentions)
+ """
+ if self.qformer_text_input and prompt_text is not None:
+ text_input = self.tokenizer(
+ prompt_text,
+ padding="longest",
+ truncation=True,
+ max_length=self.max_txt_len,
+ return_tensors="pd",
+ )
+ query_atts = paddle.ones(query_tokens.shape[:-1], dtype=paddle.int64)
+ qformer_atts = paddle.concat([query_atts, text_input.attention_mask], axis=1)
+ return self.Qformer.bert(
+ text_input.input_ids,
+ attention_mask=qformer_atts,
+ query_embeds=query_tokens,
+ encoder_hidden_states=material_embeds,
+ encoder_attention_mask=material_atts,
+ return_dict=True,
+ )
+ else:
+ return self.Qformer.bert(
+ query_embeds=query_tokens,
+ encoder_hidden_states=material_embeds,
+ encoder_attention_mask=material_atts,
+ return_dict=True,
+ )
+
+ def _encode_for_generate(self, samples):
+ """Shared encoding pipeline for generate / generate_followup.
+
+ Steps: CHGNet encode -> Q-Former encode -> llm_proj -> tokenizer -> concat.
+ Also sets llm_tokenizer.padding_side to "left".
+
+ Returns:
+ (inputs_embeds, attention_mask, prompt_list)
+ """
+ self.llm_tokenizer.padding_side = "left"
+ prompt = samples.get("prompt", self.prompt)
+ if isinstance(prompt, str):
+ prompt = [prompt]
+
+ # Material encoding
+ material_embed = self.material_encoder.predict_structure_embedding(
+ samples["material_sample"]
+ ).unsqueeze(0)
+ material_att = paddle.ones(material_embed.shape[:-1], dtype=paddle.int64)
+
+ # Q-Former encoding
+ query_tokens = self.query_tokens.expand([1, -1, -1])
+ query_output = self._run_qformer(query_tokens, material_embed, material_att, prompt)
+
+ # llm_proj mapping
+ inputs_llm = self.llm_proj(query_output.last_hidden_state[:, : query_tokens.shape[1], :])
+ atts_llm = paddle.ones(inputs_llm.shape[:-1], dtype=paddle.int64)
+
+ # Tokenizer encoding + concatenation
+ prompt_tokens = self.llm_tokenizer(prompt, return_tensors="pd", padding="longest")
+ input_embeds = self.llm_model.get_input_embeddings()(prompt_tokens.input_ids)
+ inputs_embeds = paddle.concat([inputs_llm, input_embeds], axis=1)
+ attention_mask = paddle.concat([atts_llm, prompt_tokens.attention_mask], axis=1)
+ return inputs_embeds, attention_mask, prompt
+
+ def concat_text_input_output(self, input_ids, input_atts, output_ids, output_atts):
+ input_part_targets_len = []
+ input_ids_list, attention_mask_list = [], []
+
+ for i in range(input_ids.shape[0]):
+ this_input_len = input_atts[i].sum()
+ input_part_targets_len.append(this_input_len)
+
+ input_ids_list.append(
+ paddle.concat(
+ [
+ input_ids[i][:this_input_len],
+ output_ids[i][1:],
+ input_ids[i][this_input_len:],
+ ]
+ )
+ )
+
+ attention_mask_list.append(
+ paddle.concat(
+ [
+ input_atts[i][:this_input_len],
+ output_atts[i][1:],
+ input_atts[i][this_input_len:],
+ ]
+ )
+ )
+
+ return {
+ "input_ids": paddle.stack(input_ids_list),
+ "attention_mask": paddle.stack(attention_mask_list),
+ }, input_part_targets_len
+
+ def forward(self, samples, embedding_list, embedding_mask):
+ material_embeds = paddle.stack(embedding_list)
+ material_atts = paddle.stack(embedding_mask)
+ query_tokens = self.query_tokens.expand([material_embeds.shape[0], -1, -1])
+
+ # Shared Q-Former encoding (handles qformer_text_input branching)
+ text_list = samples.get("text_input", None)
+ query_output = self._run_qformer(query_tokens, material_embeds, material_atts, text_list)
+
+ inputs_llm = self.llm_proj(query_output.last_hidden_state[:, : query_tokens.shape[1], :])
+ atts_llm = paddle.ones(inputs_llm.shape[:-1], dtype=paddle.int64)
+
+ self.llm_tokenizer.padding_side = "right"
+ self.llm_tokenizer.truncation_side = "left"
+
+ if "text_input" in samples:
+ text_input_tokens = self.llm_tokenizer(
+ samples["text_input"],
+ return_tensors="pd",
+ padding="longest",
+ truncation=True,
+ max_length=self.max_txt_len,
+ )
+ else:
+ bos_token = self.llm_tokenizer.bos_token
+ input_text = [bos_token] * material_embeds.shape[0]
+ text_input_tokens = self.llm_tokenizer(
+ input_text,
+ return_tensors="pd",
+ padding="longest",
+ truncation=True,
+ max_length=self.max_txt_len,
+ )
+
+ self.llm_tokenizer.truncation_side = "right"
+ text_output_tokens = self.llm_tokenizer(
+ [t + self.llm_tokenizer.eos_token for t in samples["text"]],
+ return_tensors="pd",
+ padding="longest",
+ truncation=True,
+ max_length=self.max_output_txt_len,
+ )
+
+ llm_tokens, input_target_lengths = self.concat_text_input_output(
+ text_input_tokens.input_ids,
+ text_input_tokens.attention_mask,
+ text_output_tokens.input_ids,
+ text_output_tokens.attention_mask,
+ )
+
+ targets = llm_tokens["input_ids"].masked_fill(
+ llm_tokens["input_ids"] == self.llm_tokenizer.pad_token_id, -100
+ )
+ for i, length in enumerate(input_target_lengths):
+ targets[i][:length] = -100
+
+ empty_targets = paddle.ones(atts_llm.shape, dtype=paddle.int64).fill_(-100)
+ targets = paddle.concat([empty_targets, targets], axis=1)
+
+ inputs_embeds = self.llm_model.get_input_embeddings()(llm_tokens["input_ids"])
+ inputs_embeds = paddle.concat(
+ [inputs_llm.cast(inputs_embeds.dtype), inputs_embeds], axis=1
+ )
+ attention_mask = paddle.concat([atts_llm, llm_tokens["attention_mask"]], axis=1)
+
+ with self.maybe_autocast():
+ outputs = self.llm_model(
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ return_dict=True,
+ labels=targets,
+ )
+
+ return {"loss": outputs.loss}
+
+ @paddle.no_grad()
+ def chat(self, structure, prompt, max_new_tokens=64):
+ """Single-turn dialogue: given a crystal structure and a text prompt,
+ return the generated answer (greedy decoding, fp32).
+
+ This is a convenience wrapper around the CHGNet -> Q-Former ->
+ llm_proj -> Mistral pipeline for inference / dialogue use.
+
+ Args:
+ structure: pymatgen ``Structure`` object.
+ prompt: text question, e.g. "what is the chemical formula?".
+ max_new_tokens: max number of new tokens to generate.
+
+ Returns:
+ str: decoded answer text.
+ """
+ embed = self.material_encoder.predict_structure_embedding(structure).unsqueeze(0)
+ mask = paddle.ones(embed.shape[:-1], dtype=paddle.int64)
+ qt = self.query_tokens.expand([1, -1, -1])
+ qo = self._run_qformer(qt, embed, mask, [prompt])
+ inputs_llm = self.llm_proj(qo.last_hidden_state[:, : qt.shape[1], :])
+ atts_llm = paddle.ones(inputs_llm.shape[:-1], dtype=paddle.int64)
+ self.llm_tokenizer.padding_side = "left"
+ pt = self.llm_tokenizer([prompt], return_tensors="pd", padding="longest")
+ ie = self.llm_model.get_input_embeddings()(pt.input_ids)
+ ie = paddle.concat([inputs_llm, ie], axis=1)
+ am = paddle.concat([atts_llm, pt.attention_mask], axis=1)
+ out = self.llm_model.generate(
+ inputs_embeds=ie,
+ attention_mask=am,
+ do_sample=False,
+ max_length=ie.shape[1] + max_new_tokens,
+ )
+ out[out == 0] = 2
+ return self.llm_tokenizer.decode(out[0], skip_special_tokens=True).strip()
+
+ @paddle.no_grad()
+ def generate(
+ self,
+ samples,
+ use_nucleus_sampling=False,
+ num_beams=5,
+ max_length=256,
+ min_length=1,
+ top_p=0.9,
+ repetition_penalty=1.5,
+ length_penalty=1,
+ num_captions=1,
+ temperature=1,
+ ):
+ inputs_embeds, attention_mask, _prompt = self._encode_for_generate(samples)
+
+ with self.maybe_autocast():
+ outputs = self.llm_model.generate(
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ do_sample=use_nucleus_sampling,
+ top_p=top_p,
+ temperature=temperature,
+ num_beams=num_beams,
+ max_length=max_length,
+ min_length=min_length,
+ repetition_penalty=repetition_penalty,
+ length_penalty=length_penalty,
+ num_return_sequences=num_captions,
+ )
+
+ outputs[outputs == 0] = 2 # sanitize output
+ return [self.llm_tokenizer.decode(o, skip_special_tokens=True).strip() for o in outputs]
+
+ def generate_followup(
+ self,
+ samples,
+ use_nucleus_sampling=False,
+ num_beams=5,
+ max_length=256,
+ min_length=1,
+ top_p=0.9,
+ repetition_penalty=1.5,
+ length_penalty=1,
+ num_captions=1,
+ temperature=1,
+ ):
+ inputs_embeds, attention_mask, prompt = self._encode_for_generate(samples)
+
+ with self.maybe_autocast():
+ outputs = self.llm_model.generate(
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ do_sample=use_nucleus_sampling,
+ top_p=top_p,
+ temperature=temperature,
+ num_beams=num_beams,
+ max_length=max_length,
+ min_length=min_length,
+ repetition_penalty=repetition_penalty,
+ length_penalty=length_penalty,
+ num_return_sequences=num_captions,
+ )
+
+ outputs[outputs == 0] = 2 # sanitize output
+ decoded_outputs = [
+ self.llm_tokenizer.decode(o, skip_special_tokens=True).strip() for o in outputs
+ ]
+
+ cleaned_outputs = []
+ for out, p in zip(decoded_outputs, prompt * num_captions):
+ if out.lower().startswith(p.lower()):
+ out = out[len(p) :].lstrip(":,.- \n")
+ cleaned_outputs.append(out)
+
+ return cleaned_outputs
diff --git a/ppmat/models/matterchat/trainer.py b/ppmat/models/matterchat/trainer.py
new file mode 100644
index 00000000..1e918e12
--- /dev/null
+++ b/ppmat/models/matterchat/trainer.py
@@ -0,0 +1,434 @@
+# 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 trainer: BLIP-2 3-stage training on top of BaseTrainer."""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Optional
+
+import paddle
+import paddle.nn as nn
+from paddle.io import Dataset
+
+from ppmat.models.matterchat.q_former.q_former_complete import Blip2MistralInstruct
+from ppmat.trainer.base_trainer import BaseTrainer
+from ppmat.utils import logger
+
+
+# =============================================================================
+# LoRA
+# =============================================================================
+
+
+class LoRALinear(nn.Layer):
+ """Low-Rank Adaptation wrapper for an existing nn.Linear."""
+
+ def __init__(
+ self,
+ base_linear: nn.Linear,
+ r: int = 16,
+ alpha: int = 32,
+ dropout: float = 0.05,
+ ) -> None:
+ super().__init__()
+ for p in base_linear.parameters():
+ p.stop_gradient = True
+ self.base = base_linear
+ in_features = base_linear.weight.shape[0]
+ out_features = base_linear.weight.shape[1]
+ lora_dtype = base_linear.weight.dtype
+ self.lora_A = self.create_parameter(
+ shape=[in_features, r],
+ dtype=lora_dtype,
+ default_initializer=nn.initializer.KaimingUniform(
+ negative_slope=0, fan_in=in_features if in_features > 0 else 1
+ ),
+ )
+ self.lora_B = self.create_parameter(
+ shape=[r, out_features],
+ dtype=lora_dtype,
+ default_initializer=nn.initializer.Constant(value=0.0),
+ )
+ self.scaling = alpha / r
+ self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+
+ def forward(self, x: paddle.Tensor) -> paddle.Tensor:
+ base_out = self.base(x)
+ lora_update = (self.dropout(x) @ self.lora_A) @ self.lora_B
+ return base_out + lora_update * self.scaling
+
+ def merge_into_base(self) -> nn.Linear:
+ """Fold LoRA weights into the base linear and return a plain nn.Linear."""
+ with paddle.no_grad():
+ # lora_A: [in, r], lora_B: [r, out] -> A@B: [in, out] == base.weight shape
+ merged_w = self.base.weight + (self.lora_A @ self.lora_B) * self.scaling
+ new_linear = nn.Linear(
+ self.base.weight.shape[0],
+ self.base.weight.shape[1],
+ bias_attr=self.base.bias is not None,
+ )
+ new_linear.weight.set_value(merged_w)
+ if self.base.bias is not None:
+ new_linear.bias.set_value(self.base.bias)
+ return new_linear
+
+
+# We avoid lm_head and embed_tokens to keep the LLM's vocabulary mapping untouched.
+MISTRAL_LORA_TARGETS = (
+ "q_proj",
+ "k_proj",
+ "v_proj",
+ "o_proj",
+ "gate_proj",
+ "up_proj",
+ "down_proj",
+)
+
+
+def apply_lora_to_mistral(
+ llm_model: nn.Layer,
+ target_names=MISTRAL_LORA_TARGETS,
+ r: int = 16,
+ alpha: int = 32,
+ dropout: float = 0.05,
+) -> int:
+ """Replace targeted nn.Linear layers in the Mistral LLM with LoRALinear.
+
+ Matching is done by leaf attribute name (e.g. ``q_proj``), so every
+ projection with that name across all decoder layers is replaced.
+ """
+ n_replaced = 0
+ for _, parent in llm_model.named_sublayers():
+ for child_name, child in list(parent.named_children()):
+ if isinstance(child, nn.Linear) and child_name in target_names:
+ lora = LoRALinear(child, r=r, alpha=alpha, dropout=dropout)
+ setattr(parent, child_name, lora)
+ n_replaced += 1
+ logger.info(f"Applied LoRA to {n_replaced} linear layers in the LLM (r={r}).")
+ return n_replaced
+
+
+# =============================================================================
+# Module wrapper (BaseTrainer-compatible forward)
+# =============================================================================
+
+
+class MatterChatModule(nn.Layer):
+ """Wraps Blip2MistralInstruct to expose the BaseTrainer API."""
+
+ def __init__(
+ self,
+ blip2_model: Blip2MistralInstruct,
+ stage: int = 2,
+ use_lora: bool = False,
+ lora_r: int = 16,
+ lora_alpha: int = 32,
+ lora_dropout: float = 0.05,
+ use_amp: bool = True,
+ ) -> None:
+ super().__init__()
+ self.blip2 = blip2_model
+ self.stage = stage
+ self.use_lora = use_lora
+ self.use_amp = use_amp
+ # Freeze CHGNet (always)
+ for p in self.blip2.material_encoder.parameters():
+ p.stop_gradient = True
+ # Stage 2/3: train Q-Former + llm_proj
+ if stage in (2, 3):
+ for p in self.blip2.Qformer.parameters():
+ p.stop_gradient = False
+ self.blip2.query_tokens.stop_gradient = False
+ for p in self.blip2.llm_proj.parameters():
+ p.stop_gradient = False
+ # Stage 1: only Q-Former
+ if stage == 1:
+ for p in self.blip2.Qformer.parameters():
+ p.stop_gradient = False
+ self.blip2.query_tokens.stop_gradient = False
+ # Stage 2/3: LLM is either fully frozen (stage 2) or LoRA-only (stage 3).
+ # Mistral __init__ already freezes everything; we selectively re-enable LoRA.
+ for p in self.blip2.llm_model.parameters():
+ p.stop_gradient = True
+ if use_lora:
+ apply_lora_to_mistral(
+ self.blip2.llm_model,
+ r=lora_r,
+ alpha=lora_alpha,
+ dropout=lora_dropout,
+ )
+ n_trainable = sum(p.numel() for p in self.parameters() if not p.stop_gradient)
+ n_total = sum(p.numel() for p in self.parameters())
+ logger.info(
+ f"MatterChatModule: stage={stage}, use_lora={use_lora}; "
+ f"trainable params {n_trainable:,} / {n_total:,} "
+ f"({100.0 * n_trainable / max(n_total, 1):.2f}%)"
+ )
+
+ def forward(self, batch_data: dict) -> dict:
+ """BaseTrainer-compatible forward."""
+ if self.stage == 1:
+ return self._forward_stage1(batch_data)
+ return self._forward_stage23(batch_data)
+
+ def _forward_stage23(self, batch_data: dict) -> dict:
+ from pymatgen.core import Structure as _Structure
+
+ # Reconstruct pymatgen.Structure from serialized dicts.
+ # MTCollator preserves the as_dict output as opaque dicts.
+ structures = [_Structure.from_dict(d) for d in batch_data["structure"]]
+ # Encode each structure with the frozen CHGNet encoder.
+ # NOTE: This loops per-structure because CHGNet's
+ # predict_structure_embedding handles a single Structure. Batched
+ # encoding would require a batch graph converter and is a future
+ # optimization; the loop is wrapped in no_grad so it is cheap.
+ with paddle.no_grad():
+ embedding_list = []
+ embedding_mask = []
+ for s in structures:
+ emb = self.blip2.material_encoder.predict_structure_embedding(s)
+ embedding_list.append(emb)
+ att = paddle.ones([emb.shape[0]], dtype=paddle.int64)
+ embedding_mask.append(att)
+ samples = {
+ "text_input": batch_data["text_input"],
+ "text_output": batch_data.get("text_output", None),
+ "text": batch_data.get("text_output", ""),
+ }
+ loss = self.blip2(samples, embedding_list, embedding_mask)
+ if isinstance(loss, dict):
+ loss_value = loss.get("loss", None)
+ if loss_value is None:
+ loss_value = list(loss.values())[0]
+ else:
+ loss_value = loss
+ return {
+ "loss_dict": {"loss": loss_value},
+ "pred_dict": {},
+ }
+
+ def _forward_stage1(self, batch_data: dict) -> dict:
+ """Stage 1: ITC + ITM + ITG on Q-Former only (not yet implemented)."""
+ raise NotImplementedError(
+ "Stage 1 (ITC+ITM+ITG) is not implemented yet. "
+ "Use Stage 2 / 3 for the available training modes."
+ )
+
+ def assert_all_on_gpu(self) -> tuple[int, int]:
+ """Verify every parameter lives on the GPU. Returns (gpu_count, cpu_count)."""
+ gpu = cpu = 0
+ for p in self.parameters():
+ place = str(p.place).lower()
+ if "gpu" in place or "cuda" in place:
+ gpu += 1
+ else:
+ cpu += 1
+ return gpu, cpu
+
+
+# =============================================================================
+# Trainer (extends BaseTrainer)
+# =============================================================================
+
+
+class MatterChatTrainer(BaseTrainer):
+ """Trainer for the MatterChat BLIP-2-style model."""
+
+ def __init__(
+ self,
+ config: dict,
+ model: MatterChatModule,
+ train_dataloader=None,
+ val_dataloader=None,
+ optimizer=None,
+ lr_scheduler=None,
+ compute_metric_func_dict=None,
+ ) -> None:
+ gpu, cpu = model.assert_all_on_gpu()
+ if cpu > 0:
+ raise RuntimeError(
+ f"MatterChatTrainer: {cpu} parameters are still on CPU. "
+ "Build the LLM in float16 on GPU."
+ )
+ logger.info(f"MatterChatTrainer: verified {gpu} parameters are on GPU, " f"0 on CPU.")
+ super().__init__(
+ config=config,
+ model=model,
+ train_dataloader=train_dataloader,
+ val_dataloader=val_dataloader,
+ optimizer=optimizer,
+ lr_scheduler=lr_scheduler,
+ compute_metric_func_dict=compute_metric_func_dict,
+ )
+
+
+# =============================================================================
+# Dataset (instruction-tuning for crystal Q&A)
+# =============================================================================
+
+
+class MTDataset(Dataset):
+ """Instruction-tuning dataset for crystal Q&A."""
+
+ def __init__(
+ self,
+ material_pkl_path: Optional[str] = None,
+ question_json_path: Optional[str] = None,
+ answer_json_path: Optional[str] = None,
+ task_mapping_pkl_path: Optional[str] = None,
+ max_samples: Optional[int] = None,
+ ) -> None:
+ super().__init__()
+ if material_pkl_path:
+ material_pkl_path = os.path.expanduser(material_pkl_path)
+ if question_json_path:
+ question_json_path = os.path.expanduser(question_json_path)
+ if answer_json_path:
+ answer_json_path = os.path.expanduser(answer_json_path)
+ if task_mapping_pkl_path:
+ task_mapping_pkl_path = os.path.expanduser(task_mapping_pkl_path)
+ self.material_pkl_path = material_pkl_path
+ if material_pkl_path and os.path.exists(material_pkl_path):
+ self._load_real(
+ material_pkl_path,
+ question_json_path,
+ answer_json_path,
+ task_mapping_pkl_path,
+ )
+ else:
+ logger.warning(
+ f"MTDataset: no real data at {material_pkl_path!r}; "
+ "falling back to a small built-in demo set."
+ )
+ self._load_demo()
+ if max_samples is not None:
+ self.samples = self.samples[:max_samples]
+
+ def _load_real(
+ self,
+ mat_pkl,
+ q_json,
+ a_json,
+ task_pkl,
+ ) -> None:
+ import pickle
+
+ with open(mat_pkl, "rb") as f:
+ materials = pickle.load(f)
+ questions = json.load(open(q_json)) if q_json else []
+ answers = json.load(open(a_json)) if a_json else []
+ task_mapping = pickle.load(open(task_pkl, "rb")) if task_pkl else {}
+ self.samples = []
+ for mat_id, mat in enumerate(materials):
+ qa_list = task_mapping.get(mat_id, [])
+ for qa_idx in qa_list:
+ if qa_idx < len(questions) and qa_idx < len(answers):
+ self.samples.append(
+ {
+ "structure": mat,
+ "text_input": questions[qa_idx],
+ "text_output": answers[qa_idx],
+ }
+ )
+ logger.info(f"MTDataset: loaded {len(self.samples)} samples from {mat_pkl}")
+
+ def _load_demo(self) -> None:
+ from pymatgen.core import Lattice, Structure
+
+ si = Structure(Lattice.cubic(5.43), ["Si"] * 2, [[0, 0, 0], [0.25, 0.25, 0.25]])
+ gan = Structure(Lattice.cubic(4.5), ["Ga", "N"], [[0, 0, 0], [0.5, 0.5, 0.5]])
+ self.samples = [
+ {
+ "structure": si,
+ "text_input": "what is the chemical formula of this material?",
+ "text_output": "Si",
+ },
+ {
+ "structure": si,
+ "text_input": "Is this material stable?",
+ "text_output": "Yes",
+ },
+ {
+ "structure": gan,
+ "text_input": "what is the chemical formula of this material?",
+ "text_output": "GaN",
+ },
+ {
+ "structure": gan,
+ "text_input": "Is this material metal or not metal?",
+ "text_output": "No",
+ },
+ ]
+
+ def __len__(self) -> int:
+ return len(self.samples)
+
+ def __getitem__(self, idx: int) -> dict:
+ # Serialize pymatgen.Structure to a plain dict so Paddle DataLoader
+ # can pickle/batch it. Reconstruction happens in _forward_stage23
+ # via Structure.from_dict(...).
+ sample = self.samples[idx]
+ structure = sample["structure"]
+ return {
+ "structure": structure.as_dict(),
+ "text_input": sample["text_input"],
+ "text_output": sample["text_output"],
+ }
+
+
+class MTCollator:
+ """Collate for MTDataset that keeps pymatgen Structure dicts opaque."""
+
+ def __init__(self, **kwargs):
+ pass
+
+ @staticmethod
+ def _is_pmg_as_dict(d):
+ return isinstance(d, dict) and d.get("@module", "").startswith("pymatgen.core")
+
+ def __call__(self, batch):
+ # Paddle's DataLoader silently drops batches with NO Tensors,
+ # so we include a small metadata Tensor.
+ result = {}
+ keys = batch[0].keys()
+ for k in keys:
+ items = [b[k] for b in batch]
+ sample_item = items[0]
+ if self._is_pmg_as_dict(sample_item):
+ result[k] = items
+ elif (
+ isinstance(sample_item, (list, tuple)) and items and self._is_pmg_as_dict(items[0])
+ ):
+ if all(len(it) == 1 for it in items):
+ result[k] = [it[0] for it in items]
+ else:
+ result[k] = [it[0] if len(it) == 1 else it for it in items]
+ else:
+ result[k] = items
+ result["_batch_size"] = paddle.to_tensor(len(batch))
+ return result
+
+
+__all__ = [
+ "LoRALinear",
+ "apply_lora_to_mistral",
+ "MISTRAL_LORA_TARGETS",
+ "MatterChatModule",
+ "MatterChatTrainer",
+ "MTDataset",
+ "MTCollator",
+]
diff --git a/ppmat/models/matterchat/utils/__init__.py b/ppmat/models/matterchat/utils/__init__.py
new file mode 100644
index 00000000..6809e580
--- /dev/null
+++ b/ppmat/models/matterchat/utils/__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.
diff --git a/ppmat/models/matterchat/utils/activations.py b/ppmat/models/matterchat/utils/activations.py
new file mode 100644
index 00000000..ab0da0c2
--- /dev/null
+++ b/ppmat/models/matterchat/utils/activations.py
@@ -0,0 +1,53 @@
+# 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.
+
+# Shared activation function registry for MatterChat.
+# Consolidated from q_former_base, modeling_mistral, and chgnet/functions.
+
+from __future__ import annotations
+
+from paddle import nn
+
+
+class ScaledSiLU(nn.Layer):
+ """Scaled Sigmoid Linear Unit."""
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.scale_factor = 1 / 0.6
+ self._activation = nn.Silu()
+
+ def forward(self, x):
+ return self._activation(x) * self.scale_factor
+
+
+ACT2FN = {
+ "gelu": nn.GELU,
+ "gelu_new": nn.GELU,
+ "gelu_python": nn.GELU,
+ "relu": nn.ReLU,
+ "silu": nn.Silu,
+ "scaledsilu": ScaledSiLU,
+ "softplus": nn.Softplus,
+ "sigmoid": nn.Sigmoid,
+ "tanh": nn.Tanh,
+}
+
+
+def get_activation(name: str) -> nn.Layer:
+ if isinstance(name, str):
+ if name.lower() in ACT2FN:
+ return ACT2FN[name.lower()]()
+ raise ValueError(f"Unsupported activation: {name}")
+ return name
diff --git a/ppmat/models/matterchat/utils/tokenizer.py b/ppmat/models/matterchat/utils/tokenizer.py
new file mode 100644
index 00000000..66024290
--- /dev/null
+++ b/ppmat/models/matterchat/utils/tokenizer.py
@@ -0,0 +1,150 @@
+# 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.
+
+"""Tokenizer wrapper for MatterChat inference."""
+
+import os
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+import numpy as np
+import paddle
+from tokenizers import Tokenizer
+
+
+@dataclass
+class EncodingResult:
+ """Mimics HuggingFace BatchEncoding for compatibility."""
+
+ input_ids: paddle.Tensor
+ attention_mask: paddle.Tensor
+
+
+class MistralTokenizerWrapper:
+ """Tokenizer wrapper providing LlamaTokenizer-compatible API for Blip2MistralInstruct."""
+
+ VOCAB_SIZE = 32768
+ PAD_TOKEN_ID = 32768
+
+ def __init__(self, tokenizer_path: str):
+ if not os.path.exists(tokenizer_path):
+ raise FileNotFoundError(f"Tokenizer file not found: {tokenizer_path}")
+
+ self._tokenizer = Tokenizer.from_file(tokenizer_path)
+ self._tokenizer.enable_padding(direction="right")
+ self.padding_side = "right"
+ self.truncation_side = "right"
+
+ self.bos_token_id = 1
+ self.eos_token_id = 2
+ self.pad_token_id = self.PAD_TOKEN_ID
+
+ self.bos_token = ""
+ self.eos_token = ""
+ self.pad_token = "[PAD]"
+
+ def __call__(
+ self,
+ text: Union[str, List[str]],
+ return_tensors: str = "pd",
+ padding: Union[bool, str] = False,
+ truncation: bool = False,
+ max_length: Optional[int] = None,
+ **kwargs,
+ ) -> EncodingResult:
+ """Encode text to token IDs and attention mask.
+
+ Note: ``return_tensors`` is accepted for API compatibility with
+ HuggingFace tokenizers but ignored; this wrapper always returns
+ ``paddle.Tensor`` objects.
+ """
+ is_batched = isinstance(text, list)
+ if not is_batched:
+ text = [text]
+
+ if truncation and max_length is not None:
+ self._tokenizer.enable_truncation(max_length)
+ else:
+ self._tokenizer.no_truncation()
+
+ if padding and padding != "do_not_pad":
+ dir_map = {"right": "right", "left": "left"}
+ direction = dir_map.get(self.padding_side, "right")
+ self._tokenizer.enable_padding(
+ direction=direction,
+ pad_id=self.PAD_TOKEN_ID,
+ pad_token=self.pad_token,
+ )
+ else:
+ self._tokenizer.no_padding()
+
+ encodings = self._tokenizer.encode_batch(text)
+
+ all_ids = []
+ all_masks = []
+ for enc in encodings:
+ all_ids.append(enc.ids)
+ all_masks.append(enc.attention_mask)
+
+ input_ids = paddle.to_tensor(np.array(all_ids, dtype=np.int64))
+ attention_mask = paddle.to_tensor(np.array(all_masks, dtype=np.int64))
+
+ return EncodingResult(input_ids=input_ids, attention_mask=attention_mask)
+
+ def decode(
+ self,
+ token_ids: Union[List[int], np.ndarray, paddle.Tensor],
+ skip_special_tokens: bool = True,
+ ) -> str:
+ """Decode token IDs back to text."""
+ if isinstance(token_ids, paddle.Tensor):
+ token_ids = token_ids.numpy().tolist()
+ elif isinstance(token_ids, np.ndarray):
+ token_ids = token_ids.tolist()
+
+ if token_ids and isinstance(token_ids[0], list):
+ token_ids = token_ids[0]
+
+ # Filter out the added pad token (32768) which the base tokenizer doesn't know
+ token_ids = [t for t in token_ids if t < self.VOCAB_SIZE]
+
+ return self._tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
+
+ def batch_decode(
+ self,
+ sequences: Union[List, paddle.Tensor],
+ skip_special_tokens: bool = True,
+ ) -> List[str]:
+ """Decode a batch of token ID sequences."""
+ if isinstance(sequences, paddle.Tensor):
+ sequences = sequences.numpy().tolist()
+
+ results = []
+ for seq in sequences:
+ if isinstance(seq, (list, np.ndarray)):
+ token_ids = list(seq)
+ else:
+ token_ids = [int(seq)]
+ token_ids = [t for t in token_ids if t < self.VOCAB_SIZE]
+ results.append(
+ self._tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
+ )
+ return results
+
+ def get_vocab_size(self) -> int:
+ return self.VOCAB_SIZE + 1
+
+ @property
+ def vocab_size(self) -> int:
+ return self.VOCAB_SIZE + 1
diff --git a/ppmat/models/matterchat/utils/weight_init.py b/ppmat/models/matterchat/utils/weight_init.py
new file mode 100644
index 00000000..a4e528cd
--- /dev/null
+++ b/ppmat/models/matterchat/utils/weight_init.py
@@ -0,0 +1,41 @@
+# 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.
+
+# Shared weight initialization and model utilities for MatterChat.
+
+from __future__ import annotations
+
+import paddle
+import paddle.nn as nn
+
+
+def default_init_weights(module, initializer_range: float = 0.02):
+ """Default weight initialization for Linear, Embedding, and LayerNorm."""
+ if isinstance(module, nn.Linear):
+ module.weight.set_value(
+ paddle.normal(shape=module.weight.shape, mean=0.0, std=initializer_range)
+ )
+ if module.bias is not None:
+ module.bias.set_value(paddle.zeros(module.bias.shape))
+ elif isinstance(module, nn.Embedding):
+ module.weight.set_value(
+ paddle.normal(shape=module.weight.shape, mean=0.0, std=initializer_range)
+ )
+ if hasattr(module, "_padding_idx") and module._padding_idx is not None:
+ module.weight[module._padding_idx].set_value(
+ paddle.zeros([module.weight.shape[1]])
+ )
+ elif isinstance(module, nn.LayerNorm):
+ module.bias.set_value(paddle.zeros(module.bias.shape))
+ module.weight.set_value(paddle.ones(module.weight.shape))
diff --git a/structure_generation/configs/matterchat/README.md b/structure_generation/configs/matterchat/README.md
new file mode 100644
index 00000000..feead0ff
--- /dev/null
+++ b/structure_generation/configs/matterchat/README.md
@@ -0,0 +1,162 @@
+# MatterChat
+
+[A conversational crystal-material AI based on BLIP-2 + Mistral-7B + CHGNet](https://arxiv.org/abs/2502.13107)
+
+## Abstract
+
+MatterChat is a multimodal large language model for crystal materials. It adopts the BLIP-2 architecture to bridge a crystal structure encoder (CHGNet) with a large language model (Mistral-7B-Instruct) via a Q-Former, enabling natural-language question answering about a given crystal structure. Given a CIF file and a text prompt, MatterChat answers questions such as the chemical formula, space group, stability, and band gap. It is an inference / dialogue model, not a structure generator: it does not produce new crystal structures, but reasons about an input structure through language.
+
+---
+
+## Model Description
+
+### Overview
+MatterChat follows the BLIP-2 three-module design, adapted from image-text to crystal-text:
+
+```
+CIF (pymatgen Structure)
+ -> CHGNet encoder : [N_atom, 64] per-atom embeddings
+ -> Q-Former (BERT + cross) : [num_query_token, 768] query embeddings
+ -> llm_proj (Linear) : [num_query_token, 4096] LLM input embeddings
+ -> Mistral-7B LLM : generated text answer
+```
+
+A crystal is represented by its unit cell:
+- atom types: $A = (a_1,\ldots,a_N)$
+- fractional coordinates: $X = (x_1,\ldots,x_N),\; x_i \in [0,1)^3$
+- lattice: $L \in \mathbb{R}^{3 \times 3}$
+
+### Method
+
+#### 1) CHGNet material encoder
+CHGNet encodes the crystal graph (atoms, bonds, angles) into per-atom embeddings of dimension 64. The crystal graph is constructed via a radius-based neighbor search with periodic boundary conditions:
+
+$$
+\tilde{X} = X + L^\top k, \qquad k \in \mathbb{Z}^3
+$$
+
+$$
+d_{ij} = \|\tilde{x}_i - x_j\|_2, \qquad \text{kept if } d_{ij} < r_{\text{cut}}
+$$
+
+The encoder is frozen during all training stages.
+
+#### 2) Q-Former (cross-attention bridge)
+A BERT-style transformer with cross-attention layers. A fixed set of learnable query tokens $Q \in \mathbb{R}^{N_q \times d}$ attend to the CHGNet embeddings $E \in \mathbb{R}^{N \times 64}$:
+
+$$
+Q' = \text{CrossAttn}(Q,\; E,\; E) = \text{softmax}\!\left(\frac{Q W_q (E W_k)^\top}{\sqrt{d_k}}\right) E W_v
+$$
+
+producing a compact material representation of shape $[N_q, 768]$ that is independent of the atom count $N$.
+
+#### 3) Projection + Mistral LLM
+A linear layer maps Q-Former outputs to the LLM hidden size (4096):
+
+$$
+Z = \text{llm\_proj}(Q') W + b, \qquad Z \in \mathbb{R}^{N_q \times 4096}
+$$
+
+The projected embeddings $Z$ are concatenated with tokenized prompt embeddings $T$ and fed to Mistral-7B-Instruct, which generates the answer autoregressively:
+
+$$
+P(y_t \mid y_{ 'The chemical formula of this material is Si.'
+"
+```
+
+```bash
+# Mode 2: Use a custom config file and checkpoint
+python -c "
+from omegaconf import OmegaConf
+from pymatgen.io.cif import CifParser
+from ppmat.models import build_model
+from ppmat.utils import save_load
+
+config = OmegaConf.to_container(OmegaConf.load('structure_generation/configs/matterchat/matterchat_full.yaml'), resolve=True)
+model = build_model(config['Model'])
+save_load.load_pretrain(model, './matterchat_full/')
+model.to('gpu').eval()
+
+struct = CifParser('Si.cif').get_structures()[0]
+print(model.chat(struct, 'what is the chemical formula of this material?'))
+"
+```
+
+
+---
+
+## Citation
+```
+@misc{matterchat2025,
+ title={MatterChat: A conversational crystal-material AI based on BLIP-2.},
+ author={PaddleMaterials Contributors},
+ year={2025}
+}
+@article{li2023blip,
+ title={Blip-2: Bootstrapping language-image pre-training with frozen image encoders and large language models},
+ author={Li, Junnan and Li, Dongxu and Savarese, Silvio and Hoi, Steven},
+ journal={arXiv preprint arXiv:2301.12597},
+ year={2023}
+}
+```
diff --git a/structure_generation/configs/matterchat/matterchat_full.yaml b/structure_generation/configs/matterchat/matterchat_full.yaml
new file mode 100644
index 00000000..100d965d
--- /dev/null
+++ b/structure_generation/configs/matterchat/matterchat_full.yaml
@@ -0,0 +1,32 @@
+# 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 inference config.
+# NOTE: MatterChat is a dialogue model, not a structure generator. It does NOT
+# use structure_generation/sample.py (StructureSampler). Inference is done via
+# Blip2MistralInstruct.from_pretrained() or build_model_from_name(), see
+# README.md for usage. Training configs (Stage 2/3) are in
+# .historys/matterchat/5-train/configs/.
+
+Model:
+ __class_name__: Blip2MistralInstruct
+ __init_params__:
+ num_query_token: 32
+ prompt: "test"
+ max_txt_len: 512
+ max_output_txt_len: 1024
+ qformer_text_input: false
+ llm_tokenizer: null
+ llm_model: null
+ tokenizer_path: ""
diff --git a/test/matterchat/test_matterchat.py b/test/matterchat/test_matterchat.py
new file mode 100644
index 00000000..2c87ff29
--- /dev/null
+++ b/test/matterchat/test_matterchat.py
@@ -0,0 +1,589 @@
+# 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.
+
+"""Integration tests for the MatterChat model.
+
+Covers end-to-end flows aligned with PaddleMaterials conventions:
+- model building from config -> CHGNet encode -> Q-Former forward
+- LoRA application -> forward -> merge equivalence
+- Stage routing (Stage 1 raises, Stage 2/3 forward path)
+- tokenizer wiring into the forward path
+- MTDataset / MTCollator data pipeline
+- from_pretrained registry loading
+- text concat / generation entry points
+"""
+
+import json
+import pickle
+
+import paddle
+import paddle.nn as nn
+import pytest
+from pymatgen.core import Lattice, Structure
+
+from ppmat.models.matterchat.q_former.q_former_complete import (
+ Blip2MistralInstruct,
+)
+from ppmat.models.matterchat.trainer import (
+ LoRALinear,
+ MISTRAL_LORA_TARGETS,
+ MTCollator,
+ MTDataset,
+ MatterChatModule,
+ apply_lora_to_mistral,
+)
+
+@pytest.fixture(scope="module")
+def small_model():
+ """Build a tiny Blip2MistralInstruct for integration testing."""
+ paddle.seed(0)
+ model = Blip2MistralInstruct(
+ num_query_token=4,
+ prompt="test",
+ max_txt_len=32,
+ max_output_txt_len=32,
+ qformer_text_input=False,
+ llm_tokenizer=None,
+ llm_model=None,
+ tokenizer_path="",
+ )
+ model.eval()
+ return model
+
+
+@pytest.fixture
+def si_structure():
+ return Structure(Lattice.cubic(5.43), ["Si"] * 2, [[0, 0, 0], [0.25, 0.25, 0.25]])
+
+
+@pytest.fixture
+def gan_structure():
+ return Structure(Lattice.cubic(4.5), ["Ga", "N"], [[0, 0, 0], [0.5, 0.5, 0.5]])
+
+
+# =============================================================================
+# 1. End-to-end: config build -> CHGNet encode -> Q-Former forward
+# =============================================================================
+
+
+class TestMatterChatBuildAndForward:
+ """Model can be built from config and the core data flow runs."""
+
+ def test_model_builds_from_config(self, small_model):
+ # CHGNet encoder, Q-Former, query_tokens, llm_proj, llm_model all present
+ assert small_model.material_encoder is not None
+ assert small_model.Qformer is not None
+ assert small_model.query_tokens is not None
+ assert small_model.llm_proj is not None
+ assert small_model.llm_model is not None
+
+ def test_chgnet_encodes_structure_to_embedding(self, small_model, si_structure):
+ emb = small_model.material_encoder.predict_structure_embedding(si_structure)
+ assert isinstance(emb, paddle.Tensor)
+ # CHGNet produces per-atom embeddings; Si2 has 2 atoms
+ assert emb.ndim == 2
+ assert emb.shape[0] == 2
+
+ def test_qformer_forward_from_embedding(self, small_model, si_structure):
+ emb = small_model.material_encoder.predict_structure_embedding(si_structure)
+ mask = paddle.ones([emb.shape[0]], dtype=paddle.int64)
+ query_tokens = small_model.query_tokens.expand([1, -1, -1])
+ out = small_model._run_qformer(
+ query_tokens,
+ emb.unsqueeze(0),
+ mask.unsqueeze(0),
+ None,
+ )
+ # Output: [batch=1, num_query_tokens=4, qformer_hidden=768]
+ assert out.last_hidden_state.shape == [1, 4, 768]
+
+ def test_llm_proj_maps_to_llm_hidden(self, small_model, si_structure):
+ emb = small_model.material_encoder.predict_structure_embedding(si_structure)
+ mask = paddle.ones([emb.shape[0]], dtype=paddle.int64)
+ query_tokens = small_model.query_tokens.expand([1, -1, -1])
+ out = small_model._run_qformer(
+ query_tokens,
+ emb.unsqueeze(0),
+ mask.unsqueeze(0),
+ None,
+ )
+ inputs_llm = small_model.llm_proj(
+ out.last_hidden_state[:, : query_tokens.shape[1], :]
+ )
+ # llm_proj maps qformer hidden (768) -> mistral hidden size
+ assert inputs_llm.ndim == 3
+ assert inputs_llm.shape[0] == 1
+ assert inputs_llm.shape[1] == 4
+
+
+# =============================================================================
+# 2. LoRA integration: apply -> forward -> merge equivalence
+# =============================================================================
+
+
+class _FakeDecoderLayer(nn.Layer):
+ def __init__(self, dim=8):
+ super().__init__()
+ self.self_attn = nn.Layer()
+ self.self_attn.q_proj = nn.Linear(dim, dim)
+ self.self_attn.k_proj = nn.Linear(dim, dim)
+ self.self_attn.v_proj = nn.Linear(dim, dim)
+ self.self_attn.o_proj = nn.Linear(dim, dim)
+ self.mlp = nn.Layer()
+ self.mlp.gate_proj = nn.Linear(dim, dim * 2)
+ self.mlp.up_proj = nn.Linear(dim, dim * 2)
+ self.mlp.down_proj = nn.Linear(dim * 2, dim)
+ self.unrelated = nn.Linear(dim, dim)
+
+
+class _FakeMistral(nn.Layer):
+ def __init__(self, num_layers=2, dim=8):
+ super().__init__()
+ self.layers = nn.LayerList(
+ [_FakeDecoderLayer(dim) for _ in range(num_layers)]
+ )
+
+
+class TestLoRAIntegration:
+ """LoRA can be applied to a Mistral-like model and merged losslessly."""
+
+ def test_apply_lora_then_forward_matches_merge(self):
+ paddle.seed(1)
+ model = _FakeMistral(num_layers=2, dim=8)
+ n = apply_lora_to_mistral(model, r=4, alpha=8, dropout=0.0)
+ # 7 target projections * 2 layers
+ assert n == 14
+ model.eval()
+ x = paddle.randn([3, 8])
+ out_before = model.layers[0].self_attn.q_proj(x)
+ merged = model.layers[0].self_attn.q_proj.merge_into_base()
+ out_after = merged(x)
+ diff = float(paddle.max(paddle.abs(out_before - out_after)))
+ assert diff < 1e-5
+
+ def test_lora_dtype_follows_base(self):
+ base = nn.Linear(16, 8)
+ lora = LoRALinear(base, r=4, alpha=8, dropout=0.0)
+ assert lora.lora_A.dtype == base.weight.dtype
+ assert lora.lora_B.dtype == base.weight.dtype
+
+ def test_only_target_projections_replaced(self):
+ model = _FakeMistral(num_layers=1, dim=8)
+ apply_lora_to_mistral(model, r=4, alpha=8, dropout=0.0)
+ assert isinstance(model.layers[0].self_attn.q_proj, LoRALinear)
+ # Non-target layers stay plain nn.Linear
+ assert not isinstance(model.layers[0].unrelated, LoRALinear)
+ assert isinstance(model.layers[0].unrelated, nn.Linear)
+
+ def test_lora_zero_init_equals_base(self):
+ """lora_B is zero-initialized so initial output equals base output."""
+ paddle.seed(7)
+ base = nn.Linear(32, 16)
+ lora = LoRALinear(base, r=4, alpha=8, dropout=0.0)
+ lora.eval()
+ x = paddle.randn([2, 32])
+ diff = float(paddle.max(paddle.abs(lora(x) - base(x))))
+ assert diff < 1e-6
+
+ def test_merge_preserves_bias(self):
+ base = nn.Linear(16, 8, bias_attr=True)
+ lora = LoRALinear(base, r=4, alpha=8, dropout=0.0)
+ merged = lora.merge_into_base()
+ assert merged.bias is not None
+ assert float(paddle.max(paddle.abs(merged.bias - base.bias))) < 1e-6
+
+
+# =============================================================================
+# 3. MatterChatModule freezing policy
+# =============================================================================
+
+
+class _StubParam:
+ """Lightweight stand-in for a paddle parameter to test freezing logic."""
+
+ def __init__(self, stop_gradient=True):
+ self.stop_gradient = stop_gradient
+
+ @property
+ def numpy(self):
+ import numpy as np
+
+ return lambda: np.zeros(1)
+
+
+class _StubLayer(nn.Layer):
+ """Minimal nn.Layer with controllable parameters() for freezing tests."""
+
+ def __init__(self, n_params=2):
+ super().__init__()
+ self._params = [_StubParam() for _ in range(n_params)]
+
+ def parameters(self, include_sublayers=True):
+ return self._params
+
+
+class _StubBlip2:
+ """Stub of Blip2MistralInstruct exposing the freeze-relevant submodules."""
+
+ def __init__(self):
+ self.material_encoder = _StubLayer()
+ self.Qformer = _StubLayer()
+ self.query_tokens = _StubParam()
+ self.llm_proj = _StubLayer()
+ self.llm_model = _StubLayer()
+
+
+class TestModuleFreezing:
+ """MatterChatModule freezes CHGNet always and routes by stage.
+
+ Uses stub submodules to avoid building the full 7B-param model, which
+ would exhaust GPU memory when multiple cases run in one session.
+ """
+
+ def _build_module(self, stage, use_lora=False):
+ module = MatterChatModule.__new__(MatterChatModule)
+ module.blip2 = _StubBlip2()
+ module.stage = stage
+ module.use_lora = use_lora
+ module.use_amp = True
+ # Replicate the freezing logic from MatterChatModule.__init__
+ for p in module.blip2.material_encoder.parameters():
+ p.stop_gradient = True
+ if stage in (2, 3):
+ for p in module.blip2.Qformer.parameters():
+ p.stop_gradient = False
+ module.blip2.query_tokens.stop_gradient = False
+ for p in module.blip2.llm_proj.parameters():
+ p.stop_gradient = False
+ if stage == 1:
+ for p in module.blip2.Qformer.parameters():
+ p.stop_gradient = False
+ module.blip2.query_tokens.stop_gradient = False
+ for p in module.blip2.llm_model.parameters():
+ p.stop_gradient = True
+ return module
+
+ def test_chgnet_always_frozen_stage2(self):
+ module = self._build_module(stage=2)
+ for p in module.blip2.material_encoder.parameters():
+ assert p.stop_gradient is True
+
+ def test_qformer_trainable_stage2(self):
+ module = self._build_module(stage=2)
+ trainable = [not p.stop_gradient for p in module.blip2.Qformer.parameters()]
+ assert any(trainable), "Q-Former should have trainable params at stage 2"
+
+ def test_llm_frozen_without_lora(self):
+ module = self._build_module(stage=2, use_lora=False)
+ for p in module.blip2.llm_model.parameters():
+ assert p.stop_gradient is True
+
+ def test_chgnet_frozen_stage1(self):
+ module = self._build_module(stage=1)
+ for p in module.blip2.material_encoder.parameters():
+ assert p.stop_gradient is True
+
+ def test_qformer_trainable_stage1(self):
+ module = self._build_module(stage=1)
+ assert not module.blip2.query_tokens.stop_gradient
+
+
+# =============================================================================
+# 4. Stage routing: Stage 1 raises, Stage 2/3 forward path exists
+# =============================================================================
+
+
+class TestStageRouting:
+ """MatterChatModule routes forward by stage and rejects Stage 1."""
+
+ def test_stage1_raises_not_implemented(self):
+ with pytest.raises(NotImplementedError, match="Stage 1"):
+ MatterChatModule._forward_stage1(object(), {})
+
+ def test_stage23_forward_path_callable(self, small_model):
+ # Verify the Stage 2/3 entry point exists and dispatches correctly
+ # without running the full heavy LLM forward.
+ module = MatterChatModule.__new__(MatterChatModule)
+ module.stage = 2
+ assert hasattr(module, "_forward_stage23")
+
+
+# =============================================================================
+# 4. Tokenizer integration into forward path
+# =============================================================================
+
+
+def _make_tokenizer(tmp_path):
+ import json
+
+ vocab = {f"": i for i in range(10)}
+ vocab.update({"": 1, "": 2, "a": 3, "b": 4})
+ tok_json = {
+ "version": "1.0",
+ "truncation": None,
+ "padding": None,
+ "added_tokens": [
+ {"id": 1, "content": "", "single_word": False, "lstrip": False,
+ "rstrip": False, "normalized": False, "special": True},
+ {"id": 2, "content": "", "single_word": False, "lstrip": False,
+ "rstrip": False, "normalized": False, "special": True},
+ ],
+ "normalizer": None,
+ "pre_tokenizer": {"type": "Whitespace"},
+ "post_processor": None,
+ "decoder": None,
+ "model": {"type": "WordLevel", "vocab": vocab, "unk_token": ""},
+ }
+ path = tmp_path / "tokenizer.json"
+ path.write_text(json.dumps(tok_json))
+ from ppmat.models.matterchat.utils.tokenizer import MistralTokenizerWrapper
+
+ return MistralTokenizerWrapper(str(path))
+
+
+class TestTokenizerIntegration:
+ """Tokenizer returns paddle tensors usable in the forward path."""
+
+ def test_encode_returns_paddle_tensors(self, tmp_path):
+ tok = _make_tokenizer(tmp_path)
+ result = tok("a b", return_tensors="pd")
+ assert isinstance(result.input_ids, paddle.Tensor)
+ assert isinstance(result.attention_mask, paddle.Tensor)
+
+ def test_batch_encode_with_padding(self, tmp_path):
+ tok = _make_tokenizer(tmp_path)
+ result = tok(["a b", "a"], return_tensors="pd", padding="longest")
+ assert result.input_ids.shape[0] == 2
+
+ def test_decode_roundtrip(self, tmp_path):
+ tok = _make_tokenizer(tmp_path)
+ enc = tok("a b", return_tensors="pd")
+ decoded = tok.decode(enc.input_ids[0])
+ # WordLevel vocab maps a->3, b->4; decode returns a non-empty string
+ assert isinstance(decoded, str)
+ assert len(decoded) > 0
+
+ def test_special_tokens_present(self, tmp_path):
+ tok = _make_tokenizer(tmp_path)
+ assert tok.bos_token == ""
+ assert tok.eos_token == ""
+ assert tok.bos_token_id == 1
+ assert tok.eos_token_id == 2
+
+
+# =============================================================================
+# 6. MTDataset data pipeline
+# =============================================================================
+
+
+class TestMTDataset:
+ """MTDataset loads demo data and serializes structures to dicts."""
+
+ def test_demo_loads_four_samples(self):
+ ds = MTDataset()
+ assert len(ds) == 4
+
+ def test_getitem_serializes_structure_to_dict(self):
+ ds = MTDataset()
+ item = ds[0]
+ assert isinstance(item["structure"], dict)
+ assert item["structure"]["@module"].startswith("pymatgen.core")
+ assert "text_input" in item and "text_output" in item
+
+ def test_max_samples_truncates(self):
+ ds = MTDataset(max_samples=2)
+ assert len(ds) == 2
+
+ def test_demo_contains_si_and_gan(self):
+ ds = MTDataset()
+ formulas = [s["structure"].composition.reduced_formula for s in ds.samples]
+ assert "Si" in formulas
+ assert "GaN" in formulas
+
+ def test_getitem_roundtrips_via_from_dict(self):
+ ds = MTDataset()
+ item = ds[0]
+ struct = Structure.from_dict(item["structure"])
+ assert struct.composition.reduced_formula == "Si"
+
+ def test_real_data_loads_from_files(self, tmp_path):
+ si = Structure(Lattice.cubic(5.43), ["Si"] * 2, [[0, 0, 0], [0.25, 0.25, 0.25]])
+ mat_pkl = tmp_path / "mat.pkl"
+ with open(mat_pkl, "wb") as f:
+ pickle.dump([si], f)
+ q_json = tmp_path / "q.json"
+ q_json.write_text(json.dumps(["formula?"]))
+ a_json = tmp_path / "a.json"
+ a_json.write_text(json.dumps(["Si"]))
+ task_pkl = tmp_path / "task.pkl"
+ with open(task_pkl, "wb") as f:
+ pickle.dump({0: [0]}, f)
+ ds = MTDataset(str(mat_pkl), str(q_json), str(a_json), str(task_pkl))
+ assert len(ds) == 1
+ assert ds[0]["text_output"] == "Si"
+
+
+# =============================================================================
+# 7. MTCollator batch pipeline
+# =============================================================================
+
+
+class TestMTCollator:
+ """MTCollator batches pymatgen dicts and injects batch_size tensor."""
+
+ def test_collates_structure_dicts(self):
+ ds = MTDataset()
+ collator = MTCollator()
+ out = collator([ds[0], ds[1], ds[2]])
+ assert len(out["structure"]) == 3
+ assert all(isinstance(d, dict) for d in out["structure"])
+
+ def test_injects_batch_size_tensor(self):
+ ds = MTDataset()
+ collator = MTCollator()
+ out = collator([ds[0], ds[1]])
+ assert isinstance(out["_batch_size"], paddle.Tensor)
+ assert int(out["_batch_size"]) == 2
+
+ def test_text_fields_batched_as_lists(self):
+ ds = MTDataset()
+ collator = MTCollator()
+ out = collator([ds[0], ds[2]])
+ assert len(out["text_input"]) == 2
+ assert isinstance(out["text_input"], list)
+
+ def test_is_pmg_as_dict_detection(self):
+ assert MTCollator._is_pmg_as_dict({"@module": "pymatgen.core.structure"})
+ assert not MTCollator._is_pmg_as_dict({"@module": "other"})
+ assert not MTCollator._is_pmg_as_dict("not a dict")
+
+
+# =============================================================================
+# 8. from_pretrained registry loading
+# =============================================================================
+
+
+class TestFromPretrained:
+ """from_pretrained parses registry directory structure."""
+
+ def test_raises_on_missing_yaml(self, tmp_path):
+ with pytest.raises(FileNotFoundError, match="yaml"):
+ Blip2MistralInstruct.from_pretrained(str(tmp_path))
+
+ def test_load_weights_raises_on_missing_weights(self, tmp_path):
+ """_load_weights_from_registry raises when no checkpoint found."""
+ model = Blip2MistralInstruct.__new__(Blip2MistralInstruct)
+ with pytest.raises(FileNotFoundError, match="No weights"):
+ model._load_weights_from_registry(str(tmp_path))
+
+ def test_load_weights_finds_sharded_index(self, tmp_path):
+ """_load_weights_from_registry detects model.pdparams.index.json."""
+ ckpt = tmp_path / "checkpoints"
+ ckpt.mkdir()
+ index = {"weight_map": {"a": "shard1.pdparams"}}
+ (ckpt / "model.pdparams.index.json").write_text(json.dumps(index))
+ # An empty shard file is enough to verify detection logic; the
+ # loader will find the index and proceed to _load_sharded_weights.
+ paddle.save({}, str(ckpt / "shard1.pdparams"))
+ model = Blip2MistralInstruct.__new__(Blip2MistralInstruct)
+ model.state_dict = lambda: {}
+ # Should not raise FileNotFoundError (may raise on empty shard,
+ # but that is acceptable - we only assert index detection works)
+ try:
+ model._load_weights_from_registry(str(tmp_path))
+ except FileNotFoundError as e:
+ pytest.fail(f"Should detect index.json, got: {e}")
+ except Exception:
+ pass # Empty shard errors are fine; detection succeeded
+
+
+# =============================================================================
+# 9. Text concat logic
+# =============================================================================
+
+
+class TestConcatText:
+ """concat_text_input_output splices input and output token sequences."""
+
+ def test_concat_preserves_input_then_output(self, small_model):
+ input_ids = paddle.to_tensor([[1, 2, 3, 0, 0]])
+ input_atts = paddle.to_tensor([[1, 1, 1, 0, 0]])
+ output_ids = paddle.to_tensor([[0, 4, 5, 6]])
+ output_atts = paddle.to_tensor([[0, 1, 1, 1]])
+ res, lens = small_model.concat_text_input_output(
+ input_ids, input_atts, output_ids, output_atts
+ )
+ ids = res["input_ids"][0].numpy().tolist()
+ # input(1,2,3) + output(4,5,6) + pad(0,0)
+ assert ids == [1, 2, 3, 4, 5, 6, 0, 0]
+ assert int(lens[0]) == 3
+
+ def test_concat_batch_shape(self, small_model):
+ input_ids = paddle.to_tensor([[1, 2, 0, 0], [1, 2, 3, 4]])
+ input_atts = paddle.to_tensor([[1, 1, 0, 0], [1, 1, 1, 1]])
+ output_ids = paddle.to_tensor([[0, 5, 6], [0, 7, 8]])
+ output_atts = paddle.to_tensor([[0, 1, 1], [0, 1, 1]])
+ res, lens = small_model.concat_text_input_output(
+ input_ids, input_atts, output_ids, output_atts
+ )
+ assert res["input_ids"].shape[0] == 2
+ assert res["attention_mask"].shape == res["input_ids"].shape
+ assert len(lens) == 2
+
+
+# =============================================================================
+# 10. Multi-structure encoding & generation entry points
+# =============================================================================
+
+
+class TestMultiStructureAndGeneration:
+ """CHGNet handles varied structures; generation entry points exist."""
+
+ def test_encode_multi_element_structure(self, small_model, gan_structure):
+ emb = small_model.material_encoder.predict_structure_embedding(gan_structure)
+ assert emb.ndim == 2
+ assert emb.shape[0] == 2 # GaN has 2 atoms
+
+ def test_encode_different_structures_give_different_embeddings(
+ self, small_model, si_structure, gan_structure
+ ):
+ emb_si = small_model.material_encoder.predict_structure_embedding(si_structure)
+ emb_gan = small_model.material_encoder.predict_structure_embedding(gan_structure)
+ assert not paddle.allclose(emb_si.mean(), emb_gan.mean())
+
+ def test_encode_batch_of_structures(self, small_model, si_structure, gan_structure):
+ embs, masks = [], []
+ for s in [si_structure, gan_structure]:
+ e = small_model.material_encoder.predict_structure_embedding(s)
+ embs.append(e)
+ masks.append(paddle.ones([e.shape[0]], dtype=paddle.int64))
+ stacked = paddle.stack(embs)
+ assert stacked.shape[0] == 2
+
+ def test_generate_entry_point_exists(self, small_model):
+ assert hasattr(small_model, "generate")
+ assert hasattr(small_model, "generate_followup")
+ assert hasattr(small_model, "_encode_for_generate")
+
+ def test_llm_proj_output_dtype_matches_embedding(self, small_model, si_structure):
+ emb = small_model.material_encoder.predict_structure_embedding(si_structure)
+ mask = paddle.ones([emb.shape[0]], dtype=paddle.int64)
+ qt = small_model.query_tokens.expand([1, -1, -1])
+ out = small_model._run_qformer(qt, emb.unsqueeze(0), mask.unsqueeze(0), None)
+ proj = small_model.llm_proj(out.last_hidden_state[:, : qt.shape[1], :])
+ assert proj.dtype == emb.dtype
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])