data_process.py

from typing import Tuple
import pandas as pd
import pickle
from sklearn.model_selection import train_test_split
import os

# for classification approach
def data_process_class(train_path: str, test_path: str, save_path: str, seed: int = 42, val_ratio: float = 0.1, abnormal_ratio: float = 0.1, train_size = None, test_size = None, normal_only: bool = False) -> Tuple[str, str, str]:
    with open(train_path, "rb") as f:
        train_data = pickle.load(f)
        train_df = pd.DataFrame(train_data)
        # Check if 'label' column contains lists
        if train_df['Label'].apply(isinstance, args=(list,)).all():
            # Find the max of each list in the 'label' column
            train_df['Label'] = train_df['Label'].apply(max)
        
        # Print label
        print(f"Labels: {train_df['Label'].unique()}")

        if normal_only:
            train_df = train_df[train_df['Label'] == 0].reset_index(drop=True)
        
        # Filter out only EventTemplate and Label columns
        train_df = train_df[['EventTemplate', 'Label']]

        # Rename Label to label
        train_df.rename(columns={'Label': 'label'}, inplace=True)

        # Make EventTemplate from list to string and convert using vocab
        # train_df['EventTemplate'] = train_df['EventTemplate'].apply(lambda x: " ".join([str(vocab.get_event(i)) for i in x]))
        train_df['EventTemplate'] = train_df['EventTemplate'].apply(lambda x: " ".join(x))

        if train_size is not None:
            # Collect all abnomal data
            abnormal_data = train_df[train_df['label'] != 0]
            # Print 
            print(f"Abnormal data: {abnormal_data.shape[0]}")
            # Collect all normal data
            normal_data = train_df[train_df['label'] == 0]
            # Print
            print(f"Normal data: {normal_data.shape[0]}")

            # Check if number of abnormal data is more than 10% of the train_size, if is, then sample only % of the train_size
            abnormal_size = int(train_size*val_ratio*abnormal_ratio + train_size*(1-val_ratio)*abnormal_ratio)
            if abnormal_data.shape[0] > abnormal_size:
                abnormal_data = abnormal_data.sample(n=abnormal_size, random_state=seed)

            abnormal_data_sampled = abnormal_data.shape[0]
            print(f"Abnormal data sampled: {abnormal_data_sampled}")
            
            # Calculate the number of train_normal_data and train_abnormal_data, val_normal_data and val_abnormal_data and print
            train_abnormal_data = int(train_size*(1-val_ratio)*abnormal_ratio)
            train_normal_data = int(train_size * (1-val_ratio) - train_abnormal_data)
            val_abnormal_data = abnormal_data_sampled - train_abnormal_data
            val_normal_data = int(train_size * val_ratio - val_abnormal_data)
            print(f"Train normal data: {train_normal_data}")
            print(f"Train abnormal data: {train_abnormal_data}")
            print(f"Val normal data: {val_normal_data}")
            print(f"Val abnormal data: {val_abnormal_data}")

            split_normal_ratio = val_normal_data / (train_normal_data + val_normal_data)
            split_abnormal_ratio = val_abnormal_data / (val_abnormal_data + train_abnormal_data)

            normal_data = normal_data.sample(n=train_normal_data + val_normal_data, random_state=seed)

            # Split train and val for normal and abnormal data
            normal_train_df, normal_val_df = train_test_split(normal_data, test_size=split_normal_ratio, random_state=seed)
            abnormal_train_df, abnormal_val_df = train_test_split(abnormal_data, test_size=split_abnormal_ratio, random_state=seed)

            # Concat normal and abnormal data for train and val
            for_train_df = pd.concat([normal_train_df, abnormal_train_df], ignore_index=True)
            val_df = pd.concat([normal_val_df, abnormal_val_df], ignore_index=True)
        else:
            # Normal
            normal_data = train_df[train_df['label'] == 0]
            # Abnormal
            abnormal_data = train_df[train_df['label'] != 0]

            # Print
            print(f"Train Normal data: {normal_data.shape[0]}")
            print(f"Train Abnormal data: {abnormal_data.shape[0]}")
            print(f"Total Train data: {train_df.shape[0]}")

            # Split
            normal_train_df, normal_val_df = train_test_split(normal_data, test_size=val_ratio, random_state=seed)
            abnormal_train_df, abnormal_val_df = train_test_split(abnormal_data, test_size=val_ratio, random_state=seed)

            # Concat normal and abnormal data for train and val
            for_train_df = pd.concat([normal_train_df, abnormal_train_df], ignore_index=True)
            val_df = pd.concat([normal_val_df, abnormal_val_df], ignore_index=True)

            # Calculate ratio of val_df to total df
            calculated_val_ratio = val_df.shape[0] / (for_train_df.shape[0] + val_df.shape[0])
            print(f"Calculated val ratio: {calculated_val_ratio}")

        # Make label from int to list
        for_train_df['label'] = for_train_df['label'].apply(lambda x: [x])
        val_df['label'] = val_df['label'].apply(lambda x: [x])
        # Print size 
        print(f"Train data: {for_train_df.shape[0]}")
        print(f"Val data: {val_df.shape[0]}")

        # Save as .json file
        final_train_save_path = save_path + '/' + train_path.split('/')[-2] + '_train.json'
        final_val_save_path = save_path + '/' + train_path.split('/')[-2] + '_val.json'

        # Make dir if none
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        for_train_df.to_json(final_train_save_path, orient='records', lines=True)
        val_df.to_json(final_val_save_path, orient='records', lines=True)


    with open(test_path, "rb") as f:
        test_data = pickle.load(f)
        test_df = pd.DataFrame(test_data)
        # Check if 'label' column contains lists
        if test_df['Label'].apply(isinstance, args=(list,)).all():
            # Find the max of each list in the 'label' column
            test_df['Label'] = test_df['Label'].apply(max)
        
        # Filter out only EventTemplate and Label columns
        test_df = test_df[['EventTemplate', 'Label']]

        # Rename Label to label
        test_df.rename(columns={'Label': 'label'}, inplace=True)

        # Make EventTemplate from list to string and convert using vocab
        # test_df['EventTemplate'] = test_df['EventTemplate'].apply(lambda x: " ".join([str(vocab.get_event(i)) for i in x]))
        test_df['EventTemplate'] = test_df['EventTemplate'].apply(lambda x: " ".join(x))
        if test_size is not None:
            # Collect all abnomal data
            abnormal_data = test_df[test_df['label'] != 0]
            # Print 
            print(f"Abnormal data: {abnormal_data.shape[0]}")
            # Collect all normal data
            normal_data = test_df[test_df['label'] == 0]
            # Print
            print(f"Normal data: {normal_data.shape[0]}")

            # Sample normal data
            normal_data = normal_data.sample(n=test_size, random_state=seed)
            # Check if number of abnormal data is more than 10% of the test_size, if is, then sample only 10% of the test_size
            if abnormal_data.shape[0] > test_size * abnormal_ratio:
                abnormal_data = abnormal_data.sample(n=int(test_size * abnormal_ratio), random_state=seed)

            abnormal_data_sampled = abnormal_data.shape[0]
            print(f"Abnormal data sampled: {abnormal_data_sampled}")
            
            # Calculate the number of test_normal_data and test_abnormal_data and print
            test_abnormal_data = int(abnormal_data_sampled)
            test_normal_data = int(test_size - test_abnormal_data)
            print(f"Test normal data: {test_normal_data}")
            print(f"Test abnormal data: {test_abnormal_data}")

            # Sample and concat for test
            test_df = pd.concat([normal_data.sample(n=int(test_normal_data), random_state=seed), abnormal_data.sample(n=int(test_abnormal_data), random_state=seed)], ignore_index=True)

        # Print normal and abnormal len
        print(f"Test Normal data: {test_df[test_df['label'] == 0].shape[0]}")
        print(f"Test Abnormal data: {test_df[test_df['label'] != 0].shape[0]}")
        # Make label from int to list
        test_df['label'] = test_df['label'].apply(lambda x: [x])
        # Print size
        print(f"Test data: {test_df.shape[0]}")

        # Make dir if none
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        # Save as .json file
        final_test_save_path  = save_path + '/' + test_path.split('/')[-2] + '_test.json'
        test_df.to_json(final_test_save_path, orient='records', lines=True)
    
    return final_train_save_path, final_val_save_path, final_test_save_path

