| import torch |
| from torch import nn |
|
|
| from torch_geometric.utils import to_dense_batch |
|
|
| from models.base.graph_encoder import GraphEncoder |
| from models.base.config import GraphPromptEncoderConfig |
|
|
|
|
| class BaseGraphPromptEncoder(nn.Module): |
| """ |
| Base class to define the graph prompt encoder for encoding graph-based structural information. |
| |
| Args: |
| config ([`GraphPromptEncoderConfig`]): The configuration of the graph prompt encoder. |
| """ |
| def __init__(self, config: GraphPromptEncoderConfig): |
| super().__init__() |
| |
| self.graph_encoder = GraphEncoder(config) |
| |
| def reshape_node_embeds( |
| self, |
| x: torch.Tensor, |
| batch_index: torch.Tensor, |
| )-> tuple[torch.Tensor, torch.Tensor]: |
| """ |
| Reshape the node embeddings with batch axis to match the input word embeddings. |
| |
| Args: |
| x (torch.Tensor): Node embeddings of shape (num_nodes, output_size). |
| batch_index (torch.Tensor): Batch indices for each node (shape: [num_nodes]). |
| |
| Returns: |
| tuple[torch.Tensor, torch.Tensor]: A tuple containing the reshaped node embeddings and the corresponding node mask. |
| """ |
| |
| |
| node_embeds, node_mask = to_dense_batch(x, batch_index, fill_value=0.0) |
| |
| return node_embeds, node_mask.long() |
|
|
