codingutf8 2022 microsoft clefourrier the huggingface inc team 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 orglicenseslicense2 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 pytorch graphormer model import math from typing import iterable iterator list optional tuple union import torch import torch nn as nn from torch nn import bcewithlogitsloss crossentropyloss mseloss from activations import act2fn from modelingoutputs import basemodeloutputwithnoattention sequenceclassifieroutput from modelingutils import pretrainedmodel from utils import logging from configurationgraphormer import graphormerconfig logger logging getloggername checkpointfordoc graphormerbasepcqm4mv1 configfordoc graphormerconfig graphormerpretrainedmodelarchivelist clefourriergraphormerbasepcqm4mv1 clefourriergraphormerbasepcqm4mv2 see all graphormer models at https huggingface comodels filtergraphormer def quantnoisemodule nn module p float blocksize int if no quantization noise don t register hook if p 0 return module supported modules if not isinstancemodule nn linear nn embedding nn conv2d raise notimplementederrormodule unsupported for quantnoise test whether module weight has the right sizes wrt blocksize isconv module weight ndim 4 2d matrix if not isconv if module weight size1 blocksize 0 raise assertionerrorinput features must be a multiple of block sizes 4d matrix else 1x1 convolutions if module kernelsize 1 1 if module inchannels blocksize 0 raise assertionerrorinput channels must be a multiple of block sizes regular convolutions else k module kernelsize0 module kernelsize1 if k blocksize 0 raise assertionerrorkernel size must be a multiple of block size def forwardprehookmod input no noise for evaluation if mod training if not isconv gather weight and sizes weight mod weight infeatures weight size1 outfeatures weight size0 split weight matrix into blocks and randomly drop selected blocks mask torch zerosinfeatures blocksize outfeatures deviceweight device mask bernoullip mask mask repeatinterleaveblocksize 1 view1 infeatures else gather weight and sizes weight mod weight inchannels mod inchannels outchannels mod outchannels split weight matrix into blocks and randomly drop selected blocks if mod kernelsize 1 1 mask torch zeros intinchannels blocksize outchannels deviceweight device mask bernoullip mask mask repeatinterleaveblocksize 1 view1 inchannels else mask torch zerosweight size0 weight size1 deviceweight device mask bernoullip mask mask unsqueeze2 unsqueeze3 repeat1 1 mod kernelsize0 mod kernelsize1 scale weights and apply mask mask mask totorch bool x bool is not currently supported in torchscript s 1 1 p mod weight data s weight maskedfillmask 0 module registerforwardprehookforwardprehook return module class layerdropmodulelistnn modulelist def initself p float modules optionaliterablenn module none super initmodules self p p def iterself iteratornn module dropoutprobs torch emptylenself uniform for i m in enumeratesuper iter if not self training or dropoutprobsi self p yield m class graphormergraphnodefeaturenn module def initself config graphormerconfig super init self numheads config numattentionheads self numatoms config numatoms self atomencoder nn embeddingconfig numatoms 1 config hiddensize paddingidxconfig padtokenid self indegreeencoder nn embedding config numindegree config hiddensize paddingidxconfig padtokenid self outdegreeencoder nn embedding config numoutdegree config hiddensize paddingidxconfig padtokenid self graphtoken nn embedding1 config hiddensize def forward self inputnodes torch longtensor indegree torch longtensor outdegree torch longtensor torch tensor ngraph nnode inputnodes size 2 nodefeature node feature graph token self atomencoderinputnodes sumdim2 ngraph nnode nhidden self indegreeencoderindegree self outdegreeencoderoutdegree graphtokenfeature self graphtoken weight unsqueeze0 repeatngraph 1 1 graphnodefeature torch catgraphtokenfeature nodefeature dim1 return graphnodefeature class graphormergraphattnbiasnn module def initself config graphormerconfig super init self numheads config numattentionheads self multihopmaxdist config multihopmaxdist we do not change edge feature embedding learning as edge embeddings are represented as a combination of the original features shortest path self edgeencoder nn embeddingconfig numedges 1 config numattentionheads paddingidx0 self edgetype config edgetype if self edgetype multihop self edgedisencoder nn embedding config numedgedis config numattentionheads config numattentionheads 1 self spatialposencoder nn embeddingconfig numspatial config numattentionheads paddingidx0 self graphtokenvirtualdistance nn embedding1 config numattentionheads def forward self inputnodes torch longtensor attnbias torch tensor spatialpos torch longtensor inputedges torch longtensor attnedgetype torch longtensor torch tensor ngraph nnode inputnodes size 2 graphattnbias attnbias clone graphattnbias graphattnbias unsqueeze1 repeat 1 self numheads 1 1 ngraph nhead nnode1 nnode1 spatial pos ngraph nnode nnode nhead ngraph nhead nnode nnode spatialposbias self spatialposencoderspatialpos permute0 3 1 2 graphattnbias 1 1 graphattnbias 1 1 spatialposbias reset spatial pos here t self graphtokenvirtualdistance weight view1 self numheads 1 graphattnbias 1 0 graphattnbias 1 0 t graphattnbias 0 graphattnbias 0 t edge feature if self edgetype multihop spatialpos spatialpos clone spatialposspatialpos 0 1 set pad to 1 set 1 to 1 inputnodes 1 to inputnodes 1 spatialpos torch wherespatialpos 1 spatialpos 1 spatialpos if self multihopmaxdist 0 spatialpos spatialpos clamp0 self multihopmaxdist inputedges inputedges self multihopmaxdist ngraph nnode nnode maxdist nhead inputedges self edgeencoderinputedges mean2 maxdist inputedges size2 edgeinputflat inputedges permute3 0 1 2 4 reshapemaxdist 1 self numheads edgeinputflat torch bmm edgeinputflat self edgedisencoder weight reshape1 self numheads self numheads maxdist inputedges edgeinputflat reshapemaxdist ngraph nnode nnode self numheads permute 1 2 3 0 4 inputedges inputedges sum2 spatialpos float unsqueeze1 permute0 3 1 2 else ngraph nnode nnode nhead ngraph nhead nnode nnode inputedges self edgeencoderattnedgetype mean2 permute0 3 1 2 graphattnbias 1 1 graphattnbias 1 1 inputedges graphattnbias graphattnbias attnbias unsqueeze1 reset return graphattnbias class graphormermultiheadattentionnn module def initself config graphormerconfig super init self embeddingdim config embeddingdim self kdim config kdim if config kdim is not none else config embeddingdim self vdim config vdim if config vdim is not none else config embeddingdim self qkvsamedim self kdim config embeddingdim and self vdim config embeddingdim self numheads config numattentionheads self attentiondropoutmodule torch nn dropoutpconfig attentiondropout inplacefalse self headdim config embeddingdim config numattentionheads if not self headdim config numattentionheads self embeddingdim raise assertionerrorthe embeddingdim must be divisible by numheads self scaling self headdim0 5 self selfattention true config selfattention if not self selfattention raise notimplementederrorthe graphormer model only supports self attention for now if self selfattention and not self qkvsamedim raise assertionerrorselfattention requires query key and value to be of the same size self kproj quantnoise nn linearself kdim config embeddingdim biasconfig bias config qnoise config qnblocksize self vproj quantnoise nn linearself vdim config embeddingdim biasconfig bias config qnoise config qnblocksize self qproj quantnoise nn linearconfig embeddingdim config embeddingdim biasconfig bias config qnoise config qnblocksize self outproj quantnoise nn linearconfig embeddingdim config embeddingdim biasconfig bias config qnoise config qnblocksize self onnxtrace false def resetparametersself if self qkvsamedim empirically observed the convergence to be much better with the scaled initialization nn init xavieruniformself kproj weight gain1 math sqrt2 nn init xavieruniformself vproj weight gain1 math sqrt2 nn init xavieruniformself qproj weight gain1 math sqrt2 else nn init xavieruniformself kproj weight nn init xavieruniformself vproj weight nn init xavieruniformself qproj weight nn init xavieruniformself outproj weight if self outproj bias is not none nn init constantself outproj bias 0 0 def forward self query torch longtensor key optionaltorch tensor value optionaltorch tensor attnbias optionaltorch tensor keypaddingmask optionaltorch tensor none needweights bool true attnmask optionaltorch tensor none beforesoftmax bool false needheadweights bool false tupletorch tensor optionaltorch tensor if needheadweights needweights true tgtlen bsz embeddingdim query size srclen tgtlen if not embeddingdim self embeddingdim raise assertionerror fthe query embedding dimension embeddingdim is not equal to the expected embeddingdim f self embeddingdim if not listquery size tgtlen bsz embeddingdim raise assertionerrorquery size incorrect in graphormer compared to model dimensions if key is not none srclen keybsz key size if not torch jit isscripting if keybsz bsz or value is none or not srclen bsz value shape 2 raise assertionerror the batch shape does not match the key or value shapes provided to the attention q self qprojquery k self kprojquery v self vprojquery q self scaling q q contiguous viewtgtlen bsz self numheads self headdim transpose0 1 if k is not none k k contiguous view1 bsz self numheads self headdim transpose0 1 if v is not none v v contiguous view1 bsz self numheads self headdim transpose0 1 if k is none or not k size1 srclen raise assertionerrorthe shape of the key generated in the attention is incorrect this is part of a workaround to get around forkjoin parallelism not supporting optional types if keypaddingmask is not none and keypaddingmask dim 0 keypaddingmask none if keypaddingmask is not none if keypaddingmask size0 bsz or keypaddingmask size1 srclen raise assertionerror the shape of the generated padding mask for the key does not match expected dimensions attnweights torch bmmq k transpose1 2 attnweights self applysparsemaskattnweights tgtlen srclen bsz if listattnweights size bsz self numheads tgtlen srclen raise assertionerrorthe attention weights generated do not match the expected dimensions if attnbias is not none attnweights attnbias viewbsz self numheads tgtlen srclen if attnmask is not none attnmask attnmask unsqueeze0 attnweights attnmask if keypaddingmask is not none don t attend to padding symbols attnweights attnweights viewbsz self numheads tgtlen srclen attnweights attnweights maskedfill keypaddingmask unsqueeze1 unsqueeze2 totorch bool floatinf attnweights attnweights viewbsz self numheads tgtlen srclen if beforesoftmax return attnweights v attnweightsfloat torch nn functional softmaxattnweights dim1 attnweights attnweightsfloat typeasattnweights attnprobs self attentiondropoutmoduleattnweights if v is none raise assertionerrorno value generated attn torch bmmattnprobs v if listattn size bsz self numheads tgtlen self headdim raise assertionerrorthe attention generated do not match the expected dimensions attn attn transpose0 1 contiguous viewtgtlen bsz embeddingdim attn torch tensor self outprojattn attnweights none if needweights attnweights attnweightsfloat contiguous viewbsz self numheads tgtlen srclen transpose1 0 if not needheadweights average attention weights over heads attnweights attnweights meandim0 return attn attnweights def applysparsemaskself attnweights torch tensor tgtlen int srclen int bsz int torch tensor return attnweights class graphormergraphencoderlayernn module def initself config graphormerconfig none super init initialize parameters self embeddingdim config embeddingdim self numattentionheads config numattentionheads self qnoise config qnoise self qnblocksize config qnblocksize self prelayernorm config prelayernorm self dropoutmodule torch nn dropoutpconfig dropout inplacefalse self activationdropoutmodule torch nn dropoutpconfig activationdropout inplacefalse initialize blocks self activationfn act2fnconfig activationfn self selfattn graphormermultiheadattentionconfig layer norm associated with the self attention layer self selfattnlayernorm nn layernormself embeddingdim self fc1 self buildfc self embeddingdim config ffnembeddingdim qnoiseconfig qnoise qnblocksizeconfig qnblocksize self fc2 self buildfc config ffnembeddingdim self embeddingdim qnoiseconfig qnoise qnblocksizeconfig qnblocksize layer norm associated with the position wise feedforward nn self finallayernorm nn layernormself embeddingdim def buildfc self inputdim int outputdim int qnoise float qnblocksize int unionnn module nn linear nn embedding nn conv2d return quantnoisenn linearinputdim outputdim qnoise qnblocksize def forward self inputnodes torch tensor selfattnbias optionaltorch tensor none selfattnmask optionaltorch tensor none selfattnpaddingmask optionaltorch tensor none tupletorch tensor optionaltorch tensor residual inputnodes if self prelayernorm inputnodes self selfattnlayernorminputnodes inputnodes attn self selfattn queryinputnodes keyinputnodes valueinputnodes attnbiasselfattnbias keypaddingmaskselfattnpaddingmask needweightsfalse attnmaskselfattnmask inputnodes self dropoutmoduleinputnodes inputnodes residual inputnodes if not self prelayernorm inputnodes self selfattnlayernorminputnodes residual inputnodes if self prelayernorm inputnodes self finallayernorminputnodes inputnodes self activationfnself fc1inputnodes inputnodes self activationdropoutmoduleinputnodes inputnodes self fc2inputnodes inputnodes self dropoutmoduleinputnodes inputnodes residual inputnodes if not self prelayernorm inputnodes self finallayernorminputnodes return inputnodes attn class graphormergraphencodernn module def initself config graphormerconfig super init self dropoutmodule torch nn dropoutpconfig dropout inplacefalse self layerdrop config layerdrop self embeddingdim config embeddingdim self applygraphormerinit config applygraphormerinit self traceable config traceable self graphnodefeature graphormergraphnodefeatureconfig self graphattnbias graphormergraphattnbiasconfig self embedscale config embedscale if config qnoise 0 self quantnoise quantnoise nn linearself embeddingdim self embeddingdim biasfalse config qnoise config qnblocksize else self quantnoise none if config encodernormalizebefore self emblayernorm nn layernormself embeddingdim else self emblayernorm none if config prelayernorm self finallayernorm nn layernormself embeddingdim if self layerdrop 0 0 self layers layerdropmodulelistpself layerdrop else self layers nn modulelist self layers extendgraphormergraphencoderlayerconfig for in rangeconfig numhiddenlayers apply initialization of model params after building the model if config freezeembeddings raise notimplementederrorfreezing embeddings is not implemented yet for layer in rangeconfig numtranslayerstofreeze m self layerslayer if m is not none for p in m parameters p requiresgrad false def forward self inputnodes torch longtensor inputedges torch longtensor attnbias torch tensor indegree torch longtensor outdegree torch longtensor spatialpos torch longtensor attnedgetype torch longtensor perturbnone laststateonly bool false tokenembeddings optionaltorch tensor none attnmask optionaltorch tensor none tupleuniontorch tensor listtorch longtensor torch tensor compute padding mask this is needed for multihead attention datax inputnodes ngraph nnode datax size 2 paddingmask datax 0 eq0 paddingmaskcls torch zerosngraph 1 devicepaddingmask device dtypepaddingmask dtype paddingmask torch catpaddingmaskcls paddingmask dim1 attnbias self graphattnbiasinputnodes attnbias spatialpos inputedges attnedgetype if tokenembeddings is not none inputnodes tokenembeddings else inputnodes self graphnodefeatureinputnodes indegree outdegree if perturb is not none inputnodes 1 perturb if self embedscale is not none inputnodes inputnodes self embedscale if self quantnoise is not none inputnodes self quantnoiseinputnodes if self emblayernorm is not none inputnodes self emblayernorminputnodes inputnodes self dropoutmoduleinputnodes inputnodes inputnodes transpose0 1 innerstates if not laststateonly innerstates appendinputnodes for layer in self layers inputnodes layer inputnodes selfattnpaddingmaskpaddingmask selfattnmaskattnmask selfattnbiasattnbias if not laststateonly innerstates appendinputnodes graphrep inputnodes0 if laststateonly innerstates inputnodes if self traceable return torch stackinnerstates graphrep else return innerstates graphrep class graphormerdecoderheadnn module def initself embeddingdim int numclasses int super init an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models with fsdp module params will be on cuda so we cast them back to cpu so that the rng is consistent with and without fsdp initialize the weights specific to the graphormer model initialize the weights we might be missing part of the linear init dependant on the layer num the graphormer model is a graphencoder model it goes from a graph to its representation if you want to use the model for a downstream classification task use graphormerforgraphclassification instead for any other downstream task feel free to add a new class or combine this model with a downstream model of your choice following the example in graphormerforgraphclassification remove head is set to true during finetuning last inner state then revert batch and graph len project masked tokens only project back to size of vocabulary maximum output length supported by the encoder return self maxnodes class graphormerforgraphclassificationgraphormerpretrainedmodel def initself config graphormerconfig super initconfig self encoder graphormermodelconfig self embeddingdim config embeddingdim self numclasses config numclasses self classifier graphormerdecoderheadself embeddingdim self numclasses self isencoderdecoder true initialize weights and apply final processing self postinit def forward self inputnodes torch longtensor inputedges torch longtensor attnbias torch tensor indegree torch longtensor outdegree torch longtensor spatialpos torch longtensor attnedgetype torch longtensor labels optionaltorch longtensor none returndict optionalbool none unused uniontupletorch tensor sequenceclassifieroutput returndict returndict if returndict is not none else self config usereturndict encoderoutputs self encoder inputnodes inputedges attnbias indegree outdegree spatialpos attnedgetype returndicttrue outputs hiddenstates encoderoutputslasthiddenstate encoderoutputshiddenstates headoutputs self classifieroutputs logits headoutputs 0 contiguous loss none if labels is not none mask torch isnanlabels if self numclasses 1 regression lossfct mseloss loss lossfctlogitsmask squeeze labelsmask squeeze float elif self numclasses 1 and lenlabels shape 1 one task classification lossfct crossentropyloss loss lossfctlogitsmask view1 self numclasses labelsmask view1 else binary multitask classification lossfct bcewithlogitslossreductionsum loss lossfctlogitsmask labelsmask if not returndict return tuplex for x in loss logits hiddenstates if x is not none return sequenceclassifieroutputlossloss logitslogits hiddenstateshiddenstates attentionsnone coding utf 8 2022 microsoft clefourrier the huggingface inc team 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 pytorch graphormer model see all graphormer models at https huggingface co models filter graphormer from https github com facebookresearch fairseq blob dd0079bde7f678b0cd0715cbd0ae68d661b7226d fairseq modules quant_noise py wraps modules and applies quantization noise to the weights for subsequent quantization with iterative product quantization as described in training with quantization noise for extreme model compression args module nn module p amount of quantization noise block_size size of the blocks for subsequent quantization with ipq remarks module weights must have the right sizes wrt the block size only linear embedding and conv2d modules are supported for the moment for more detail on how to quantize by blocks with convolutional weights see and the bit goes down revisiting the quantization of neural networks we implement the simplest form of noise here as stated in the paper which consists in randomly dropping blocks if no quantization noise don t register hook supported modules test whether module weight has the right sizes wrt block_size 2d matrix 4d matrix 1x1 convolutions regular convolutions no noise for evaluation gather weight and sizes split weight matrix into blocks and randomly drop selected blocks gather weight and sizes split weight matrix into blocks and randomly drop selected blocks scale weights and apply mask x bool is not currently supported in torchscript from https github com facebookresearch fairseq blob dd0079bde7f678b0cd0715cbd0ae68d661b7226d fairseq modules layer_drop py a layerdrop implementation based on torch nn modulelist layerdrop as described in https arxiv org abs 1909 11556 we refresh the choice of which layers to drop every time we iterate over the layerdropmodulelist instance during evaluation we always iterate over all layers usage python layers layerdroplist p 0 5 modules layer1 layer2 layer3 for layer in layers this might iterate over layers 1 and 3 x layer x for layer in layers this might iterate over all layers x layer x for layer in layers this might not iterate over any layers x layer x args p float probability of dropping out each layer modules iterable optional an iterable of modules to add compute node features for each node in the graph node feature graph token n_graph n_node n_hidden compute attention bias for each head we do not change edge feature embedding learning as edge embeddings are represented as a combination of the original features shortest path n_graph n_head n_node 1 n_node 1 spatial pos n_graph n_node n_node n_head n_graph n_head n_node n_node reset spatial pos here edge feature set pad to 1 set 1 to 1 input_nodes 1 to input_nodes 1 n_graph n_node n_node max_dist n_head n_graph n_node n_node n_head n_graph n_head n_node n_node reset multi headed attention see attention is all you need for more details config self_attention empirically observed the convergence to be much better with the scaled initialization args key_padding_mask bytetorch tensor optional mask to exclude keys that are pads of shape batch src_len where padding elements are indicated by 1s need_weights bool optional return the attention weights averaged over heads default false attn_mask bytetorch tensor optional typically used to implement causal attention where the mask prevents the attention from looking forward in time default none before_softmax bool optional return the raw attention weights and values before the attention softmax need_head_weights bool optional return the attention weights for each head implies need_weights default return the average attention weights over all heads this is part of a workaround to get around fork join parallelism not supporting optional types don t attend to padding symbols average attention weights over heads initialize parameters initialize blocks layer norm associated with the self attention layer layer norm associated with the position wise feed forward nn nn layernorm is applied either before or after the self attention ffn modules similar to the original transformer implementation apply initialization of model params after building the model compute padding mask this is needed for multi head attention num_classes should be 1 for regression or the number of classes for classification an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models with fsdp module params will be on cuda so we cast them back to cpu so that the rng is consistent with and without fsdp initialize the weights specific to the graphormer model initialize the weights we might be missing part of the linear init dependant on the layer num the graphormer model is a graph encoder model it goes from a graph to its representation if you want to use the model for a downstream classification task use graphormerforgraphclassification instead for any other downstream task feel free to add a new class or combine this model with a downstream model of your choice following the example in graphormerforgraphclassification remove head is set to true during fine tuning last inner state then revert batch and graph len project masked tokens only project back to size of vocabulary maximum output length supported by the encoder this model can be used for graph level classification or regression tasks it can be trained on regression by setting config num_classes to 1 there should be one float type label per graph one task classification by setting config num_classes to the number of classes there should be one integer label per graph binary multi task classification by setting config num_classes to the number of labels there should be a list of integer labels for each graph initialize weights and apply final processing regression one task classification binary multi task classification
import math
from typing import Iterable, Iterator, List, Optional, Tuple, Union

