code/extra_score_sources_ablation.py

"""Ablate extra non-LightGCN score sources for the post95 stacker."""

from __future__ import annotations

import argparse
import importlib.util
import pickle as pkl
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import precision_recall_curve, roc_auc_score
from sklearn.model_selection import GroupKFold, StratifiedKFold


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 best_f1(y: np.ndarray, s: np.ndarray):
    p, r, t = precision_recall_curve(y, s)
    f = 2 * p * r / (p + r + 1e-12)
    i = int(np.argmax(f))
    th = float(t[i]) if i < len(t) else 0.5
    return float(f[i]), th, float(roc_auc_score(y, s)), float(p[i]), float(r[i])


def rank01(x: np.ndarray) -> np.ndarray:
    order = np.argsort(x, kind="mergesort")
    out = np.empty(len(x), dtype=np.float32)
    out[order] = np.linspace(0.0, 1.0, len(x), dtype=np.float32)
    return out


def zscore(x: np.ndarray) -> np.ndarray:
    return ((x - x.mean()) / (x.std() + 1e-8)).astype(np.float32)


def score_to_features(scores: np.ndarray, prefix: str, pairs: np.ndarray) -> tuple[np.ndarray, list[str]]:
    author_rank = np.zeros(len(scores), dtype=np.float32)
    df = pd.DataFrame({"idx": np.arange(len(scores)), "author": pairs[:, 0], "score": scores})
    for _, g in df.groupby("author", sort=False):
        idx = g["idx"].to_numpy()
        order = np.argsort(g["score"].to_numpy(), kind="mergesort")
        vals = np.linspace(0, 1, len(idx), dtype=np.float32) if len(idx) > 1 else np.array([1.0], dtype=np.float32)
        author_rank[idx[order]] = vals
    X = np.column_stack([scores.astype(np.float32), zscore(scores), rank01(scores), author_rank]).astype(np.float32)
    return X, [prefix, f"{prefix}_z", f"{prefix}_rank", f"{prefix}_author_rank"]


def fit_lgb_oof(X: np.ndarray, y: np.ndarray, seed: int, n_splits: int) -> np.ndarray:
    oof = np.zeros(len(y), dtype=np.float32)
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
    for fold, (tr, va) in enumerate(skf.split(X, y), start=1):
        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=seed + fold,
        )
        clf.fit(X[tr], y[tr])
        oof[va] = clf.predict_proba(X[va])[:, 1]
    return oof


def content_mean_score(root: Path, pairs: np.ndarray, builder) -> np.ndarray:
    cache = root / "validation_runs" / "feature_cache"
    cache.mkdir(parents=True, exist_ok=True)
    path = cache / f"content_mean_{len(pairs)}_{int(pairs[:,0].sum())}_{int(pairs[:,1].sum())}.npy"
    if path.exists():
        return np.load(path)
    with (root / "data_and_docs" / "feature.pkl").open("rb") as f:
        feat = pkl.load(f).numpy().astype(np.float32)
    feat /= np.linalg.norm(feat, axis=1, keepdims=True) + 1e-8
    author_mean = np.zeros((builder.num_authors, feat.shape[1]), dtype=np.float32)
    for a in range(builder.num_authors):
        hist = list(builder.author_papers[a])
        if hist:
            v = feat[np.asarray(hist, dtype=np.int64)].mean(axis=0)
            author_mean[a] = v / (np.linalg.norm(v) + 1e-8)
    out = np.sum(author_mean[pairs[:, 0]] * feat[pairs[:, 1]], axis=1).astype(np.float32)
    np.save(path, out)
    return out


class MF(nn.Module):
    def __init__(self, n_author: int, n_paper: int, dim: int):
        super().__init__()
        self.a = nn.Embedding(n_author, dim)
        self.p = nn.Embedding(n_paper, dim)
        self.ab = nn.Embedding(n_author, 1)
        self.pb = nn.Embedding(n_paper, 1)
        nn.init.normal_(self.a.weight, std=0.05)
        nn.init.normal_(self.p.weight, std=0.05)
        nn.init.zeros_(self.ab.weight)
        nn.init.zeros_(self.pb.weight)

    def score(self, pairs):
        return (self.a(pairs[:, 0]) * self.p(pairs[:, 1])).sum(-1) + self.ab(pairs[:, 0]).squeeze(-1) + self.pb(pairs[:, 1]).squeeze(-1)


def train_mf_bpr_score(root: Path, train_refs: pd.DataFrame, val_pairs: pd.DataFrame, out_dir: Path, device: str, seed: int, dim: int = 256, epochs: int = 220) -> np.ndarray:
    out_path = out_dir / f"val_mf_bpr_s{seed}_d{dim}.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)
    val = val_pairs[["source", "target"]].to_numpy(np.int64)
    y = val_pairs["label"].to_numpy(np.int8)
    train_set = set(map(tuple, train.tolist()))
    model = 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)
    val_t = torch.as_tensor(val, dtype=torch.long, device=device)
    batch_size = 65536
    best = (-1.0, None)
    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 or ep == epochs - 1:
            with torch.no_grad():
                scores = []
                for st in range(0, len(val), 131072):
                    scores.append(model.score(val_t[st : st + 131072]).detach().cpu().numpy())
                scores = np.concatenate(scores).astype(np.float32)
            f1, th, auc, _, _ = best_f1(y, scores)
            if f1 > best[0]:
                best = (f1, scores.copy())
            print(f"mf epoch={ep+1:03d} loss={loss.item():.4f} f1={f1:.6f} th={th:.6f} auc={auc:.6f}")
    np.save(out_path, best[1])
    return best[1]


