updated model

dataset/ogbn_link_pred_dataset.py

@@ -1,9 +1,11 @@
 import os
+import torch.nn.functional as F
+import random
+from torch_sparse import SparseTensor
 import pandas as pd
 import torch
 from ogb.nodeproppred import PygNodePropPredDataset
 from torch_geometric.transforms import RandomLinkSplit
-from torch_geometric.loader import LinkNeighborLoader
 from torch_geometric.data import Data
 
 import requests
@@ -23,9 +25,8 @@ class OGBNLinkPredDataset:
         self._download_abstracts()
         self.corpus = self._load_corpus()
 
-        self.train_data, self.val_data, self.test_data = self._split_data(
-            val_size, test_size
-        )
+        self.val_size = val_size
+        self.test_size = test_size
 
     def _download_abstracts(self):
         target_dir = os.path.join(self.root, "mapping")
@@ -38,22 +39,17 @@ class OGBNLinkPredDataset:
             os.makedirs(target_dir, exist_ok=True)
 
             try:
-                print(f"Downloading from {url}...")
                 response = requests.get(url, stream=True)
                 response.raise_for_status()
                 with open(gz_path, "wb") as f:
                     for chunk in response.iter_content(chunk_size=8192):
                         f.write(chunk)
-                print(f"File downloaded to: {gz_path}")
 
-                print(f"Decompressing {gz_path}...")
-                with gzip.open(gz_path, 'rb') as f_in:
-                    with open(tsv_path, 'wb') as f_out:
+                with gzip.open(gz_path, "rb") as f_in:
+                    with open(tsv_path, "wb") as f_out:
                         shutil.copyfileobj(f_in, f_out)
-                print(f"File extracted to: {tsv_path}")
 
                 os.remove(gz_path)
-                print(f"Removed temporary file: {gz_path}")
 
             except requests.exceptions.RequestException as e:
                 print(f"Error downloading file: {e}")
@@ -80,22 +76,226 @@ class OGBNLinkPredDataset:
                 + "\n "
                 + df_text_aligned["abstract"].fillna("")
             ).tolist()
-            print(f"Corpus created with {len(corpus)} documents.")
             return corpus
         except FileNotFoundError:
             print("Error: titleabs.tsv not found. Could not create corpus.")
             return []
 
-    def _split_data(self, val_size: float, test_size: float) -> tuple[Data, Data, Data]:
+    def get_splits(self) -> tuple[Data, Data, Data]:
         transform = RandomLinkSplit(
-            num_val=val_size,
-            num_test=test_size,
+            num_val=self.val_size,
+            num_test=self.test_size,
             is_undirected=False,
-            add_negative_train_samples=False,
+            add_negative_train_samples=True,
+            neg_sampling_ratio=1.0,
         )
         train_split, val_split, test_split = transform(self.data)
-        print("Data successfully split into train, validation, and test sets.")
         return train_split, val_split, test_split
 
+
+class OGBNLinkPredNegDataset(OGBNLinkPredDataset):
+    """Degree similar hard negatives sampling"""
+
+    def __init__(
+        self, root_dir: str = "data", val_size: float = 0.1, test_size: float = 0.2
+    ):
+        super().__init__(root_dir, val_size, test_size)
+        self.degree_tol = 0
+
     def get_splits(self) -> tuple[Data, Data, Data]:
