# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="JWEh39ZL1p3Z"
# # Setup

# + colab={"base_uri": "https://localhost:8080/"} id="Fgj88POfmDYv" outputId="4aaed3f1-5a00-40ce-bc4f-ae9c8d867383"
from google.colab import drive
drive.mount('/content/drive')
# %cd "/content/drive/MyDrive/Colab Notebooks/NLP/project/"


# + [markdown] id="sB_ZhA2G2zbP"
# # Load training data

# + id="beVtVqF4hte8"
from torch.utils.data import IterableDataset, DataLoader
import os
import sys
import torch
from pathlib import Path
import gc
import numpy as np

module_path = "ut-mit-news-classify/NYT/"
if module_path not in sys.path:
    sys.path.append(module_path)
import utils


class ConcatenatingDataset(IterableDataset):
    def __init__(self, ensemble_dataset, gpt_dataset):
#       assert len(ensemble_dataset) == len(gpt_dataset)
        self.ensemble_dataset = ensemble_dataset
        self.gpt_dataset = gpt_dataset

    def __len__(self):
        return len(self.gpt_dataset) if len(self.gpt_dataset) < len(self.ensemble_dataset) else len(self.ensemble_dataset)

    def __iter__(self):       
        for ensemble_X, gpt in zip(self.ensemble_dataset, self.gpt_dataset):
            yield torch.cat((ensemble_X, gpt[0]), 0), gpt[1]

            
class ChunkDataset(IterableDataset):
    def __init__(self, file_paths):
        self.file_paths = file_paths
        self.total_length = None

    def __len__(self):
        if self.total_length is not None:
            return self.total_length
        
        self.total_length = 0
        for fp in self.file_paths:
            print(f'Loading for len {fp}')
            self.total_length += len(torch.load(fp).X)
            gc.collect()
        return self.total_length

    def __iter__(self):
        for fp in self.file_paths:
            print(f'Loading for iter {fp}')
            dataset = torch.load(fp)
            gc.collect()
            if hasattr(dataset, 'y'):
                for x, y in zip(dataset.X, dataset.y):
                    yield x, y
            else:
                for x in dataset.X:
                    yield x



# + id="lxaAq4SJTmmZ"
gpt_vectorized_chunks_path = Path('gpt2-vectors/train')
gpt_file_paths = os.listdir(gpt_vectorized_chunks_path)
sorted_gpt_filenames = sorted(gpt_file_paths, key=lambda fn: int(fn.split('chunk')[1].split('of')[0]))
sorted_gpt_filepaths = [gpt_vectorized_chunks_path / Path(p) for p in sorted_gpt_filenames]


# + id="BOKl3tDudyhU"
ensemble_vectorized_chunks_path = 'vectorized-fixed/train'
ensemble_file_paths = os.listdir(ensemble_vectorized_chunks_path)
sorted_ensemble_filenames = sorted(ensemble_file_paths, key=lambda fn: int(fn.split('chunk')[1].split('of')[0]))
sorted_ensemble_filepaths = [ensemble_vectorized_chunks_path / Path(p) for p in sorted_ensemble_filenames]


# + id="nf7jKT4ZeA6k" colab={"base_uri": "https://localhost:8080/"} outputId="2c02fe32-3993-4e9b-a4a3-fdc7c133be86"
concat_dataset = ConcatenatingDataset(
    ChunkDataset(sorted_ensemble_filepaths),
    ChunkDataset(sorted_gpt_filepaths))

train_X, train_y = [], []
gc.collect()
for x, y in concat_dataset:
    train_X.append(x.numpy())
    train_y.append(y.numpy())

train_X = np.asarray(train_X)
train_y = np.asarray(train_y)
gc.collect()

# + colab={"base_uri": "https://localhost:8080/"} id="O0RjHYdi_Cd-" outputId="214f3052-218e-4a4a-9d68-cf4884d3beb5"
train_X.shape, train_y.shape

# + colab={"base_uri": "https://localhost:8080/"} id="jVdZPNTC-K4I" outputId="03fbde4f-1a91-423e-e3aa-498805933732"
del concat_dataset
gc.collect()

# + [markdown] id="oReOHtUJ2LFS"
# # Configure and start Wandb job

# + id="6IzK2sYUISVU"
# %%capture
pip install --upgrade wandb

