{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":75254,"databundleVersionId":8264048,"sourceType":"competition"}],"dockerImageVersionId":30698,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nif torch.cuda.is_available():\n# GPU enabled\n    !pip install  dgl==1.1.3 -f https://data.dgl.ai/wheels/cu121/repo.html\n    !pip install  dglgo -f https://data.dgl.ai/wheels-test/repo.html\nelse:\n# just use CPU\n    !pip install  dgl -f https://data.dgl.ai/wheels/repo.html\n    !pip install  dglgo -f https://data.dgl.ai/wheels-test/repo.html","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-05-06T12:30:25.459991Z","iopub.execute_input":"2024-05-06T12:30:25.460871Z","iopub.status.idle":"2024-05-06T12:31:18.844308Z","shell.execute_reply.started":"2024-05-06T12:30:25.460834Z","shell.execute_reply":"2024-05-06T12:31:18.843146Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import dgl.nn as dglnn\nimport dgl.function as fn\nfrom dgl.nn.pytorch import MetaPath2Vec\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\nimport tqdm","metadata":{"execution":{"iopub.status.busy":"2024-05-06T12:31:23.205817Z","iopub.execute_input":"2024-05-06T12:31:23.206801Z","iopub.status.idle":"2024-05-06T12:31:23.977874Z","shell.execute_reply.started":"2024-05-06T12:31:23.206757Z","shell.execute_reply":"2024-05-06T12:31:23.977089Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport pickle as pkl\nimport numpy as np\n\ndef read_txt(file):\n    res_list = list()\n    with open(file, \"r\") as f:\n        line_list = f.readlines()\n    for line in line_list:\n        res_list.append(list(map(int, line.strip().split(' '))))      \n    return res_list\n\nbase_path = \"/kaggle/input/cs3319-01-project-2-recommendation-2024-spring/\"\n\ncite_file = \"paper_file_ann.txt\"\ntrain_ref_file = \"bipartite_train_ann.txt\"\ntest_ref_file = \"bipartite_test_ann.txt\"\ncoauthor_file = \"author_file_ann.txt\"\nfeature_file = \"feature.pkl\"\n\ncitation = read_txt(os.path.join(base_path, cite_file))\nexisting_refs = read_txt(os.path.join(base_path, train_ref_file))\nrefs_to_pred = read_txt(os.path.join(base_path, test_ref_file))\ncoauthor = read_txt(os.path.join(base_path, coauthor_file))\n\nfeature_file = os.path.join(base_path, feature_file)\nwith open(feature_file, 'rb') as f:\n      paper_feature = pkl.load(f)\n        \n\n\nprint(\"Number of citation edges: {}\\n\\\nNumber of existing references: {}\\n\\\nNumber of author-paper pairs to predict: {}\\n\\\nNumber of coauthor edges: {}\\n\\\nShape of paper features: {}\"\n.format(len(citation), len(existing_refs), len(refs_to_pred), len(coauthor), paper_feature.shape))","metadata":{"execution":{"iopub.status.busy":"2024-05-06T12:31:26.641552Z","iopub.execute_input":"2024-05-06T12:31:26.641928Z","iopub.status.idle":"2024-05-06T12:31:34.904360Z","shell.execute_reply.started":"2024-05-06T12:31:26.641896Z","shell.execute_reply":"2024-05-06T12:31:34.903419Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\n\ncite_edges = pd.DataFrame(citation, columns=['source', 'target'])\ncite_edges = cite_edges.set_index(\n    \"c-\" + cite_edges.index.astype(str)\n)\n\nref_edges = pd.DataFrame(existing_refs, columns=['source', 'target'])\nref_edges = ref_edges.set_index(\n    \"r-\" + ref_edges.index.astype(str)\n)\n\ncoauthor_edges = pd.DataFrame(coauthor, columns=['source', 'target'])\ncoauthor_edges = coauthor_edges.set_index(\n    \"a-\" + coauthor_edges.index.astype(str)\n)\n\ncite_edges.head()","metadata":{"execution":{"iopub.status.busy":"2024-05-06T12:31:38.466094Z","iopub.execute_input":"2024-05-06T12:31:38.467006Z","iopub.status.idle":"2024-05-06T12:31:40.523008Z","shell.execute_reply.started":"2024-05-06T12:31:38.466972Z","shell.execute_reply":"2024-05-06T12:31:40.521965Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"node_tmp = pd.concat([cite_edges.loc[:, 'source'], cite_edges.loc[:, 'target'], ref_edges.loc[:, 'target']])\nnode_papers = pd.DataFrame(index=pd.unique(node_tmp))\n\nnode_tmp = pd.concat([ref_edges['source'], coauthor_edges['source'], coauthor_edges['target']])\nnode_authors = pd.DataFrame(index=pd.unique(node_tmp))\n\nprint(\"Number of paper nodes: {}, number of author nodes: {}\".format(len(node_papers), len(node_authors)))","metadata":{"execution":{"iopub.status.busy":"2024-05-06T12:31:43.513652Z","iopub.execute_input":"2024-05-06T12:31:43.514330Z","iopub.status.idle":"2024-05-06T12:31:43.579335Z","shell.execute_reply.started":"2024-05-06T12:31:43.514289Z","shell.execute_reply":"2024-05-06T12:31:43.578377Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_refs = ref_edges.sample(frac=0.9,random_state=0,axis=0)\ntest_true_refs = ref_edges[~ref_edges.index.isin(train_refs.index)]\ntest_true_refs.loc[:, 'label'] = 1\n\nfalse_source = node_authors.sample(frac=test_true_refs.shape[0]/node_authors.shape[0],random_state=0,replace=True,axis=0)\nfalse_target = node_papers.sample(frac=test_true_refs.shape[0]/node_papers.shape[0],random_state=0,replace=True,axis=0)\nfalse_source = false_source.reset_index()\nfalse_target = false_target.reset_index()\ntest_false_refs = pd.concat([false_source, false_target], axis=1)\ntest_false_refs.columns = ['source', 'target']\ntest_false_refs = test_false_refs[test_false_refs.isin(ref_edges) == False]\ntest_false_refs.loc[:, 'label'] = 0\n\ntest_refs = pd.concat([test_true_refs, test_false_refs.iloc[:min(len(false_source), len(false_target))]])\ntest_refs = test_refs.sample(frac=1,random_state=0,axis=0)\ntest_refs.head()","metadata":{"execution":{"iopub.status.busy":"2024-05-06T12:31:49.518017Z","iopub.execute_input":"2024-05-06T12:31:49.518409Z","iopub.status.idle":"2024-05-06T12:31:50.121489Z","shell.execute_reply.started":"2024-05-06T12:31:49.518373Z","shell.execute_reply":"2024-05-06T12:31:50.120232Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import dgl\n\nos.environ['DGLBACKEND'] = 'pytorch'\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\ntrain_ref_tensor = torch.from_numpy(train_refs.values)\ncite_tensor = torch.from_numpy(cite_edges.values)\ncoauthor_tensor = torch.from_numpy(coauthor_edges.values)\nrel_list = [('author', 'ref', 'paper'), ('paper', 'cite', 'paper'), ('author', 'coauthor', 'author'), ('paper', 'beref', 'author')]\ngraph_data = {\n   rel_list[0]: (train_ref_tensor[:, 0], train_ref_tensor[:, 1]),\n   rel_list[1]: (torch.cat([cite_tensor[:, 0], cite_tensor[:, 1]]), torch.cat([cite_tensor[:, 1], cite_tensor[:, 0]])),\n   rel_list[2]: (torch.cat([coauthor_tensor[:, 0], coauthor_tensor[:, 1]]), torch.cat([coauthor_tensor[:, 1], coauthor_tensor[:, 0]])),\n   rel_list[3]: (train_ref_tensor[:, 1], train_ref_tensor[:, 0])\n}\nhetero_graph = dgl.heterograph(graph_data)\nnode_features = {'author': torch.rand(6611, 512), 'paper': paper_feature}\nhetero_graph.ndata['features'] = node_features\nhetero_graph = hetero_graph.to(device)\nhetero_graph","metadata":{"execution":{"iopub.status.busy":"2024-05-06T12:31:52.935304Z","iopub.execute_input":"2024-05-06T12:31:52.935714Z","iopub.status.idle":"2024-05-06T12:31:53.143174Z","shell.execute_reply.started":"2024-05-06T12:31:52.935682Z","shell.execute_reply":"2024-05-06T12:31:53.142056Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = MetaPath2Vec(hetero_graph, [rel_list[1][1], rel_list[3][1], rel_list[0][1]], window_size=1).to(device)\ndataloader = DataLoader(torch.arange(hetero_graph.num_nodes('paper')), batch_size=8192,\n                        shuffle=True, collate_fn=model.sample)\nopt = torch.optim.SparseAdam(model.parameters(), lr=0.025)\n\nfor epoch in range(100):\n    with tqdm.tqdm(dataloader) as tq:\n        for (pos_u, pos_v, neg_v) in tq:\n            pos_u = pos_u.to(device)\n            pos_v = pos_v.to(device)\n            neg_v = neg_v.to(device)\n\n            loss = model(pos_u, pos_v, neg_v)\n\n            opt.zero_grad()\n            loss.backward()\n            opt.step()\n            tq.set_postfix({'loss': '%.03f' % loss.item()}, refresh=False)","metadata":{"execution":{"iopub.status.busy":"2024-04-28T07:53:31.206775Z","iopub.execute_input":"2024-04-28T07:53:31.207154Z","iopub.status.idle":"2024-04-28T07:57:27.566623Z","shell.execute_reply.started":"2024-04-28T07:53:31.207117Z","shell.execute_reply":"2024-04-28T07:57:27.565610Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with torch.no_grad():\n    user_nids = torch.LongTensor(model.local_to_global_nid['author']).to(device)\n    item_nids = torch.LongTensor(model.local_to_global_nid['paper']).to(device)\n    node_embeddings = {}\n    node_embeddings['author'] = model.node_embed(user_nids).detach().cpu()\n    node_embeddings['paper'] = model.node_embed(item_nids).detach().cpu()\n    print(node_embeddings['author'].shape, node_embeddings['paper'].shape)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class StochasticTwoLayerRGCN(nn.Module):\n    def __init__(self, in_feat, hidden_feat, out_feat, rel_names):\n        super().__init__()\n        self.conv1 = dglnn.HeteroGraphConv({\n                rel : dglnn.GraphConv(in_feat, hidden_feat, norm='right')\n                for rel in rel_names\n            })\n        self.conv2 = dglnn.HeteroGraphConv({\n                rel : dglnn.GraphConv(hidden_feat, out_feat, norm='right')\n                for rel in rel_names\n            })\n\n    def forward(self, blocks, x):\n        x = self.conv1(blocks[0], x)\n        x = self.conv2(blocks[1], x)\n        return x\n\nclass ScorePredictor(nn.Module):\n    def forward(self, edge_subgraph, h):\n        with edge_subgraph.local_scope():\n            edge_subgraph.ndata['h'] = h\n            for etype in edge_subgraph.canonical_etypes:\n                edge_subgraph.apply_edges(fn.u_dot_v('h', 'h', 'score'), etype=etype)\n            return edge_subgraph.edata['score']\n\nclass Model(nn.Module):\n    def __init__(self, in_features, hidden_features, out_features, etypes):\n        super().__init__()\n        self.rgcn = StochasticTwoLayerRGCN(\n            in_features, hidden_features, out_features, etypes)\n        self.pred = ScorePredictor()\n\n    def forward(self, positive_graph, negative_graph, blocks, x):\n        x = self.rgcn(blocks, x)\n        pos_score = self.pred(positive_graph, x)\n        neg_score = self.pred(negative_graph, x)\n        return pos_score, neg_score\n\ndef compute_loss(pos_score, neg_score, etype):\n    n_edges = pos_score[etype].shape[0]\n    return (1 - pos_score[etype].unsqueeze(1) + neg_score[etype].view(n_edges, -1)).clamp(min=0).mean()\n\nsampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)\ntrain_eid_dict = {\n    etype: hetero_graph.edges(etype=etype, form='eid')\n    for etype in hetero_graph.etypes}\n# train_eid_dict['device'] = device\n\nsampler = dgl.dataloading.NeighborSampler([4, 4])\nsampler = dgl.dataloading.as_edge_prediction_sampler(\n    sampler, negative_sampler=dgl.dataloading.negative_sampler.Uniform(5))\ndataloader = dgl.dataloading.DataLoader(\n    # The following arguments are specific to DataLoader.\n    hetero_graph,                                  # The graph\n    train_eid_dict,  # The edges to iterate over\n    sampler,                                # The neighbor sampler\n    device=device,                          # Put the MFGs on CPU or GPU\n    # The following arguments are inherited from PyTorch DataLoader.\n    batch_size=32768,    # Batch size\n    shuffle=True,       # Whether to shuffle the nodes for every epoch\n    drop_last=False,    # Whether to drop the last incomplete batch\n    num_workers=0,       # Number of sampler processes\n)\n\nmodel = Model(512, 64, 10, rel_list).to(device)\nopt = torch.optim.Adam(model.parameters())\nfor epoch in range(10):\n    with tqdm.tqdm(dataloader) as tq:\n        for step, (input_nodes, positive_graph, negative_graph, blocks) in enumerate(tq):\n            input_features = blocks[0].srcdata['features']\n            pos_score, neg_score = model(positive_graph, negative_graph, blocks, input_features)\n            score = torch.cat([pos_score[rel_list[0]], neg_score[rel_list[0]]])\n            label = torch.cat([torch.ones_like(pos_score[rel_list[0]]), torch.zeros_like(neg_score[rel_list[0]])])\n            loss = F.binary_cross_entropy_with_logits(score, label)\n            loss = compute_loss(pos_score, neg_score, rel_list[0])\n            opt.zero_grad()\n            loss.backward()\n            opt.step()\n            tq.set_postfix({'loss': '%.03f' % loss.item()}, refresh=False)","metadata":{"execution":{"iopub.status.busy":"2024-05-06T12:32:06.532306Z","iopub.execute_input":"2024-05-06T12:32:06.532637Z","iopub.status.idle":"2024-05-06T12:34:17.909505Z","shell.execute_reply.started":"2024-05-06T12:32:06.532613Z","shell.execute_reply":"2024-05-06T12:34:17.908442Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with torch.no_grad():\n    node_features_gpu = {k: v.to(device) for k, v in node_features.items()}\n    node_embeddings = model.rgcn([hetero_graph.to(device), hetero_graph.to(device)], node_features_gpu)\n    node_embeddings = {k: v.detach().cpu() for k, v in node_embeddings.items()}\n    print(node_embeddings['author'].shape, node_embeddings['paper'].shape)","metadata":{"execution":{"iopub.status.busy":"2024-05-06T12:36:07.124476Z","iopub.execute_input":"2024-05-06T12:36:07.125492Z","iopub.status.idle":"2024-05-06T12:36:07.182966Z","shell.execute_reply.started":"2024-05-06T12:36:07.125459Z","shell.execute_reply":"2024-05-06T12:36:07.182024Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from numpy import dot\nfrom numpy.linalg import norm\n\n# Use Cosine Similarity to measure the probability of recommandation\ndef cos_sim(a, b):\n    cos_sim = np.sum(a * b, axis = 1) / (norm(a, axis = 1) * norm(b, axis = 1))\n    return cos_sim\n\ntest_arr = np.array(test_refs.values)\nres = cos_sim(np.array(node_embeddings['author'][test_arr[:, 0]]), np.array(node_embeddings['paper'][test_arr[:, 1]]))\n\n# Generate predict labels\nlbl_true = test_refs.label.to_numpy()\nlbl_true = lbl_true.flatten()\n\nimport matplotlib.pyplot as plt\nfrom sklearn.metrics import (precision_recall_curve, PrecisionRecallDisplay)\n\n# ROC curve\nprecision, recall, _ = precision_recall_curve(lbl_true, np.array(res))\ndisp = PrecisionRecallDisplay(precision=precision, recall=recall)\ndisp.plot()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-05-06T12:36:10.344493Z","iopub.execute_input":"2024-05-06T12:36:10.344858Z","iopub.status.idle":"2024-05-06T12:36:11.095831Z","shell.execute_reply.started":"2024-05-06T12:36:10.344830Z","shell.execute_reply":"2024-05-06T12:36:11.094722Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.metrics import f1_score\n\nlbl_pred = np.array(res)\nthreshold = 0.5\nlbl_pred[lbl_pred >= threshold] = 1\nlbl_pred[lbl_pred < threshold] = 0\n# Test scores\nprint(\"F1-Score: {:.3f}\".format(f1_score(lbl_true, lbl_pred)))","metadata":{"execution":{"iopub.status.busy":"2024-05-06T12:36:18.124810Z","iopub.execute_input":"2024-05-06T12:36:18.125479Z","iopub.status.idle":"2024-05-06T12:36:18.152757Z","shell.execute_reply.started":"2024-05-06T12:36:18.125445Z","shell.execute_reply":"2024-05-06T12:36:18.151872Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Output your prediction\ntest_arr = np.array(refs_to_pred)\nres = cos_sim(np.array(node_embeddings['author'][test_arr[:, 0]]), np.array(node_embeddings['paper'][test_arr[:, 1]]))\nres[res >= threshold] = 1\nres[res < threshold] = 0\ndata = []\nfor index, p in enumerate(list(res)):\n    tp = [index, str(int(p))]\n    data.append(tp)\n\ndf = pd.DataFrame(data, columns=[\"Index\", \"Predicted\"], dtype=object)\ndf.to_csv('/kaggle/working/SampleSubmission.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2024-05-06T12:36:21.760186Z","iopub.execute_input":"2024-05-06T12:36:21.760862Z","iopub.status.idle":"2024-05-06T12:36:29.871991Z","shell.execute_reply.started":"2024-05-06T12:36:21.760824Z","shell.execute_reply":"2024-05-06T12:36:29.871009Z"},"trusted":true},"execution_count":null,"outputs":[]}]}