-        return self.train_data, self.val_data, self.test_data
+        transform = RandomLinkSplit(
+            num_val=self.val_size,
+            num_test=self.test_size,
+            is_undirected=False,
+            add_negative_train_samples=False,
+            neg_sampling_ratio=0.0,
+        )
+        train_split, val_split, test_split = transform(self.data)
+
+        print("Generating hard negatives...")
+        adj_matrix = SparseTensor.from_edge_index(
+            train_split.edge_index,  # only from train_split
+            sparse_sizes=(self.data.num_nodes, self.data.num_nodes),
+        )
+        self.degrees = adj_matrix.sum(dim=0).to(torch.long)
+        # to prevent creating negative edges that are positive in other split
+        self.all_edge_set = set(zip(*self.data.edge_index.tolist()))
+        train_split = self._add_balanced_negs(train_split)
+        val_split = self._add_balanced_negs(val_split)
+        test_split = self._add_balanced_negs(test_split)
+        return train_split, val_split, test_split
+
+    def _add_balanced_negs(self, split_data):
+        assert (split_data.edge_label == 1).all(), "Expected only positive edges"
+
+        pos_edges = split_data.edge_label_index
+        pos_list = pos_edges.t().tolist()
+        num_negs = pos_edges.size(1)
+
+        negs = []
+        for _ in range(num_negs):
+            u, v_orig = random.choice(pos_list)
+            target_deg = int(self.degrees[v_orig])
+
+            found = False
+            for _ in range(20):
+                w = random.randrange(self.data.num_nodes)
+                if (
+                    (u, w) not in self.all_edge_set
+                    and w != u
+                    and abs(int(self.degrees[w]) - target_deg) <= self.degree_tol
+                ):
+                    negs.append((u, w))
+                    found = True
+                    break
+
+            if not found:
+                while True:
+                    w = random.randrange(self.data.num_nodes)
+                    if (u, w) not in self.all_edge_set and w != u:
+                        negs.append((u, w))
+                        break
+
+        neg_edges = torch.tensor(negs, dtype=torch.long).t()
+        N = pos_edges.size(1)
+
+        split_data.edge_label_index = torch.cat([pos_edges, neg_edges], dim=1)
+        split_data.edge_label = torch.cat(
+            [
+                torch.ones(N, dtype=torch.long, device=pos_edges.device),
+                torch.zeros(N, dtype=torch.long, device=pos_edges.device),
+            ]
+        )
+
+        return split_data
+
+
+# class OGBNLinkPredNegDataset2(OGBNLinkPredDataset):
+#     """Degree and semantically similar hard negatives sampling"""
+#
+#     def __init__(self, root_dir="data", val_size=0.1, test_size=0.2):
+#         super().__init__(root_dir, val_size, test_size)
+#
+#     def get_splits(self) -> tuple[Data, Data, Data]:
+#         transform = RandomLinkSplit(
+#             num_val=self.val_size,
+#             num_test=self.test_size,
+#             is_undirected=False,
+#             add_negative_train_samples=False,
+#             neg_sampling_ratio=0.0,
+#         )
+#         train_split, val_split, test_split = transform(self.data)
+#
+#         print("Generating semantic hard negatives...")
+#         train_split = self._add_balanced_negs(train_split)
+#         val_split = self._add_balanced_negs(val_split)
+#         test_split = self._add_balanced_negs(test_split)
+#         return train_split, val_split, test_split
+#
+#     def _add_balanced_negs(self, split_data):
+#         assert (split_data.edge_label == 1).all(), "Expected only positive edges"
+#
+#         BS = 1_000
+#         B = self.data.x.to("cuda", dtype=torch.bfloat16)  # (num_nodes, dim)
+#         B = F.normalize(B, p=2, dim=1)
+#         K = 100
+#
+#         pos_edges = split_data.edge_label_index
+#         adj_matrix = SparseTensor.from_edge_index(
+#             split_data.edge_index,
+#             sparse_sizes=(self.data.num_nodes, self.data.num_nodes),
+#         )
+#         degrees = adj_matrix.sum(dim=0).to("cuda")
+#
+#         topk_val = torch.empty((BS, K), dtype=torch.bfloat16, device="cuda")
+#         topk_idx = torch.empty((BS, K), dtype=torch.int64, device="cuda")
+#
+#         neg_edges = []
+#
+#         for i in range(0, pos_edges.shape[1], BS):
+#             batch_end = min(i + BS, pos_edges.shape[1])
+#             src_idx = pos_edges[0, i:batch_end]  # (batch_size,)
+#             dst_idx = pos_edges[1, i:batch_end]  # (batch_size,)
+#
+#             A = B[src_idx]  # (batch_size, dim)
+#
+#             with torch.autocast("cuda", dtype=torch.bfloat16):
+#                 sim = torch.mm(A, B.t())  # equivalent to cos-sim
+#
+#                 # mask for similarity with itself and existing edges
+#                 sim[torch.arange(len(A)), dst_idx] = -1
+#                 sim[torch.arange(len(A)), src_idx] = -1
+#                 # TODO: exclude edges from val&test sets
+#
+#                 torch.topk(sim, K, out=(topk_val, topk_idx))
+#                 topk_idx2 = topk_idx[: len(A)]
+#
+#                 # sample degree matched negs
+#                 topk_deg = degrees[topk_idx2]
+#                 src_deg = degrees[src_idx]
+#
+#                 deg_diffs = torch.abs(topk_deg - src_deg.unsqueeze(1))
+#                 closest_idx = torch.argmin(deg_diffs, dim=1)  # (batch_size,)
+#                 sampled_negs = topk_idx[
+#                     torch.arange(len(A), device="cuda"), closest_idx
+#                 ]
+#                 neg_edges.append(sampled_negs)
+#
+#         neg_dsts = torch.cat(neg_edges, dim=0).to("cpu")
+#         neg_edge_index = torch.stack([pos_edges[0].cpu(), neg_dsts], dim=0)
+#         edge_label_index = torch.cat([pos_edges.cpu(), neg_edge_index], dim=1)
+#         edge_label = torch.cat(
+#             [split_data.edge_label, torch.zeros(neg_dsts.shape[0])], dim=0
+#         )
+#         assert edge_label.shape[0] == edge_label_index.shape[1], (
+#             "Label and index shape mismatch"
+#         )
+#         assert len(neg_dsts) == pos_edges.shape[1], (
+#             "Expected same amount of positive and negative edges"
+#         )
+#         return Data(
+#             x=split_data.x,
+#             edge_index=edge_label,
+#             edge_label_index=edge_label_index,
+#             edge_label=edge_label,
+#         )
+
+
+if __name__ == "__main__":
+    dataset = OGBNLinkPredNegDataset()
+    train, val, test = dataset.get_splits()
+
+    def extract_pos_neg_edges(split):
+        pos = split.edge_label_index[:, split.edge_label == 1]
+        neg = split.edge_label_index[:, split.edge_label == 0]
+        return pos, neg
+
+    for name, split in [("train", train), ("val", val), ("test", test)]:
+        assert split.edge_label_index.shape[0] == 2, (
+            f"{name}: edge_label_index must have 2 rows"
+        )
+        assert split.edge_label_index.shape[1] == split.edge_label.shape[0], (
+            f"{name}: label/index shape mismatch"
+        )
+        assert torch.all(0 <= split.edge_label) and torch.all(split.edge_label <= 1), (
+            f"{name}: labels not 0/1"
+        )
+
+        pos, neg = extract_pos_neg_edges(split)
+        assert pos.size(1) == neg.size(1), f"{name}: pos/neg count mismatch"
+
+        pos_set = set(tuple(e) for e in pos.t().tolist())
+        neg_set = set(tuple(e) for e in neg.t().tolist())
+        assert pos_set.isdisjoint(neg_set), f"{name}: pos/neg overlap"
+
+        assert all(u != v for u, v in pos_set), f"{name}: pos self-loops"
+        assert all(u != v for u, v in neg_set), f"{name}: neg self-loops"
+
+        assert len(pos_set) == pos.size(1), f"{name}: pos duplicates"
+        assert len(neg_set) == neg.size(1), f"{name}: neg duplicates"
+
+        assert pos.size(1) / neg.size(1) == 1.0 if neg.size(1) > 0 else True, (
+            f"{name}: ratio not 1.0"
+        )
+
+    print("All asserts passed!")

