inference.py

#!/usr/bin/env python3
"""
Noise inference for DiffuQKT.

This script evaluates robustness under 16 perturbation settings:
1) clean baseline
2) drop_0.1 ... drop_0.5
3) c2w_0.1 ... c2w_0.5   (correct -> wrong)
4) w2c_0.1 ... w2c_0.5   (wrong -> correct)

It can also export the summary to TXT and/or CSV.
"""

import argparse
import csv
import math
import os
import random

import numpy as np
import pandas as pd
import torch
from sklearn import metrics
from tqdm import tqdm

from load_data import getLoader
from model import DiffuQKT


DATASET_CONFIG = {
    "assist09": {
        "folder": "ASSIST09",
        "ques_skill": "ques_skill.csv",
        "train_question": "train_question.txt",
        "test_question": "test_question.txt",
        "train_skill": "train_skill.txt",
        "test_skill": "test_skill.txt",
    },
    "assist17": {
        "folder": "ASSIST17",
        "ques_skill": "ques_skill.csv",
        "train_question": "train_question.txt",
        "test_question": "test_question.txt",
        "train_skill": "train_skill.txt",
        "test_skill": "test_skill.txt",
    },
    "xes3g5m": {
        "folder": "XES3G5M",
        "ques_skill": "ques_skill.csv",
        "train_question": "train_question.txt",
        "test_question": "test_question.txt",
        "train_skill": "train_skill.txt",
        "test_skill": "test_skill.txt",
    },
    "statics2011": {
        "folder": "STATICS2011",
        "ques_skill": "ques_skill.csv",
        "train_question": "train_question.txt",
        "test_question": "test_question.txt",
        "train_skill": "train_skill.txt",
        "test_skill": "test_skill.txt",
    },
}


def parse_args():
    parser = argparse.ArgumentParser(description="Noise Inference for DiffuQKT")
    parser.add_argument("--model_path", type=str, required=True, help="Path to model weights")
    parser.add_argument(
        "--dataset",
        type=str,
        required=True,
        choices=sorted(DATASET_CONFIG.keys()),
        help="Dataset name",
    )
    parser.add_argument(
        "--data_dir",
        type=str,
        default="",
        help="Base data directory. Auto-detects ./data when omitted.",
    )
    parser.add_argument("--split", type=str, default="test", choices=["train", "test"])
    parser.add_argument("--device", type=str, default="cuda:0")
    parser.add_argument("--batch_size", type=int, default=512)
    parser.add_argument("--max_seq", type=int, default=200)
    parser.add_argument("--min_seq", type=int, default=3)
    parser.add_argument("--seed", type=int, default=33)

    parser.add_argument("--d", type=int, default=128)
    parser.add_argument("--dropout", type=float, default=0.4)
    parser.add_argument("--head", type=int, default=2)
    parser.add_argument("--beta_start", type=float, default=0.0001)
    parser.add_argument("--beta_end", type=float, default=0.02)
    parser.add_argument("--diff_num_step", type=int, default=100)
    parser.add_argument("--lamda_1", type=float, default=0.01)
    parser.add_argument("--lamda_2", type=float, default=0.001)

    parser.add_argument(
        "--noise_levels",
        type=float,
        nargs="+",
        default=[0.0, 0.1, 0.2, 0.3, 0.4, 0.5],
        help="Noise levels (flip/drop ratios) to test",
    )
    parser.add_argument("--output_csv", type=str, default="", help="Optional CSV export path")
    parser.add_argument("--output_txt", type=str, default="", help="Optional TXT export path")
    return parser.parse_args()


def set_seed(seed: int):
    os.environ["PYTHONHASHSEED"] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    try:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    except Exception:
        pass


def resolve_dataset_paths(data_dir: str, dataset: str):
    if dataset not in DATASET_CONFIG:
        raise ValueError(f"Unknown dataset: {dataset}")

    config = DATASET_CONFIG[dataset]
    script_dir = os.path.dirname(os.path.abspath(__file__))
    candidate_roots = []

    if data_dir:
        candidate_roots.append(os.path.abspath(data_dir))
    else:
        candidate_roots.extend(
            [
                os.path.join(script_dir, "data"),
                os.path.join(script_dir, "./data"),
                os.path.join(os.path.dirname(script_dir), "data"),
                "data",
                "../data",
            ]
        )

    for root in candidate_roots:
        dataset_root = os.path.join(os.path.abspath(root), config["folder"])
        ques_skill_path = os.path.join(dataset_root, config["ques_skill"])
        if os.path.isfile(ques_skill_path):
            return {
                "root": dataset_root,
                "ques_skill_path": ques_skill_path,
                "train_path": os.path.join(dataset_root, config["train_question"]),
                "test_path": os.path.join(dataset_root, config["test_question"]),
                "train_skill_path": os.path.join(dataset_root, config["train_skill"]),
                "test_skill_path": os.path.join(dataset_root, config["test_skill"]),
            }

    raise FileNotFoundError(
        f"Could not find dataset '{dataset}' under any of: "
        + ", ".join(os.path.abspath(path) for path in candidate_roots)
    )