def train_ranker_oof(X: np.ndarray, y: np.ndarray, pairs: np.ndarray, seed: int, out_dir: Path) -> np.ndarray:
    out_path = out_dir / "val_lgbmranker_oof.npy"
    if out_path.exists():
        return np.load(out_path)
    oof = np.zeros(len(y), dtype=np.float32)
    gkf = GroupKFold(n_splits=5)
    groups = pairs[:, 0]
    for fold, (tr, va) in enumerate(gkf.split(X, y, groups=groups), start=1):
        tr_order = np.lexsort((np.arange(len(tr)), pairs[tr, 0]))
        va_order = np.lexsort((np.arange(len(va)), pairs[va, 0]))
        tr_idx = tr[tr_order]
        va_idx = va[va_order]
        tr_group = pd.Series(pairs[tr_idx, 0]).value_counts(sort=False).to_numpy()
        ranker = lgb.LGBMRanker(
            objective="lambdarank",
            metric="ndcg",
            n_estimators=700,
            learning_rate=0.03,
            num_leaves=31,
            subsample=0.9,
            colsample_bytree=0.9,
            reg_lambda=10.0,
            min_child_samples=60,
            random_state=seed + fold,
            verbose=-1,
        )
        ranker.fit(X[tr_idx], y[tr_idx], group=tr_group)
        oof[va_idx] = ranker.predict(X[va_idx]).astype(np.float32)
        print(f"ranker fold={fold} f1={best_f1(y[va_idx], oof[va_idx])[0]:.6f}")
    np.save(out_path, oof)
    return oof


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("--n-splits", type=int, default=5)
    parser.add_argument("--skip-mf", action="store_true")
    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")
    train_refs, val_pairs = lgcn.make_notebook_style_split(root, args.split_seed, 0.9)
    pairs = val_pairs[["source", "target"]].to_numpy(np.int64)
    y = val_pairs["label"].to_numpy(np.int8)
    main = np.load(args.main_val_score_file).astype(np.float32)
    out_dir = root / "validation_runs" / f"dynamic_seed{args.split_seed}" / "extra_score_sources"
    out_dir.mkdir(parents=True, exist_ok=True)

    print("building post95 base features")
    builder = stack.ExplicitGraphFeatures(root, train_refs)
    X_hand = builder.transform(pairs)
    X = np.column_stack(
        [
            stack.add_rank_features(pairs, main),
            X_hand,
            post.negative_evidence_features(X_hand, main),
            gen.topk_content_similarity_fast(root, pairs, 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 = np.column_stack([X, gen.variant_feature_matrix(post, [np.load(p).astype(np.float32) for p in selected])]).astype(np.float32)

    rows = []
    base_oof = fit_lgb_oof(X, y, args.seed, args.n_splits)
    f1, th, auc, p, r = best_f1(y, base_oof)
    rows.append({"stage": "post95_lgbm_baseline", "f1": f1, "threshold": th, "auc": auc, "precision": p, "recall": r, "n_features": X.shape[1]})
    np.save(out_dir / "post95_lgbm_baseline_oof.npy", base_oof)

    extra_blocks = []
    extra_names = []

    print("adding pure content mean-cos score")
    content = content_mean_score(root, pairs, builder)
    Xc, names = score_to_features(content, "content_mean_cos", pairs)
    extra_blocks.append(Xc)
    extra_names.extend(names)
    X_cur = np.column_stack([X, *extra_blocks]).astype(np.float32)
    oof = fit_lgb_oof(X_cur, y, args.seed + 10, args.n_splits)
    f1, th, auc, p, r = best_f1(y, oof)
    rows.append({"stage": "+content_mean_cos", "f1": f1, "threshold": th, "auc": auc, "precision": p, "recall": r, "n_features": X_cur.shape[1]})
    np.save(out_dir / "content_mean_cos_stack_oof.npy", oof)

    if not args.skip_mf:
        print("training/adding BPR-MF score")
        mf = train_mf_bpr_score(root, train_refs, val_pairs, out_dir, args.device, args.seed)
        Xm, names = score_to_features(mf, "mf_bpr", pairs)
        extra_blocks.append(Xm)
        extra_names.extend(names)
        X_cur = np.column_stack([X, *extra_blocks]).astype(np.float32)
        oof = fit_lgb_oof(X_cur, y, args.seed + 20, args.n_splits)
        f1, th, auc, p, r = best_f1(y, oof)
        rows.append({"stage": "+bpr_mf", "f1": f1, "threshold": th, "auc": auc, "precision": p, "recall": r, "n_features": X_cur.shape[1]})
        np.save(out_dir / "bpr_mf_stack_oof.npy", oof)

    print("training/adding author-group LGBMRanker OOF score")
    ranker_scores = train_ranker_oof(X, y, pairs, args.seed, out_dir)
    Xr, names = score_to_features(ranker_scores, "lgbmranker_author_oof", pairs)
    extra_blocks.append(Xr)
    extra_names.extend(names)
    X_cur = np.column_stack([X, *extra_blocks]).astype(np.float32)
    oof = fit_lgb_oof(X_cur, y, args.seed + 30, args.n_splits)
    f1, th, auc, p, r = best_f1(y, oof)
    rows.append({"stage": "+lgbmranker_author", "f1": f1, "threshold": th, "auc": auc, "precision": p, "recall": r, "n_features": X_cur.shape[1]})
    np.save(out_dir / "lgbmranker_stack_oof.npy", oof)

    pd.Series(extra_names).to_csv(out_dir / "extra_feature_names.csv", index=False)
    result = pd.DataFrame(rows).sort_values("f1", ascending=False)
    result.to_csv(out_dir / "extra_score_ablation.csv", index=False)
    print(result.to_string(index=False))


if __name__ == "__main__":
    main()