<jupyter_start><jupyter_text><jupyter_code>!pip install dgl-cu110 
import dgl
import torch
import torch.nn as nn
import torch.nn.functional as F
import itertools
import numpy as np
import scipy.sparse as sp
%matplotlib inline
!pip install scikit-plot
import scikitplot as skplt
import dgl.data

dataset = dgl.data.CoraGraphDataset()
g = dataset[0]
g = dgl.add_self_loop(g)
# Split edge set for training and testing
print(g.edges())
u, v = g.edges()
print(len(u))
print(len(v))
eids = np.arange(g.number_of_edges())
print(eids)
eids = np.random.permutation(eids)
print(len(eids))
test_size = int(len(eids) * 0.1)
train_size = g.number_of_edges() - test_size
test_pos_u, test_pos_v = u[eids[:test_size]], v[eids[:test_size]]
train_pos_u, train_pos_v = u[eids[test_size:]], v[eids[test_size:]]

# Find all negative edges and split them for training and testing
adj = sp.coo_matrix((np.ones(len(u)), (u.numpy(), v.numpy())))
print(adj.todense())
adj_neg = 1 - adj.todense() - np.eye(g.number_of_nodes())
print(adj_neg)
neg_u, neg_v = np.where(adj_neg != 0)
print(len(neg_u))#, neg_v)
print(g.number_of_edges())
neg_eids = np.random.choice(len(neg_u), g.number_of_edges())
test_neg_u, test_neg_v = neg_u[neg_eids[:test_size]], neg_v[neg_eids[:test_size]]
train_neg_u, train_neg_v = neg_u[neg_eids[test_size:]], neg_v[neg_eids[test_size:]]
train_g = dgl.remove_edges(g, eids[:test_size])
from dgl.nn import SAGEConv

# ----------- 2. create model -------------- #
# build a two-layer GraphSAGE model
class GraphSAGE(nn.Module):
    def __init__(self, in_feats, h_feats):
        super(GraphSAGE, self).__init__()
        self.conv1 = SAGEConv(in_feats, h_feats, 'mean')
        self.conv2 = SAGEConv(h_feats, h_feats, 'mean')

    def forward(self, g, in_feat):
        h = self.conv1(g, in_feat)
        h = F.relu(h)
        h = self.conv2(g, h)
        return h
train_pos_g = dgl.graph((train_pos_u, train_pos_v), num_nodes=g.number_of_nodes())
train_neg_g = dgl.graph((train_neg_u, train_neg_v), num_nodes=g.number_of_nodes())

test_pos_g = dgl.graph((test_pos_u, test_pos_v), num_nodes=g.number_of_nodes())
test_neg_g = dgl.graph((test_neg_u, test_neg_v), num_nodes=g.number_of_nodes())
import dgl.function as fn
import pandas as pd

class DotPredictor(nn.Module):
    def forward(self, g, h):
        with g.local_scope():
            g.ndata['h'] = h
            # Compute a new edge feature named 'score' by a dot-product between the
            # source node feature 'h' and destination node feature 'h'.
            g.apply_edges(fn.u_dot_v('h', 'h', 'score'))
            # u_dot_v returns a 1-element vector for each edge so you need to squeeze it.
            return g.edata['score'][:, 0]
#model = GraphSAGE(train_g.ndata['feat'].shape[1], 16)
# You can replace DotPredictor with MLPPredictor.
#pred = MLPPredictor(16)
from sklearn.metrics import accuracy_score
def train(train_g,train_pos_g,train_neg_g, model):
    pred = DotPredictor()

    train=[]
    test=[]

    def compute_loss(pos_score, neg_score):
        scores = torch.cat([pos_score, neg_score])
        labels = torch.cat([torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])])
        #print(scores)
        #print(labels)
        return F.binary_cross_entropy_with_logits(scores, labels)

    def compute_auc(pos_score, neg_score):
        scores = torch.cat([pos_score, neg_score]).numpy()
        
        labels = torch.cat(
            [torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])]).numpy()
        return roc_auc_score(labels, scores)
    def accuracy(pos_score, neg_score):
        scores = torch.cat([pos_score, neg_score]).detach().numpy()
        scores = (scores>0.5)
        labels = torch.cat(
            [torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])]).numpy()
        print(scores)    
        return accuracy_score(labels, scores)    
    # ----------- 3. set up loss and optimizer -------------- #
    # in this case, loss will in training loop
    optimizer = torch.optim.Adam(itertools.chain(model.parameters(), pred.parameters()), lr=0.01)

    # ----------- 4. training -------------------------------- #
    all_logits = []
    for e in range(100):
        # forward
        h = model(train_g, train_g.ndata['feat'])
        pos_train_score = pred(train_pos_g, h)
        print('pos_train_score',pos_train_score)
        neg_train_score = pred(train_neg_g, h)
        print('neg_train_score',neg_train_score)
        loss = compute_loss(pos_train_score, neg_train_score)
        # train_auc = compute_auc(pos_train_score, neg_train_score)
        # pos_test_score = pred(test_pos_g, h)
        # neg_test_score = pred(test_neg_g, h)
        # test_loss = compute_loss(pos_test_score, neg_test_score)
        # #test_auc = compute_auc(pos_test_score, neg_test_score)
        # train.append(loss.item())
        # test.append(test_loss.item())
        

        # backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if e % 5 == 0:
            print('In epoch {}, loss: {}'.format(e, loss))

    # ----------- 5. check results ------------------------ #
    # df= pd.DataFrame()
    # df['test_loss']= test
    # df['train_loss'] = train
    
    # print(df)
    # ax = df.plot.line(rot=0,ylim=(0,1),figsize=[4,4],title='comparative_accuracy_plot')#,colormap='flare')
    # #ax = bar_data.plot.bar(rot=0)
    # #ax.label_outer('comparative_accuracy_bar_plot')
    # fig=ax.get_figure()
    # fig.savefig(path)
    # ax.set_xlabel("iteration")
    # ax.set_ylabel("loss")
    from sklearn.metrics import roc_auc_score
    with torch.no_grad():
        pos_score = pred(test_pos_g, h)
        neg_score = pred(test_neg_g, h)
        print(neg_score)
        print('AUC', compute_auc(pos_score, neg_score))
        print('Auccuracy', accuracy(pos_score, neg_score))
        #confusion_matrix(pos_score, neg_score)


    
