random_flip_inference.py

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

This script evaluates a trained model under answer-flip noise only:
0.0, 0.1, 0.2, 0.3, 0.4, 0.5 by default.
"""

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 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="Random Flip 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="Random flip ratios to test",
    )
    parser.add_argument(
        "--output_csv",
        type=str,
        default="",
        help="CSV output path for results; defaults to this script directory when not specified",
    )
    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(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 resolve_output_csv_path(output_csv: str, dataset: str, split: str):
    if output_csv:
        return output_csv

    script_dir = os.path.dirname(os.path.abspath(__file__))
    filename = f"flip_results_{dataset}_{split}.csv"
    return os.path.join(script_dir, filename)


def perturb_answers_random_flip(last_ans, next_ans, mask, noise_level: float):
    if float(noise_level) <= 0.0:
        return last_ans.clone(), next_ans.clone()

    noisy_last_ans = last_ans.clone().long()
    noisy_next_ans = next_ans.clone().float()

    batch_size = noisy_last_ans.shape[0]
    seq_width = noisy_last_ans.shape[1] + 1
    device = noisy_last_ans.device

    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_answers = torch.cat(
            [noisy_last_ans[b, first_idx].view(1), noisy_next_ans[b, valid_idx].long()],
            dim=0,
        ).clone()

        actual_len = int(full_answers.shape[0])
        if actual_len <= 0:
            continue

        candidate = np.arange(actual_len)
        num_to_flip = max(1, int(candidate.size * float(noise_level)))
        num_to_flip = min(num_to_flip, int(candidate.size))
        flip_indices = np.random.choice(candidate, size=num_to_flip, replace=False)
        flip_indices = torch.as_tensor(flip_indices, dtype=torch.long, device=device)
        full_answers.index_put_((flip_indices,), 1 - full_answers.index_select(0, flip_indices))

        rebuilt_last = torch.zeros((seq_width - 1,), dtype=torch.long, device=device)
        rebuilt_next = torch.zeros((seq_width - 1,), dtype=torch.float32, device=device)

        if actual_len > 1:
            start = (seq_width - 1) - (actual_len - 1)
            rebuilt_last[start:] = full_answers[:-1]
            rebuilt_next[start:] = full_answers[1:].float()

        noisy_last_ans[b] = rebuilt_last
        noisy_next_ans[b] = rebuilt_next

    return noisy_last_ans, noisy_next_ans


def run_inference_with_random_flip(model, data_loader, noise_level: float, seed: int):
    set_seed(seed)
    model.eval()
    y_true_all = []
    y_pred_all = []
    device = next(model.parameters()).device

    with torch.no_grad():
        for data in data_loader:
            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_ans, noisy_next_ans = perturb_answers_random_flip(
                last_ans=last_ans,
                next_ans=next_ans,
                mask=mask,
                noise_level=noise_level,
            )
            predict, _, _ = model(
                last_problem,
                last_skill,
                noisy_last_ans,
                next_problem,
                next_skill,
                noisy_next_ans,
                mask,
                is_train=False,
            )

            next_predict = torch.masked_select(predict, mask)
            next_true = torch.masked_select(next_ans, 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(["noise_level", "auc", "acc", "rmse", "num_samples"])
        for result in results_list:
            writer.writerow(
                [
                    result["noise_level"],
                    f"{result['metrics']['auc']:.6f}",
                    f"{result['metrics']['acc']:.6f}",
                    f"{result['metrics']['rmse']:.6f}",
                    result["num_samples"],
                ]
            )


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("=" * 72)
    print(f"{'Noise Level':<12} {'AUC':<14} {'ACC':<14} {'RMSE':<14} {'Samples':<8}")
    print("=" * 72)

    results_list = []
    for noise_level in sorted(float(x) for x in args.noise_levels):
        y_true, y_pred = run_inference_with_random_flip(
            model=model,
            data_loader=data_loader,
            noise_level=noise_level,
            seed=args.seed,
        )
        metrics_row = calc_metrics(y_true, y_pred)
        result = {
            "noise_level": noise_level,
            "metrics": metrics_row,
            "num_samples": int(y_true.size),
        }
        results_list.append(result)
        print(
            f"{noise_level:<12.1f} "
            f"{metrics_row['auc']:<14.6f} "
            f"{metrics_row['acc']:<14.6f} "
            f"{metrics_row['rmse']:<14.6f} "
            f"{y_true.size:<8}"
        )

    print("=" * 72)

    output_csv_path = resolve_output_csv_path(args.output_csv, args.dataset, args.split)
    write_results_csv(output_csv_path, results_list)
    print(f"Results saved to: {output_csv_path}")


if __name__ == "__main__":
    main()