galis_app.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 import streamlit as st
-
+import torch.nn.functional as F
 from predictor.link_predictor import (
     prepare_system,
     get_citation_predictions,
@@ -9,7 +9,7 @@ from predictor.link_predictor import (
 )
 from llm.related_work_generator import generate_related_work
 
-MODEL_PATH = Path("predictor/model.pth")
+MODEL_PATH = Path("model.pth")
 
 
 @st.cache_resource
@@ -94,8 +94,9 @@ def app():
                     new_vector = abstract_to_vector(
                         abstract_input, abstract_title, st_model
                     )
+
                     probabilities = get_citation_predictions(
-                        vector=new_vector,
+                        vector=F.normalize(new_vector.view(1, -1), p=2, dim=1),
                         model=gcn_model,
                         z_all=z_all,
                         num_nodes=dataset.data.num_nodes,
@@ -112,7 +113,9 @@ def app():
 
                 with related_work_placeholder.container():
                     with st.spinner("Generating related work section..."):
-                        related_work = generate_related_work(st.session_state.references)
+                        related_work = generate_related_work(
+                            st.session_state.references
+                        )
                         st.session_state.related_work = related_work
 
     if st.session_state.references:

model/cos-sim.py

@@ -0,0 +1,104 @@
+from pathlib import Path
+import numpy as np
+import torch
+import torch.nn.functional as F
+from sklearn.metrics import roc_auc_score, average_precision_score
+from sentence_transformers import SentenceTransformer
+import argparse
+from dataset.ogbn_link_pred_dataset import OGBNLinkPredDataset, OGBNLinkPredNegDataset
+
+BATCH_SIZE_EDGES = 100_000  # edge batching for scoring
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--custom-neg", action=argparse.BooleanOptionalAction, default=False
+    )
+    parser.add_argument(
+        "--bert-embed", action=argparse.BooleanOptionalAction, default=False
+    )
+    return parser.parse_args()
+
+
+@torch.no_grad()
+def eval_edges_cos(global_emb, edge_index, edge_label, batch_size=BATCH_SIZE_EDGES):
+    # edge_index shape: [2, M] with GLOBAL node ids; edge_label: [M] in {0,1}
+    assert edge_index.dim() == 2 and edge_index.size(0) == 2
+    assert edge_index.max() < global_emb.size(0), "Edge node id out of range."
+    assert (edge_label == 0).any() and (edge_label == 1).any(), "Need both classes."
+
+    scores_list, labels_list = [], []
+    M = edge_index.size(1)
+    for i in range(0, M, batch_size):
+        j = min(i + batch_size, M)
+        src = edge_index[0, i:j].to(global_emb.device)
+        dst = edge_index[1, i:j].to(global_emb.device)
+        scores = (global_emb[src] * global_emb[dst]).sum(
+            dim=1
+        )  # cosine (L2-normalized)
+        scores_list.append(scores.float().cpu().numpy())
+        labels_list.append(edge_label[i:j].cpu().numpy())
+    y_scores = np.concatenate(scores_list)
+    y_true = np.concatenate(labels_list)
+    roc = roc_auc_score(y_true, y_scores)
+    ap = average_precision_score(y_true, y_scores)
+    return roc, ap
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    USE_CUSTOM_NEG = args.custom_neg
+    USE_BERT_EMBED = args.bert_embed
+    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # --- Load dataset + frozen embeddings ---
+    if USE_CUSTOM_NEG:
+        print("using hard negatives")
+        dataset = OGBNLinkPredNegDataset(val_size=0.1, test_size=0.2)
+    else:
+        print("using random negatives")
+        dataset = OGBNLinkPredDataset(val_size=0.1, test_size=0.2)
+    if USE_BERT_EMBED:
+        print("using BERT embeds")
+        if Path("model/embeddings.pth").exists():
+            emb = torch.load("model/embeddings.pth", map_location=DEVICE)
+        else:
+            st = SentenceTransformer("bongsoo/kpf-sbert-128d-v1", device=DEVICE)
+            emb = st.encode(
+                dataset.corpus, convert_to_tensor=True, show_progress_bar=True
+            )
+            Path("model").mkdir(parents=True, exist_ok=True)
+            torch.save(emb, "model/embeddings.pth")
+        emb = F.normalize(emb.to(DEVICE), p=2, dim=1)
+    else:
+        print("using skipgram embeds")
+        emb = dataset.data.x
+
+    train_data, val_data, test_data = dataset.get_splits()
+
+    # Sanity checks
+    for split_name, data in [
+        ("train", train_data),
+        ("val", val_data),
+        ("test", test_data),
+    ]:
+        assert data.edge_label_index.size(1) == data.edge_label.size(0), (
+            f"{split_name} size mismatch"
+        )
+        assert (data.edge_label == 0).any() and (data.edge_label == 1).any(), (
+            f"{split_name} lacks negatives"
+        )
+        assert data.edge_label_index.max() < emb.size(0), (
+            f"{split_name} has node ids >= num_nodes"
+        )
+
+    val_roc, val_ap = eval_edges_cos(
+        emb, val_data.edge_label_index, val_data.edge_label
+    )
+    test_roc, test_ap = eval_edges_cos(
+        emb, test_data.edge_label_index, test_data.edge_label
+    )
+
+    print(f"Val ROC-AUC:  {val_roc:.4f}, Val AP:  {val_ap:.4f}")
+    print(f"Test ROC-AUC: {test_roc:.4f}, Test AP: {test_ap:.4f}")

model/mlp.py

@@ -0,0 +1,137 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from sklearn.metrics import roc_auc_score, average_precision_score
+import numpy as np
+from dataset.ogbn_link_pred_dataset import (
+    OGBNLinkPredDataset,
+    OGBNLinkPredNegDataset,
+    # OGBNLinkPredNegDataset2,
+)
+from pathlib import Path
+from sentence_transformers import SentenceTransformer
+import argparse
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+BATCH_SIZE = 2048
+NUM_EPOCHS = 50
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--custom-neg", action=argparse.BooleanOptionalAction, default=False
+    )
+    parser.add_argument(
+        "--bert-embed", action=argparse.BooleanOptionalAction, default=False
+    )
+    return parser.parse_args()
+
+
+# --- Feature builder ---
+def edge_features(emb, ei):
+    u, v = ei
+    eu, ev = emb[u], emb[v]
+    return torch.cat([eu * ev, torch.abs(eu - ev)], dim=1)
+
+
+# --- Simple MLP ---
+class PairMLP(nn.Module):
+    def __init__(self, in_dim, hidden=256):
+        super().__init__()
+        self.fc1 = nn.Linear(in_dim, hidden)
+        self.fc2 = nn.Linear(hidden, 1)
+
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        return self.fc2(x).squeeze(-1)
+
+
+# --- Training loop ---
+def run_epoch(data, train=True):
+    model.train(train)
+    total_loss = 0
+    idx = (
+        torch.randperm(data.edge_label.size(0))
+        if train
+        else torch.arange(data.edge_label.size(0))
+    )
+    for i in range(0, len(idx), BATCH_SIZE):
+        batch_end = min(i + BATCH_SIZE, data.edge_label.size(0))
+        batch_idx = idx[i:batch_end]
+        feats = edge_features(emb, data.edge_label_index[:, batch_idx]).to(DEVICE)
+        labels = data.edge_label[batch_idx].float().to(DEVICE)
+        scores = model(feats)
+        loss = F.binary_cross_entropy_with_logits(scores, labels)
+        if train:
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+        total_loss += loss.item() * len(batch_idx)
+    return total_loss / len(idx)
+
+
+@torch.no_grad()
+def evaluate(data):
+    scores_all, labels_all = [], []
+    for i in range(0, data.edge_label.size(0), BATCH_SIZE):
+        batch_end = min(i + BATCH_SIZE, data.edge_label.size(0))
+        feats = edge_features(emb, data.edge_label_index[:, i:batch_end]).to(DEVICE)
+        labels = data.edge_label[i : i + BATCH_SIZE]
+        scores = torch.sigmoid(model(feats)).cpu().numpy()
+        scores_all.append(scores)
+        labels_all.append(labels.numpy())
+    y_scores = np.concatenate(scores_all)
+    y_true = np.concatenate(labels_all)
+    return roc_auc_score(y_true, y_scores), average_precision_score(y_true, y_scores)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    USE_CUSTOM_NEG = args.custom_neg
+    USE_BERT_EMBED = args.bert_embed
+
+    # --- Load dataset + frozen embeddings ---
+    if USE_CUSTOM_NEG:
+        print("using hard negatives")
+        dataset = OGBNLinkPredNegDataset(val_size=0.1, test_size=0.2)
+    else:
+        print("using random negatives")
+        dataset = OGBNLinkPredDataset(val_size=0.1, test_size=0.2)
+    if USE_BERT_EMBED:
+        print("using BERT embeds")
+        if Path("model/embeddings.pth").exists():
+            emb = torch.load("model/embeddings.pth", map_location=DEVICE)
+        else:
+            st = SentenceTransformer("bongsoo/kpf-sbert-128d-v1", device=DEVICE)
+            emb = st.encode(
+                dataset.corpus, convert_to_tensor=True, show_progress_bar=True
+            )
+            Path("model").mkdir(parents=True, exist_ok=True)
+            torch.save(emb, "model/embeddings.pth")
+        emb = emb.to(DEVICE)
+    else:
+        print("using skipgram embeds")
+        emb = dataset.data.x
+
+    train_data, val_data, test_data = dataset.get_splits()
+
+    model = PairMLP(emb.size(1) * 2).to(DEVICE)
+    opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
+
+    # --- Training ---
+    best_roc, best_ap = 0.0, 0.0
+    for epoch in range(NUM_EPOCHS):
+        loss = run_epoch(train_data, train=True)
+        val_roc, val_ap = evaluate(val_data)
+        if val_roc > best_roc:
+            torch.save(
+                model.state_dict(), f"model_roc{str(val_roc)[:4].replace('.', '_')}.pth"
+            )
+        print(
+            f"Epoch {epoch + 1} | Loss {loss:.4f} | Val ROC {val_roc:.4f} | Val AP {val_ap:.4f}"
+        )
+
+    # --- Final test ---
+    test_roc, test_ap = evaluate(test_data)
+    print(f"Test ROC {test_roc:.4f} | Test AP {test_ap:.4f}")

