pipeline.py

import os
import re
import json
import pandas as pd
import numpy as np
from tabulate import tabulate

import torch

from dataloader.attacker_data import AttackerData
from dataloader.detector_data import DetectorData
from dataloader.recommender_data import RecommenderData
from models.attacker.Attacker_Factory import AttackModelFactory
from models.detector.Detector_Factory import DetectionModelFactory
from models.recommender.Recommeder_Factory import RecommendationModelFactory
from utils.config_utils import ConfigManager, print_args_info
from utils.common_utils import log_exists, update_log


def pipeline_init_cm(args):
    def override_global_config(config, args, exclude = ("scenario",)):
        for k, v in vars(args).items():
            if k in exclude:
                continue
            if hasattr(config, k):
                print(f"[Override] {k} = {v}")
                setattr(config, k, v)
            else:
                print(f"[Error] Global config has no key `{k}` — aborting.")
                exit(1)

    cm = ConfigManager(scenario = args.scenario)
    assert len(cm.scenario_config.attack_models) == len(
        cm.scenario_config.helper_models), "[ERROR] attack_models and helper_models must have the same length!"
    override_global_config(cm.global_config, args)
    return cm


def pipeline_injection(cm, flag):
    # injection
    if flag == True:
        print(f"++" * 50)
        print("\033[92m[Flag] Injection Flag is ON!!!\033[0m")
        print(f"++" * 50)
        for dataset in cm.scenario_config.attack_targets:
            for position in cm.scenario_config.attack_targets[dataset]:
                for target in cm.scenario_config.attack_targets[dataset][position]:
                    for i in range(len(cm.scenario_config.attack_models)):
                        # Configurations for A SINGLE injected instance.
                        print(f"==" * 50)

                        # attack args
                        attack_args = cm.generate_attack_injection_args(dataset = dataset, targets = target, i = i)
                        print_args_info(args = attack_args, text = "[INFO] Attack Configuration:")

                        # helper args
                        helper_args = cm.generate_helper_injection_args(dataset = dataset, i = i)
                        print_args_info(args = helper_args, text = "[INFO] Helper Configuration:")

                        # save config
                        if not hasattr(cm.global_config, "save_dir") or cm.global_config.save_dir is None:
                            print("[INFO] Skipping save: save_dir is not set in global.yaml.")
                            dir_name = None
                            target_str = None
                        else:
                            attack_model_str = "-".join(attack_args.model)
                            target_str = "-".join("_".join(map(str, sorted(set(group)))) for group in attack_args.targets)
                            dir_name = os.path.join(
                                cm.global_config.save_dir,
                                cm.scenario_config.name,
                                dataset,
                                attack_model_str,
                                position,
                            )

                            # check saved files
                            rating_path = os.path.join(dir_name, f"{target_str}_rating.txt")
                            label_path = os.path.join(dir_name, f"{target_str}_label.txt")
                            if os.path.exists(rating_path) and os.path.exists(label_path):
                                print(f"[INFO] Skipping injection: already exists at '{dir_name}'")
                                continue

                        # attack injection
                        print("[INFO] Attack Injection:")
                        data = AttackerData(attack_args.dataset)
                        for j in range(len(attack_args.model)):
                            args = cm.generate_single_injection_args(attack_args, j)
                            p_df, p_label = AttackModelFactory(data, args)
                            data.merge_sysnthetic_files(rating = p_df, label = p_label)
                        data.update()

                        # helper injection
                        print("[INFO] Helper Injection:")
                        for j in range(len(helper_args.model)):
                            args = cm.generate_single_injection_args(helper_args, j)
                            p_df, p_label = AttackModelFactory(data, args)
                            data.merge_sysnthetic_files(rating = p_df, label = p_label)
                        data.update()

                        # save
                        if all([dir_name, target_str]):
                            data.save(dir_name, target_str)

                        print(f"==" * 50)
    else:
        print(f"++" * 50)
    print("\033[91m[Flag] Injection Flag is OFF!!!\033[0m")
    print(f"++" * 50)