import torch
import torch.nn as nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACT2FN
from ...modeling_outputs import (
    BaseModelOutputWithNoAttention,
    SequenceClassifierOutput,
)
from ...modeling_utils import PreTrainedModel
from ...utils import logging
from .configuration_graphormer import GraphormerConfig


logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "graphormer-base-pcqm4mv1"
_CONFIG_FOR_DOC = "GraphormerConfig"


GRAPHORMER_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "clefourrier/graphormer-base-pcqm4mv1",
    "clefourrier/graphormer-base-pcqm4mv2",
    
]


def quant_noise(module: nn.Module, p: float, block_size: int):
    

    
    if p <= 0:
        return module

    
    if not isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d)):
        raise NotImplementedError("Module unsupported for quant_noise.")

    
    is_conv = module.weight.ndim == 4

    
    if not is_conv:
        if module.weight.size(1) % block_size != 0:
            raise AssertionError("Input features must be a multiple of block sizes")

    
    else:
        
        if module.kernel_size == (1, 1):
            if module.in_channels % block_size != 0:
                raise AssertionError("Input channels must be a multiple of block sizes")
        
        else:
            k = module.kernel_size[0] * module.kernel_size[1]
            if k % block_size != 0:
                raise AssertionError("Kernel size must be a multiple of block size")

    def _forward_pre_hook(mod, input):
        
        if mod.training:
            if not is_conv:
                
                weight = mod.weight
                in_features = weight.size(1)
                out_features = weight.size(0)

                
                mask = torch.zeros(in_features // block_size * out_features, device=weight.device)
                mask.bernoulli_(p)
                mask = mask.repeat_interleave(block_size, -1).view(-1, in_features)

            else:
                
                weight = mod.weight
                in_channels = mod.in_channels
                out_channels = mod.out_channels

                
                if mod.kernel_size == (1, 1):
                    mask = torch.zeros(
                        int(in_channels // block_size * out_channels),
                        device=weight.device,
                    )
                    mask.bernoulli_(p)
                    mask = mask.repeat_interleave(block_size, -1).view(-1, in_channels)
                else:
                    mask = torch.zeros(weight.size(0), weight.size(1), device=weight.device)
                    mask.bernoulli_(p)
                    mask = mask.unsqueeze(2).unsqueeze(3).repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1])

            
            mask = mask.to(torch.bool)  
            s = 1 / (1 - p)
            mod.weight.data = s * weight.masked_fill(mask, 0)

    module.register_forward_pre_hook(_forward_pre_hook)
    return module