# + colab={"base_uri": "https://localhost:8080/"} id="pf-4hN5NFbFM" outputId="d9c79562-c0f5-4d16-a26d-0955126e88fc"
# !wandb login 

# + colab={"base_uri": "https://localhost:8080/", "height": 136} id="LIecSECbFg3U" outputId="0b55b445-2f50-45d1-8204-0fa1d707f0ae"
import wandb

# Start a new run
run = wandb.init(project='ensemble-plus-gpt2', entity='ut-mit-news-classify')

#hyperparams
epochs = 200
patience = 50
batch_size = 256
learning_rate = 1e-5
validation_split = 0.2

#reduce learning rate callback params
factor = 0.2
reduce_lr_patience = 5
min_lr = 0.001

# Save hyperparameters
config = wandb.config
config.batch_size = batch_size
config.epochs = epochs
config.early_stopping_patience = patience
config.learning_rate = learning_rate
config.validation_split = validation_split

config.reduce_lr = f'factor-{factor}__patience-{reduce_lr_patience}__min_lr-{min_lr}'


# + [markdown] id="cZ8Jp3HH11Vo"
# # Metrics functions for training with Keras

# + id="plCDxgG99TvJ"
from keras import backend as K

def recall(y_true, y_pred):
    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
    possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))
    recall = true_positives / (possible_positives + K.epsilon())
    return recall

def precision(y_true, y_pred):
    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
    predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))
    precision = true_positives / (predicted_positives + K.epsilon())
    return precision

def f1(y_true, y_pred):
    p = precision(y_true, y_pred)
    r = recall(y_true, y_pred)
    return 2*((p*r)/(p+r+K.epsilon()))



# + [markdown] id="yUt--FSj2A94"
# # Create and Compile Model

# + colab={"base_uri": "https://localhost:8080/"} id="POZNPfIO8gVS" outputId="66ce2091-9ff7-4c81-ec9a-4ec073d7f086"
import keras
from keras import layers

model = keras.Sequential()
model.add(layers.Dense(1400, activation="relu", input_shape=(1768,)))
model.add(layers.Dense(900, activation="relu", input_shape=(1768,)))
model.add(layers.Dense(538, activation="sigmoid", input_shape=(1768,)))

model.summary()

# + id="MjMcel-Mw-x4"
best_val_loss_model_filename = 'ensemble-gpt-vectors-with-filtered-and-cut-off-ends-298984-articles.h5'


# + id="O5FRM9hpW4JZ"
from wandb.keras import WandbCallback
from keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau


mc = ModelCheckpoint(best_val_loss_model_filename, monitor='val_loss', mode='min')
es = EarlyStopping(monitor='val_loss', mode='min', verbose=1, patience=patience)
reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=factor, patience=reduce_lr_patience, min_lr=min_lr)
callbacks = [mc, es, reduce_lr, WandbCallback()]

opt = keras.optimizers.Adam(learning_rate=learning_rate)

model.compile(loss='binary_crossentropy', optimizer=opt, metrics=['accuracy', recall, precision, f1])


# + [markdown] id="ZpDTuczX254e"
# # Train!

# + colab={"base_uri": "https://localhost:8080/"} id="eBfBEZcv_TnW" outputId="8a251ee5-9172-470a-88d2-a9afc5718a8c"

model.fit(train_X, train_y, epochs=epochs, callbacks=callbacks, validation_split=validation_split)


# + colab={"base_uri": "https://localhost:8080/"} id="gy-SMi-VnZWn" outputId="3b7e600c-8919-4ec5-864f-ced4df1fc656"
config.train_set_shape = str(train_X.shape)

del train_X
del train_y
gc.collect()

# + [markdown] id="lOdOr9Y36j8Z"
# # Save model to Wandb

# + id="VbGcydaR-3tH" colab={"base_uri": "https://localhost:8080/"} outputId="3bc251a6-c6da-4ab8-b47b-6737834d1adb"
artifact = wandb.Artifact('ensemble-plus-gpt2-bigger-layers-full-train-filtered-and-cut-off-ends', type='keras-model')
artifact.add_file(best_val_loss_model_filename)
run.log_artifact(artifact)

# + [markdown] id="bMenWG9W6qVA"
# # Load model and test on test set

# + id="8i8ILR02H6Pi"