def _pick_indices(candidates: np.ndarray, noise_level: float):
    if candidates.size == 0:
        return np.array([], dtype=np.int64)
    num_to_change = max(1, int(candidates.size * float(noise_level)))
    num_to_change = min(num_to_change, int(candidates.size))
    return np.random.choice(candidates, size=num_to_change, replace=False)


def _rebuild_from_full_sequences(
    full_problem,
    full_skill,
    full_ans_input,
    full_ans_target,
    max_seq_len,
    device,
):
    seq_len = int(max(0, full_problem.shape[0] - 1))
    last_problem = torch.zeros((max_seq_len,), dtype=torch.long, device=device)
    last_skill = torch.zeros((max_seq_len,), dtype=torch.long, device=device)
    last_ans = torch.zeros((max_seq_len,), dtype=torch.long, device=device)
    next_problem = torch.zeros((max_seq_len,), dtype=torch.long, device=device)
    next_skill = torch.zeros((max_seq_len,), dtype=torch.long, device=device)
    next_ans_input = torch.zeros((max_seq_len,), dtype=torch.float32, device=device)
    next_ans_target = torch.zeros((max_seq_len,), dtype=torch.float32, device=device)
    mask = torch.zeros((max_seq_len,), dtype=torch.bool, device=device)

    if seq_len <= 0:
        return last_problem, last_skill, last_ans, next_problem, next_skill, next_ans_input, next_ans_target, mask

    start = max_seq_len - seq_len
    last_problem[start:] = full_problem[:-1]
    last_skill[start:] = full_skill[:-1]
    last_ans[start:] = full_ans_input[:-1]
    next_problem[start:] = full_problem[1:]
    next_skill[start:] = full_skill[1:]
    next_ans_input[start:] = full_ans_input[1:].float()
    next_ans_target[start:] = full_ans_target[1:].float()
    mask[start:] = True
    return last_problem, last_skill, last_ans, next_problem, next_skill, next_ans_input, next_ans_target, mask


def perturb_batch(
    last_problem,
    last_skill,
    last_ans,
    next_problem,
    next_skill,
    next_ans,
    mask,
    noise_level,
    noise_mode,
):
    noise_mode = str(noise_mode).lower().strip()
    if noise_mode == "clean" or float(noise_level) <= 0.0:
        return (
            last_problem,
            last_skill,
            last_ans,
            next_problem,
            next_skill,
            next_ans,
            next_ans,
            mask,
        )

    batch_size, seq_width = last_problem.shape
    device = last_problem.device

    out_last_problem = torch.zeros_like(last_problem)
    out_last_skill = torch.zeros_like(last_skill)
    out_last_ans = torch.zeros_like(last_ans)
    out_next_problem = torch.zeros_like(next_problem)
    out_next_skill = torch.zeros_like(next_skill)
    out_next_ans_input = torch.zeros_like(next_ans)
    out_next_ans_target = torch.zeros_like(next_ans)
    out_mask = torch.zeros_like(mask, dtype=torch.bool)

    for b in range(batch_size):
        valid_idx = torch.nonzero(mask[b], as_tuple=True)[0]
        if valid_idx.numel() == 0:
            continue

        first_idx = int(valid_idx[0].item())
        full_problem = torch.cat(
            [last_problem[b, first_idx].view(1), next_problem[b, valid_idx]],
            dim=0,
        ).clone()
        full_skill = torch.cat(
            [last_skill[b, first_idx].view(1), next_skill[b, valid_idx]],
            dim=0,
        ).clone()
        full_ans_original = torch.cat(
            [last_ans[b, first_idx].view(1), next_ans[b, valid_idx].long()],
            dim=0,
        ).clone()
        full_ans_input = full_ans_original.clone()
        full_ans_target = full_ans_original.clone()

        if noise_mode == "drop":
            seq_len = int(full_problem.shape[0])
            if seq_len <= 2:
                keep = np.ones((seq_len,), dtype=bool)
            else:
                num_to_drop = max(1, int(seq_len * float(noise_level)))
                num_to_drop = min(num_to_drop, seq_len - 2)
                drop_idx = np.random.choice(seq_len, size=num_to_drop, replace=False)
                keep = np.ones((seq_len,), dtype=bool)
                keep[drop_idx] = False
            keep_idx = torch.as_tensor(np.flatnonzero(keep), dtype=torch.long, device=device)
            full_problem = full_problem.index_select(0, keep_idx)
            full_skill = full_skill.index_select(0, keep_idx)
            full_ans_input = full_ans_input.index_select(0, keep_idx)
            full_ans_target = full_ans_target.index_select(0, keep_idx)
        else:
            answers_np = full_ans_input.detach().cpu().numpy()
            if noise_mode == "flip":
                candidates = np.arange(answers_np.shape[0])
            elif noise_mode == "c2w":
                candidates = np.flatnonzero(answers_np > 0)
            elif noise_mode == "w2c":
                candidates = np.flatnonzero(answers_np <= 0)
            else:
                raise ValueError(f"Unsupported noise_mode: {noise_mode}")

            change_idx = _pick_indices(candidates, noise_level)
            if change_idx.size > 0:
                change_idx_tensor = torch.as_tensor(change_idx, dtype=torch.long, device=device)
                full_ans_input.index_put_((change_idx_tensor,), 1 - full_ans_input.index_select(0, change_idx_tensor))

        rebuilt = _rebuild_from_full_sequences(
            full_problem=full_problem,
            full_skill=full_skill,
            full_ans_input=full_ans_input.long(),
            full_ans_target=full_ans_target.long(),
            max_seq_len=seq_width,
            device=device,
        )
        (
            out_last_problem[b],
            out_last_skill[b],
            out_last_ans[b],
            out_next_problem[b],
            out_next_skill[b],
            out_next_ans_input[b],
            out_next_ans_target[b],
            out_mask[b],
        ) = rebuilt

    return (
        out_last_problem,
        out_last_skill,
        out_last_ans,
        out_next_problem,
        out_next_skill,
        out_next_ans_input,
        out_next_ans_target,
        out_mask,
    )