class LayerDropModuleList(nn.ModuleList):
    

    def __init__(self, p: float, modules: Optional[Iterable[nn.Module]] = None):
        super().__init__(modules)
        self.p = p

    def __iter__(self) -> Iterator[nn.Module]:
        dropout_probs = torch.empty(len(self)).uniform_()
        for i, m in enumerate(super().__iter__()):
            if not self.training or (dropout_probs[i] > self.p):
                yield m


class GraphormerGraphNodeFeature(nn.Module):
    

    def __init__(self, config: GraphormerConfig):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.num_atoms = config.num_atoms

        self.atom_encoder = nn.Embedding(config.num_atoms + 1, config.hidden_size, padding_idx=config.pad_token_id)
        self.in_degree_encoder = nn.Embedding(
            config.num_in_degree, config.hidden_size, padding_idx=config.pad_token_id
        )
        self.out_degree_encoder = nn.Embedding(
            config.num_out_degree, config.hidden_size, padding_idx=config.pad_token_id
        )

        self.graph_token = nn.Embedding(1, config.hidden_size)

    def forward(
        self,
        input_nodes: torch.LongTensor,
        in_degree: torch.LongTensor,
        out_degree: torch.LongTensor,
    ) -> torch.Tensor:
        n_graph, n_node = input_nodes.size()[:2]

        node_feature = (  
            self.atom_encoder(input_nodes).sum(dim=-2)  
            + self.in_degree_encoder(in_degree)
            + self.out_degree_encoder(out_degree)
        )

        graph_token_feature = self.graph_token.weight.unsqueeze(0).repeat(n_graph, 1, 1)

        graph_node_feature = torch.cat([graph_token_feature, node_feature], dim=1)

        return graph_node_feature


