prepare_dataset.py

import os
from typing import List

import numpy as np
import pandas as pd
import scipy.sparse

from basic_function import extract_id_from_file_name, load_dict, get_root_path


def get_outside_train_features():
    train = pd.read_csv(os.path.join(get_root_path(), "features", "safe_type_train.csv"))
    train.rename(columns={"id": "file_name"}, inplace=True)
    full_features = pd.read_csv(os.path.join(get_root_path(), "features", "outside.csv"), index_col=0)
    full_features["file_name"] = full_features["file_name"].map(lambda x: extract_id_from_file_name(x))

    test_name_list = load_dict(os.path.join(get_root_path(), "features", "test_name_list"))
    test_data = pd.DataFrame(columns=["file_name"], data=np.array(test_name_list))
    test_data["file_name"] = test_data["file_name"].map(lambda x: extract_id_from_file_name(x))

    # merge
    train_data = pd.merge(train, full_features, "left", on="file_name")
    test_data = pd.merge(test_data, full_features, "left", on="file_name")

    label = train_data["safe_type"]

    train_data.drop(columns=["safe_type"], inplace=True)

    return train_data, label, test_data


def load_clustering_statics_files():
    full_features = pd.read_csv(os.path.join(get_root_path(), "features", "outside_stage2.csv"), index_col=0)
    full_features["file_name"] = full_features["file_name"].map(lambda x: extract_id_from_file_name(x))
    return full_features


def load_ft_features(feature_files=None):
    if feature_files is None:
        feature_files = {"black": "black_features.csv", "white": "white_features.csv", "test": "test_features.csv"}
    train = pd.read_csv(os.path.join(get_root_path(), "features", "safe_type_train.csv"))
    train.rename(columns={"id": "file_name"}, inplace=True)

    black_features = pd.read_csv(os.path.join(get_root_path(), "features", feature_files["black"]))
    white_features = pd.read_csv(os.path.join(get_root_path(), "features", feature_files["white"]))
    full_features = pd.concat([black_features, white_features])
    full_features["file_name"] = full_features["file_name"].map(lambda x: extract_id_from_file_name(x))

    # load test data
    test_data = pd.read_csv(os.path.join(get_root_path(), "features", feature_files["test"]))
    test_data["file_name"] = test_data["file_name"].map(lambda x: extract_id_from_file_name(x))

    # merge
    train_dat = pd.merge(train, full_features, "inner", on="file_name")

    label = train_dat["safe_type"]
    train_dat.drop(columns=["safe_type"], inplace=True)
    return train_dat, label, test_data


def load_runtime_features():
    train = pd.read_csv(os.path.join(get_root_path(), "features", "safe_type_train.csv"))
    train.rename(columns={"id": "file_name"}, inplace=True)

    full_features = pd.read_csv(os.path.join("features", "train_used_time_feauture.csv"))
    full_features["file_name"] = full_features["file_name"].map(lambda x: extract_id_from_file_name(x))

    # load test data
    test_data = pd.read_csv(os.path.join("features", "test_used_time_feauture.csv"))
    test_data["file_name"] = test_data["file_name"].map(lambda x: extract_id_from_file_name(x))

    # merge
    train_dat = pd.merge(train, full_features, "inner", on="file_name")

    label = train_dat["safe_type"]
    train_dat.drop(columns=["safe_type"], inplace=True)
    return train_dat, label, test_data


def load_depth_three_features():
    return load_ft_features({"black": "black_features_depth_3.csv", "white": "white_features_depth_3.csv",
                             "test": "test_features_depth_3.csv"})


def load_nn_features():
    train = pd.read_csv(os.path.join(get_root_path(), "features", "safe_type_train.csv"))
    train.rename(columns={"id": "file_name"}, inplace=True)

    train_features = pd.read_csv(os.path.join(get_root_path(), "features", "train_nn.csv"))
    train_features["file_name"] = train_features["file_name"].map(lambda x: extract_id_from_file_name(x))

    # load test data
    test_data = pd.read_csv(os.path.join(get_root_path(), "features", "test_nn.csv"))
    test_data["file_name"] = test_data["file_name"].map(lambda x: extract_id_from_file_name(x))

    # merge
    train_dat = pd.merge(train, train_features, "inner", on="file_name")

    label = train_dat["safe_type"]
    train_dat.drop(columns=["safe_type"], inplace=True)
    return train_dat, label, test_data