def run_inference_with_noise(model, data_loader, noise_level: float, seed: int, noise_mode: str = "clean"):
    set_seed(seed)
    model.eval()
    y_true_all = []
    y_pred_all = []
    device = next(model.parameters()).device

    with torch.no_grad():
        for data in tqdm(data_loader, desc=f"{noise_mode}_{noise_level:.1f}", leave=False):
            last_problem, last_skill, last_ans, next_problem, next_skill, next_ans, mask = data
            last_problem = last_problem.to(device)
            last_skill = last_skill.to(device)
            last_ans = last_ans.to(device)
            next_problem = next_problem.to(device)
            next_skill = next_skill.to(device)
            next_ans = next_ans.to(device)
            mask = mask.to(device)
            (
                noisy_last_problem,
                noisy_last_skill,
                noisy_last_ans,
                noisy_next_problem,
                noisy_next_skill,
                noisy_next_ans_input,
                noisy_next_ans_target,
                noisy_mask,
            ) = perturb_batch(
                last_problem,
                last_skill,
                last_ans,
                next_problem,
                next_skill,
                next_ans,
                mask,
                noise_level,
                noise_mode,
            )

            predict, _, _ = model(
                noisy_last_problem,
                noisy_last_skill,
                noisy_last_ans,
                noisy_next_problem,
                noisy_next_skill,
                noisy_next_ans_input,
                noisy_mask,
                is_train=False,
            )

            next_predict = torch.masked_select(predict, noisy_mask)
            next_true = torch.masked_select(noisy_next_ans_target, noisy_mask)
            y_true_all.append(next_true.detach().cpu().contiguous().view(-1).numpy())
            y_pred_all.append(next_predict.detach().cpu().contiguous().view(-1).numpy())

    y_true = np.concatenate(y_true_all, axis=0) if y_true_all else np.array([])
    y_pred = np.concatenate(y_pred_all, axis=0) if y_pred_all else np.array([])
    return y_true, y_pred


def calc_metrics(y_true, y_pred):
    out = {"auc": float("nan"), "acc": float("nan"), "rmse": float("nan")}
    if y_true.size == 0:
        return out

    try:
        out["auc"] = float(metrics.roc_auc_score(y_true, y_pred))
    except Exception:
        pass

    try:
        out["acc"] = float(((y_pred >= 0.5) == (y_true >= 0.5)).mean())
    except Exception:
        pass

    try:
        out["rmse"] = float(math.sqrt(np.mean((y_pred - y_true) ** 2)))
    except Exception:
        pass

    return out


def write_results_csv(path, results_list):
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    with open(path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["group", "mode", "noise_level", "auc", "acc", "rmse", "num_samples"])
        for result in results_list:
            writer.writerow(
                [
                    result["group"],
                    result["mode"],
                    result["noise_level"],
                    f"{result['metrics']['auc']:.6f}",
                    f"{result['metrics']['acc']:.6f}",
                    f"{result['metrics']['rmse']:.6f}",
                    result["num_samples"],
                ]
            )


