k3-node 1.0.0__py3-none-any.whl

k3_node/rag/verbalizer.py

@@ -0,0 +1,162 @@
+"""Verbalization and LLM prompt formatting for GraphRAG."""
+
+from typing import Dict, List, Literal, Optional, Sequence, Tuple, Union
+import numpy as np
+from keras import ops
+
+from k3_node.rag.subgraph import SubgraphResult
+
+
+def subgraph_to_triples(
+    subgraph: SubgraphResult,
+    id_to_entity: Optional[Dict[int, str]] = None,
+    id_to_relation: Optional[Dict[int, str]] = None,
+) -> List[Tuple[str, str, str]]:
+    """Convert a SubgraphResult into a list of (head, relation, tail) string triples.
+
+    Args:
+        subgraph: The extracted SubgraphResult.
+        id_to_entity: Optional map from original entity integer ID to entity string.
+        id_to_relation: Optional map from relation integer ID to relation string.
+
+    Returns:
+        List of `(head, relation, tail)` string triples.
+    """
+    if subgraph.num_edges == 0:
+        return []
+
+    edge_index_np = np.asarray(ops.convert_to_numpy(subgraph.edge_index)).astype(np.int64)
+    row, col = edge_index_np[0], edge_index_np[1]
+
+    # Map subgraph indices back to original node IDs if mapping/nodes available
+    if subgraph.nodes is not None and len(subgraph.nodes) > 0:
+        orig_row = subgraph.nodes[row]
+        orig_col = subgraph.nodes[col]
+    else:
+        orig_row, orig_col = row, col
+
+    edge_type_np = None
+    if subgraph.edge_type is not None:
+        edge_type_np = np.asarray(ops.convert_to_numpy(subgraph.edge_type)).astype(np.int64)
+
+    triples = []
+    for i in range(len(row)):
+        h_id = int(orig_row[i])
+        t_id = int(orig_col[i])
+        h_name = id_to_entity.get(h_id, f"Node_{h_id}") if id_to_entity else f"Node_{h_id}"
+        t_name = id_to_entity.get(t_id, f"Node_{t_id}") if id_to_entity else f"Node_{t_id}"
+
+        if edge_type_np is not None:
+            r_id = int(edge_type_np[i])
+            r_name = id_to_relation.get(r_id, f"rel_{r_id}") if id_to_relation else f"rel_{r_id}"
+        else:
+            r_name = "connected_to"
+
+        triples.append((h_name, r_name, t_name))
+
+    return triples
+
+
+def verbalize_subgraph(
+    subgraph: SubgraphResult,
+    id_to_entity: Optional[Dict[int, str]] = None,
+    id_to_relation: Optional[Dict[int, str]] = None,
+    format_style: Literal["triples", "markdown", "natural"] = "markdown",
+    max_triples: Optional[int] = 50,
+) -> str:
+    """Verbalize an extracted subgraph into textual knowledge for LLM prompt augmentation.
+
+    Args:
+        subgraph: Extracted SubgraphResult around retrieved entities.
+        id_to_entity: Optional dictionary mapping node ID to entity name.
+        id_to_relation: Optional dictionary mapping relation ID to relation name.
+        format_style:
+            - `"triples"`: List of `(Head, Relation, Tail)` text triples.
+            - `"markdown"`: Markdown bullet list with facts.
+            - `"natural"`: Natural language sentences (`"Head relation Tail."`).
+        max_triples: Maximum number of facts to include in the context.
+
+    Returns:
+        Formatted textual context string ready to be injected into an LLM prompt.
+    """
+    triples = subgraph_to_triples(subgraph, id_to_entity, id_to_relation)
+    if not triples:
+        return "No relevant knowledge graph facts retrieved."
+
+    if max_triples is not None and len(triples) > max_triples:
+        triples = triples[:max_triples]
+
+    if format_style == "triples":
+        lines = [f"({h}, {r}, {t})" for h, r, t in triples]
+        return "\n".join(lines)
+
+    elif format_style == "natural":
+        lines = []
+        for h, r, t in triples:
+            # Clean relation string (replace underscores with spaces)
+            rel_str = r.replace("_", " ")
+            lines.append(f"{h} {rel_str} {t}.")
+        return " ".join(lines)
+
+    else:  # "markdown"
+        lines = ["### Retrieved Knowledge Graph Facts:"]
+        for h, r, t in triples:
+            lines.append(f"- **{h}** — *{r}* -> **{t}**")
+        return "\n".join(lines)
+
+
+def format_llm_prompt(
+    query: str,
+    context: str,
+    system_prompt: Optional[str] = None,
+    model_family: Literal["llama3", "mistral", "chatml", "standard"] = "llama3",
+) -> str:
+    """Format query and retrieved KG context into prompt templates for LLMs.
+
+    Supported model families include:
+    - `"llama3"`: Meta Llama 3 / 3.1 instruct template.
+    - `"mistral"`: Mistral / Mixtral instruct template.
+    - `"chatml"`: OpenAI / Qwen ChatML template.
+    - `"standard"`: General markdown system/user format.
+
+    Args:
+        query: The user's input question or instruction.
+        context: The verbalized knowledge graph context.
+        system_prompt: Optional system prompt to instruct the LLM.
+        model_family: Target model prompt format. (default: "llama3")
+
+    Returns:
+        Formatted prompt string.
+    """
+    if system_prompt is None:
+        system_prompt = (
+            "You are an expert assistant augmented with a Knowledge Graph. "
+            "Use the provided Knowledge Graph Context to accurately answer the user's question."
+        )
+
+    if model_family == "llama3":
+        return (
+            "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"
+            f"{system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n"
+            f"Knowledge Graph Context:\n{context}\n\n"
+            f"Question: {query}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        )
+    elif model_family == "mistral":
+        return (
+            f"<s>[INST] {system_prompt}\n\n"
+            f"Knowledge Graph Context:\n{context}\n\n"
+            f"Question: {query} [/INST]"
+        )
+    elif model_family == "chatml":
+        return (
+            f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
+            f"<|im_start|>user\nKnowledge Graph Context:\n{context}\n\nQuestion: {query}<|im_end|>\n"
+            f"<|im_start|>assistant\n"
+        )
+    else:  # "standard"
+        return (
+            f"System: {system_prompt}\n\n"
+            f"Knowledge Graph Context:\n{context}\n\n"
+            f"User Question: {query}\n\n"
+            f"Assistant:"
+        )

k3_node/tasks/__init__.py

@@ -0,0 +1,19 @@
+"""High-level Task APIs for K3-Node."""
+
+from k3_node.tasks.base import BaseTask
+from k3_node.tasks.backbone_resolver import resolve_backbone
+from k3_node.tasks.node_classification import NodeClassifier
+from k3_node.tasks.node_regression import NodeRegressor
+from k3_node.tasks.graph_classification import GraphClassifier
+from k3_node.tasks.graph_regression import GraphRegressor
+from k3_node.tasks.link_prediction import LinkPredictor
+
+__all__ = [
+    "BaseTask",
+    "resolve_backbone",
+    "NodeClassifier",
+    "NodeRegressor",
+    "GraphClassifier",
+    "GraphRegressor",
+    "LinkPredictor",
+]

k3_node/tasks/backbone_resolver.py

@@ -0,0 +1,125 @@
+"""Backbone resolver mapping string identifiers to K3-Node models."""
+
+from typing import Any, Dict, Optional, Union
+import keras
+
+from k3_node import models
+
+
+BACKBONE_REGISTRY = {
+    "gcn": models.GCN,
+    "gat": models.GAT,
+    "sage": models.GraphSAGE,
+    "graphsage": models.GraphSAGE,
+    "gin": models.GIN,
+    "pna": models.PNA,
+    "edge_cnn": models.EdgeCNN,
+    "edgecnn": models.EdgeCNN,
+    "mlp": models.MLP,
+    "linkx": models.LINKX,
+    "pmlp": models.PMLP,
+    "sgformer": models.SGFormer,
+    "polynormer": models.Polynormer,
+    "schnet": models.SchNet,
+    "dimenet": models.DimeNet,
+    "dimenet++": models.DimeNetPlusPlus,
+    "dimenetplusplus": models.DimeNetPlusPlus,
+    "attentive_fp": models.AttentiveFP,
+    "attentivefp": models.AttentiveFP,
+    "graph_unet": models.GraphUNet,
+}
+
+
+def resolve_backbone(
+    backbone: Union[str, keras.Model, Any],
+    in_channels: int,
+    out_channels: int,
+    hidden_channels: int = 64,
+    num_layers: int = 2,
+    dropout: float = 0.5,
+    **kwargs,
+) -> keras.Model:
+    r"""Resolves a backbone string or instance into a compiled/callable Keras model."""
+    if isinstance(backbone, keras.Model):
+        return backbone
+
+    if not isinstance(backbone, str):
+        raise TypeError(
+            f"Expected backbone to be a string name or keras.Model instance, got {type(backbone)}"
+        )
+
+    name = backbone.lower().strip()
+    if name not in BACKBONE_REGISTRY:
+        available = ", ".join(sorted(BACKBONE_REGISTRY.keys()))
+        raise ValueError(f"Unknown backbone '{backbone}'. Available backbones: {available}")
+
+    cls = BACKBONE_REGISTRY[name]
+
+    # Handle models with specialized constructors
+    if name in ("schnet",):
+        return cls(hidden_channels=hidden_channels, num_filters=hidden_channels, num_interactions=num_layers, **kwargs)
+    elif name in ("dimenet", "dimenet++", "dimenetplusplus"):
+        return cls(hidden_channels=hidden_channels, out_channels=out_channels, num_blocks=num_layers, **kwargs)
+    elif name in ("attentive_fp", "attentivefp"):
+        edge_dim = kwargs.pop("edge_dim", in_channels)
+        num_timesteps = kwargs.pop("num_timesteps", 2)
+        return cls(
+            in_channels=in_channels,
+            hidden_channels=hidden_channels,
+            out_channels=out_channels,
+            edge_dim=edge_dim,
+            num_layers=num_layers,
+            num_timesteps=num_timesteps,
+            dropout=dropout,
+            **kwargs,
+        )
+    elif name in ("mlp",):
+        channel_list = [in_channels] + [hidden_channels] * (num_layers - 1) + [out_channels]
+        return cls(channel_list=channel_list, dropout=dropout, **kwargs)
+    elif name in ("linkx",):
+        num_nodes = kwargs.pop("num_nodes", 1000)
+        return cls(
+            num_nodes=num_nodes,
+            in_channels=in_channels,
+            hidden_channels=hidden_channels,
+            out_channels=out_channels,
+            num_layers=num_layers,
+            dropout=dropout,
+            **kwargs,
+        )
+    elif name in ("sgformer",):
+        return cls(
+            in_channels=in_channels,
+            hidden_channels=hidden_channels,
+            out_channels=out_channels,
+            num_layers=num_layers,
+            dropout=dropout,
+            **kwargs,
+        )
+    elif name in ("polynormer",):
+        return cls(
+            in_channels=in_channels,
+            hidden_channels=hidden_channels,
+            out_channels=out_channels,
+            num_layers=num_layers,
+            dropout=dropout,
+            **kwargs,
+        )
+    elif name in ("graph_unet",):
+        return cls(
+            in_channels=in_channels,
+            hidden_channels=hidden_channels,
+            out_channels=out_channels,
+            depth=num_layers,
+            **kwargs,
+        )
+    else:
+        # Standard BasicGNN (GCN, GAT, SAGE, GIN, PNA, EdgeCNN)
+        return cls(
+            in_channels=in_channels,
+            hidden_channels=hidden_channels,
+            num_layers=num_layers,
+            out_channels=out_channels,
+            dropout=dropout,
+            **kwargs,
+        )

k3_node/tasks/base.py

@@ -0,0 +1,67 @@
+"""Base task abstraction for high-level K3-Node estimators."""
+
+from typing import Any, Dict, List, Optional, Tuple, Union
+import keras
+from keras import ops
+
+from k3_node.data import BaseData
+from k3_node.hub.hub_mixin import K3NodeHubMixin
+
+
+class BaseTask(K3NodeHubMixin):
+    r"""Abstract base task estimator providing common training, evaluation,
+    and serialization workflows.
+    """
+
+    def __init__(self, model: Optional[keras.Model] = None):
+        self.model = model
+        self._is_compiled = False
+
+    def compile(
+        self,
+        optimizer: Optional[Union[str, keras.optimizers.Optimizer]] = None,
+        loss: Optional[Any] = None,
+        metrics: Optional[List[Any]] = None,
+        **kwargs,
+    ):
+        r"""Configures the task model for training."""
+        if self.model is None:
+            raise RuntimeError("Model has not been initialized. Call fit() or construct with a model first.")
+
+        opt = optimizer or keras.optimizers.Adam(learning_rate=0.01)
+        self.model.compile(optimizer=opt, loss=loss, metrics=metrics, **kwargs)
+        self._is_compiled = True
+        return self
+
+    def summary(self):
+        r"""Prints a string summary of the underlying neural network."""
+        if self.model is not None:
+            return self.model.summary()
+        print("Model has not been initialized yet.")
+
+    def save(self, filepath: str):
+        r"""Saves the underlying model weights."""
+        if self.model is not None:
+            self.model.save(filepath)
+        else:
+            raise RuntimeError("Cannot save an uninitialized model.")
+
+    @classmethod
+    def load(cls, filepath: str, **kwargs):
+        r"""Loads a saved task model from disk."""
+        model = keras.models.load_model(filepath, **kwargs)
+        instance = cls(model=model)
+        instance._is_compiled = True
+        return instance
+
+    def _extract_inputs(self, data: Any):
+        r"""Extracts input tensors from a Data object, tuple, or dictionary."""
+        if hasattr(data, "inputs"):
+            return data.inputs
+        elif isinstance(data, (tuple, list)):
+            return data
+        elif hasattr(data, "x") and hasattr(data, "edge_index"):
+            if hasattr(data, "edge_attr") and data.edge_attr is not None:
+                return (data.x, data.edge_index, data.edge_attr)
+            return (data.x, data.edge_index)
+        return data

k3_node/tasks/graph_classification.py

@@ -0,0 +1,270 @@
+"""High-level Graph Classification Task."""
+
+from typing import Any, Dict, List, Optional, Union
+import keras
+from keras import layers, ops
+import numpy as np
+
+from k3_node.tasks.base import BaseTask
+from k3_node.tasks.backbone_resolver import resolve_backbone
+from k3_node.layers import pool as k3_pool
+from k3_node.loader import DataLoader
+
+
+class GraphClassificationModel(keras.Model):
+    r"""Internal wrapper combining a node-level GNN backbone, a global readout
+    pooling operation, and a final classification dense head.
+    """
+
+    def __init__(
+        self,
+        backbone: keras.Model,
+        pooling: str = "mean",
+        hidden_channels: int = 64,
+        num_classes: int = 2,
+        dropout: float = 0.5,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.pooling = pooling
+        self.dropout = layers.Dropout(dropout) if dropout > 0 else None
+        self.head = layers.Dense(num_classes)
+        self.num_graphs = None
+
+    def call(self, inputs, training=False):
+        if isinstance(inputs, (tuple, list)):
+            x, edge_index = inputs[0], inputs[1]
+            batch = inputs[2] if len(inputs) > 2 else None
+            size = inputs[3] if len(inputs) > 3 else self.num_graphs
+        else:
+            x = inputs
+            edge_index = getattr(x, "edge_index", None)
+            batch = getattr(x, "batch", None)
+            size = getattr(x, "num_graphs", self.num_graphs)
+            x = getattr(x, "x", x)
+
+        if batch is None:
+            batch = ops.zeros((ops.shape(x)[0],), dtype="int64")
+
+        h = self.backbone((x, edge_index), training=training)
+
+        if self.pooling in ("mean", "global_mean_pool"):
+            g = k3_pool.global_mean_pool(h, batch, size=size)
+        elif self.pooling in ("add", "sum", "global_add_pool"):
+            g = k3_pool.global_add_pool(h, batch, size=size)
+        elif self.pooling in ("max", "global_max_pool"):
+            g = k3_pool.global_max_pool(h, batch, size=size)
+        else:
+            g = k3_pool.global_mean_pool(h, batch, size=size)
+
+        if self.dropout is not None:
+            g = self.dropout(g, training=training)
+        return self.head(g)
+
+
+class GraphClassifier(BaseTask):
+    r"""High-level estimator for graph classification tasks (e.g., molecular property,
+    bioinformatics, social graph classification).
+
+    Args:
+        backbone: GNN architecture (``"gin"``, ``"gcn"``, ``"gat"``, ``"sage"``,
+            ``"pna"``, etc.) or a custom :class:`keras.Model`. (default: ``"gin"``)
+        in_channels (int, optional): Size of input node features.
+        hidden_channels (int, optional): Dimensionality of hidden node features. (default: ``64``)
+        num_classes (int, optional): Number of graph classes.
+        num_layers (int, optional): Number of GNN layers. (default: ``3``)
+        pooling (str, optional): Readout pooling (``"mean"``, ``"add"``, ``"max"``). (default: ``"mean"``)
+        dropout (float, optional): Dropout probability. (default: ``0.5``)
+        **backbone_kwargs: Additional arguments forwarded to the backbone constructor.
+    """
+
+    def __init__(
+        self,
+        backbone: Union[str, keras.Model] = "gin",
+        in_channels: Optional[int] = None,
+        hidden_channels: int = 64,
+        num_classes: Optional[int] = None,
+        num_layers: int = 3,
+        pooling: str = "mean",
+        dropout: float = 0.5,
+        **backbone_kwargs,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.num_classes = num_classes
+        self.num_layers = num_layers
+        self.pooling = pooling
+        self.dropout = dropout
+        self.backbone_kwargs = backbone_kwargs
+
+    def _init_model(self, sample_data: Any, dataset: Optional[Any] = None):
+        in_c = self.in_channels
+        if in_c is None:
+            if hasattr(sample_data, "num_node_features") and sample_data.num_node_features > 0:
+                in_c = sample_data.num_node_features
+            elif hasattr(sample_data, "num_features") and sample_data.num_features > 0:
+                in_c = sample_data.num_features
+            elif hasattr(sample_data, "x") and sample_data.x is not None:
+                in_c = int(ops.shape(sample_data.x)[-1])
+            else:
+                raise ValueError("Could not infer in_channels.")
+
+        out_c = self.num_classes
+        if out_c is None:
+            if dataset is not None and hasattr(dataset, "num_classes") and dataset.num_classes is not None:
+                out_c = dataset.num_classes
+            elif dataset is not None and isinstance(dataset, (list, tuple)):
+                max_y = 0
+                for g in dataset[:100]:
+                    if hasattr(g, "y") and g.y is not None:
+                        max_y = max(max_y, int(ops.convert_to_numpy(ops.max(g.y))))
+                out_c = max_y + 1
+            elif hasattr(sample_data, "num_classes") and sample_data.num_classes is not None and sample_data.num_classes > 1:
+                out_c = sample_data.num_classes
+            elif hasattr(sample_data, "y") and sample_data.y is not None:
+                out_c = int(ops.convert_to_numpy(ops.max(sample_data.y))) + 1
+            else:
+                out_c = 2  # default binary
+
+        out_c = max(int(out_c), 2)
+
+        self.in_channels = in_c
+        self.num_classes = out_c
+
+        gnn = resolve_backbone(
+            self.backbone,
+            in_channels=in_c,
+            out_channels=self.hidden_channels,
+            hidden_channels=self.hidden_channels,
+            num_layers=self.num_layers,
+            dropout=self.dropout,
+            **self.backbone_kwargs,
+        )
+
+        self.model = GraphClassificationModel(
+            backbone=gnn,
+            pooling=self.pooling,
+            hidden_channels=self.hidden_channels,
+            num_classes=out_c,
+            dropout=self.dropout,
+        )
+
+    def fit(
+        self,
+        dataset: Any,
+        epochs: int = 20,
+        lr: float = 0.01,
+        batch_size: int = 32,
+        shuffle: bool = True,
+        verbose: int = 1,
+        callbacks: Optional[List[Any]] = None,
+    ):
+        r"""Trains the graph classifier."""
+        if not isinstance(dataset, DataLoader):
+            loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
+            sample = dataset[0]
+        else:
+            loader = dataset
+            sample = next(iter(loader))
+
+        if self.model is None:
+            self._init_model(sample, dataset=dataset)
+
+        if not self._is_compiled:
+            self.model.compile(
+                optimizer=keras.optimizers.Adam(learning_rate=lr),
+                loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+                metrics=[keras.metrics.SparseCategoricalAccuracy(name="acc")],
+            )
+            self._is_compiled = True
+
+        history = {"loss": [], "acc": []}
+        for epoch in range(epochs):
+            batch_losses = []
+            batch_accs = []
+            for batch in loader:
+                x = ops.convert_to_tensor(batch.x, dtype="float32")
+                edge_index = ops.convert_to_tensor(batch.edge_index, dtype="int64")
+                batch_vec = ops.convert_to_tensor(batch.batch, dtype="int64")
+                y = ops.convert_to_tensor(batch.y, dtype="int64")
+                num_g = int(ops.shape(y)[0])
+                if hasattr(self.model, "num_graphs"):
+                    self.model.num_graphs = num_g
+
+                if not self.model.built:
+                    y_pred = self.model((x, edge_index, batch_vec), training=False)
+                    self.model.built = True
+                    if hasattr(self.model, "_compile_loss") and self.model._compile_loss is not None:
+                        self.model._compile_loss.build(y, y_pred)
+                    if hasattr(self.model, "_compile_metrics") and self.model._compile_metrics is not None:
+                        self.model._compile_metrics.build(y, y_pred)
+                    if self.model.optimizer is not None and not self.model.optimizer.built:
+                        self.model.optimizer.build(self.model.trainable_variables)
+
+                res = self.model.train_on_batch((x, edge_index, batch_vec), y)
+                if isinstance(res, (list, tuple)):
+                    batch_losses.append(float(res[0]))
+                    if len(res) > 1:
+                        batch_accs.append(float(res[1]))
+                else:
+                    batch_losses.append(float(res))
+
+            avg_loss = float(np.mean(batch_losses)) if batch_losses else 0.0
+            avg_acc = float(np.mean(batch_accs)) if batch_accs else 0.0
+            history["loss"].append(avg_loss)
+            history["acc"].append(avg_acc)
+            if verbose:
+                print(f"Epoch {epoch + 1}/{epochs} - loss: {avg_loss:.4f} - acc: {avg_acc:.4f}")
+
+        return history
+
+    def predict_proba(self, dataset_or_loader: Any, batch_size: int = 32):
+        r"""Predicts class probabilities for graphs."""
+        if not isinstance(dataset_or_loader, DataLoader):
+            loader = DataLoader(dataset_or_loader, batch_size=batch_size, shuffle=False)
+        else:
+            loader = dataset_or_loader
+
+        probs = []
+        for batch in loader:
+            x = ops.convert_to_tensor(batch.x, dtype="float32")
+            edge_index = ops.convert_to_tensor(batch.edge_index, dtype="int64")
+            batch_vec = ops.convert_to_tensor(batch.batch, dtype="int64")
+            num_g = int(ops.convert_to_numpy(ops.max(batch_vec))) + 1 if ops.shape(batch_vec)[0] > 0 else 1
+            if hasattr(self.model, "num_graphs"):
+                self.model.num_graphs = num_g
+            logits = self.model((x, edge_index, batch_vec), training=False)
+            probs.append(ops.softmax(logits, axis=-1))
+        return ops.concatenate(probs, axis=0)
+
+    def predict(self, dataset_or_loader: Any, batch_size: int = 32):
+        r"""Predicts discrete class labels for graphs."""
+        probs = self.predict_proba(dataset_or_loader, batch_size=batch_size)
+        return ops.argmax(probs, axis=-1)
+
+    def evaluate(self, dataset_or_loader: Any, batch_size: int = 32) -> Dict[str, float]:
+        r"""Evaluates classification accuracy on the dataset."""
+        if not isinstance(dataset_or_loader, DataLoader):
+            loader = DataLoader(dataset_or_loader, batch_size=batch_size, shuffle=False)
+        else:
+            loader = dataset_or_loader
+
+        correct = 0
+        total = 0
+        for batch in loader:
+            x = ops.convert_to_tensor(batch.x, dtype="float32")
+            edge_index = ops.convert_to_tensor(batch.edge_index, dtype="int64")
+            batch_vec = ops.convert_to_tensor(batch.batch, dtype="int64")
+            num_g = int(ops.convert_to_numpy(ops.max(batch_vec))) + 1 if ops.shape(batch_vec)[0] > 0 else 1
+            if hasattr(self.model, "num_graphs"):
+                self.model.num_graphs = num_g
+            logits = self.model((x, edge_index, batch_vec), training=False)
+            pred = ops.argmax(logits, axis=-1)
+            pred_np = ops.convert_to_numpy(ops.cast(pred, "int64"))
+            y_np = ops.convert_to_numpy(ops.cast(batch.y, "int64"))
+            correct += int((pred_np == y_np).sum())
+            total += int(y_np.shape[0])
+
+        return {"accuracy": float(correct / max(total, 1))}