from pytorch_lightning import Trainer

from pytorch_lightning.loggers import TensorBoardLogger
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.plugins import DDPPlugin

from KoSBERT.lightning_modules_dataloader import KoSBertDataModule
from KoSBERT.lightning_modules_train import KoSBertTrainModule

network = KoSBertTrainModule()
datamodule = KoSBertDataModule(data_dir = './dataset/')


logger = TensorBoardLogger(
                           save_dir="SBert Logs",
                           name="regression",
                           default_hp_metric=False,
                           )

checkpoint_callback = ModelCheckpoint(
                                     monitor='val_loss',
                                     dirpath='CKPT_DIR',
                                     filename='KoSBERT_regression',
                                     mode='min'
                                     )

early_stop_callback = EarlyStopping(
                                    monitor='val_loss',
                                    min_delta=1e-4,
                                    patience=20,
                                    verbose=True,
                                    mode='min'
                                    )


trainer = Trainer(max_epochs=200,
                  gpus=2,
                  accelerator='ddp',
                  callbacks=[early_stop_callback,
                             checkpoint_callback],
                  plugins=DDPPlugin(find_unused_parameters=True),
                  logger=logger
                  )

if __name__ == '__main__':
    trainer.fit(network, datamodule=datamodule)