class GraphormerGraphAttnBias(nn.Module):
    

    def __init__(self, config: GraphormerConfig):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.multi_hop_max_dist = config.multi_hop_max_dist

        
        
        self.edge_encoder = nn.Embedding(config.num_edges + 1, config.num_attention_heads, padding_idx=0)

        self.edge_type = config.edge_type
        if self.edge_type == "multi_hop":
            self.edge_dis_encoder = nn.Embedding(
                config.num_edge_dis * config.num_attention_heads * config.num_attention_heads,
                1,
            )

        self.spatial_pos_encoder = nn.Embedding(config.num_spatial, config.num_attention_heads, padding_idx=0)

        self.graph_token_virtual_distance = nn.Embedding(1, config.num_attention_heads)

    def forward(
        self,
        input_nodes: torch.LongTensor,
        attn_bias: torch.Tensor,
        spatial_pos: torch.LongTensor,
        input_edges: torch.LongTensor,
        attn_edge_type: torch.LongTensor,
    ) -> torch.Tensor:
        n_graph, n_node = input_nodes.size()[:2]
        graph_attn_bias = attn_bias.clone()
        graph_attn_bias = graph_attn_bias.unsqueeze(1).repeat(
            1, self.num_heads, 1, 1
        )  

        
        
        spatial_pos_bias = self.spatial_pos_encoder(spatial_pos).permute(0, 3, 1, 2)
        graph_attn_bias[:, :, 1:, 1:] = graph_attn_bias[:, :, 1:, 1:] + spatial_pos_bias

        
        t = self.graph_token_virtual_distance.weight.view(1, self.num_heads, 1)
        graph_attn_bias[:, :, 1:, 0] = graph_attn_bias[:, :, 1:, 0] + t
        graph_attn_bias[:, :, 0, :] = graph_attn_bias[:, :, 0, :] + t

        
        if self.edge_type == "multi_hop":
            spatial_pos_ = spatial_pos.clone()

            spatial_pos_[spatial_pos_ == 0] = 1  
            
            spatial_pos_ = torch.where(spatial_pos_ > 1, spatial_pos_ - 1, spatial_pos_)
            if self.multi_hop_max_dist > 0:
                spatial_pos_ = spatial_pos_.clamp(0, self.multi_hop_max_dist)
                input_edges = input_edges[:, :, :, : self.multi_hop_max_dist, :]
            

            input_edges = self.edge_encoder(input_edges).mean(-2)
            max_dist = input_edges.size(-2)
            edge_input_flat = input_edges.permute(3, 0, 1, 2, 4).reshape(max_dist, -1, self.num_heads)
            edge_input_flat = torch.bmm(
                edge_input_flat,
                self.edge_dis_encoder.weight.reshape(-1, self.num_heads, self.num_heads)[:max_dist, :, :],
            )
            input_edges = edge_input_flat.reshape(max_dist, n_graph, n_node, n_node, self.num_heads).permute(
                1, 2, 3, 0, 4
            )
            input_edges = (input_edges.sum(-2) / (spatial_pos_.float().unsqueeze(-1))).permute(0, 3, 1, 2)
        else:
            
            input_edges = self.edge_encoder(attn_edge_type).mean(-2).permute(0, 3, 1, 2)

        graph_attn_bias[:, :, 1:, 1:] = graph_attn_bias[:, :, 1:, 1:] + input_edges
        graph_attn_bias = graph_attn_bias + attn_bias.unsqueeze(1)  

        return graph_attn_bias


