code/generate_extra_bprmf_submission.py

"""Generate post95 submissions with content-mean and BPR-MF score features."""

from __future__ import annotations

import argparse
import importlib.util
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F


def load_module(name: str, path: Path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def train_mf_test_scores(extra, root: Path, train_refs: pd.DataFrame, test_pairs: np.ndarray, out_dir: Path, device: str, seed: int, dim: int, epochs: int) -> np.ndarray:
    out_path = out_dir / f"test_mf_bpr_dynamic_s{seed}_d{dim}_e{epochs}.npy"
    if out_path.exists():
        return np.load(out_path)
    torch.manual_seed(seed)
    np.random.seed(seed)
    rng = np.random.default_rng(seed)
    train = train_refs[["source", "target"]].to_numpy(np.int64)
    train_set = set(map(tuple, train.tolist()))
    model = extra.MF(6611, 79937, dim).to(torch.device(device))
    opt = torch.optim.AdamW(model.parameters(), lr=0.01, weight_decay=1e-6)
    train_t = torch.as_tensor(train, dtype=torch.long, device=device)
    batch_size = 65536
    for ep in range(epochs):
        idx = torch.randint(0, train_t.size(0), (batch_size,), device=device)
        pos = train_t[idx]
        neg_np = np.empty((batch_size, 2), dtype=np.int64)
        authors = pos[:, 0].detach().cpu().numpy()
        filled = 0
        while filled < batch_size:
            papers = rng.integers(0, 79937, size=batch_size - filled)
            for a, p in zip(authors[filled:], papers):
                if (int(a), int(p)) not in train_set:
                    neg_np[filled] = (a, p)
                    filled += 1
                    if filled >= batch_size:
                        break
        neg = torch.as_tensor(neg_np, dtype=torch.long, device=device)
        loss = -F.logsigmoid(model.score(pos) - model.score(neg)).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
        if (ep + 1) % 20 == 0:
            print(f"mf-test epoch={ep+1:03d} loss={loss.item():.4f}")
    test_t = torch.as_tensor(test_pairs, dtype=torch.long, device=device)
    scores = []
    with torch.no_grad():
        for st in range(0, len(test_pairs), 131072):
            scores.append(model.score(test_t[st : st + 131072]).detach().cpu().numpy())
    scores = np.concatenate(scores).astype(np.float32)
    np.save(out_path, scores)
    return scores


def make_subs(root: Path, out_dir: Path, score: np.ndarray, ratios: list[float]) -> None:
    known = np.load(root / "cached_scores" / "test_known_mask.npy").astype(bool)
    for ratio in ratios:
        pred = np.zeros(len(score), dtype=np.int8)
        pred[np.argsort(score)[-int(round(len(score) * ratio)):]] = 1
        pred[known] = 1
        path = out_dir / f"submission_post95_content_mf_lgb_r{ratio:.3f}.csv"
        pd.DataFrame({"Index": np.arange(len(pred), dtype=np.int64), "Predicted": pred}).to_csv(path, index=False)
        print(path, int(pred.sum()), float(pred.mean()))
    pred = (score >= 0.5).astype(np.int8)
    pred[known] = 1
    path = out_dir / "submission_post95_content_mf_lgb_score_ge0.500.csv"
    pd.DataFrame({"Index": np.arange(len(pred), dtype=np.int64), "Predicted": pred}).to_csv(path, index=False)
    print(path, int(pred.sum()), float(pred.mean()))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--package-root", type=Path, default=Path(__file__).resolve().parents[1])
    parser.add_argument("--split-seed", type=int, default=202)
    parser.add_argument("--main-val-score-file", type=Path, required=True)
    parser.add_argument("--device", default="cuda:0" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--seed", type=int, default=202)
    parser.add_argument("--mf-dim", type=int, default=256)
    parser.add_argument("--mf-epochs", type=int, default=220)
    parser.add_argument("--ratios", nargs="*", type=float, default=[0.498, 0.500, 0.502, 0.504, 0.505])
    args = parser.parse_args()

    root = args.package_root
    stack = load_module("stack", root / "code" / "stack_rank_calibration.py")
    lgcn = load_module("lgcn", root / "code" / "train_val_lgcn_ensemble.py")
    post = load_module("post", root / "code" / "post95_ablation.py")
    gen = load_module("gen", root / "code" / "generate_post95_submission.py")
    extra = load_module("extra", root / "code" / "extra_score_sources_ablation.py")

    out_dir = root / "validation_runs" / f"dynamic_seed{args.split_seed}" / "extra_bprmf_submission"
    out_dir.mkdir(parents=True, exist_ok=True)

    train_refs, val_pairs = lgcn.make_notebook_style_split(root, args.split_seed, 0.9)
    val_pairs_arr = val_pairs[["source", "target"]].to_numpy(np.int64)
    y = val_pairs["label"].to_numpy(np.int8)
    main_val = np.load(args.main_val_score_file).astype(np.float32)

    val_builder = stack.ExplicitGraphFeatures(root, train_refs)
    Xh = val_builder.transform(val_pairs_arr)
    X_val = np.column_stack(
        [
            stack.add_rank_features(val_pairs_arr, main_val),
            Xh,
            post.negative_evidence_features(Xh, main_val),
            gen.topk_content_similarity_fast(root, val_pairs_arr, val_builder),
        ]
    ).astype(np.float32)
    selected = [Path(x.strip()) for x in (root / "validation_runs" / f"dynamic_seed{args.split_seed}" / "post95_submission" / "selected_variant_val_scores.txt").read_text().splitlines() if x.strip()]
    X_val = np.column_stack([X_val, gen.variant_feature_matrix(post, [np.load(p).astype(np.float32) for p in selected])]).astype(np.float32)
    content_val = extra.content_mean_score(root, val_pairs_arr, val_builder)
    mf_val = np.load(root / "validation_runs" / f"dynamic_seed{args.split_seed}" / "extra_score_sources" / f"val_mf_bpr_s{args.seed}_d{args.mf_dim}.npy").astype(np.float32)
    Xc, _ = extra.score_to_features(content_val, "content_mean_cos", val_pairs_arr)
    Xm, _ = extra.score_to_features(mf_val, "mf_bpr", val_pairs_arr)
    X_val = np.column_stack([X_val, Xc, Xm]).astype(np.float32)
    print("fit LightGBM", X_val.shape)
    clf = lgb.LGBMClassifier(
        n_estimators=1200,
        learning_rate=0.025,
        num_leaves=31,
        subsample=0.9,
        colsample_bytree=0.9,
        reg_lambda=5.0,
        min_child_samples=80,
        objective="binary",
        verbose=-1,
        random_state=args.seed,
    )
    clf.fit(X_val, y)

    test_pairs = np.array(gen.read_txt(root / "data_and_docs" / "bipartite_test_ann.txt"), dtype=np.int64)
    main_test = np.load(root / "validation_runs" / f"dynamic_seed{args.split_seed}" / "post95_test_scores" / "dyn202_l2d512_bpr_bigbatch_more" / "scores" / "test_vanilla_ensemble_mean.npy").astype(np.float32)
    full_refs = pd.DataFrame(gen.read_txt(root / "data_and_docs" / "bipartite_train_ann.txt"), columns=["source", "target"])
    test_builder = stack.ExplicitGraphFeatures(root, full_refs)
    Xht = test_builder.transform(test_pairs)
    X_test = np.column_stack(
        [
            stack.add_rank_features(test_pairs, main_test),
            Xht,
            post.negative_evidence_features(Xht, main_test),
            gen.topk_content_similarity_fast(root, test_pairs, test_builder),
        ]
    ).astype(np.float32)
    test_scores = []
    for p in selected:
        rel = p.resolve().relative_to(root / "validation_runs" / f"dynamic_seed{args.split_seed}")
        tp = root / "validation_runs" / f"dynamic_seed{args.split_seed}" / "post95_test_scores" / rel.parent / rel.name.replace("val_", "test_", 1)
        test_scores.append(np.load(tp).astype(np.float32))
    X_test = np.column_stack([X_test, gen.variant_feature_matrix(post, test_scores)]).astype(np.float32)
    content_test = extra.content_mean_score(root, test_pairs, test_builder)
    mf_test = train_mf_test_scores(extra, root, train_refs, test_pairs, out_dir, args.device, args.seed, args.mf_dim, args.mf_epochs)
    Xct, _ = extra.score_to_features(content_test, "content_mean_cos", test_pairs)
    Xmt, _ = extra.score_to_features(mf_test, "mf_bpr", test_pairs)
    X_test = np.column_stack([X_test, Xct, Xmt]).astype(np.float32)
    print("predict", X_test.shape)
    pred_score = clf.predict_proba(X_test)[:, 1].astype(np.float32)
    np.save(out_dir / "test_post95_content_mf_lgb_pred.npy", pred_score)
    make_subs(root, out_dir, pred_score, args.ratios)


if __name__ == "__main__":
    main()