def pipeline_detection(cm, flag):
    # detection
    if flag == True:
        print(f"++" * 50)
        print("\033[92m[Flag] Detection Flag is ON!!!\033[0m")
        print(f"++" * 50)
        results_file = os.path.join(cm.global_config.save_dir, cm.scenario_config.name, "detector_results.json")
        if os.path.exists(results_file):
            with open(results_file, "r") as f:
                log = json.load(f)
        else:
            log = {}
        for dataset in cm.scenario_config.attack_targets:
            for position in cm.scenario_config.attack_targets[dataset]:
                for target in cm.scenario_config.attack_targets[dataset][position]:
                    for i in range(len(cm.scenario_config.attack_models)):
                        # last dir
                        attack_args = cm.generate_attack_injection_args(dataset = dataset, targets = target, i = i)
                        attack_model_str = "-".join(attack_args.model)
                        dir_name = os.path.join(
                            cm.global_config.save_dir,
                            cm.scenario_config.name,
                            dataset,
                            attack_model_str,
                            position,
                        )

                        # for every injected instance
                        print(f"==" * 50)
                        target_str = "-".join("_".join(map(str, sorted(set(group)))) for group in attack_args.targets)
                        log_keys_prefix = [cm.scenario_config.name, dataset, attack_model_str, position, target_str]

                        # data preparation
                        data_args = cm.generate_detector_data_args(dataset = dataset, i = i)
                        data_args.poisoned_trn_path = os.path.join(dir_name, f"{target_str}_rating.txt")
                        data_args.label_path = os.path.join(dir_name, f"{target_str}_label.txt")
                        data_args.targets = set(re.findall(r'\d+', target_str))
                        data = DetectorData(data_args)

                        # Detection models
                        for detector in cm.scenario_config.detector_models:
                            log_keys = log_keys_prefix + [detector]
                            if log_exists(log, log_keys):
                                print(f"[Skipped {detector}] Already Evaluated. {'/'.join(log_keys)}")
                                continue

                            detector_args = cm.generate_model_args(dataset = dataset, model = 'detector', model_name = detector)
                            conf_matrix, report = DetectionModelFactory(data, detector_args)
                            result = {
                                "conf_matrix": conf_matrix.tolist() if hasattr(conf_matrix, 'tolist') else str(conf_matrix),
                                "report": report
                            }  # results

                            torch.cuda.empty_cache()
                            torch.cuda.ipc_collect()

                            update_log(log, log_keys, result)
                            with open(results_file, "w") as f:
                                json.dump(log, f, indent = 2)

                        print(f"==" * 50)
    else:
        print(f"++" * 50)
        print("\033[91m[Flag] Detection Flag is OFF!!!\033[0m")
        print(f"++" * 50)


def pipeline_assessor(cm, flag):
    # assessor (RS Impact and Stealthiness)
    if flag == True:
        print(f"++" * 50)
        print("\033[92m[Flag] Assessor Flag is ON!!!\033[0m")
        print(f"++" * 50)
        results_file = os.path.join(cm.global_config.save_dir, cm.scenario_config.name, "assessor_results.json")
        if os.path.exists(results_file):
            with open(results_file, "r") as f:
                log = json.load(f)
        else:
            log = {}
        for dataset in cm.scenario_config.attack_targets:
            for position in cm.scenario_config.attack_targets[dataset]:
                for target in cm.scenario_config.attack_targets[dataset][position]:
                    for i in range(len(cm.scenario_config.attack_models)):
                        # last dir
                        attack_args = cm.generate_attack_injection_args(dataset = dataset, targets = target, i = i)
                        attack_model_str = "-".join(attack_args.model)
                        dir_name = os.path.join(
                            cm.global_config.save_dir,
                            cm.scenario_config.name,
                            dataset,
                            attack_model_str,
                            position,
                        )

                        # for every injected instance
                        print(f"==" * 50)
                        target_str = "-".join("_".join(map(str, sorted(set(group)))) for group in attack_args.targets)
                        log_keys_prefix = [cm.scenario_config.name, dataset, attack_model_str, position, target_str]

                        # data preparation
                        data_args = cm.generate_recommender_data_args(dataset = dataset, i = i)
                        data_args.poisoned_trn_path = os.path.join(dir_name, f"{target_str}_rating.txt")
                        data_args.label_path = os.path.join(dir_name, f"{target_str}_label.txt")
                        data_args.targets = set(re.findall(r'\d+', target_str))
                        data = RecommenderData(data_args)

                        # RS Impact
                        for rs in cm.scenario_config.assessor_rs_models:
                            log_keys = log_keys_prefix + [rs]
                            if log_exists(log, log_keys):
                                print(f"[Skipped {rs}] Already Assessed. {'/'.join(log_keys)}")
                                continue

                            rs_args = cm.generate_model_args(dataset = dataset, model = 'recommender', model_name = rs)
                            score_summary, impact_summary = RecommendationModelFactory(data, rs_args)
                            result = {
                                "PS": impact_summary,
                                "Summary": score_summary
                            }

                            torch.cuda.empty_cache()
                            torch.cuda.ipc_collect()

                            update_log(log, log_keys, result)
                            with open(results_file, "w") as f:
                                json.dump(log, f, indent = 2)

                        # RS Stealthiness
                        log_keys = log_keys_prefix + ['Distance']
                        if log_exists(log, log_keys):
                            print(f"[Skipped JSD] Already Assessed. {'/'.join(log_keys)}")
                        else:
                            jsd, tvd = data.calculate_JSD_distance(cm.scenario_config.assessor_dis_pos_rating,
                                                                   cm.scenario_config.assessor_dis_neg_rating)
                            result = {
                                "JSD": jsd,
                                "TVD": tvd
                            }
                            update_log(log, log_keys, result)
                            with open(results_file, "w") as f:
                                json.dump(log, f, indent = 2)

                        print(f"==" * 50)
    else:
        print(f"++" * 50)
        print("\033[91m[Flag] Assessor Flag is OFF!!!\033[0m")
        print(f"++" * 50)