class GraphormerMultiheadAttention(nn.Module):
    

    def __init__(self, config: GraphormerConfig):
        super().__init__()
        self.embedding_dim = config.embedding_dim
        self.kdim = config.kdim if config.kdim is not None else config.embedding_dim
        self.vdim = config.vdim if config.vdim is not None else config.embedding_dim
        self.qkv_same_dim = self.kdim == config.embedding_dim and self.vdim == config.embedding_dim

        self.num_heads = config.num_attention_heads
        self.attention_dropout_module = torch.nn.Dropout(p=config.attention_dropout, inplace=False)

        self.head_dim = config.embedding_dim // config.num_attention_heads
        if not (self.head_dim * config.num_attention_heads == self.embedding_dim):
            raise AssertionError("The embedding_dim must be divisible by num_heads.")
        self.scaling = self.head_dim**-0.5

        self.self_attention = True  
        if not (self.self_attention):
            raise NotImplementedError("The Graphormer model only supports self attention for now.")
        if self.self_attention and not self.qkv_same_dim:
            raise AssertionError("Self-attention requires query, key and value to be of the same size.")

        self.k_proj = quant_noise(
            nn.Linear(self.kdim, config.embedding_dim, bias=config.bias),
            config.q_noise,
            config.qn_block_size,
        )
        self.v_proj = quant_noise(
            nn.Linear(self.vdim, config.embedding_dim, bias=config.bias),
            config.q_noise,
            config.qn_block_size,
        )
        self.q_proj = quant_noise(
            nn.Linear(config.embedding_dim, config.embedding_dim, bias=config.bias),
            config.q_noise,
            config.qn_block_size,
        )

        self.out_proj = quant_noise(
            nn.Linear(config.embedding_dim, config.embedding_dim, bias=config.bias),
            config.q_noise,
            config.qn_block_size,
        )

        self.onnx_trace = False

    def reset_parameters(self):
        if self.qkv_same_dim:
            
            
            nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
            nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
            nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
        else:
            nn.init.xavier_uniform_(self.k_proj.weight)
            nn.init.xavier_uniform_(self.v_proj.weight)
            nn.init.xavier_uniform_(self.q_proj.weight)

        nn.init.xavier_uniform_(self.out_proj.weight)
        if self.out_proj.bias is not None:
            nn.init.constant_(self.out_proj.bias, 0.0)

    def forward(
        self,
        query: torch.LongTensor,
        key: Optional[torch.Tensor],
        value: Optional[torch.Tensor],
        attn_bias: Optional[torch.Tensor],
        key_padding_mask: Optional[torch.Tensor] = None,
        need_weights: bool = True,
        attn_mask: Optional[torch.Tensor] = None,
        before_softmax: bool = False,
        need_head_weights: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        
        if need_head_weights:
            need_weights = True

        tgt_len, bsz, embedding_dim = query.size()
        src_len = tgt_len
        if not (embedding_dim == self.embedding_dim):
            raise AssertionError(
                f"The query embedding dimension {embedding_dim} is not equal to the expected embedding_dim"
                f" {self.embedding_dim}."
            )
        if not (list(query.size()) == [tgt_len, bsz, embedding_dim]):
            raise AssertionError("Query size incorrect in Graphormer, compared to model dimensions.")

        if key is not None:
            src_len, key_bsz, _ = key.size()
            if not torch.jit.is_scripting():
                if (key_bsz != bsz) or (value is None) or not (src_len, bsz == value.shape[:2]):
                    raise AssertionError(
                        "The batch shape does not match the key or value shapes provided to the attention."
                    )

        q = self.q_proj(query)
        k = self.k_proj(query)
        v = self.v_proj(query)

        q *= self.scaling

        q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        if k is not None:
            k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        if v is not None:
            v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)

        if (k is None) or not (k.size(1) == src_len):
            raise AssertionError("The shape of the key generated in the attention is incorrect")

        
        
        if key_padding_mask is not None and key_padding_mask.dim() == 0:
            key_padding_mask = None

        if key_padding_mask is not None:
            if key_padding_mask.size(0) != bsz or key_padding_mask.size(1) != src_len:
                raise AssertionError(
                    "The shape of the generated padding mask for the key does not match expected dimensions."
                )
        attn_weights = torch.bmm(q, k.transpose(1, 2))
        attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)

        if list(attn_weights.size()) != [bsz * self.num_heads, tgt_len, src_len]:
            raise AssertionError("The attention weights generated do not match the expected dimensions.")

        if attn_bias is not None:
            attn_weights += attn_bias.view(bsz * self.num_heads, tgt_len, src_len)

        if attn_mask is not None:
            attn_mask = attn_mask.unsqueeze(0)
            attn_weights += attn_mask

        if key_padding_mask is not None:
            
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf")
            )
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        if before_softmax:
            return attn_weights, v

        attn_weights_float = torch.nn.functional.softmax(attn_weights, dim=-1)
        attn_weights = attn_weights_float.type_as(attn_weights)
        attn_probs = self.attention_dropout_module(attn_weights)

        if v is None:
            raise AssertionError("No value generated")
        attn = torch.bmm(attn_probs, v)
        if list(attn.size()) != [bsz * self.num_heads, tgt_len, self.head_dim]:
            raise AssertionError("The attention generated do not match the expected dimensions.")

        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embedding_dim)
        attn: torch.Tensor = self.out_proj(attn)

        attn_weights = None
        if need_weights:
            attn_weights = attn_weights_float.contiguous().view(bsz, self.num_heads, tgt_len, src_len).transpose(1, 0)
            if not need_head_weights:
                
                attn_weights = attn_weights.mean(dim=0)

        return attn, attn_weights

    def apply_sparse_mask(self, attn_weights: torch.Tensor, tgt_len: int, src_len: int, bsz: int) -> torch.Tensor:
        return attn_weights