predictor/link_predictor.py

@@ -1,10 +1,10 @@
-from pathlib import Path
 import torch
+import torch.nn.functional as F
+from dataset.ogbn_link_pred_dataset import OGBNLinkPredDataset
+from pathlib import Path
 import structlog
-
 from sentence_transformers import SentenceTransformer
-from model.simple_gcn_model import SimpleGCN
-from dataset.ogbn_link_pred_dataset import OGBNLinkPredDataset
+from model.mlp import edge_features, PairMLP
 
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -24,50 +24,33 @@ def abstract_to_vector(
     text = title + "\n" + abstract_text
     with torch.no_grad():
         vector = st_model.encode(text, convert_to_tensor=True, device=DEVICE)
-    return vector.unsqueeze(0)
+    return vector
 
 
 def get_citation_predictions(
-    vector: torch.Tensor, model: SimpleGCN, z_all: torch.Tensor, num_nodes: int
+    vector: torch.Tensor,
+    model: PairMLP,
+    z_all: torch.Tensor,
+    num_nodes: int,
 ) -> torch.Tensor:
     model.eval()
-    with torch.no_grad():
-        empty_edge_index = torch.empty(2, 0, dtype=torch.long, device=DEVICE)
-        h1_new = model.conv1(vector, edge_index=empty_edge_index).relu()
-        z_new = model.conv2(h1_new, edge_index=empty_edge_index)
-
-    new_node_idx = num_nodes
-    row = torch.full((num_nodes,), fill_value=new_node_idx, device=DEVICE)
-    col = torch.arange(num_nodes, device=DEVICE)
-    edge_label_index_to_check = torch.stack([row, col], dim=0)
-
-    z_combined = torch.cat([z_all, z_new], dim=0)
 
     with torch.no_grad():
-        logits = model.decode(z_combined, edge_label_index_to_check)
-
-    return torch.sigmoid(logits)
+        combined_embeddings = torch.cat([vector.view(1, -1), z_all], dim=0)
+        edge_index = torch.tensor([[0] * num_nodes, list(range(1, num_nodes + 1))]).to(
+            DEVICE
+        )
+        feat = edge_features(combined_embeddings, edge_index).to(DEVICE)
+        scores = torch.sigmoid(model(feat))
+        return scores.squeeze()
 
 
 def format_top_k_predictions(
-    probs: torch.Tensor, dataset: OGBNLinkPredDataset, top_k=10., show_prob=False
+    probs: torch.Tensor, dataset: OGBNLinkPredDataset, top_k=10, show_prob=False
 ) -> str:
-    """
-    Formats the top K predictions into a single string for display.
-
-    Args:
-        probs (torch.Tensor): The tensor of probabilities for all potential links.
-        dataset (OGBNLinkPredDataset): The dataset object containing the corpus.
-        top_k (int): The number of top predictions to format.
-
-    Returns:
-        str: A formatted string with the top K predictions.
-    """
     probs = probs.cpu()
     top_probs, top_indices = torch.topk(probs, k=top_k)
-
     output_lines = []
-
     header = f"Top {top_k} Citation Predictions:"
     output_lines.append(header)
 
@@ -86,14 +69,9 @@ def format_top_k_predictions(
 
 
 def prepare_system(model_path: Path):
-    """
-    Performs all one-time, expensive operations to prepare the system.
-    Initializes models, loads data, and pre-calculates embeddings using structured logging.
-    """
     logger.info("system_preparation.start")
 
     dataset = OGBNLinkPredDataset()
-    data = dataset.data.to(DEVICE)
     logger.info("dataset.load.success")
 
     model_name = "bongsoo/kpf-sbert-128d-v1"
@@ -103,54 +81,69 @@ def prepare_system(model_path: Path):
     st_model = SentenceTransformer(model_name, device=DEVICE)
     logger.info("model.load.success", model_type="SentenceTransformer")
 
-    gcn_model = SimpleGCN(
-        in_channels=dataset.num_features, hidden_channels=256, out_channels=128
-    ).to(DEVICE)
+    # Load corpus embeddings
+    if Path("model/embeddings.pth").exists():
+        corpus_embeddings = torch.load("model/embeddings.pth", map_location=DEVICE)
+        logger.info("embeddings.load.success")
+    else:
+        logger.info("embeddings.calculation.start")
+        corpus_embeddings = st_model.encode(
+            dataset.corpus, convert_to_tensor=True, show_progress_bar=True
+        )
+        Path("model").mkdir(parents=True, exist_ok=True)
+        torch.save(corpus_embeddings, "model/embeddings.pth")
+        logger.info("embeddings.calculation.success")
+
+    corpus_embeddings = F.normalize(corpus_embeddings.to(DEVICE), p=2, dim=1)
+
+    # Initialize PairMLP
+    embedding_dim = corpus_embeddings.size(1)
+    pair_mlp = PairMLP(embedding_dim * 2).to(DEVICE)
 
     if model_path.exists():
-        gcn_model.load_state_dict(torch.load(model_path, map_location=DEVICE))
-        logger.info("model.load.success", model_type="GCN", path=str(model_path))
+        pair_mlp.load_state_dict(torch.load(model_path, map_location=DEVICE))
+        logger.info("model.load.success", model_type="PairMLP", path=str(model_path))
     else:
         logger.warning(
             "model.load.failure",
-            model_type="GCN",
+            model_type="PairMLP",
             path=str(model_path),
             reason="File not found, using random weights.",
         )
-    gcn_model.eval()
 
-    logger.info("embeddings.calculation.start", embedding_name="z_all")
-    with torch.no_grad():
-        z_all = gcn_model(data.x, data.edge_index)
+    pair_mlp.eval()
 
     logger.info(
         "embeddings.calculation.success",
-        embedding_name="z_all",
-        shape=list(z_all.shape),
+        embedding_name="corpus_embeddings",
+        shape=list(corpus_embeddings.shape),
     )
-
     logger.info("system_preparation.finish", status="ready_for_predictions")
-    return gcn_model, st_model, dataset, z_all
+
+    return pair_mlp, st_model, dataset, corpus_embeddings
 
 
 if __name__ == "__main__":
     MODEL_PATH = Path("model.pth")
-
-    gcn_model, st_model, dataset, z_all = prepare_system(MODEL_PATH)
+    pair_model, st_model, dataset, corpus_embeddings = prepare_system(MODEL_PATH)
 
     my_title = "A Survey of Graph Neural Networks for Link Prediction"
-    my_abstract = """Link predictor is a critical task in graph analysis. "
-                   "In this paper, we review various GNN architectures like GCN and GraphSAGE for predicting edges.
-                   """
+    my_abstract = """Link prediction is a critical task in graph analysis. 
+                   In this paper, we review various GNN architectures like GCN and GraphSAGE for predicting edges."""
 
     new_vector = abstract_to_vector(my_title, my_abstract, st_model)
+    new_vector = F.normalize(
+        new_vector.view(1, -1), p=2, dim=1
+    )  # Normalize like corpus embeddings
 
     probabilities = get_citation_predictions(
         vector=new_vector,
-        model=gcn_model,
-        z_all=z_all,
+        model=pair_model,
+        z_all=corpus_embeddings,
         num_nodes=dataset.data.num_nodes,
     )
 
-    references = format_top_k_predictions(probabilities, dataset, top_k=5)
+    references = format_top_k_predictions(
+        probabilities, dataset, top_k=5, show_prob=True
+    )
     print(references)