# for generative approach
def data_process_gen(train_path: str, test_path: str, save_path: str, seed: int = 42, val_ratio: float = 0.1, abnormal_ratio: float = 0.5, train_size = None, test_size = None, normal_only: bool = False) -> Tuple[str, str, str]:
    with open(train_path, "rb") as f:
        train_data = pickle.load(f)
        train_df = pd.DataFrame(train_data)
        # Check if 'label' column contains lists
        if train_df['Label'].apply(isinstance, args=(list,)).all():
            # Find the max of each list in the 'label' column
            train_df['Label'] = train_df['Label'].apply(max)
        if normal_only:
            train_df = train_df[train_df['Label'] == 0].reset_index(drop=True)
        
        # Filter out only EventTemplate and Label columns
        train_df = train_df[['EventTemplate', 'Label']]

        # Rename Label to label
        train_df.rename(columns={'Label': 'answer'}, inplace=True)

        # Make EventTemplate from list to string
        train_df['EventTemplate'] = train_df['EventTemplate'].apply(lambda x: ' '.join(["[L] " + i for i in x]))

        if train_size is not None:
            # Collect all abnomal data
            abnormal_data = train_df[train_df['answer'] != 0]
            # Print
            print(f"Abnormal data: {abnormal_data.shape[0]}")
            # Collect all normal data
            normal_data = train_df[train_df['answer'] == 0]
            # Print
            print(f"Normal data: {normal_data.shape[0]}")

            # Check if number of abnormal data is more than 10% of the train_size, if is, then sample only 10% of the train_size
            abnormal_size = int( train_size*val_ratio*abnormal_ratio + train_size*(1-val_ratio)*abnormal_ratio)
            if abnormal_data.shape[0] > abnormal_size:
                abnormal_data = abnormal_data.sample(n=abnormal_size, random_state=seed)

            abnormal_data_sampled = abnormal_data.shape[0]
            print(f"Abnormal data sampled: {abnormal_data_sampled}")
            
            # Calculate the number of train_normal_data and train_abnormal_data, val_normal_data and val_abnormal_data and print
            train_abnormal_data = int(train_size*(1-val_ratio)*abnormal_ratio)
            train_normal_data = int(train_size * (1-val_ratio) - train_abnormal_data)
            val_abnormal_data = 1 - train_abnormal_data
            val_normal_data = int(train_size * val_ratio - val_abnormal_data)
            print(f"Train normal data: {train_normal_data}")
            print(f"Train abnormal data: {train_abnormal_data}")
            print(f"Val normal data: {val_normal_data}")
            print(f"Val abnormal data: {val_abnormal_data}")

            split_normal_ratio = val_normal_data / (train_normal_data + val_normal_data)
            split_abnormal_ratio = val_abnormal_data / (val_abnormal_data + train_abnormal_data)

            normal_data = normal_data.sample(n=train_normal_data + val_normal_data, random_state=seed)

            # Split train and val for normal and abnormal data
            normal_train_df, normal_val_df = train_test_split(normal_data, test_size=split_normal_ratio, random_state=seed)
            abnormal_train_df, abnormal_val_df = train_test_split(abnormal_data, test_size=split_abnormal_ratio, random_state=seed)

            # Concat normal and abnormal data for train and val
            for_train_df = pd.concat([normal_train_df, abnormal_train_df], ignore_index=True)
            val_df = pd.concat([normal_val_df, abnormal_val_df], ignore_index=True)
        else:
            # Normal
            normal_data = train_df[train_df['answer'] == 0]
            # Abnormal
            abnormal_data = train_df[train_df['answer'] != 0]

            # Print
            print(f"Train Normal data: {normal_data.shape[0]}")
            print(f"Train Abnormal data: {abnormal_data.shape[0]}")
            print(f"Total Train data: {train_df.shape[0]}")

            # Split
            normal_train_df, normal_val_df = train_test_split(normal_data, test_size=val_ratio, random_state=seed)
            abnormal_train_df, abnormal_val_df = train_test_split(abnormal_data, test_size=val_ratio, random_state=seed)

            # Concat normal and abnormal data for train and val
            for_train_df = pd.concat([normal_train_df, abnormal_train_df], ignore_index=True)
            val_df = pd.concat([normal_val_df, abnormal_val_df], ignore_index=True)

            # Calculate ratio of val_df to total df
            calculated_val_ratio = val_df.shape[0] / (for_train_df.shape[0] + val_df.shape[0])
            print(f"Calculated val ratio: {calculated_val_ratio}")
        
        # Make label from int to list
        for_train_df['answer'] = for_train_df['answer'].apply(lambda x: "A1" if x == 0 else "B1")
        val_df['answer'] = val_df['answer'].apply(lambda x: "A1" if x == 0 else "B1")

        # Make dir if none
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        
        # Save as .json file
        final_train_save_path = save_path + '/' + train_path.split('/')[-2] + '_train.json'
        final_val_save_path = save_path + '/' + train_path.split('/')[-2] + '_val.json'

        for_train_df.to_json(final_train_save_path, orient='records', lines=True)
        val_df.to_json(final_val_save_path, orient='records', lines=True)


    with open(test_path, "rb") as f:
        test_data = pickle.load(f)
        test_df = pd.DataFrame(test_data)
        # Check if 'label' column contains lists
        if test_df['Label'].apply(isinstance, args=(list,)).all():
            # Find the max of each list in the 'label' column
            test_df['Label'] = test_df['Label'].apply(max)
        
        # Filter out only EventTemplate and Label columns
        test_df = test_df[['EventTemplate', 'Label']]

        # Rename Label to label
        test_df.rename(columns={'Label': 'answer'}, inplace=True)

        # Make EventTemplate from list to string
        test_df['EventTemplate'] = test_df['EventTemplate'].apply(lambda x: ' '.join(["[L] " + i for i in x]))

        if test_size is not None:
            # Collect all abnomal data
            abnormal_data = test_df[test_df['answer'] != 0]
            # Print 
            print(f"Abnormal data: {abnormal_data.shape[0]}")
            # Collect all normal data
            normal_data = test_df[test_df['answer'] == 0]
            # Print
            print(f"Normal data: {normal_data.shape[0]}")

            # Sample normal data
            normal_data = normal_data.sample(n=test_size, random_state=seed)
            # Check if number of abnormal data is more than 10% of the test_size, if is, then sample only 10% of the test_size
            if abnormal_data.shape[0] > test_size * abnormal_ratio:
                abnormal_data = abnormal_data.sample(n=int(test_size * abnormal_ratio), random_state=seed)

            abnormal_data_sampled = abnormal_data.shape[0]
            print(f"Abnormal data sampled: {abnormal_data_sampled}")
            
            # Calculate the number of test_normal_data and test_abnormal_data and print
            test_abnormal_data = int(abnormal_data_sampled)
            test_normal_data = int(test_size - test_abnormal_data)
            print(f"Test normal data: {test_normal_data}")
            print(f"Test abnormal data: {test_abnormal_data}")

            # Sample and concat for test
            test_df = pd.concat([normal_data.sample(n=int(test_normal_data), random_state=seed), abnormal_data.sample(n=int(test_abnormal_data), random_state=seed)], ignore_index=True)

                # Print normal and abnormal len
        print(f"Test Normal data: {test_df[test_df['answer'] == 0].shape[0]}")
        print(f"Test Abnormal data: {test_df[test_df['answer'] != 0].shape[0]}")

        # Make label from int to list
        test_df['answer'] = test_df['answer'].apply(lambda x: "A1" if x == 0 else "B1")

        # Make dir if none
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        # Save as .json file
        final_test_save_path  = save_path + '/' + test_path.split('/')[-2] + '_test.json'
        test_df.to_json(final_test_save_path, orient='records', lines=True)

    return final_train_save_path, final_val_save_path, final_test_save_path