"""Generate full-test dot-score ensemble submissions from saved full checkpoints.""" from __future__ import annotations import argparse import importlib.util import pickle as pkl from pathlib import Path import numpy as np import pandas as pd import torch import torch.nn as nn from torch_geometric.data import HeteroData EDGE_TYPES = [ ("author", "ref", "paper"), ("paper", "beref", "author"), ("paper", "cite", "paper"), ("author", "coauthor", "author"), ] def read_txt(path: Path): return [list(map(int, line.strip().split())) for line in path.open()] def log_norm(x): x = np.log1p(x) return (x - x.mean()) / (x.std() + 1e-8) class LightGCNLayer(nn.Module): def forward(self, x_dict, edge_index_dict): agg_dict = {node_type: [] for node_type in x_dict} for et in EDGE_TYPES: if et not in edge_index_dict: continue st, _, dt = et src, dst = edge_index_dict[et] sx = x_dict[st] agg = sx.new_zeros((x_dict[dt].size(0), sx.size(-1))) deg = sx.new_zeros((x_dict[dt].size(0), 1)) agg.index_add_(0, dst, sx[src]) deg.index_add_(0, dst, torch.ones((dst.numel(), 1), dtype=sx.dtype, device=sx.device)) agg_dict[dt].append(agg / deg.clamp(min=1.0)) return {nt: sum(v) / len(v) if v else x_dict[nt] for nt, v in agg_dict.items()} class LightGCN(nn.Module): def __init__(self, n_author, feat_dim, dim, layers=4): super().__init__() self.author_emb = nn.Embedding(n_author, dim) self.paper_proj = nn.Linear(feat_dim, dim) self.layers = nn.ModuleList([LightGCNLayer() for _ in range(layers)]) self.num_layers = layers def encode(self, data): x = {"author": self.author_emb.weight, "paper": self.paper_proj(data["paper"].x)} all_x = [x] for layer in self.layers: x = layer(x, data.edge_index_dict) all_x.append(x) w = 1.0 / len(all_x) return {nt: sum(w * xx[nt] for xx in all_x) for nt in x} def build(root: Path, device): data_dir = root / "data_and_docs" refs = read_txt(data_dir / "bipartite_train_ann.txt") test = read_txt(data_dir / "bipartite_test_ann.txt") cite = read_txt(data_dir / "paper_file_ann.txt") coa = read_txt(data_dir / "author_file_ann.txt") with (data_dir / "feature.pkl").open("rb") as f: feat = pkl.load(f).numpy().astype(np.float32) n_paper = 79937 ref_deg = np.zeros(n_paper, np.float32) cout = np.zeros(n_paper, np.float32) cin = np.zeros(n_paper, np.float32) for _, p in refs: ref_deg[p] += 1 for s, t in cite: cout[s] += 1 cin[t] += 1 deg = np.stack([log_norm(ref_deg), log_norm(cout), log_norm(cin)], axis=-1) paper_x = np.concatenate([feat, deg], axis=1) paper_x = (paper_x - paper_x.mean(0)) / (paper_x.std(0) + 1e-8) rt = torch.as_tensor(np.array(refs), dtype=torch.long) ct = torch.as_tensor(np.array(cite), dtype=torch.long) co = torch.as_tensor(np.array(coa), dtype=torch.long) data = HeteroData() data["author"].num_nodes = 6611 data["paper"].num_nodes = n_paper data["paper"].x = torch.as_tensor(paper_x, dtype=torch.float) data["author", "ref", "paper"].edge_index = rt.t().contiguous() data["paper", "beref", "author"].edge_index = rt[:, [1, 0]].t().contiguous() data["paper", "cite", "paper"].edge_index = torch.cat([ct, ct[:, [1, 0]]], 0).t().contiguous() data["author", "coauthor", "author"].edge_index = torch.cat([co, co[:, [1, 0]]], 0).t().contiguous() return data.to(device), np.array(test, dtype=np.int64), refs, paper_x.shape[1] @torch.no_grad() def predict_dot(model, data, pairs, batch_size): z = model.encode(data) a = z["author"].detach().cpu().numpy() p = z["paper"].detach().cpu().numpy() out = [] for st in range(0, len(pairs), batch_size): b = pairs[st : st + batch_size] out.append(np.sum(a[b[:, 0]] * p[b[:, 1]], axis=1).astype(np.float32)) return np.concatenate(out) def rank01(x): order = np.argsort(x, kind="mergesort") r = np.empty(len(x), dtype=np.float32) r[order] = np.linspace(0, 1, len(x), dtype=np.float32) return r def write(scores, known, out_dir, prefix, ratios): forced = scores.copy() forced[known] = np.inf order = np.argsort(forced)[::-1] for ratio in ratios: k = int(round(len(scores) * ratio)) pred = np.zeros(len(scores), dtype=np.int8) pred[order[:k]] = 1 df = pd.DataFrame({"Index": np.arange(len(pred)), "Predicted": pred.astype(str)}) path = out_dir / f"{prefix}_r{ratio:.3f}.csv" df.to_csv(path, index=False) print(path, int(pred.sum()), float(pred.mean())) def main(): parser = argparse.ArgumentParser() parser.add_argument("--package-root", type=Path, default=Path(__file__).resolve().parents[1]) parser.add_argument("--device", default="cuda:0" if torch.cuda.is_available() else "cpu") parser.add_argument("--batch-size", type=int, default=65536) parser.add_argument("--ratios", nargs="*", type=float, default=[0.505, 0.515, 0.521, 0.530, 0.540]) args = parser.parse_args() root = args.package_root device = torch.device(args.device) data, pairs, refs, feat_dim = build(root, device) train_set = set(map(tuple, refs)) known = np.array([tuple(x) in train_set for x in pairs]) ckpts = [ "model_best_s0_d512.pt", "model_best_s0_d384.pt", "model_lgcn_s23.pt", "model_lgcn_s0.pt", "model_lgcn_s100.pt", "model_lgcn_s77.pt", "model_lgcn_s42.pt", "model_lgcn_s2024.pt", ] score_dir = root / "cached_scores" / "dot_full" out_dir = root / "submissions" / "dot_full" score_dir.mkdir(parents=True, exist_ok=True) out_dir.mkdir(parents=True, exist_ok=True) scores = [] for name in ckpts: path = root / "checkpoints" / "extra_models" / name cache = score_dir / f"{path.stem}_dot.npy" if cache.exists(): s = np.load(cache) else: state = torch.load(path, map_location=device) dim = state["author_emb.weight"].shape[1] model = LightGCN(6611, feat_dim, dim, 4).to(device) model.load_state_dict(state) s = predict_dot(model, data, pairs, args.batch_size) np.save(cache, s) del model torch.cuda.empty_cache() print(name, s.mean(), s.std()) scores.append(s) zmean = np.mean([(s - s.mean()) / (s.std() + 1e-8) for s in scores[:4]], axis=0) rmean = np.mean([rank01(s) for s in scores[:4]], axis=0) write(zmean, known, out_dir, "sub_dot_top4_z", args.ratios) write(rmean, known, out_dir, "sub_dot_top4_rank", args.ratios) if __name__ == "__main__": main()