def write_results_txt(path, results_list, model_path: str, dataset: str, split: str):
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        f.write(f"Model: {model_path}\n")
        f.write(f"Dataset: {dataset}\n")
        f.write(f"Split: {split}\n")
        f.write("=" * 96 + "\n")
        f.write(f"{'Group':<18} {'Mode':<10} {'Level':<8} {'AUC':<14} {'ACC':<14} {'RMSE':<14} {'Samples':<8}\n")
        f.write("=" * 96 + "\n")
        for result in results_list:
            metrics_row = result["metrics"]
            f.write(
                f"{result['group']:<18} "
                f"{result['mode']:<10} "
                f"{result['noise_level']:<8.1f} "
                f"{metrics_row['auc']:<14.6f} "
                f"{metrics_row['acc']:<14.6f} "
                f"{metrics_row['rmse']:<14.6f} "
                f"{result['num_samples']:<8}\n"
            )


def main():
    args = parse_args()
    set_seed(int(args.seed))

    if not os.path.isfile(args.model_path):
        raise FileNotFoundError(f"Model not found: {args.model_path}")

    dataset_paths = resolve_dataset_paths(args.data_dir, args.dataset)
    print(f"Using dataset root: {dataset_paths['root']}")
    print(f"Looking for dataset: {args.dataset}")

    requested_device = torch.device(args.device)
    if requested_device.type == "cuda" and not torch.cuda.is_available():
        print("CUDA is not available, falling back to CPU.")
        device = torch.device("cpu")
    else:
        device = requested_device

    try:
        state_dict = torch.load(args.model_path, map_location=device, weights_only=True)
    except TypeError:
        state_dict = torch.load(args.model_path, map_location=device)

    ques_skill_df = pd.read_csv(dataset_paths["ques_skill_path"])
    skill_max = int(ques_skill_df["skill_id"].max()) + 2
    pro_max = int(ques_skill_df["problem_id"].max()) + 2

    model = DiffuQKT(
        pro_max=pro_max,
        skill_max=skill_max,
        d=args.d,
        p=args.dropout,
        head=args.head,
        beta_start=args.beta_start,
        beta_end=args.beta_end,
        diff_num_step=args.diff_num_step,
        lamda_1=args.lamda_1,
        lamda_2=args.lamda_2,
    ).to(device)
    model.load_state_dict(state_dict, strict=True)
    model.eval()

    question_key = f"{args.split}_path"
    skill_key = f"{args.split}_skill_path"
    data_loader = getLoader(
        pro_max,
        dataset_paths[question_key],
        dataset_paths[skill_key],
        args.batch_size,
        False,
        args.min_seq,
        args.max_seq,
    )

    print(f"Model: {args.model_path}")
    print(f"Dataset: {args.dataset}")
    print(f"Split: {args.split}")
    print("=" * 88)
    print(f"{'Group':<18} {'Level':<8} {'AUC':<14} {'ACC':<14} {'RMSE':<14} {'Samples':<8}")
    print("=" * 88)

    positive_levels = [float(x) for x in sorted(args.noise_levels) if float(x) > 0.0]
    eval_plan = [("clean", "clean", 0.0)]
    eval_plan.extend((f"drop_{lvl:.1f}", "drop", lvl) for lvl in positive_levels)
    eval_plan.extend((f"c2w_{lvl:.1f}", "c2w", lvl) for lvl in positive_levels)
    eval_plan.extend((f"w2c_{lvl:.1f}", "w2c", lvl) for lvl in positive_levels)

    results_list = []
    for group_name, noise_mode, noise_level in eval_plan:
        y_true, y_pred = run_inference_with_noise(
            model=model,
            data_loader=data_loader,
            noise_level=noise_level,
            seed=args.seed,
            noise_mode=noise_mode,
        )
        metrics_row = calc_metrics(y_true, y_pred)
        result = {
            "group": group_name,
            "mode": noise_mode,
            "noise_level": noise_level,
            "metrics": metrics_row,
            "num_samples": int(y_true.size),
        }
        results_list.append(result)
        print(
            f"{group_name:<18} {noise_level:<8.1f} "
            f"{metrics_row['auc']:<14.6f} {metrics_row['acc']:<14.6f} "
            f"{metrics_row['rmse']:<14.6f} {y_true.size:<8}"
        )

    print("=" * 88)

    if args.output_csv:
        write_results_csv(args.output_csv, results_list)
        print(f"Results saved to: {args.output_csv}")

    if args.output_txt:
        write_results_txt(args.output_txt, results_list, args.model_path, args.dataset, args.split)
        print(f"TXT report saved to: {args.output_txt}")


if __name__ == "__main__":
    main()