Models

model_manager

class openddi.models.model_manager.model_manager(args: <module 'argparse' from '/home/docs/.asdf/installs/python/3.10.17/lib/python3.10/argparse.py'>)

Bases: object

Model manager for handling different model types based on configuration.

This class maps model names to their corresponding model classes and provides functionality to load the appropriate model with proper parameters.

__init__(args: <module 'argparse' from '/home/docs/.asdf/installs/python/3.10.17/lib/python3.10/argparse.py'>)

Initialize the model manager.

Parameters:

args – Command line arguments containing model specification and parameters

load_model()

Load the model corresponding to the specified model type.

Returns:

Initialized model instance on the appropriate device

Raises:

ValueError – If num_classes is not properly set

load_origin_model(ddi_dataset)

Load the original model based on the specified model type.

Parameters:

ddi_dataset – Dataset object containing DDI data and statistics.

Returns:

Initialized model on the appropriate device.

Return type:

torch.nn.Module

CASTER

class openddi.models.CASTER.CASTER(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Bases: Module

CASTER baseline (tensor decomposition/bilinear scoring style).

Architecture: - X -> drug embedding e_i - Each relation k learns a vector r_k - logits_k(i,j) = < (e_i ⊙ r_k), e_j >

Signature construction consistent with existing model_manager: - (feature, hidden1, hidden2, num_relations, num_classes, dropout) - Where hidden2 is used as the final embedding dimension d.

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Initialize CASTER model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension

• hidden2 – Final embedding dimension (d)

• num_relations – Number of relation types

• num_classes – Number of output classes

• dropout – Dropout rate

bind_graph(data_graph)

Bind graph data to the model.

Parameters:

data_graph – Graph data containing node features

forward(graph_or_none, idx_batch)

Forward pass of the model.

Parameters:

• graph_or_none – Can be None (already bound) or Data containing x

• idx_batch – (i_idx, j_idx, y) batch indices

Returns:

[B, K] output logits

Return type:

logits

ConvLSTM

class openddi.models.ConvLSTM.ConvLSTM(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.2, timesteps=1)

Bases: Module

ConvLSTM model combining convolutional and LSTM layers for sequence processing.

This model uses 1D convolutional layers followed by LSTM layers to process sequential data, with global pooling and fully connected layers for classification.

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.2, timesteps=1)

Initialize ConvLSTM model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension for conv and LSTM

• hidden2 – Second hidden layer dimension for conv and LSTM

• num_relations – Number of relation types (unused in current implementation)

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

• timesteps – Number of timesteps for sequence processing

forward(data_o, idx)

Forward pass of the ConvLSTM model.

Parameters:

• data_o – Graph data object containing node features, edge indices, and edge types

• idx – Batch indices containing pairs of nodes to process

Returns:

Output logits for the given node pairs

DDIMDL

class openddi.models.DDIMDL.DDIMDL(features: list, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.5, pooling_ratio: float = 0.3)

Bases: Module

DDIMDL model for multi-modal drug-drug interaction prediction.

This model processes multiple feature modalities independently through separate MLPs and combines their outputs through averaging for final prediction.

__init__(features: list, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.5, pooling_ratio: float = 0.3)

Initialize DDIMDL model.

Parameters:

• features – List of feature dimensions for each modality

• hidden1 – First hidden layer dimension for MLPs

• hidden2 – Second hidden layer dimension for MLPs

• num_relations – Number of relation types (unused in current implementation)

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

• pooling_ratio – Pooling ratio parameter (used as dropout in current implementation)

forward(data_o, idx)

Forward pass of the DDIMDL model.

Parameters:

• data_o – Graph data object containing node features, edge indices, and edge types

• idx – Batch indices containing pairs of nodes to process

Returns:

Averaged output logits from all modalities

DDKG