model = keras.models.load_model(best_val_loss_model_filename, custom_objects={"f1": f1, "recall": recall, "precision": precision})
model.save('ensemble-gpt-vectors-with-filtered-and-cut-off-ends-all-train-articles.h5')

# + colab={"base_uri": "https://localhost:8080/"} id="YXDBPAfwK6oL" outputId="c8ebea8a-ac8d-4a89-d3a2-e51b0f3bb985"
gpt_vectorized_chunks_path = Path('gpt2-vectors/test')
gpt_file_paths = os.listdir(gpt_vectorized_chunks_path)
sorted_gpt_filenames = sorted(gpt_file_paths, key=lambda fn: int(fn.split('chunk')[1].split('of')[0]))
sorted_gpt_filepaths = [gpt_vectorized_chunks_path / Path(p) for p in sorted_gpt_filenames]

ensemble_vectorized_chunks_path = 'vectorized-fixed/test'
ensemble_file_paths = os.listdir(ensemble_vectorized_chunks_path)
sorted_ensemble_filenames = sorted(ensemble_file_paths, key=lambda fn: int(fn.split('chunk')[1].split('of')[0]))
sorted_ensemble_filepaths = [ensemble_vectorized_chunks_path / Path(p) for p in sorted_ensemble_filenames]

concat_dataset = ConcatenatingDataset(
    ChunkDataset(sorted_ensemble_filepaths),
    ChunkDataset(sorted_gpt_filepaths))

test_X, test_y = [], []
gc.collect()
for x, y in concat_dataset:
    test_X.append(x.numpy())
    test_y.append(y.numpy())

test_X = np.asarray(test_X)
test_y = np.asarray(test_y)
gc.collect()

# + id="Y4S2TqtV-njA"
config.test_set_shape = str(test_X.shape)


# + id="dSfr0gjuUFgt" colab={"base_uri": "https://localhost:8080/"} outputId="fa2ef055-e8cb-4f56-947b-85c3b405d70a"
# %%time
test_preds = model.predict(test_X)

# + id="cEW8VsMvubLc"
import torch

from sklearn.metrics import accuracy_score
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score
from sklearn.metrics import f1_score

def multi_label_accuracy(correct_labels, predicted_labels, average):
    predicted_labels = torch.round(torch.tensor(predicted_labels)).numpy()
    accuracy = accuracy_score(correct_labels, predicted_labels)
    precision = precision_score(correct_labels, predicted_labels, average=average, zero_division=0)
    recall = recall_score(correct_labels, predicted_labels, average=average, zero_division=0)
    f_1_score = f1_score(correct_labels, predicted_labels, average=average, zero_division=0)
    
    return accuracy, precision, recall, f_1_score



# + id="sQNLDf5cAud3" colab={"base_uri": "https://localhost:8080/"} outputId="3aa360b2-cf35-4afc-c762-137fcefc92f3"
for average in ['weighted', 'micro', 'macro', 'samples']:
    accuracy, precision, recall, f_score = multi_label_accuracy(test_y, test_preds, average)
    print(average + ' ' + str(round(f_score,3)))
    metrics = { 
            "test_precision"+'_'+ average: precision, 
            "test_f1"+'_'+ average: f_score, 
            "test_recall"+'_'+ average: recall,
            }
    wandb.log(metrics)
    
wandb.log({ 
        "test_accuracy": accuracy,
        })

# + id="28ZIvUxtaRSP" colab={"base_uri": "https://localhost:8080/", "height": 1000, "referenced_widgets": ["2888fe3a64224cc6b67fbc855f0f0e89", "40b85a686d814512958520d1a1a326c9", "e2b93c9cc95d42f39ea4202162c92c70", "dfda9ecb01d2454cb112a9e2659b360b", "fef96ae5bd5547aaac0044ba91969847", "bf5c16d1b9b6461b80ab9b896bc80a30", "75622757a53c458899c0979235ec7b1e", "<KEY>"]} outputId="50ff6ee2-8b3c-4e97-f567-854131a7a2f2"
run.finish()
len(test_X)

# + colab={"base_uri": "https://localhost:8080/"} id="Dv8Va6EkoahM" outputId="df0a5b51-b895-4f35-e0fa-2c0da54a836b"
del test_y
del test_X
gc.collect()

# + id="UtYU1CW3zQUE"


StarcoderdataJupyterScriptsDedupFiltered