def load_tfidf_features(suffix, type_name=""):
    black = scipy.sparse.load_npz(
        os.path.join(get_root_path(), "features", "black" + suffix + type_name + ".npz")).toarray()
    white = scipy.sparse.load_npz(
        os.path.join(get_root_path(), "features", "white" + suffix + type_name + ".npz")).toarray()
    test = scipy.sparse.load_npz(
        os.path.join(get_root_path(), "features", "test" + suffix + type_name + ".npz")).toarray()

    black_l = np.ones((black.shape[0],))
    white_l = np.zeros((white.shape[0],))
    train_data = pd.DataFrame(np.concatenate((black, white), axis=0))

    label = pd.DataFrame(np.concatenate((black_l, white_l), axis=0))

    test_df = pd.DataFrame(test)

    black_name_list = load_dict(os.path.join(get_root_path(), "features", "black_name_list" + suffix + type_name))
    white_name_list = load_dict(os.path.join(get_root_path(), "features", "white_name_list" + suffix + type_name))
    train_name_list = np.concatenate((black_name_list, white_name_list), axis=0)

    test_name_list = load_dict(os.path.join(get_root_path(), "features", "test_name_list" + suffix + type_name))

    train_data["file_name"] = train_name_list
    train_data["file_name"] = train_data["file_name"].map(lambda x: extract_id_from_file_name(x))

    test_df["file_name"] = test_name_list
    test_df["file_name"] = test_df["file_name"].map(lambda x: extract_id_from_file_name(x))

    return train_data, label, test_df


def load_autoencoder_features():
    features = np.load("train_nn.npy")
    print(features.shape)
    name = np.load("file_name_list_stage2.npy")
    label = np.load("label_nn.npy")
    features = pd.DataFrame(data=features)
    features["file_name"] = name
    features["file_name"] = features["file_name"].map(lambda x: extract_id_from_file_name(x))

    """
    useless test_df. just read to avoid error
    """
    test = scipy.sparse.load_npz(os.path.join(get_root_path(), "features", "test.npz")).toarray()
    test_df = pd.DataFrame(test)
    test_name_list = load_dict(os.path.join(get_root_path(), "features", "test_name_list"))

    test_df["file_name"] = test_name_list
    test_df["file_name"] = test_df["file_name"].map(lambda x: extract_id_from_file_name(x))
    return features, label, test_df


def load_stage2_tf_idf(suffix, type_name=""):
    stage2 = scipy.sparse.load_npz(
        os.path.join(get_root_path(), "features", "stage2" + suffix + type_name + ".npz")).toarray()

    train_data = pd.DataFrame(stage2)

    stage2_name_list = load_dict(os.path.join(get_root_path(), "features", "stage2_name_list" + suffix + type_name))

    train_data["file_name"] = stage2_name_list
    train_data["file_name"] = train_data["file_name"].map(lambda x: extract_id_from_file_name(x))

    return train_data


def load_nn_stage2_features():
    nn_features = np.load("nn_features.npy")
    name = np.load("file_name_list_stage2.npy")

    train_data = pd.DataFrame(nn_features)
    train_data["file_name"] = name
    train_data["file_name"] = train_data["file_name"].map(lambda x: extract_id_from_file_name(x))
    return train_data


def load_tianchi_tf_idf():
    stage2 = scipy.sparse.load_npz(
        os.path.join(get_root_path(), "features", "tianchi" + ".npz")).toarray()

    train_data = pd.DataFrame(stage2)

    stage2_name_list = load_dict(os.path.join(get_root_path(), "features", "tianchi_name_list"))

    train_data["file_name"] = stage2_name_list
    train_data["file_name"] = train_data["file_name"].map(lambda x: extract_id_from_file_name(x))
    tianchi = pd.read_csv("security_train.csv")[["label", "file_id"]].drop_duplicates()
    tianchi = tianchi.rename(columns={"file_id": "file_name"})
    full = pd.merge(train_data, tianchi, how="left", on="file_name")
    label = full["label"]

    return train_data, label


def load_tfidf_sparse_features(suffix):
    black = scipy.sparse.load_npz(os.path.join(get_root_path(), "black" + suffix + ".npz"))
    white = scipy.sparse.load_npz(os.path.join(get_root_path(), "white" + suffix + ".npz"))
    test = scipy.sparse.load_npz(os.path.join(get_root_path(), "test" + suffix + ".npz"))

    white_file_id = load_dict("white_name_list")
    black_file_id = load_dict("black_name_list")

    black_l = np.ones((black.shape[0],))
    white_l = np.zeros((white.shape[0],))
    train_data = scipy.sparse.vstack([black, white])
    label = pd.DataFrame(np.concatenate((black_l, white_l), axis=0))

    test_df = test
    file_id = load_dict(os.path.join(get_root_path(), "test_name_list" + suffix))
    return train_data, label, test_df, file_id, np.array(black_file_id + white_file_id)


def merge_features(features: List):
    train_data, label, test_data = features.pop(0)
    train_data["label"] = label

    for i in range(len(features)):
        train_data = pd.merge(train_data, features[i][0], how="left", on="file_name")
        test_data = pd.merge(test_data, features[i][2], how="left", on="file_name")

    label = train_data["label"]
    train_data.drop(columns=["label"], inplace=True)

    return train_data, label, test_data


def drop_id(features: List):
    features[0].drop(columns=["file_name"], inplace=True)
    features[2].drop(columns=["file_name"], inplace=True)
    return features


if __name__ == '__main__':
    train_data = load_stage2_tf_idf("1000")