text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.autograd import Variable
  import time
  from layers import SpGraphAttentionLayer, ConvKB
  from config import Config
  
  args = Config()
  args.load_config()
  device = "cuda" if args.cuda else "cpu"
  
  class SpGAT(nn.Module):
      def __init__(self, num_nodes, nfeat, nhid, relation_dim, dropout, alpha, nheads):
          """
              Sparse version of GAT
              nfeat -> Entity Input Embedding dimensions
              nhid  -> Entity Output Embedding dimensions
              relation_dim -> Relation Embedding dimensions
              num_nodes -> number of nodes in the Graph
              nheads -> Used for Multihead attention
  
          """
          super(SpGAT, self).__init__()
          self.dropout = dropout
          self.dropout_layer = nn.Dropout(self.dropout)
          self.attentions = [SpGraphAttentionLayer(num_nodes, nfeat,
                                                   nhid,
                                                   relation_dim,
                                                   dropout=dropout,
                                                   alpha=alpha,
                                                   concat=True)
                             for _ in range(nheads)]
  
          for i, attention in enumerate(self.attentions):
              self.add_module('attention_{}'.format(i), attention)
  
          # W matrix to convert h_input to h_output dimension
          self.W = nn.Parameter(torch.zeros(size=(relation_dim, nheads * nhid)))
          nn.init.xavier_uniform_(self.W.data, gain=1.414)
  
          self.out_att = SpGraphAttentionLayer(num_nodes, nhid * nheads,
                                               nheads * nhid, nheads * nhid,
                                               dropout=dropout,
                                               alpha=alpha,
                                               concat=False
                                               )
  
      def forward(self, Corpus_, batch_inputs, entity_embeddings, relation_embed,
                  edge_list, edge_type, edge_embed, edge_list_nhop, edge_type_nhop):
          x = entity_embeddings
  
          edge_embed_nhop = relation_embed[
              edge_type_nhop[:, 0]] + relation_embed[edge_type_nhop[:, 1]]
  
          x = torch.cat([att(x, edge_list, edge_embed, edge_list_nhop, edge_embed_nhop)
                         for att in self.attentions], dim=1)
          x = self.dropout_layer(x)
  
          out_relation_1 = relation_embed.mm(self.W)
  
          edge_embed = out_relation_1[edge_type]
          edge_embed_nhop = out_relation_1[
              edge_type_nhop[:, 0]] + out_relation_1[edge_type_nhop[:, 1]]
  
          x = F.elu(self.out_att(x, edge_list, edge_embed,
                                 edge_list_nhop, edge_embed_nhop))
          return x, out_relation_1
  
  
  class SpKBGATModified(nn.Module):
      def __init__(self, initial_entity_emb, initial_relation_emb, entity_out_dim, relation_out_dim,
                   drop_GAT, alpha, nheads_GAT):
          '''Sparse version of KBGAT
          entity_in_dim -> Entity Input Embedding dimensions
          entity_out_dim  -> Entity Output Embedding dimensions, passed as a list
          num_relation -> number of unique relations
          relation_dim -> Relation Embedding dimensions
          num_nodes -> number of nodes in the Graph
          nheads_GAT -> Used for Multihead attention, passed as a list '''
  
          super().__init__()
  
          self.num_nodes = initial_entity_emb.shape[0]
          self.entity_in_dim = initial_entity_emb.shape[1]
          self.entity_out_dim_1 = entity_out_dim[0]
          self.nheads_GAT_1 = nheads_GAT[0]
          self.entity_out_dim_2 = entity_out_dim[1]
          self.nheads_GAT_2 = nheads_GAT[1]
  
          # Properties of Relations
          self.num_relation = initial_relation_emb.shape[0]
          self.relation_dim = initial_relation_emb.shape[1]
          self.relation_out_dim_1 = relation_out_dim[0]
  
          self.drop_GAT = drop_GAT
          self.alpha = alpha      # For leaky relu
  
          self.final_entity_embeddings = nn.Parameter(
              torch.randn(self.num_nodes, self.entity_out_dim_1 * self.nheads_GAT_1))
  
          self.final_relation_embeddings = nn.Parameter(
              torch.randn(self.num_relation, self.entity_out_dim_1 * self.nheads_GAT_1))
  
          self.entity_embeddings = nn.Parameter(initial_entity_emb)
          self.relation_embeddings = nn.Parameter(initial_relation_emb)
  
          self.sparse_gat_1 = SpGAT(self.num_nodes, self.entity_in_dim, self.entity_out_dim_1, self.relation_dim,
                                    self.drop_GAT, self.alpha, self.nheads_GAT_1)
  
          self.W_entities = nn.Parameter(torch.zeros(
              size=(self.entity_in_dim, self.entity_out_dim_1 * self.nheads_GAT_1)))
          nn.init.xavier_uniform_(self.W_entities.data, gain=1.414)
  
      def forward(self, Corpus_, adj, batch_inputs, train_indices_nhop):
          # getting edge list
          edge_list = adj[0].to(device)
          edge_type = adj[1].to(device)
  
          edge_list_nhop = torch.cat(
              (train_indices_nhop[:, 3].unsqueeze(-1), train_indices_nhop[:, 0].unsqueeze(-1)), dim=1).t()
          edge_type_nhop = torch.cat(
              [train_indices_nhop[:, 1].unsqueeze(-1), train_indices_nhop[:, 2].unsqueeze(-1)], dim=1)
  
          if args.cuda:
              edge_list_nhop = edge_list_nhop.to(device)
              edge_type_nhop = edge_type_nhop.to(device)
  
          edge_embed = self.relation_embeddings[edge_type]
  
          start = time.time()
  
          self.entity_embeddings.data = F.normalize(
              self.entity_embeddings.data, p=2, dim=1).detach()
  
          # self.relation_embeddings.data = F.normalize(
          #     self.relation_embeddings.data, p=2, dim=1)
  
          out_entity_1, out_relation_1 = self.sparse_gat_1(
              Corpus_, batch_inputs, self.entity_embeddings, self.relation_embeddings,
              edge_list, edge_type, edge_embed, edge_list_nhop, edge_type_nhop)
  
          mask_indices = torch.unique(batch_inputs[:, 2])
          mask = torch.zeros(self.entity_embeddings.shape[0])
          mask[mask_indices] = 1.0
          entities_upgraded = self.entity_embeddings.mm(self.W_entities).to(device)
  		
          if args.cuda:
              mask_indices = mask_indices.to(device)
              mask = mask.to(device)
              out_entity_1 = out_entity_1.to(device)
  
          out_entity_1 = entities_upgraded + mask.unsqueeze(-1).expand_as(out_entity_1) * out_entity_1
          out_entity_1 = F.normalize(out_entity_1, p=2, dim=1)
  
          self.final_entity_embeddings.data = out_entity_1.data
          self.final_relation_embeddings.data = out_relation_1.data
  
          return out_entity_1, out_relation_1
  
  
  class SpKBGATConvOnly(nn.Module):
      def __init__(self, initial_entity_emb, initial_relation_emb, entity_out_dim, relation_out_dim,
                   drop_GAT, drop_conv, alpha, alpha_conv, nheads_GAT, conv_out_channels):
          '''Sparse version of KBGAT
          entity_in_dim -> Entity Input Embedding dimensions
          entity_out_dim  -> Entity Output Embedding dimensions, passed as a list
          num_relation -> number of unique relations
          relation_dim -> Relation Embedding dimensions
          num_nodes -> number of nodes in the Graph
          nheads_GAT -> Used for Multihead attention, passed as a list '''
  
          super().__init__()
  
          self.num_nodes = initial_entity_emb.shape[0]
          self.entity_in_dim = initial_entity_emb.shape[1]
          self.entity_out_dim_1 = entity_out_dim[0]
          self.nheads_GAT_1 = nheads_GAT[0]
          self.entity_out_dim_2 = entity_out_dim[1]
          self.nheads_GAT_2 = nheads_GAT[1]
  
          # Properties of Relations
          self.num_relation = initial_relation_emb.shape[0]
          self.relation_dim = initial_relation_emb.shape[1]
          self.relation_out_dim_1 = relation_out_dim[0]
  
          self.drop_GAT = drop_GAT
          self.drop_conv = drop_conv
          self.alpha = alpha      # For leaky relu
          self.alpha_conv = alpha_conv
          self.conv_out_channels = conv_out_channels
  
          self.final_entity_embeddings = nn.Parameter(
              torch.randn(self.num_nodes, self.entity_out_dim_1 * self.nheads_GAT_1))
  
          self.final_relation_embeddings = nn.Parameter(
              torch.randn(self.num_relation, self.entity_out_dim_1 * self.nheads_GAT_1))
  
          self.convKB = ConvKB(self.entity_out_dim_1 * self.nheads_GAT_1, 3, 1,
                               self.conv_out_channels, self.drop_conv, self.alpha_conv)
  
      def forward(self, Corpus_, adj, batch_inputs):
          conv_input = torch.cat((self.final_entity_embeddings[batch_inputs[:, 0], :].unsqueeze(1), self.final_relation_embeddings[
              batch_inputs[:, 1]].unsqueeze(1), self.final_entity_embeddings[batch_inputs[:, 2], :].unsqueeze(1)), dim=1)
          out_conv = self.convKB(conv_input)
          return out_conv
  
      def batch_test(self, batch_inputs):
          conv_input = torch.cat((self.final_entity_embeddings[batch_inputs[:, 0], :].unsqueeze(1), self.final_relation_embeddings[
              batch_inputs[:, 1]].unsqueeze(1), self.final_entity_embeddings[batch_inputs[:, 2], :].unsqueeze(1)), dim=1)
          out_conv = self.convKB(conv_input)
          return out_conv
  