class openddi.models.DDKG.DDKG(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Bases: Module

DDKG baseline model for drug-drug interaction prediction.

Architecture: - Linear dimensionality reduction: feature → hid1 - Two-layer RGCNConv with num_bases for parameter reduction in multi-relation graphs - Checkpointing (without use_reentrant parameter) - Pair features → MLP → logits

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Initialize DDKG model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension

• hidden2 – Second hidden layer dimension

• num_relations – Number of relation types

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

bind_graph(data_graph)

Bind graph data to the model for caching.

Parameters:

data_graph – Graph data object containing node features, edge indices, and edge types

forward(graph_or_none, idx_batch)

Forward pass of the DDKG model.

Parameters:

• graph_or_none – Can be None (already bound) or Data containing graph information

• idx_batch – (i_idx, j_idx, y) batch indices

Returns:

[B, K] output logits

Return type:

logits

DeepDDI

class openddi.models.DeepDDI.DeepDDI(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Bases: Module

DeepDDI baseline model for drug-drug interaction prediction.

Architecture:

• No graph convolution, directly uses node feature matrix X

• Pair representation: [x_i || x_j || absolute difference || element-wise product]

• MLP -> logits (supports single/multi-label classification)

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Initialize DeepDDI model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension

• hidden2 – Second hidden layer dimension

• num_relations – Number of relation types (unused in current implementation)

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

bind_graph(data_graph)

Bind graph data to the model.

Parameters:

data_graph – Graph data object (only uses X, not edge_index)

forward(graph_or_none, idx_batch)

Forward pass of the DeepDDI model.

Parameters:

• graph_or_none – Can be None (already bound) or Data containing node features

• idx_batch – (i_idx, j_idx, y) batch indices

Returns:

[B, K] output logits

Return type:

logits

DSNDDI

class openddi.models.DSNDDI.DSNDDI(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Bases: Module

Dual-View Structure-aware Neural Network for Drug-Drug Interaction prediction.

This model integrates both local drug embeddings and global graph structure to predict DDI types.

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Initialize the DSNDDI model.

Parameters:

• feature – Input feature dimension

• hidden1 – Hidden dimension for local and global encoders

• hidden2 – Hidden dimension for DDI prediction MLP

• num_relations – Number of relation types in the graph

• num_classes – Number of output classes (DDI types)

• dropout – Dropout rate

forward(data_o, idx)

Forward pass of the model.

Parameters:

• data_o – Graph data object containing node features, edge indices and types

• idx – Tuple of (head_indices, tail_indices) for drug pairs

Returns:

Predicted logits for each DDI type [batch_size, num_classes]

Return type:

logits

ExDDI

class openddi.models.ExDDI.FeatureMask(in_dim: int)

Bases: Module

Feature explainability layer: learns a weight (0-1) for each dimension.

Applies masking in the reduced hid1 space to save GPU memory.

__init__(in_dim: int)

Initialize FeatureMask layer.

Parameters:

in_dim – Input dimension for the mask

forward(x: Tensor) → Tensor

Apply feature mask to input tensor.

Parameters:

x – Input tensor [N, d]

Returns:

Masked tensor [N, d]

class openddi.models.ExDDI.ExDDI(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Bases: Module

ExDDI baseline model for explainable drug-drug interaction prediction.

Architecture: - Linear dimensionality reduction + FeatureMask (hid1 dimension) - Single-layer GraphSAGE (hid1→hid2) - Checkpointing (without use_reentrant parameter) - Pair representation → MLP → logits

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Initialize ExDDI model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension

• hidden2 – Second hidden layer dimension

• num_relations – Number of relation types (unused in current implementation)

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

bind_graph(data_graph)

Bind graph data to the model.

Parameters:

data_graph – Graph data object containing node features and edge indices

forward(graph_or_none, idx_batch)

Forward pass of the ExDDI model.

Parameters:

• graph_or_none – Can be None (already bound) or Data containing graph information

• idx_batch – (i_idx, j_idx, y) batch indices

Returns:

Output logits

Return type:

logits

GOGNN

class openddi.models.GOGNN.GOGNN(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.5, pooling_ratio: float = 0.5)

Bases: Module

GOGNN model for drug-drug interaction prediction using relational graph convolutional networks.

This model uses RGCNConv for graph encoding and MLP for pair classification.

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.5, pooling_ratio: float = 0.5)

Initialize GOGNN model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension (unused in current implementation)

• hidden2 – Second hidden layer dimension (unused in current implementation)

• num_relations – Number of relation types for RGCN

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

• pooling_ratio – Pooling ratio parameter (unused in current implementation)

forward(data_o, idx)

Forward pass of the GOGNN model.

Parameters:

• data_o – Graph data object containing node features, edge indices, and edge types

• idx – Batch indices containing pairs of nodes to process

Returns:

Output logits for the given node pairs

KGNN

class openddi.models.KGNN.KGNN(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Bases: Module

KGNN model for knowledge graph-based drug-drug interaction prediction.

This model uses RGCN layers for graph encoding and linear layer for pair classification.

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Initialize KGNN model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension for RGCN and linear projection

• hidden2 – Second hidden layer dimension for RGCN

• num_relations – Number of relation types for RGCN

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

forward(data_o, idx)

Forward pass of the KGNN model.

Parameters:

• data_o – Graph data object containing node features, edge indices, and edge types

• idx – Batch indices containing pairs of nodes to process

Returns:

Output logits for the given node pairs

LaGAT

class openddi.models.LaGAT.TBA(config)

Bases: Module

Target-Based Attention module for attention mechanism in graph networks.

__init__(config)

Initialize TBA module.

Parameters:

config – Configuration object containing model parameters

forward(inputs)

Forward pass of TBA module.

Parameters:

inputs – Tuple of (drug_embed, neigh_embed) drug_embed: (batch, dim), neigh_embed: (batch, n_neighbor^(hop+1), dim)

Returns:

Tuple of (neighbor_embed, attention_weights)

class openddi.models.LaGAT.NeighAggregator(activation='relu', l2_weight=1e-05, name='neigh_aggregator')

Bases: Module

Neighborhood aggregator for combining neighbor embeddings.

__init__(activation='relu', l2_weight=1e-05, name='neigh_aggregator')

Initialize NeighAggregator.

Parameters:

• activation – Activation function type (‘relu’ or ‘tanh’)

• l2_weight – L2 regularization weight

• name – Module name

build(ent_embed_dim, neighbor_embed_dim)

Build linear layer dynamically.

Parameters:

• ent_embed_dim – Entity embedding dimension

• neighbor_embed_dim – Neighbor embedding dimension

forward(inputs)

Forward pass of NeighAggregator.

Parameters:

inputs – Tuple of (entity, neighbor) entity: (batch, 1, dim), neighbor: (batch, n_shape, dim)

Returns:

Aggregated neighbor embeddings

class openddi.models.LaGAT.GetReceptiveField(config)

Bases: Module

Module for getting receptive field in graph networks.

__init__(config)

Initialize GetReceptiveField module.

Parameters:

config – Configuration object containing model parameters

forward(x)

Forward pass to get receptive field.

Parameters:

x – Input tensor

Returns:

List containing neighbor entities and relations

class openddi.models.LaGAT.SqueezeLayer

Bases: Module

Simple squeeze layer to remove dimension of size 1.

__init__()

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Forward pass of SqueezeLayer.

Parameters:

x – Input tensor

Returns:

Squeezed tensor

class openddi.models.LaGAT.LaGAT(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Bases: Module

LaGAT model for drug-drug interaction prediction with hierarchical feature concatenation.

This model uses RGCN layers and hierarchical feature concatenation for enhanced representation.

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Initialize LaGAT model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension for RGCN and linear projection

• hidden2 – Second hidden layer dimension for RGCN

• num_relations – Number of relation types for RGCN

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

forward(data_o, idx)

Forward pass of the LaGAT model.

Parameters:

• data_o – Graph data object containing node features, edge indices, and edge types

• idx – Batch indices containing pairs of nodes to process

Returns:

Output logits for the given node pairs

MIRACLE

class openddi.models.MIRACLE.InteractionPredictor(dg, hidden, k)

Bases: Module

Interaction predictor module for MIRACLE model.

__init__(dg, hidden, k)

Initialize InteractionPredictor.

Parameters:

• dg – Input dimension

• hidden – Hidden layer dimension

• k – Output dimension (number of classes)

forward(l)

Forward pass of InteractionPredictor.

Parameters:

l – Input tensor

Returns:

Predicted interaction scores

class openddi.models.MIRACLE.MIRACLE(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3, pooling_ratio: float = 0.5)

Bases: Module

MIRACLE model for drug-drug interaction prediction.

This model combines GCN and RGCN layers with interaction prediction.

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3, pooling_ratio: float = 0.5)

Initialize MIRACLE model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension for GCN

• hidden2 – Second hidden layer dimension for RGCN

• num_relations – Number of relation types for RGCN

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

• pooling_ratio – Pooling ratio parameter (unused in current implementation)

forward(data_o, idx)

Forward pass of the MIRACLE model.

Parameters:

• data_o – Graph data object containing node features, edge indices, and edge types

• idx – Batch indices containing pairs of nodes to process

Returns:

Predicted interaction scores

MKGFENN

class openddi.models.MKGFENN.FENN(modal_dims, out_dim, dropout=0.0)

Bases: Module

Feature-level fusion with gating mechanism.

For each modality m: h_m = ReLU(W_m x_m), g_m = sigmoid(wg_m^T x_m + b_m) Fusion: h = Σ_m g_m · h_m

__init__(modal_dims, out_dim, dropout=0.0)

Initialize FENN module.

Parameters:

• modal_dims – List of dimensions for each modality

• out_dim – Output dimension

• dropout – Dropout rate

forward(x, splits)

Forward pass of FENN module.

Parameters:

• x – Input tensor [N, sum(d_m)]

• splits – Split indices [0, d1, d1+d2, …, sum]

Returns:

Fused feature representation

class openddi.models.MKGFENN.MKGFENN(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.5)

Bases: Module

MKG-FENN baseline model.

Architecture: - FENN (multi-modal fusion) → RGCN → Pair-MLP → logits - Signature construction consistent with model_manager

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.5)

Initialize MKGFENN model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension

• hidden2 – Second hidden layer dimension

• num_relations – Number of relation types

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

set_modal_splits(split_str: str)

Set modal splits at runtime (consistent with CSV concatenation order).

Parameters:

split_str – Comma-separated string of modal dimensions, e.g., “1024,768,256,128”

bind_graph(data_graph)

Bind graph data to the model.

Parameters:

data_graph – Graph data object

forward(graph_or_none, idx_batch)

Forward pass of the MKGFENN model.

Parameters:

• graph_or_none – Can be None (already bound) or Data containing graph information

• idx_batch – (i_idx, j_idx, y) batch indices

Returns:

Output logits

Return type:

logits

MMDGDTI

class openddi.models.MMDGDTI.ModalAttnFusion(modal_dims, out_dim, dropout=0.0)

Bases: Module

Multi-modal attention fusion.

For each modality: x_m -> proj_m(x_m) in R^d Scoring: s_m = w_m^T x_m Attention: α = softmax([s_m]); h = sum_m α_m * ReLU(proj_m(x_m))

__init__(modal_dims, out_dim, dropout=0.0)

Initialize ModalAttnFusion.

Parameters:

• modal_dims – List of dimensions for each modality

• out_dim – Output dimension

• dropout – Dropout rate

forward(x, splits)

Forward pass of ModalAttnFusion.

Parameters:

• x – Input tensor [N, sum(d_m)]

• splits – Split indices [0, d1, d1+d2, …, sum]

Returns:

Tuple of (fused features, attention weights)

class openddi.models.MMDGDTI.MMDGDTI(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Bases: Module

MMDG-DTI baseline (adapted for DDI).

Architecture: - Modal attention fusion -> GraphSAGE (two layers) -> Pair representation -> MLP classification - Signature construction consistent with model_manager

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Initialize MMDGDTI model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension

• hidden2 – Second hidden layer dimension

• num_relations – Number of relation types (unused in current implementation)

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

set_modal_splits(split_str: str)

Set modal splits at runtime (consistent with CSV concatenation order).

Parameters:

split_str – Comma-separated string of modal dimensions

bind_graph(data_graph)

Bind graph data to the model.

Parameters:

data_graph – Graph data object

forward(graph_or_none, idx_batch)

Forward pass of the MMDGDTI model.

Parameters:

• graph_or_none – Can be None (already bound) or Data containing graph information

• idx_batch – (i_idx, j_idx, y) batch indices

Returns:

Output logits

Return type:

logits

MRCGNN

openddi.models.MRCGNN.reset_parameters(w)

Reset parameters using uniform distribution.

class openddi.models.MRCGNN.Discriminator(hidden_dim: int)

Bases: Module

Discriminator module for contrastive learning.

__init__(hidden_dim: int)

Initialize Discriminator.

Parameters:

hidden_dim – Hidden dimension size

forward(c, h_pl, h_mi, s_bias1=None, s_bias2=None)

Forward pass of Discriminator.

Parameters:

• c – Context vector

• h_pl – Positive samples

• h_mi – Negative samples

• s_bias1 – Optional bias for positive samples

• s_bias2 – Optional bias for negative samples

Returns:

Discriminator logits

class openddi.models.MRCGNN.AvgReadout(*args: Any, **kwargs: Any)

Bases: Module

Average readout module for graph pooling.

forward(seq, msk=None)

Forward pass of AvgReadout.

Parameters:

• seq – Input sequence

• msk – Optional mask

Returns:

Averaged sequence

class openddi.models.MRCGNN.MRCGNN(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Bases: Module

MRCGNN model for multi-relational graph neural networks with contrastive learning.

Parameters:

• num_relations – Number of relation types for RGCN (= number of multi-class categories)

• num_classes – Final prediction category number (recommended to be same as num_relations)

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Initialize MRCGNN model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension

• hidden2 – Second hidden layer dimension

• num_relations – Number of relation types

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

forward(data_o, data_s, data_a, idx)

Forward pass of the MRCGNN model.

Parameters:

• data_o – Original graph data

• data_s – Structural data

• data_a – Adversarial data

• idx – Batch indices

Returns:

Tuple of (logits, contrastive_loss1, contrastive_loss2, node_embeddings)

MUFFIN

class openddi.models.MUFFIN.MUFFIN(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3, entity_dim=128)

Bases: Module

MUFFIN model for drug-drug interaction prediction using convolutional networks.

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3, entity_dim=128)

Initialize MUFFIN model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension

• hidden2 – Second hidden layer dimension

• num_relations – Number of relation types (unused in current implementation)

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

• entity_dim – Entity embedding dimension

forward(data_o, idx)

Forward pass of the MUFFIN model.

Parameters:

• data_o – Graph data object containing node features

• idx – Batch indices containing pairs of nodes to process

Returns:

Output logits for the given node pairs

MVA

class openddi.models.MVA.MVA(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Bases: Module

MVA model for multi-view drug-drug interaction prediction.

This model combines GCN subgraph processing with attention-based fusion.

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Initialize MVA model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension for GCN

• hidden2 – Second hidden layer dimension for GCN

• num_relations – Number of relation types (unused in current implementation)

• num_classes – Number of output classes

• dropout – Dropout rate for regularization

forward(data_o, idx)

Forward pass of the MVA model.

Parameters:

• data_o – Graph data object containing node features, edge indices, and edge types

• idx – Batch indices containing pairs of nodes to process

Returns:

Output logits for the given node pairs

class openddi.models.MVA.AFF(channels=128, r=4)

Bases: Module

Attentional Feature Fusion module.

__init__(channels=128, r=4)

Initialize AFF module.

Parameters:

• channels – Number of channels

• r – Reduction ratio

forward(x, y)

Forward pass of AFF module.

Parameters:

• x – First input feature

• y – Second input feature

Returns:

Fused feature

class openddi.models.MVA.LayerNorm(hidden_size, variance_epsilon=1e-12)

Bases: Module

Layer normalization module.

__init__(hidden_size, variance_epsilon=1e-12)

Initialize LayerNorm.

Parameters:

• hidden_size – Hidden size dimension

• variance_epsilon – Epsilon for numerical stability

forward(x)

Forward pass of LayerNorm.

Parameters:

x – Input tensor

Returns:

Normalized tensor

class openddi.models.MVA.Embeddings(vocab_size, hidden_size, max_position_size, dropout_rate)

Bases: Module

Construct embeddings from protein/target and position embeddings.

__init__(vocab_size, hidden_size, max_position_size, dropout_rate)

Initialize Embeddings module.

Parameters:

• vocab_size – Vocabulary size

• hidden_size – Hidden size dimension

• max_position_size – Maximum position size

• dropout_rate – Dropout rate

forward(input_ids)

Forward pass of Embeddings.

Parameters:

input_ids – Input token IDs

Returns:

Embedded representations

class openddi.models.MVA.SelfAttention(hidden_size, num_attention_heads, attention_probs_dropout_prob)

Bases: Module

Self-attention module.

__init__(hidden_size, num_attention_heads, attention_probs_dropout_prob)

Initialize SelfAttention.

Parameters:

• hidden_size – Hidden size dimension

• num_attention_heads – Number of attention heads

• attention_probs_dropout_prob – Attention dropout probability

transpose_for_scores(x)

Transpose input for multi-head attention.

forward(hidden_states, attention_mask)

Forward pass of SelfAttention.

Parameters:

• hidden_states – Input hidden states

• attention_mask – Attention mask

Returns:

Context layer after attention

class openddi.models.MVA.SelfOutput(hidden_size, hidden_dropout_prob)

Bases: Module

Self-output module for attention.

__init__(hidden_size, hidden_dropout_prob)

Initialize SelfOutput.

Parameters:

• hidden_size – Hidden size dimension

• hidden_dropout_prob – Hidden dropout probability

forward(hidden_states, input_tensor)

Forward pass of SelfOutput.

Parameters:

• hidden_states – Hidden states from attention

• input_tensor – Original input tensor

Returns:

Output after residual connection and normalization

class openddi.models.MVA.Attention(hidden_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob)

Bases: Module

Complete attention module with self-attention and output.

__init__(hidden_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob)

Initialize Attention module.

Parameters:

• hidden_size – Hidden size dimension

• num_attention_heads – Number of attention heads

• attention_probs_dropout_prob – Attention dropout probability

• hidden_dropout_prob – Hidden dropout probability

forward(input_tensor, attention_mask)

Forward pass of Attention module.

Parameters:

• input_tensor – Input tensor

• attention_mask – Attention mask

Returns:

Attention output

class openddi.models.MVA.Intermediate(hidden_size, intermediate_size)

Bases: Module

Intermediate module for transformer.

__init__(hidden_size, intermediate_size)

Initialize Intermediate module.

Parameters:

• hidden_size – Hidden size dimension

• intermediate_size – Intermediate size dimension

forward(hidden_states)

Forward pass of Intermediate module.

Parameters:

hidden_states – Input hidden states

Returns:

Intermediate representations

class openddi.models.MVA.Output(intermediate_size, hidden_size, hidden_dropout_prob)

Bases: Module

Output module for transformer.

__init__(intermediate_size, hidden_size, hidden_dropout_prob)

Initialize Output module.

Parameters:

• intermediate_size – Intermediate size dimension

• hidden_size – Hidden size dimension

• hidden_dropout_prob – Hidden dropout probability

forward(hidden_states, input_tensor)

Forward pass of Output module.

Parameters:

• hidden_states – Intermediate hidden states

• input_tensor – Original input tensor

Returns:

Output after residual connection

class openddi.models.MVA.Encoder(hidden_size, intermediate_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob)

Bases: Module

Single encoder layer for transformer.

__init__(hidden_size, intermediate_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob)

Initialize Encoder layer.

Parameters:

• hidden_size – Hidden size dimension

• intermediate_size – Intermediate size dimension

• num_attention_heads – Number of attention heads

• attention_probs_dropout_prob – Attention dropout probability

• hidden_dropout_prob – Hidden dropout probability

forward(hidden_states, attention_mask)

Forward pass of Encoder layer.

Parameters:

• hidden_states – Input hidden states

• attention_mask – Attention mask

Returns:

Layer output

class openddi.models.MVA.Encoder_MultipleLayers(n_layer, hidden_size, intermediate_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob)

Bases: Module

Multiple encoder layers for transformer.

__init__(n_layer, hidden_size, intermediate_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob)

Initialize multiple encoder layers.

Parameters:

• n_layer – Number of layers

• hidden_size – Hidden size dimension

• intermediate_size – Intermediate size dimension

• num_attention_heads – Number of attention heads

• attention_probs_dropout_prob – Attention dropout probability

• hidden_dropout_prob – Hidden dropout probability

forward(hidden_states, attention_mask, output_all_encoded_layers=True)

Forward pass of multiple encoder layers.

Parameters:

• hidden_states – Input hidden states

• attention_mask – Attention mask

• output_all_encoded_layers – Whether to output all layers

Returns:

Final hidden states

PHGLDDI

class openddi.models.PHGLDDI.GCN_Bottom(hidden=512, feature=300)

Bases: Module

Bottom GCN module for hierarchical graph processing.

__init__(hidden=512, feature=300)

Initialize GCN_Bottom module.

Parameters:

• hidden – Hidden dimension size

• feature – Input feature dimension

forward(x, edge_index)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class openddi.models.PHGLDDI.GIN_Top(fea, hid, hidden=256, train_eps=True)

Bases: Module

Top GIN module for hierarchical graph processing.

__init__(fea, hid, hidden=256, train_eps=True)

Initialize GIN_Top module.

Parameters:

• fea – Input feature dimension

• hid – Hidden dimension size

• hidden – Output hidden dimension size

• train_eps – Whether to train epsilon parameter

reset_parameters()

Reset model parameters.

forward(x, edge_index)

Forward pass of GIN_Top.

Parameters:

• x – Node features

• edge_index – Graph edge indices

Returns:

Processed node representations

class openddi.models.PHGLDDI.PHGLDDI(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int)

Bases: Module

PHGLDDI model for hierarchical graph learning in drug-drug interaction prediction.

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int)

Initialize PHGLDDI model.

Parameters:

• feature – Input feature dimension

• hidden1 – First hidden layer dimension

• hidden2 – Second hidden layer dimension

• num_relations – Number of relation types (unused in current implementation)

• num_classes – Number of output classes

forward(graph_or_none, idx_batch)

Forward pass of the PHGLDDI model.

Parameters:

• graph_or_none – Graph data object containing node features and edge indices

• idx_batch – Batch indices containing pairs of nodes to process

Returns:

Output logits for the given node pairs

SumGNN

class openddi.models.SumGNN.SumGNN(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3, proj_dim: int | None = None, num_bases: int | None = None)

Bases: Module

SumGNN model with two RGCN layers.

Projects high-dimensional multimodal features to lower dimensions before message passing. Pair representation = concat(h1_i, h1_j, h2_i, h2_j) -> fully connected output logits.

Parameters:

• feature (int) – Original multimodal feature dimension (e.g., 2860)

• hidden1 (int) – First hidden layer dimension

• hidden2 (int) – Second hidden layer dimension

• num_relations (int) – Number of relation types in the graph

• num_classes (int) – Number of output classes

• dropout (float) – Dropout rate, default 0.3

• proj_dim (int, optional) – Projection dimension for input features

• num_bases (int, optional) – Number of bases for RGCN weight decomposition

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3, proj_dim: int | None = None, num_bases: int | None = None)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

drug_feat(emb)

Compatibility interface for old code; not used.

forward(data_o, idx_batch)

Forward pass of the model.

Parameters:

• data_o – Graph data object containing node features, edge indices and types

• idx_batch – Batch indices for drug pairs

Returns:

Output logits for the batch

Return type:

torch.Tensor

TIGER

class openddi.models.TIGER.TIGER(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Bases: Module

TIGER (Implementation version): Graph encoding + pair representation + MLP classification.

Maintains compatible signature with existing model_manager: (feature, hidden1, hidden2, num_relations, num_classes, dropout)

Parameters:

• feature (int) – Input feature dimension

• hidden1 (int) – First hidden layer dimension

• hidden2 (int) – Second hidden layer dimension

• num_relations (int) – Number of relation types (for compatibility)

• num_classes (int) – Number of output classes

• dropout (float) – Dropout rate, default 0.3

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

bind_graph(data_graph)

Bind graph data to model for caching.

Parameters:

data_graph – Graph data to cache

forward(graph_or_none, idx_batch)

Forward pass of the model.

Parameters:

• graph_or_none – torch_geometric.data.Data or None (uses cached graph if None)

• idx_batch – (i_idx, j_idx, y) - compatible with Base_multi_dataset / Base_multilabel_dataset

Returns:

Output logits [B, K]

Return type:

torch.Tensor

ZeroDDI

class openddi.models.ZeroDDI.ZeroDDI(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Bases: Module

ZeroDDI: Maps dual representation z_ij and event semantic prototypes U_k to the same unit sphere. logits = (z · U^T) / tau

Enhancements:

• Dynamic prototypes: U = normalize(sem_proj(S_raw)) computed in real-time during forward pass (with gradient)

• Semantic projection sem_proj participates in classification loss backpropagation

• Feature stability: LayerNorm + Normalize

• Stable initial temperature value tau=0.2 (learnable)

Parameters:

• feature (int) – Input feature dimension

• hidden1 (int) – First hidden layer dimension

• hidden2 (int) – Second hidden layer dimension

• num_relations (int) – Number of relation types

• num_classes (int) – Number of output classes

• dropout (float) – Dropout rate, default 0.3

__init__(feature: int, hidden1: int, hidden2: int, num_relations: int, num_classes: int, dropout: float = 0.3)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

update_event_U(event_sem: Tensor)

Only responsible for caching raw event semantics to self.S_raw and (if necessary) initializing sem_proj. Note: Does not generate U here (to avoid no_grad preventing sem_proj gradient flow).

Parameters:

event_sem (torch.Tensor) – Event semantic tensor [K, d_e]

bind_graph(data_graph)

Bind graph data to model for caching.

Parameters:

data_graph – Graph data to cache

forward(graph_or_none, idx_batch)

Forward pass of the model.

Parameters:

• graph_or_none – torch_geometric.data.Data or None (uses cached graph if None)

• idx_batch – Batch indices for drug pairs

Returns:

(logits [B,K], z [B,hid2])

Return type:

tuple

Submodules

gognn

class openddi.models.gognn.gognn.GoGNN(args, num_features, nhid, ddi_nhid, pooling_ratio, dropout_ratio, num_rel)

Bases: Module

GoGNN model for DDI prediction using graph neural networks.

Features: - Molecular graph processing with GCN layers and SAGPooling - DDI graph processing for drug-drug interaction prediction - Supports both multiclass and multilabel classification

Parameters:

• args – Configuration arguments.

• num_features – Input feature dimension.

• nhid – Hidden dimension for molecular GNN.

• ddi_nhid – Hidden dimension for DDI GNN.

• pooling_ratio – Ratio for SAGPooling layers.

• dropout_ratio – Dropout rate.

• num_rel – Number of relation types for classification.

__init__(args, num_features, nhid, ddi_nhid, pooling_ratio, dropout_ratio, num_rel)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(data)

Forward pass of GoGNN model.

Parameters:

data – Tuple containing (data_list, ddi_edge_index, ddi_edge_attr or None, …) - data_list: List of molecular graph Data objects - ddi_edge_index: Edge indices for DDI graph

Returns:

Logits for DDI edge classification.

Return type:

torch.Tensor

loss(logits, labels)

Compute supervised loss for DDI edge classification.

Parameters:

• logits – Model predictions.

• labels – Ground truth labels.

Returns:

Loss value.

Return type:

torch.Tensor

Note

Uses BCEWithLogitsLoss for multilabel tasks, CrossEntropyLoss for multiclass.

muffin

class openddi.models.muffin.muffin.MUFFIN(args, entity_dim, structure_dim, num_rel)

Bases: Module

MUFFIN model for multi-modal drug-drug interaction prediction.

Features: - Dual embedding fusion (entity and structure embeddings) - Cross-modal attention mechanisms - CNN-based feature extraction - Supports both multiclass and multilabel classification

Parameters:

• args – Configuration arguments.

• entity_dim – Dimension of entity embeddings.

• structure_dim – Dimension of structure embeddings.

• num_rel – Number of relation types for classification.

__init__(args, entity_dim, structure_dim, num_rel)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

generate_fusion_feature(batch_data)

Generate fused features from entity and structure embeddings.

Parameters:

batch_data – Tuple containing (entity_embed_pre, structure_embed_pre, …)

Returns:

Fused feature representation.

Return type:

torch.Tensor

forward(batch_data)

Forward pass of MUFFIN model.

Parameters:

batch_data – Tuple containing (entity_embed, structure_embed, ddi_edge_index, labels)

Returns:

Logits for DDI prediction.

Return type:

torch.Tensor

loss(logits, labels)

Compute supervised loss for DDI edge classification.

Parameters:

• logits – Model predictions.

• labels – Ground truth labels.

Returns:

Loss value.

Return type:

torch.Tensor

mva

class openddi.models.mva.mva.MVA(args, gcn_in_features, gcn_out_features, num_rel, bias=True)

Bases: Module

Multi-View Attention model for DDI prediction.

Features: - Graph convolutional networks for molecular structure - Transformer-based sequence encoders for SMILES - Attention-based feature fusion - Supports both multiclass and multilabel classification

Parameters:

• args – Configuration arguments.

• gcn_in_features – Input dimension for GCN.

• gcn_out_features – Output dimension for GCN.

• num_rel – Number of relation types.

• bias – Whether to use bias in GCN (default=True).

__init__(args, gcn_in_features, gcn_out_features, num_rel, bias=True)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

aggregate_message_1(nodes, node_neighbours, edges, mask)

Aggregate messages for first view (to be implemented).

aggregate_message_2(nodes, node_neighbours, edges, mask)

Aggregate messages for second view (to be implemented).

update_1(nodes, messages)

Update node representations for first view (to be implemented).

update_2(nodes, messages)

Update node representations for second view (to be implemented).

readout_1(hidden_nodes, input_nodes, node_mask)

Readout function for first view (to be implemented).

readout_2(hidden_nodes, input_nodes, node_mask)

Readout function for second view (to be implemented).

readout(input_nodes, node_mask)

General readout function (to be implemented).

final_layer(out)

Final layer processing (to be implemented).

reset_parameters()

Reset model parameters using uniform initialization.

forward(data)

Forward pass of MVA model.

Parameters:

data – Tuple containing molecular graph and sequence data.

Returns:

Logits for DDI prediction.

Return type:

torch.Tensor

loss(logits, labels)

Compute supervised loss for DDI edge classification.

Parameters:

• logits – Model predictions.

• labels – Ground truth labels.

Returns:

Loss value.

Return type:

torch.Tensor

class openddi.models.mva.mva.AFF(channels=128, r=4)

Bases: Module

Attention Feature Fusion module for multi-view feature integration.

Parameters:

• channels – Input channel dimension (default=128).

• r – Reduction ratio for intermediate channels (default=4).

__init__(channels=128, r=4)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x, y)

Forward pass for attention feature fusion.

Parameters:

• x – First feature tensor.

• y – Second feature tensor.

Returns:

Fused feature tensor.

Return type:

torch.Tensor

class openddi.models.mva.mva.LayerNorm(hidden_size, variance_epsilon=1e-12)

Bases: Module

Layer normalization module.

Parameters:

• hidden_size – Hidden dimension size.

• variance_epsilon – Small epsilon for numerical stability (default=1e-12).

__init__(hidden_size, variance_epsilon=1e-12)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Apply layer normalization.

Parameters:

x – Input tensor.

Returns:

Normalized tensor.

Return type:

torch.Tensor

class openddi.models.mva.mva.Embeddings(vocab_size, hidden_size, max_position_size, dropout_rate)

Bases: Module

Construct embeddings from sequence tokens and positional information.

Parameters:

• vocab_size – Vocabulary size.

• hidden_size – Hidden dimension size.

• max_position_size – Maximum sequence length.

• dropout_rate – Dropout rate.

__init__(vocab_size, hidden_size, max_position_size, dropout_rate)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(input_ids)

Create embeddings for input tokens.

Parameters:

input_ids – Token indices.

Returns:

Combined token and position embeddings.

Return type:

torch.Tensor

class openddi.models.mva.mva.SelfAttention(hidden_size, num_attention_heads, attention_probs_dropout_prob)

Bases: Module

Self-attention module.

Parameters:

• hidden_size – Hidden dimension size.

• num_attention_heads – Number of attention heads.

• attention_probs_dropout_prob – Dropout probability for attention.

__init__(hidden_size, num_attention_heads, attention_probs_dropout_prob)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

transpose_for_scores(x)

Reshape tensor for multi-head attention.

Parameters:

x – Input tensor.

Returns:

Reshaped tensor for attention computation.

Return type:

torch.Tensor

forward(hidden_states, attention_mask)

Compute self-attention.

Parameters:

• hidden_states – Input hidden states.

• attention_mask – Attention mask.

Returns:

Attention output.

Return type:

torch.Tensor

class openddi.models.mva.mva.SelfOutput(hidden_size, hidden_dropout_prob)

Bases: Module

Output module for self-attention with residual connection.

Parameters:

• hidden_size – Hidden dimension size.

• hidden_dropout_prob – Dropout probability.

__init__(hidden_size, hidden_dropout_prob)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden_states, input_tensor)

Apply output transformation with residual connection.

Parameters:

• hidden_states – Attention output.

• input_tensor – Original input.

Returns:

Output with residual connection.

Return type:

torch.Tensor

class openddi.models.mva.mva.Attention(hidden_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob)

Bases: Module

Complete attention module with self-attention and output projection.

Parameters:

• hidden_size – Hidden dimension size.

• num_attention_heads – Number of attention heads.

• attention_probs_dropout_prob – Dropout probability for attention.

• hidden_dropout_prob – Dropout probability for hidden layers.

__init__(hidden_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(input_tensor, attention_mask)

Apply attention mechanism.

Parameters:

• input_tensor – Input tensor.

• attention_mask – Attention mask.

Returns:

Attention output.

Return type:

torch.Tensor

class openddi.models.mva.mva.Intermediate(hidden_size, intermediate_size)

Bases: Module

Intermediate feed-forward module.

Parameters:

• hidden_size – Input hidden dimension.

• intermediate_size – Intermediate dimension size.

__init__(hidden_size, intermediate_size)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden_states)

Apply intermediate transformation.

Parameters:

hidden_states – Input hidden states.

Returns:

Transformed tensor.

Return type:

torch.Tensor

class openddi.models.mva.mva.Output(intermediate_size, hidden_size, hidden_dropout_prob)

Bases: Module

Output module with residual connection.

Parameters:

• intermediate_size – Intermediate dimension size.

• hidden_size – Output hidden dimension.

• hidden_dropout_prob – Dropout probability.

__init__(intermediate_size, hidden_size, hidden_dropout_prob)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden_states, input_tensor)

Apply output transformation with residual connection.

Parameters:

• hidden_states – Intermediate tensor.

• input_tensor – Original input.

Returns:

Output with residual connection.

Return type:

torch.Tensor

class openddi.models.mva.mva.Encoder(hidden_size, intermediate_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob)

Bases: Module

Transformer encoder block with attention and feed-forward layers.

Parameters:

• hidden_size – Hidden dimension size.

• intermediate_size – Intermediate dimension size.

• num_attention_heads – Number of attention heads.

• attention_probs_dropout_prob – Dropout probability for attention.

• hidden_dropout_prob – Dropout probability for hidden layers.

__init__(hidden_size, intermediate_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden_states, attention_mask)

Apply transformer encoder block.

Parameters:

• hidden_states – Input hidden states.

• attention_mask – Attention mask.

Returns:

Encoded tensor.

Return type:

torch.Tensor

class openddi.models.mva.mva.Encoder_MultipleLayers(n_layer, hidden_size, intermediate_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob)

Bases: Module

Multi-layer transformer encoder.

Parameters:

• n_layer – Number of encoder layers.

• hidden_size – Hidden dimension size.

• intermediate_size – Intermediate dimension size.

• num_attention_heads – Number of attention heads.

• attention_probs_dropout_prob – Dropout probability for attention.

• hidden_dropout_prob – Dropout probability for hidden layers.

__init__(n_layer, hidden_size, intermediate_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden_states, attention_mask, output_all_encoded_layers=True)

Apply multiple transformer encoder layers.

Parameters:

• hidden_states – Input hidden states.

• attention_mask – Attention mask.

• output_all_encoded_layers – Whether to output all layers (default=True).

Returns:

Encoded tensor.

Return type:

torch.Tensor

tiger

class openddi.models.tiger.tiger.NodeFeatures(degree, feature_num, embedding_dim, layer=2, type='graph')

Bases: Module

Node feature encoder with degree information.

Parameters:

• degree – Maximum node degree.

• feature_num – Number of input features.

• embedding_dim – Output embedding dimension.

• layer – Number of layers for initialization (default=2).

• type – Type of features (‘graph’ or ‘node’) (default=’graph’).

__init__(degree, feature_num, embedding_dim, layer=2, type='graph')

Initialize internal Module state, shared by both nn.Module and ScriptModule.

reset_parameters()

Reset model parameters.

forward(data)

Encode node features with degree information.

Parameters:

data – Graph data object.

Returns:

Encoded node features.

Return type:

torch.Tensor

class openddi.models.tiger.tiger.TIGER(max_layer=6, num_features_drug=78, num_nodes=200, num_relations_mol=10, num_relations_graph=10, output_dim=64, max_degree_graph=100, max_degree_node=100, sub_coeff=0.2, mi_coeff=0.5, dropout=0.2, device='cuda', num_rel=None, args=None)

Bases: Module

TIGER model for knowledge graph-enhanced DDI prediction.

Features: - Molecular graph transformer - Knowledge graph transformer - Mutual information maximization - Multi-view representation learning

Parameters:

• max_layer – Maximum number of transformer layers (default=6).

• num_features_drug – Number of drug features (default=78).

• num_nodes – Number of nodes in knowledge graph (default=200).

• num_relations_mol – Number of molecular relation types (default=10).

• num_relations_graph – Number of graph relation types (default=10).

• output_dim – Output dimension (default=64).

• max_degree_graph – Maximum degree in molecular graph (default=100).

• max_degree_node – Maximum degree in knowledge graph (default=100).

• sub_coeff – Subgraph coefficient for loss (default=0.2).

• mi_coeff – Mutual information coefficient for loss (default=0.5).

• dropout – Dropout rate (default=0.2).

• device – Computation device (default=’cuda’).

• num_rel – Number of relation types for classification.

• args – Configuration arguments.

__init__(max_layer=6, num_features_drug=78, num_nodes=200, num_relations_mol=10, num_relations_graph=10, output_dim=64, max_degree_graph=100, max_degree_node=100, sub_coeff=0.2, mi_coeff=0.5, dropout=0.2, device='cuda', num_rel=None, args=None)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

to(device)

Move model to specified device.

Parameters:

device – Target device.

Returns:

Model on target device.

Return type:

TIGER

loss(pred, label)

Compute supervised loss for DDI prediction.

Parameters:

• pred – Model predictions.

• label – Ground truth labels.

Returns:

Combined loss value.

Return type:

torch.Tensor

reset_parameters()

Reset all model parameters.

forward(data)

Forward pass of TIGER model.

Parameters:

data – Tuple containing (drug1_mol, drug1_subgraph, drug2_mol, drug2_subgraph)

Returns:

DDI prediction scores.

Return type:

torch.Tensor

MI(graph_embeddings, sub_embeddings)

Compute mutual information between graph and subgraph embeddings.

Parameters:

• graph_embeddings – Graph-level embeddings.

• sub_embeddings – Subgraph-level embeddings.

Returns:

Discriminator logits.

Return type:

torch.Tensor

loss_MI(logits)

Compute mutual information loss.

Parameters:

logits – Discriminator logits.

Returns:

Binary cross-entropy loss.

Return type:

torch.Tensor

save(path)

Save model state to file.

Parameters:

path – Directory path for saving.

Returns:

Path to saved model file.

Return type:

str

class openddi.models.tiger.tiger.Discriminator(n_h)

Bases: Module

Discriminator for mutual information estimation.

Parameters:

n_h – Hidden dimension size.

__init__(n_h)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

weights_init(m)

Initialize discriminator weights.

Parameters:

m – Module to initialize.

forward(c, h_pl, h_mi, s_bias1=None, s_bias2=None)

Forward pass for discriminator.

Parameters:

• c – Context embeddings.

• h_pl – Positive embeddings.

• h_mi – Negative embeddings.

• s_bias1 – Optional bias for positive scores.

• s_bias2 – Optional bias for negative scores.

Returns:

Discriminator logits.

Return type:

torch.Tensor

openddi.models.tiger.tiger.init_params(module, layers=2)

Initialize module parameters.

Parameters:

• module – Module to initialize.

• layers – Number of layers for initialization scaling (default=2).

GraphTransformer

class openddi.models.tiger.GraphTransformer.GraphTransformerEncode(num_heads, in_dim, dim_forward, rel_encoder, spatial_encoder, dropout)

Bases: Module

Graph Transformer encoder layer.

Parameters:

• num_heads – Number of attention heads.

• in_dim – Input dimension.

• dim_forward – Feed-forward dimension.

• rel_encoder – Relation encoder module.

• spatial_encoder – Spatial encoder module.

• dropout – Dropout rate.

__init__(num_heads, in_dim, dim_forward, rel_encoder, spatial_encoder, dropout)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

reset_parameters()

Reset model parameters.

forward(feature, sp_edge_index, sp_value, edge_rel)

Forward pass for graph transformer encoder.

Parameters:

• feature – Node features.

• sp_edge_index – Shortest path edge indices.

• sp_value – Shortest path distances.

• edge_rel – Edge relation types.

Returns:

(output features, attention weights)

Return type:

tuple

class openddi.models.tiger.GraphTransformer.SpatialEncoding(dim_model)

Bases: Module

Spatial encoding module for graph distances.

Parameters:

dim_model – Dimension of spatial encoding.

__init__(dim_model)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

reset_parameters()

Reset model parameters.

forward(lap)

Encode spatial distances.

Parameters:

lap – Shortest path distances.

Returns:

Spatial encoding.

Return type:

torch.Tensor

class openddi.models.tiger.GraphTransformer.MultiheadAttention(*args: Any, **kwargs: Any)

Bases: MessagePassing

Multi-head attention for graph transformer.

Parameters:

• dim_model – Model dimension.

• num_heads – Number of attention heads.

• rel_encoder – Relation encoder.

• spatial_encoder – Spatial encoder.

• **kwargs – Additional MessagePassing arguments.

__init__(dim_model, num_heads, rel_encoder, spatial_encoder, **kwargs)

reset_parameters()

Reset model parameters.

softmax_kernel_transformation(data, is_query, projection_matrix=None, numerical_stabilizer=1e-06)

Apply softmax kernel transformation for efficient attention.

Parameters:

• data – Input data.

• is_query – Whether input is query.

• projection_matrix – Projection matrix.

• numerical_stabilizer – Numerical stability term.

Returns:

Transformed data.

Return type:

torch.Tensor

denominator(qs, ks)

Compute attention denominator for normalization.

Parameters:

• qs – Query embeddings.

• ks – Key embeddings.

Returns:

Normalization denominator.

Return type:

torch.Tensor

forward(x, sp_edge_index, sp_value, edge_rel)

Forward pass for multi-head attention.

Parameters:

• x – Node features.

• sp_edge_index – Shortest path edge indices.

• sp_value – Shortest path distances.

• edge_rel – Edge relation types.

Returns:

(attention output, attention weights)

Return type:

tuple

class openddi.models.tiger.GraphTransformer.GraphTransformer(layer_num=3, embedding_dim=64, num_heads=4, num_rel=10, dropout=0.2, type='graph')

Bases: Module

Graph Transformer model for molecular representation learning.

Parameters:

• layer_num – Number of transformer layers (default=3).

• embedding_dim – Embedding dimension (default=64).

• num_heads – Number of attention heads (default=4).

• num_rel – Number of relation types (default=10).

• dropout – Dropout rate (default=0.2).

• type – Type of representation (‘graph’ or ‘node’) (default=’graph’).

__init__(layer_num=3, embedding_dim=64, num_heads=4, num_rel=10, dropout=0.2, type='graph')

Initialize internal Module state, shared by both nn.Module and ScriptModule.

reset_parameters()

Reset model parameters.

forward(feature, data)

Forward pass for graph transformer.

Parameters:

• feature – Input node features.

• data – Graph data object.

Returns:

(graph representation, subgraph representations, attention weights)

Return type:

tuple