{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceType":"competition","sourceId":141936,"databundleVersionId":17147212}],"dockerImageVersionId":31328,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### It is recommended to use GPU T4$\\times$2 to run the code.","metadata":{}},{"cell_type":"code","source":"import sys\nprint(\"Python:\", sys.version)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-05-09T08:31:10.195835Z","iopub.execute_input":"2026-05-09T08:31:10.196115Z","iopub.status.idle":"2026-05-09T08:31:10.204322Z","shell.execute_reply.started":"2026-05-09T08:31:10.196089Z","shell.execute_reply":"2026-05-09T08:31:10.203336Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install -q torch-geometric","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T08:31:24.287425Z","iopub.execute_input":"2026-05-09T08:31:24.287705Z","iopub.status.idle":"2026-05-09T08:31:30.148123Z","shell.execute_reply.started":"2026-05-09T08:31:24.287682Z","shell.execute_reply":"2026-05-09T08:31:30.147444Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport pickle as pkl\nimport random\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport tqdm\n\nfrom torch_geometric.data import HeteroData\n\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\nprint('torch:', torch.__version__)\nprint('torch cuda:', torch.version.cuda)\nprint('cuda available:', torch.cuda.is_available())\nprint('device:', device)\n\ndef set_seed(seed=0):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\n\nset_seed(0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T08:31:40.296897Z","iopub.execute_input":"2026-05-09T08:31:40.297193Z","iopub.status.idle":"2026-05-09T08:31:56.824501Z","shell.execute_reply.started":"2026-05-09T08:31:40.297165Z","shell.execute_reply":"2026-05-09T08:31:56.823739Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nimport os\nimport pickle as pkl\nimport numpy as np\n\n\ndef read_txt(file):\n    res_list = []\n    with open(file, \"r\") as f:\n        for line in f:\n            res_list.append(list(map(int, line.strip().split())))\n    return res_list\n\nbase_path = \"/kaggle/input/competitions/cs-3319-02-project-2-recommendation-systems-2026-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\nwith open(os.path.join(base_path, feature_file), 'rb') as f:\n    paper_feature = pkl.load(f)\n\nprint(\"base_path:\", base_path)\nprint(f\"Number of citation edges: {len(citation)}\")\nprint(f\"Number of existing references: {len(existing_refs)}\")\nprint(f\"Number of author-paper pairs to predict: {len(refs_to_pred)}\")\nprint(f\"Number of coauthor edges: {len(coauthor)}\")\nprint(f\"Shape of paper features: {paper_feature.shape}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T08:32:00.193224Z","iopub.execute_input":"2026-05-09T08:32:00.19449Z","iopub.status.idle":"2026-05-09T08:32:05.891321Z","shell.execute_reply.started":"2026-05-09T08:32:00.194439Z","shell.execute_reply":"2026-05-09T08:32:05.890584Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T08:32:12.05479Z","iopub.execute_input":"2026-05-09T08:32:12.055745Z","iopub.status.idle":"2026-05-09T08:32:12.720801Z","shell.execute_reply.started":"2026-05-09T08:32:12.055693Z","shell.execute_reply":"2026-05-09T08:32:12.72019Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T08:32:21.193435Z","iopub.execute_input":"2026-05-09T08:32:21.193741Z","iopub.status.idle":"2026-05-09T08:32:21.246163Z","shell.execute_reply.started":"2026-05-09T08:32:21.193711Z","shell.execute_reply":"2026-05-09T08:32:21.245539Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Build train/test edges.\n\ntrain_refs = ref_edges.sample(frac=0.9, random_state=0, axis=0)\ntest_true_refs = ref_edges[~ref_edges.index.isin(train_refs.index)].copy()\ntest_true_refs.loc[:, 'label'] = 1\n\nexisting_ref_set = set(map(tuple, ref_edges[['source', 'target']].to_numpy().tolist()))\nnum_test_pos = len(test_true_refs)\nauthor_ids = node_authors.index.to_numpy(dtype=np.int64)\npaper_ids = node_papers.index.to_numpy(dtype=np.int64)\n\nneg_pairs = []\nrng = np.random.default_rng(0)\nwhile len(neg_pairs) < num_test_pos:\n    src = int(rng.choice(author_ids))\n    dst = int(rng.choice(paper_ids))\n    if (src, dst) not in existing_ref_set:\n        neg_pairs.append((src, dst))\n\ntest_false_refs = pd.DataFrame(neg_pairs, columns=['source', 'target'])\ntest_false_refs.loc[:, 'label'] = 0\n\ntest_refs = pd.concat([test_true_refs, test_false_refs], ignore_index=True)\ntest_refs = test_refs.sample(frac=1, random_state=0, axis=0).reset_index(drop=True)\ntest_refs.head()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T08:32:30.271602Z","iopub.execute_input":"2026-05-09T08:32:30.272843Z","iopub.status.idle":"2026-05-09T08:32:32.173608Z","shell.execute_reply.started":"2026-05-09T08:32:30.272799Z","shell.execute_reply":"2026-05-09T08:32:32.172627Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_ref_tensor = torch.as_tensor(train_refs[['source', 'target']].to_numpy(), dtype=torch.long)\ncite_tensor = torch.as_tensor(cite_edges[['source', 'target']].to_numpy(), dtype=torch.long)\ncoauthor_tensor = torch.as_tensor(coauthor_edges[['source', 'target']].to_numpy(), dtype=torch.long)\n\ntest_ref_arr = np.array(refs_to_pred, dtype=np.int64) if len(refs_to_pred) > 0 else np.zeros((0, 2), dtype=np.int64)\nnum_authors = int(max(\n    ref_edges['source'].max(),\n    coauthor_edges['source'].max(),\n    coauthor_edges['target'].max(),\n    test_ref_arr[:, 0].max() if len(test_ref_arr) else 0,\n) + 1)\nnum_papers = int(max(\n    cite_edges['source'].max(),\n    cite_edges['target'].max(),\n    ref_edges['target'].max(),\n    test_ref_arr[:, 1].max() if len(test_ref_arr) else 0,\n    paper_feature.shape[0] - 1,\n) + 1)\n\npaper_x = torch.as_tensor(paper_feature, dtype=torch.float)\nif paper_x.size(0) < num_papers:\n    pad = torch.zeros(num_papers - paper_x.size(0), paper_x.size(1), dtype=paper_x.dtype)\n    paper_x = torch.cat([paper_x, pad], dim=0)\nelif paper_x.size(0) > num_papers:\n    paper_x = paper_x[:num_papers]\n\ndata = HeteroData()\ndata['author'].num_nodes = num_authors\ndata['paper'].x = paper_x\ndata['paper'].num_nodes = num_papers\n\n# author -> paper reference edges\ndata['author', 'ref', 'paper'].edge_index = train_ref_tensor.t().contiguous()\n# reverse relation for message passing from paper back to author\ndata['paper', 'beref', 'author'].edge_index = train_ref_tensor[:, [1, 0]].t().contiguous()\n\n# citation and coauthor edges are treated as undirected by adding reverse edges,\n# matching the original DGL code.\ndata['paper', 'cite', 'paper'].edge_index = torch.cat([\n    cite_tensor,\n    cite_tensor[:, [1, 0]],\n], dim=0).t().contiguous()\ndata['author', 'coauthor', 'author'].edge_index = torch.cat([\n    coauthor_tensor,\n    coauthor_tensor[:, [1, 0]],\n], dim=0).t().contiguous()\n\ndata = data.to(device)\nprint(data)\nprint('metadata:', data.metadata())\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class HeteroMeanConv(nn.Module):\n    \"\"\"A lightweight hetero mean-aggregation layer implemented with pure PyTorch.\n\n    The layer consumes PyG's HeteroData / edge_index_dict format but avoids compiled\n    PyG extensions. For efficiency, it first averages source node features at each\n    destination node and then applies relation-specific linear projections.\n    \"\"\"\n    def __init__(self, metadata, in_dims, out_dim):\n        super().__init__()\n        node_types, edge_types = metadata\n        self.node_types = list(node_types)\n        self.edge_types = list(edge_types)\n        self.rel_lins = nn.ModuleDict({\n            self._key(edge_type): nn.Linear(in_dims[edge_type[0]], out_dim, bias=False)\n            for edge_type in self.edge_types\n        })\n        self.self_lins = nn.ModuleDict({\n            node_type: nn.Linear(in_dims[node_type], out_dim)\n            for node_type in self.node_types\n        })\n\n    @staticmethod\n    def _key(edge_type):\n        return '__'.join(edge_type)\n\n    def reset_parameters(self):\n        for layer in self.rel_lins.values():\n            layer.reset_parameters()\n        for layer in self.self_lins.values():\n            layer.reset_parameters()\n\n    def forward(self, x_dict, edge_index_dict, num_nodes_dict):\n        out_dict = {\n            node_type: self.self_lins[node_type](x_dict[node_type])\n            for node_type in self.node_types\n        }\n        rel_count = {node_type: 1 for node_type in self.node_types}  # self term\n\n        for edge_type, edge_index in edge_index_dict.items():\n            src_type, _, dst_type = edge_type\n            src, dst = edge_index\n            src_x = x_dict[src_type]\n            agg = src_x.new_zeros((num_nodes_dict[dst_type], src_x.size(-1)))\n            deg = src_x.new_zeros((num_nodes_dict[dst_type], 1))\n            agg.index_add_(0, dst, src_x[src])\n            deg.index_add_(\n                0,\n                dst,\n                torch.ones((dst.numel(), 1), dtype=src_x.dtype, device=src_x.device),\n            )\n            agg = agg / deg.clamp(min=1.0)\n            out_dict[dst_type] = out_dict[dst_type] + self.rel_lins[self._key(edge_type)](agg)\n            rel_count[dst_type] += 1\n\n        return {\n            node_type: out_dict[node_type] / rel_count[node_type]\n            for node_type in self.node_types\n        }\n\n\nclass HeteroRecommender(nn.Module):\n    def __init__(self, metadata, paper_in_dim, hidden_dim=64, out_dim=10, author_in_dim=512):\n        super().__init__()\n        self.author_emb = nn.Embedding(num_authors, author_in_dim)\n        self.paper_lin = nn.Linear(paper_in_dim, author_in_dim)\n        self.num_nodes_dict = {\n            'author': num_authors,\n            'paper': num_papers,\n        }\n\n        self.conv1 = HeteroMeanConv(\n            metadata,\n            in_dims={'author': author_in_dim, 'paper': author_in_dim},\n            out_dim=hidden_dim,\n        )\n        self.conv2 = HeteroMeanConv(\n            metadata,\n            in_dims={'author': hidden_dim, 'paper': hidden_dim},\n            out_dim=out_dim,\n        )\n\n        self.reset_parameters()\n\n    def reset_parameters(self):\n        nn.init.xavier_uniform_(self.author_emb.weight)\n        self.paper_lin.reset_parameters()\n        self.conv1.reset_parameters()\n        self.conv2.reset_parameters()\n\n    def encode(self, data):\n        x_dict = {\n            'author': self.author_emb.weight,\n            'paper': self.paper_lin(data['paper'].x),\n        }\n        x_dict = self.conv1(x_dict, data.edge_index_dict, self.num_nodes_dict)\n        x_dict = {k: F.relu(v) for k, v in x_dict.items()}\n        x_dict = self.conv2(x_dict, data.edge_index_dict, self.num_nodes_dict)\n        return x_dict\n\n    def decode(self, z_dict, edge_label_index):\n        src, dst = edge_label_index\n        return (z_dict['author'][src] * z_dict['paper'][dst]).sum(dim=-1)\n\n\ndef sample_negative_edges(num_samples, num_authors, num_papers, existing_edges, device):\n    # Ensures sampled pairs are not known positive refs.\n    neg_edges = []\n    while len(neg_edges) < num_samples:\n        need = num_samples - len(neg_edges)\n        src = torch.randint(0, num_authors, (need * 2,), device='cpu')\n        dst = torch.randint(0, num_papers, (need * 2,), device='cpu')\n        for s, d in zip(src.tolist(), dst.tolist()):\n            if (s, d) not in existing_edges:\n                neg_edges.append((s, d))\n                if len(neg_edges) == num_samples:\n                    break\n    return torch.tensor(neg_edges, dtype=torch.long, device=device).t().contiguous()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = HeteroRecommender(\n    data.metadata(),\n    paper_in_dim=data['paper'].x.size(-1),\n    hidden_dim=64,\n    out_dim=10,\n    author_in_dim=512,\n).to(device)\n\noptimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=1e-5)\n\npos_edge_index = data['author', 'ref', 'paper'].edge_index\nexisting_train_set = set(map(tuple, train_refs[['source', 'target']].to_numpy().tolist()))\n\n# Each epoch performs one full-graph message-passing pass, then trains on a sampled\n# set of positive/negative author-paper edges.\ntrain_edge_batch_size = min(32768, pos_edge_index.size(1))\nnum_epochs = 60\n\nfor epoch in range(num_epochs):\n    model.train()\n    optimizer.zero_grad()\n\n    batch_perm = torch.randperm(pos_edge_index.size(1), device=device)[:train_edge_batch_size]\n    pos_batch = pos_edge_index[:, batch_perm]\n    neg_batch = sample_negative_edges(\n        pos_batch.size(1), num_authors, num_papers, existing_train_set, device\n    )\n\n    z_dict = model.encode(data)\n    pos_score = model.decode(z_dict, pos_batch)\n    neg_score = model.decode(z_dict, neg_batch)\n\n    loss = (1.0 - pos_score + neg_score).clamp(min=0).mean()\n    loss.backward()\n    optimizer.step()\n\n    print(f'Epoch {epoch:03d}, loss={loss.item():.4f}')\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"with torch.no_grad():\n    model.eval()\n    if 'z_dict' in globals():\n        node_embeddings = {k: v.detach().cpu() for k, v in z_dict.items()}\n    else:\n        node_embeddings = model.encode(data)\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":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from numpy.linalg import norm\n\n# Use cosine similarity to measure recommendation probability.\ndef cos_sim(a, b, eps=1e-12):\n    return np.sum(a * b, axis=1) / (norm(a, axis=1) * norm(b, axis=1) + eps)\n\ntest_arr = test_refs[['source', 'target']].to_numpy(dtype=np.int64)\nres = cos_sim(\n    node_embeddings['author'][test_arr[:, 0]].numpy(),\n    node_embeddings['paper'][test_arr[:, 1]].numpy(),\n)\n\nlbl_true = test_refs['label'].to_numpy().flatten()\n\nimport matplotlib.pyplot as plt\nfrom sklearn.metrics import precision_recall_curve, PrecisionRecallDisplay\n\nprecision, recall, _ = precision_recall_curve(lbl_true, np.array(res))\ndisp = PrecisionRecallDisplay(precision=precision, recall=recall)\ndisp.plot()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.metrics import f1_score\n\nlbl_pred = np.array(res).copy()\nthreshold = 0.1\nlbl_pred[lbl_pred >= threshold] = 1\nlbl_pred[lbl_pred < threshold] = 0\n\nprint('F1-Score: {:.3f}'.format(f1_score(lbl_true, lbl_pred)))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Output your prediction\nif 'threshold' not in globals():\n    threshold = 0.5\nif 'cos_sim' not in globals():\n    from numpy.linalg import norm\n    def cos_sim(a, b, eps=1e-12):\n        return np.sum(a * b, axis=1) / (norm(a, axis=1) * norm(b, axis=1) + eps)\n\noutput_path = '/kaggle/working/Submission.csv'\n\ntest_arr = np.array(refs_to_pred, dtype=np.int64)\nres = cos_sim(\n    node_embeddings['author'][test_arr[:, 0]].numpy(),\n    node_embeddings['paper'][test_arr[:, 1]].numpy(),\n)\nres[res >= threshold] = 1\nres[res < threshold] = 0\n\ndata_out = []\nfor index, p in enumerate(list(res)):\n    data_out.append([index, str(int(p))])\n\ndf = pd.DataFrame(data_out, columns=['Index', 'Predicted'], dtype=object)\ndf.to_csv(output_path, index=False)\nprint('Saved to:', output_path)\ndf.head()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}