def pipeline_evaluation(cm, flag):
    if flag == True:
        print(f"++" * 50)
        print("\033[92m[Flag] Evaluator Flag is ON!!!\033[0m")
        print(f"++" * 50)

        base_dir = os.path.join(cm.global_config.save_dir, cm.scenario_config.name)
        detector_path = os.path.join(base_dir, "detector_results.json")
        assessor_path = os.path.join(base_dir, "assessor_results.json")
        with open(detector_path, "r") as f:
            detector_log = json.load(f)

        args = cm.generate_evaluation_args()

        detection_records = []
        for dataset, injection_sets in detector_log[cm.scenario_config.name].items():
            for injection_setting, positions in injection_sets.items():
                for position, targets in positions.items():
                    for target, methods in targets.items():
                        for method, results in methods.items():
                            micro_results = eval(args.eval_detection_metric_call)
                            target_numbers = re.findall(r'\d+', target)
                            unique_numbers = set(target_numbers)
                            mode = "ST" if len(unique_numbers) == 1 else "MT"
                            detection_records.append({
                                "dataset": dataset,
                                "attack": injection_setting,
                                "position": position,
                                "target": target,
                                "mode": mode,
                                "detection model": method,
                                args.eval_detection_metric: micro_results,
                            })
        detection_df = pd.DataFrame(detection_records)
        detection_df["setting"] = detection_df[args.eval_setting_details].agg(" | ".join, axis = 1)
        pivot_table = pd.pivot_table(
            detection_df,
            index = "detection model",
            columns = "setting",
            values = args.eval_detection_metric,
            aggfunc = "mean"
        )

        mean = np.nanmean(pivot_table, axis=1)
        std = np.nanstd(pivot_table, axis=1)
        w = args.evaluator_w
        r_score = w * mean + (1 - w) * (1 - std)
        pivot_table["E-score"] = mean
        pivot_table["R-score"] = r_score
        pivot_table = pivot_table.loc[args.detection_models]
        print(f"==" * 50)
        print("[INFO] Detection Results")
        print(tabulate(pivot_table.round(3), headers = 'keys', tablefmt = 'pretty'))
        print(f"==" * 50)

        with open(assessor_path, "r") as f:
            assessor_log = json.load(f)

        assessor_records = []
        for dataset, injection_sets in assessor_log[cm.scenario_config.name].items():
            for injection_setting, positions in injection_sets.items():
                for position, targets in positions.items():
                    for target, tmps in targets.items():
                        dis = tmps['Distance'][args.eval_distance_metric]
                        for rs in args.recommender_models:
                            micro_results = eval(args.eval_impact_metric_call)
                            target_numbers = re.findall(r'\d+', target)
                            unique_numbers = set(target_numbers)
                            mode = "ST" if len(unique_numbers) == 1 else "MT"
                            assessor_records.append({
                                "dataset": dataset,
                                "attack": injection_setting,
                                "position": position,
                                "target": target,
                                "mode": mode,
                                "recommendation model": rs,
                                args.eval_distance_metric: dis,
                                args.eval_impact_metric: micro_results

                            })
        assessor_df = pd.DataFrame(assessor_records)

        assessor_df["setting"] = assessor_df[args.eval_setting_details].agg(" | ".join, axis=1)
        model_ps = (
            assessor_df.groupby(["setting", "recommendation model"])[args.eval_impact_metric]
            .mean()
            .unstack()
            .add_suffix(f"_{args.eval_impact_metric}_mean")
        )
        agg_stats = (
            assessor_df.groupby("setting")[[args.eval_distance_metric]]
            .mean()
            .rename(columns={args.eval_distance_metric: f"{args.eval_distance_metric}_mean"})
        )
        summary_table = pd.concat([model_ps, agg_stats], axis = 1)
        print(f"==" * 50)
        print("[INFO] Attack Results")
        print(tabulate(summary_table.round(4), headers = 'keys', tablefmt = 'pretty'))
        print(f"==" * 50)
    else:
        print(f"++" * 50)
        print("\033[91m[Flag] Evaluator Flag is OFF!!!\033[0m")
        print(f"++" * 50)