model = GraphSAGE(train_g.ndata['feat'].shape[1], 16)
train(train_g,train_pos_g,train_neg_g, model)
from dgl.nn import GraphConv
g = dgl.add_self_loop(g)
class GCN(nn.Module):
    def __init__(self, in_feats, h_feats):
        super(GCN, self).__init__()
        self.conv1 = GraphConv(in_feats, h_feats, allow_zero_in_degree=True)
        self.conv2 = GraphConv(h_feats, h_feats, allow_zero_in_degree=True)
        
    def forward(self, g, in_feat):
        h = self.conv1(g, in_feat)
        h = F.relu(h)
        h = self.conv2(g, h)
        return h

     
    #     class gin_conv(nn.Module):
    # def __init__(self, in_feats, h_feats):
    #     super(gin_conv, self).__init__()
    #     self.conv1 =  GINConv(in_feats, h_feats, 'max')
    #     self.conv2 = GINConv(h_feats, h_feats, 'max')

    # def forward(self, g, in_feat):
    #     h = self.conv1(g, in_feat)
    #     h = F.relu(h)
    #     h = self.conv2(g, h)
    #     return h

# Create the model with given dimensions
model2 = GCN(train_g.ndata['feat'].shape[1], 16)
train(train_g,train_pos_g,train_neg_g, model2)
from dgl.nn.pytorch.conv import GATConv

class GAT(torch.nn.Module):
    def __init__(self, in_feats,hidden_dim,h_feats, num_heads):
        super(GAT, self).__init__()
        self.layer1 = GATConv(in_feats, hidden_dim, num_heads, allow_zero_in_degree=True)
        self.layer2 = GATConv(hidden_dim * num_heads, h_feats, 1, allow_zero_in_degree=True)
    def forward(self, g, h):
        h = self.layer1(g, h)
        h = h.view(-1, h.size(1) * h.size(2))
        h = F.elu(h)
        h = self.layer2(g, h)
        h = h.squeeze() 
        return h

model3 =GAT(train_g.ndata['feat'].shape[1], 16,30,2)
train(train_g,train_pos_g,train_neg_g, model3)<jupyter_output>In epoch 0, loss: 0.6927238702774048
In epoch 5, loss: 0.6795674562454224
In epoch 10, loss: 0.6202573180198669
In epoch 15, loss: 0.5692955255508423
In epoch 20, loss: 0.5319256782531738
In epoch 25, loss: 0.5131011605262756
In epoch 30, loss: 0.4980435073375702
In epoch 35, loss: 0.478251188993454
In epoch 40, loss: 0.4710206985473633
In epoch 45, loss: 0.4587682783603668
In epoch 50, loss: 0.44706839323043823
In epoch 55, loss: 0.43710261583328247
In epoch 60, loss: 0.42780038714408875
In epoch 65, loss: 0.4187321364879608
In epoch 70, loss: 0.4091523289680481
In epoch 75, loss: 0.3995278477668762
In epoch 80, loss: 0.3898409307003021
In epoch 85, loss: 0.38005444407463074
In epoch 90, loss: 0.3705344498157501
In epoch 95, loss: 0.36109089851379395
tensor([ 0.7568,  0.2325, -2.2273,  ..., -0.0855, -0.4884, -2.2936])
AUC 0.9234130477808946
[ True  True  True ... False False False]
Auccuracy 0.802790346907994
{'fraction_non_alphanumeric': 0.1450372786579683, 'fraction_numerical': 0.060810810810810814, 'mean_word_length': 3.271144278606965, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '26922', 'license': 'Apache-2.0', 'n_tokens_mistral': 3587, 'n_tokens_neox': 3255, 'n_words': 655, 'path': 'link_prediction.ipynb', 'repo_name': 'ananyas168/GNN', 'size': '8584'}