<filename>examples/marklearn/malware_detection_dawn.py
import string
import random
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import RidgeClassifier, LogisticRegression
from sklearn.svm import SVC
import pandas as pd
from utils import test_watermark_sklearn
from sklearn.metrics import accuracy_score

from mlmodelwatermarking.verification import verify
from sklearn.model_selection import train_test_split
from warnings import simplefilter
from math import floor
import numpy as np

from mlmodelwatermarking.marklearn import Trainer
from mlmodelwatermarking import TrainingWMArgs

simplefilter(action='ignore', category=UserWarning)


def default_key(length: int):
    elements = string.ascii_uppercase + string.digits
    return ''.join(random.choices(elements, k=length))


if __name__ == '__main__':
    # Loading classification data
    df = pd.read_csv('./data/malware_detection.csv')
    X = df.drop(["Label"], axis=1)
    y = df['Label'].replace('malicious', 0)\
                   .replace('non-malicious', 1).values

    # Random Forest Classifier
    print('RANDOM FOREST CLASSIFIER\n')
    base_model = RandomForestClassifier(max_depth=1000, random_state=42)

    X_train, X_test, y_train, y_test = train_test_split(X,
                                                        y,
                                                        test_size=0.1,
                                                        random_state=42)

    # Train a watermarked model
    print('Training watermarked model')
    number_labels = len(np.unique([floor(k) for k in y_train]))
    args = TrainingWMArgs(
                    nbr_classes=number_labels,
                    key_dawn=default_key(255),
                    probability_dawn=0.01,
                    trigger_technique='dawn',
                    metric='accuracy')

    wm_model = Trainer(RandomForestClassifier(max_depth=1000, random_state=42),
                       args=args)
    ownership = wm_model.fit(X_train, y_train)

    verification = verify(
                    ownership['labels'],
                    wm_model.predict(ownership['inputs']),
                    number_labels=number_labels,
                    metric='accuracy',
                    dawn=True)

    print(verification)