class GraphormerGraphEncoderLayer(nn.Module):
    def __init__(self, config: GraphormerConfig) -> None:
        super().__init__()

        
        self.embedding_dim = config.embedding_dim
        self.num_attention_heads = config.num_attention_heads
        self.q_noise = config.q_noise
        self.qn_block_size = config.qn_block_size
        self.pre_layernorm = config.pre_layernorm

        self.dropout_module = torch.nn.Dropout(p=config.dropout, inplace=False)

        self.activation_dropout_module = torch.nn.Dropout(p=config.activation_dropout, inplace=False)

        
        self.activation_fn = ACT2FN[config.activation_fn]
        self.self_attn = GraphormerMultiheadAttention(config)

        
        self.self_attn_layer_norm = nn.LayerNorm(self.embedding_dim)

        self.fc1 = self.build_fc(
            self.embedding_dim,
            config.ffn_embedding_dim,
            q_noise=config.q_noise,
            qn_block_size=config.qn_block_size,
        )
        self.fc2 = self.build_fc(
            config.ffn_embedding_dim,
            self.embedding_dim,
            q_noise=config.q_noise,
            qn_block_size=config.qn_block_size,
        )

        
        self.final_layer_norm = nn.LayerNorm(self.embedding_dim)

    def build_fc(
        self, input_dim: int, output_dim: int, q_noise: float, qn_block_size: int
    ) -> Union[nn.Module, nn.Linear, nn.Embedding, nn.Conv2d]:
        return quant_noise(nn.Linear(input_dim, output_dim), q_noise, qn_block_size)

    def forward(
        self,
        input_nodes: torch.Tensor,
        self_attn_bias: Optional[torch.Tensor] = None,
        self_attn_mask: Optional[torch.Tensor] = None,
        self_attn_padding_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        
        residual = input_nodes
        if self.pre_layernorm:
            input_nodes = self.self_attn_layer_norm(input_nodes)

        input_nodes, attn = self.self_attn(
            query=input_nodes,
            key=input_nodes,
            value=input_nodes,
            attn_bias=self_attn_bias,
            key_padding_mask=self_attn_padding_mask,
            need_weights=False,
            attn_mask=self_attn_mask,
        )
        input_nodes = self.dropout_module(input_nodes)
        input_nodes = residual + input_nodes
        if not self.pre_layernorm:
            input_nodes = self.self_attn_layer_norm(input_nodes)

        residual = input_nodes
        if self.pre_layernorm:
            input_nodes = self.final_layer_norm(input_nodes)
        input_nodes = self.activation_fn(self.fc1(input_nodes))
        input_nodes = self.activation_dropout_module(input_nodes)
        input_nodes = self.fc2(input_nodes)
        input_nodes = self.dropout_module(input_nodes)
        input_nodes = residual + input_nodes
        if not self.pre_layernorm:
            input_nodes = self.final_layer_norm(input_nodes)

        return input_nodes, attn


class GraphormerGraphEncoder(nn.Module):
    def __init__(self, config: GraphormerConfig):
        super().__init__()

        self.dropout_module = torch.nn.Dropout(p=config.dropout, inplace=False)
        self.layerdrop = config.layerdrop
        self.embedding_dim = config.embedding_dim
        self.apply_graphormer_init = config.apply_graphormer_init
        self.traceable = config.traceable

        self.graph_node_feature = GraphormerGraphNodeFeature(config)
        self.graph_attn_bias = GraphormerGraphAttnBias(config)

        self.embed_scale = config.embed_scale

        if config.q_noise > 0:
            self.quant_noise = quant_noise(
                nn.Linear(self.embedding_dim, self.embedding_dim, bias=False),
                config.q_noise,
                config.qn_block_size,
            )
        else:
            self.quant_noise = None

        if config.encoder_normalize_before:
            self.emb_layer_norm = nn.LayerNorm(self.embedding_dim)
        else:
            self.emb_layer_norm = None

        if config.pre_layernorm:
            self.final_layer_norm = nn.LayerNorm(self.embedding_dim)

        if self.layerdrop > 0.0:
            self.layers = LayerDropModuleList(p=self.layerdrop)
        else:
            self.layers = nn.ModuleList([])
        self.layers.extend([GraphormerGraphEncoderLayer(config) for _ in range(config.num_hidden_layers)])

        
        if config.freeze_embeddings:
            raise NotImplementedError("Freezing embeddings is not implemented yet.")

        for layer in range(config.num_trans_layers_to_freeze):
            m = self.layers[layer]
            if m is not None:
                for p in m.parameters():
                    p.requires_grad = False

    def forward(
        self,
        input_nodes: torch.LongTensor,
        input_edges: torch.LongTensor,
        attn_bias: torch.Tensor,
        in_degree: torch.LongTensor,
        out_degree: torch.LongTensor,
        spatial_pos: torch.LongTensor,
        attn_edge_type: torch.LongTensor,
        perturb=None,
        last_state_only: bool = False,
        token_embeddings: Optional[torch.Tensor] = None,
        attn_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[Union[torch.Tensor, List[torch.LongTensor]], torch.Tensor]:
        
        data_x = input_nodes
        n_graph, n_node = data_x.size()[:2]
        padding_mask = (data_x[:, :, 0]).eq(0)
        padding_mask_cls = torch.zeros(n_graph, 1, device=padding_mask.device, dtype=padding_mask.dtype)
        padding_mask = torch.cat((padding_mask_cls, padding_mask), dim=1)

        attn_bias = self.graph_attn_bias(input_nodes, attn_bias, spatial_pos, input_edges, attn_edge_type)

        if token_embeddings is not None:
            input_nodes = token_embeddings
        else:
            input_nodes = self.graph_node_feature(input_nodes, in_degree, out_degree)

        if perturb is not None:
            input_nodes[:, 1:, :] += perturb

        if self.embed_scale is not None:
            input_nodes = input_nodes * self.embed_scale

        if self.quant_noise is not None:
            input_nodes = self.quant_noise(input_nodes)

        if self.emb_layer_norm is not None:
            input_nodes = self.emb_layer_norm(input_nodes)

        input_nodes = self.dropout_module(input_nodes)

        input_nodes = input_nodes.transpose(0, 1)

        inner_states = []
        if not last_state_only:
            inner_states.append(input_nodes)

        for layer in self.layers:
            input_nodes, _ = layer(
                input_nodes,
                self_attn_padding_mask=padding_mask,
                self_attn_mask=attn_mask,
                self_attn_bias=attn_bias,
            )
            if not last_state_only:
                inner_states.append(input_nodes)

        graph_rep = input_nodes[0, :, :]

        if last_state_only:
            inner_states = [input_nodes]

        if self.traceable:
            return torch.stack(inner_states), graph_rep
        else:
            return inner_states, graph_rep


class GraphormerDecoderHead(nn.Module):
    def __init__(self, embedding_dim: int, num_classes: int):
        super().__init__()
        
        self.lm_output_learned_bias = nn.Parameter(torch.zeros(1))
        self.classifier = nn.Linear(embedding_dim, num_classes, bias=False)
        self.num_classes = num_classes

    def forward(self, input_nodes: torch.Tensor, **unused) -> torch.Tensor:
        input_nodes = self.classifier(input_nodes)
        input_nodes = input_nodes + self.lm_output_learned_bias
        return input_nodes


class GraphormerPreTrainedModel(PreTrainedModel):
    

    config_class = GraphormerConfig
    base_model_prefix = "graphormer"
    main_input_name_nodes = "input_nodes"
    main_input_name_edges = "input_edges"

    def normal_(self, data: torch.Tensor):
        
        
        data.copy_(data.cpu().normal_(mean=0.0, std=0.02).to(data.device))

    def init_graphormer_params(self, module: Union[nn.Linear, nn.Embedding, GraphormerMultiheadAttention]):
        
        if isinstance(module, nn.Linear):
            self.normal_(module.weight.data)
            if module.bias is not None:
                module.bias.data.zero_()
        if isinstance(module, nn.Embedding):
            self.normal_(module.weight.data)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        if isinstance(module, GraphormerMultiheadAttention):
            self.normal_(module.q_proj.weight.data)
            self.normal_(module.k_proj.weight.data)
            self.normal_(module.v_proj.weight.data)

    def _init_weights(
        self,
        module: Union[
            nn.Linear, nn.Conv2d, nn.Embedding, nn.LayerNorm, GraphormerMultiheadAttention, GraphormerGraphEncoder
        ],
    ):
        
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, GraphormerMultiheadAttention):
            module.q_proj.weight.data.normal_(mean=0.0, std=0.02)
            module.k_proj.weight.data.normal_(mean=0.0, std=0.02)
            module.v_proj.weight.data.normal_(mean=0.0, std=0.02)
            module.reset_parameters()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        elif isinstance(module, GraphormerGraphEncoder):
            if module.apply_graphormer_init:
                module.apply(self.init_graphormer_params)

        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


class GraphormerModel(GraphormerPreTrainedModel):
    

    def __init__(self, config: GraphormerConfig):
        super().__init__(config)
        self.max_nodes = config.max_nodes

        self.graph_encoder = GraphormerGraphEncoder(config)

        self.share_input_output_embed = config.share_input_output_embed
        self.lm_output_learned_bias = None

        
        self.load_softmax = not getattr(config, "remove_head", False)

        self.lm_head_transform_weight = nn.Linear(config.embedding_dim, config.embedding_dim)
        self.activation_fn = ACT2FN[config.activation_fn]
        self.layer_norm = nn.LayerNorm(config.embedding_dim)

        self.post_init()

    def reset_output_layer_parameters(self):
        self.lm_output_learned_bias = nn.Parameter(torch.zeros(1))

    def forward(
        self,
        input_nodes: torch.LongTensor,
        input_edges: torch.LongTensor,
        attn_bias: torch.Tensor,
        in_degree: torch.LongTensor,
        out_degree: torch.LongTensor,
        spatial_pos: torch.LongTensor,
        attn_edge_type: torch.LongTensor,
        perturb: Optional[torch.FloatTensor] = None,
        masked_tokens: None = None,
        return_dict: Optional[bool] = None,
        **unused,
    ) -> Union[Tuple[torch.LongTensor], BaseModelOutputWithNoAttention]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        inner_states, graph_rep = self.graph_encoder(
            input_nodes, input_edges, attn_bias, in_degree, out_degree, spatial_pos, attn_edge_type, perturb=perturb
        )

        
        input_nodes = inner_states[-1].transpose(0, 1)

        
        if masked_tokens is not None:
            raise NotImplementedError

        input_nodes = self.layer_norm(self.activation_fn(self.lm_head_transform_weight(input_nodes)))

        
        if self.share_input_output_embed and hasattr(self.graph_encoder.embed_tokens, "weight"):
            input_nodes = torch.nn.functional.linear(input_nodes, self.graph_encoder.embed_tokens.weight)

        if not return_dict:
            return tuple(x for x in [input_nodes, inner_states] if x is not None)
        return BaseModelOutputWithNoAttention(last_hidden_state=input_nodes, hidden_states=inner_states)

    def max_nodes(self):
        
        return self.max_nodes


class GraphormerForGraphClassification(GraphormerPreTrainedModel):
    

    def __init__(self, config: GraphormerConfig):
        super().__init__(config)
        self.encoder = GraphormerModel(config)
        self.embedding_dim = config.embedding_dim
        self.num_classes = config.num_classes
        self.classifier = GraphormerDecoderHead(self.embedding_dim, self.num_classes)
        self.is_encoder_decoder = True

        
        self.post_init()

    def forward(
        self,
        input_nodes: torch.LongTensor,
        input_edges: torch.LongTensor,
        attn_bias: torch.Tensor,
        in_degree: torch.LongTensor,
        out_degree: torch.LongTensor,
        spatial_pos: torch.LongTensor,
        attn_edge_type: torch.LongTensor,
        labels: Optional[torch.LongTensor] = None,
        return_dict: Optional[bool] = None,
        **unused,
    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        encoder_outputs = self.encoder(
            input_nodes,
            input_edges,
            attn_bias,
            in_degree,
            out_degree,
            spatial_pos,
            attn_edge_type,
            return_dict=True,
        )
        outputs, hidden_states = encoder_outputs["last_hidden_state"], encoder_outputs["hidden_states"]

        head_outputs = self.classifier(outputs)
        logits = head_outputs[:, 0, :].contiguous()

        loss = None
        if labels is not None:
            mask = ~torch.isnan(labels)

            if self.num_classes == 1:  
                loss_fct = MSELoss()
                loss = loss_fct(logits[mask].squeeze(), labels[mask].squeeze().float())
            elif self.num_classes > 1 and len(labels.shape) == 1:  
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits[mask].view(-1, self.num_classes), labels[mask].view(-1))
            else:  
                loss_fct = BCEWithLogitsLoss(reduction="sum")
                loss = loss_fct(logits[mask], labels[mask])

        if not return_dict:
            return tuple(x for x in [loss, logits, hidden_states] if x is not None)
        return SequenceClassifierOutput(loss=loss, logits=logits, hidden_states=hidden_states, attentions=None)