Spaces:
Sleeping
Sleeping
| 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() | |
| 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}") | |