whyujjwalraj/mechgraph-checkpoints

MechGraph Checkpoints

Pre-trained checkpoints for MechGraph: Neuro-Symbolic Molecular Mechanism Prediction.

Checkpoints

File	Description
neuromatch_best.pt	NeuroMatch encoder with order embeddings for subgraph retrieval
alignment_best.pt	Alignment stage checkpoint (graph-LLM alignment)

Usage

import torch
from mechgraph.models.graph_encoder import GINEncoder, NeuroMatch
from mechgraph.models.projector import Graph2TokenProjector

# Load NeuroMatch
checkpoint = torch.load("checkpoints/neuromatch_best.pt", map_location="cpu")
encoder = GINEncoder(input_dim=9, hidden_dim=256, output_dim=256)
encoder.load_state_dict(checkpoint["encoder_state_dict"])

neuromatch = NeuroMatch(encoder)
neuromatch.load_state_dict(checkpoint["neuromatch_state_dict"])

Model Architecture

• GIN Encoder: 5-layer Graph Isomorphism Network
• NeuroMatch: Order embeddings for subgraph retrieval
• Graph2Token: Cross-attention projector (8 query tokens)
• LLM: Microsoft Phi-2 (2.7B parameters)

Training

Trained using three-stage pipeline:

1. Stage 1 (NeuroMatch): Subgraph retrieval with max-margin loss
2. Stage 2 (Alignment): Graph-LLM feature alignment
3. Stage 3 (Instruction): Mechanism prediction fine-tuning

License

Apache 2.0