models/clearance_ft/modules/clearance.py

import copy
import torch
import torch.nn as nn

from transformers import RobertaModel, BertModel, BertConfig
import pytorch_lightning as pl

from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error

from clearance_ft.modules.models import transformerDecoder, transformerEncoder
from clearance_ft.utils.emetric import get_cindex, get_rm2

from clearance_ft.modules.LayerModule import make_sequential
from clearance_ft.utils.preprocess import inverse_data


class clearanceMLPModel(pl.LightningModule):
    def __init__(self, lr, dropout, model_name, feature_length,
                 act_func, df_testColData=None):
        super().__init__()
        ## -- learning rate setting -- ##
        self.lr = lr
        self.df_testColData = df_testColData

        ## -- SMILES Attention model -- ##
        self.model = RobertaModel.from_pretrained(model_name)
        
        ## -- SMILES Feature Attention model -- ##
        if feature_length != 0:
            config = BertConfig.from_pretrained("bert-base-cased")
            config.max_length = feature_length
            self.feature_model = BertModel(config)
            feature_hidden_size = config.hidden_size
        else:
            feature_hidden_size = feature_length

        ## -- clearance prediction model -- ##
        smiles_dim  = self.model.config.hidden_size + feature_hidden_size

        self.output_decoder1 = make_sequential(smiles_dim, 512, act_func, dropout)
        self.output_decoder2 = make_sequential(512, 256, act_func, dropout)
        self.output_decoder3 = make_sequential(256, 32, act_func, dropout)
        self.out_layers = nn.Linear(32, 1)
        
        ## -- optimizer function -- ##
        self.criterior = torch.nn.SmoothL1Loss()
        self.save_hyperparameters()

    def forward(self, smiles_input, features):
        outputs = self.model(smiles_input['input_ids'], smiles_input['attention_mask'])
        
        if len(features) != 0:         
            features[0] = features[0].unsqueeze(dim=2)
            features_output = self.feature_model(inputs_embeds=features[0])
            sequence_embedding = torch.mean(features_output.last_hidden_state, dim=1)
            outs = torch.cat((outputs.last_hidden_state[:,0,:], sequence_embedding), dim=1)
        else:
            outs = outputs.last_hidden_state[:,0,:]

        outs = self.output_decoder1(outs)
        outs = self.output_decoder2(outs)
        outs = self.output_decoder3(outs)
        outs = self.out_layers(outs)
        
        outs = outs.squeeze(dim=1)    

        return outs

    def training_step(self, batch):
        smiles_input, labels, *features = batch
        if len(features) != 0:
            features[0] = features[0].type(self.feature_model.dtype)

        outputs = self(smiles_input, features)

        loss = self.criterior(outputs, labels)
        self.log("train_loss", loss)
        
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        smiles_input, labels, *features = batch
        if len(features) != 0:
            features[0] = features[0].type(self.feature_model.dtype)

        outputs = self(smiles_input, features)

        loss = self.criterior(outputs, labels)
        self.log("val_loss", loss, on_step=False, on_epoch=True, logger=True)

        return {"outputs": outputs, "labels": labels}

    def validation_step_end(self, outputs):
        return {"outputs": outputs["outputs"], "labels": outputs["labels"]}

    def validation_epoch_end(self, outputs):
        preds = torch.as_tensor(torch.cat([output['outputs'] for output in outputs], dim=0))
        labels = torch.as_tensor(torch.cat([output['labels'] for output in outputs], dim=0))

        y_pred = preds.detach().cpu().numpy()
        y_label = labels.detach().cpu().numpy()

        MSE_score = mean_squared_error(y_label, y_pred)
        MAE_score = mean_absolute_error(y_label, y_pred)
        R2_score = r2_score(y_label, y_pred)
        
        rm2_score = get_rm2(y_label, y_pred)
        ci_score = get_cindex(y_label, y_pred)

        self.valid_log = {'rmse': MSE_score, 'MAE': MAE_score, 'r2': R2_score, 'rm2': rm2_score, 'ci':ci_score}

        self.log("valid_MSE", MSE_score, on_step=False, on_epoch=True, logger=True)
        self.log("valid_MAE", MAE_score, on_step=False, on_epoch=True, logger=True)
        self.log("valid_R2", R2_score, on_step=False, on_epoch=True, logger=True)
        self.log("valid_rm2", rm2_score, on_step=False, on_epoch=True, logger=True)
        self.log("valid_CI", ci_score, on_step=False, on_epoch=True, logger=True)

    def test_step(self, batch, batch_idx):
        smiles_input, labels, *features = batch
        if len(features) != 0:
            features[0] = features[0].type(self.feature_model.dtype)

        outputs = self(smiles_input, features)

        loss = self.criterior(outputs, labels)

        self.log("test_loss", loss, on_step=False, on_epoch=True, logger=True)
        return {"outputs": outputs,  "labels": labels}

    def test_step_end(self, outputs):
        return {"outputs": outputs["outputs"], "labels": outputs["labels"]}

    def test_epoch_end(self, outputs):
        preds = torch.as_tensor(torch.cat([output['outputs'] for output in outputs], dim=0))
        labels = torch.as_tensor(torch.cat([output['labels'] for output in outputs], dim=0))

        y_pred = preds.detach().cpu().numpy()
        # y_label = labels.detach().cpu().numpy()
        
        test_scaler = self.trainer.datamodule.test_scaler

        df_predData = copy.deepcopy(self.df_testColData)
        df_predData["Clint"] = y_pred

        df_pred = inverse_data(df_predData, test_scaler)
        df_label = inverse_data(self.df_testColData, test_scaler)

        pred, labels = df_pred['Clint'], df_label['Clint']

        MSE_score = mean_squared_error(labels, pred)
        MAE_score = mean_absolute_error(labels, pred)
        R2_score = r2_score(labels, pred)
        
        rm2_score = get_rm2(labels, pred)
        ci_score = get_cindex(labels, pred)

        self.test_result = {'preds': pred, 'labels': labels}
        self.test_log = {'rmse': MSE_score, 'MAE': MAE_score, 'r2': R2_score, 'rm2': rm2_score, 'ci':ci_score}

        self.log("test_MSE", MSE_score, on_step=False, on_epoch=True, logger=True)
        self.log("test_MAE", MAE_score, on_step=False, on_epoch=True, logger=True)
        self.log("test_R2", R2_score, on_step=False, on_epoch=True, logger=True)
        self.log("test_rm2", rm2_score, on_step=False, on_epoch=True, logger=True)
        self.log("test_CI", ci_score, on_step=False, on_epoch=True, logger=True)

    def predict_step(self, batch, batch_idx):
        smiles_input, *features = batch
        if len(features) != 0:
            features[0] = features[0].type(self.feature_model.dtype)

        outputs = self(smiles_input, features)
        return outputs.detach().cpu().numpy().tolist()

    
    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(filter(lambda p:p.requires_grad, self.parameters()), lr=self.lr)
        return optimizer


class clearanceEncoderModel(pl.LightningModule):
    def __init__(self, lr, dropout, model_name, feature_length,
                 df_trainColData, df_testColData, 
                 TransformerHeads = None, TransformerLayers = None, TransformerActFunc = None):
        super().__init__()
        ## -- learning rate setting -- ##
        self.lr = lr
        self.df_trainColData = df_trainColData
        self.df_testColData = df_testColData

        ## -- SMILES Attention model -- ##
        self.model = RobertaModel.from_pretrained(model_name)
        
        ## -- SMILES Feature Attention model -- ##
        config = BertConfig.from_pretrained("bert-base-cased")
        config.max_length = feature_length
        config.hidden_size = self.model.config.hidden_size
        self.feature_model = BertModel(config)
        
        ## -- clearance prediction model -- ##
        self.transformerEncoder = transformerEncoder(self.model.config.hidden_size)
        self.output_layer = nn.Sequential(nn.Linear(self.model.config.hidden_size, 256),
                                          nn.Linear(256, 1))
        self.criterior = torch.nn.SmoothL1Loss()
        self.save_hyperparameters()

    def forward(self, smiles_input, features):
        outputs = self.model(smiles_input['input_ids'], smiles_input['attention_mask'])
        
        if len(features) != 0:         
            features[0] = features[0].unsqueeze(dim=2)
            features_output = self.feature_model(inputs_embeds=features[0])
            outs = torch.cat((outputs.last_hidden_state[:,0,:].unsqueeze(dim=1), features_output.last_hidden_state), dim=1)
        else:
            outs = outputs.last_hidden_state[:,0,:]

        # features = features.unsqueeze(dim=2)
        # features_output = self.feature_model(inputs_embeds=features)

        # outputs = self.model(smiles_input['input_ids'], smiles_input['attention_mask'])
        # outputs = outputs.last_hidden_state[:,0,:]
        
        # encoder_src = torch.cat((outputs.unsqueeze(dim=1), features_output.last_hidden_state), dim=1)
        smiles_sequence, _ = self.transformerEncoder(outs)

        outs = self.output_layer(smiles_sequence)
        outs = outs.squeeze(dim=1)    

        return outs

    def training_step(self, batch):
        smiles_input, labels, *features = batch
        if len(features) != 0:
            features[0] = features[0].type(self.feature_model.dtype)

        outputs = self(smiles_input, features)

        loss = self.criterior(outputs, labels)
        self.log("train_loss", loss)
        
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        smiles_input, labels, *features = batch
        if len(features) != 0:
            features[0] = features[0].type(self.feature_model.dtype)

        outputs = self(smiles_input, features)

        loss = self.criterior(outputs, labels)
        self.log("val_loss", loss, on_step=False, on_epoch=True, logger=True)

        return {"outputs": outputs, "labels": labels}

    def validation_step_end(self, outputs):
        return {"outputs": outputs["outputs"], "labels": outputs["labels"]}

    def validation_epoch_end(self, outputs):
        preds = torch.as_tensor(torch.cat([output['outputs'] for output in outputs], dim=0))
        labels = torch.as_tensor(torch.cat([output['labels'] for output in outputs], dim=0))

        y_pred = preds.detach().cpu().numpy()
        y_label = labels.detach().cpu().numpy()

        MSE_score = mean_squared_error(y_label, y_pred)
        MAE_score = mean_absolute_error(y_label, y_pred)
        R2_score = r2_score(y_label, y_pred)

        rm2_score = get_rm2(y_label, y_pred)
        ci_score = get_cindex(y_label, y_pred)

        self.valid_log = {'rmse': MSE_score, 'MAE': MAE_score, 'r2': R2_score, 'rm2': rm2_score, 'ci':ci_score}

        self.log("valid_MSE", MSE_score, on_step=False, on_epoch=True, logger=True)
        self.log("valid_MAE", MAE_score, on_step=False, on_epoch=True, logger=True)
        self.log("valid_R2", R2_score, on_step=False, on_epoch=True, logger=True)
        self.log("valid_rm2", rm2_score, on_step=False, on_epoch=True, logger=True)
        self.log("valid_CI", ci_score, on_step=False, on_epoch=True, logger=True)

    def test_step(self, batch, batch_idx):
        smiles_input, labels, *features = batch
        if len(features) != 0:
            features[0] = features[0].type(self.feature_model.dtype)

        outputs = self(smiles_input, features)

        loss = self.criterior(outputs, labels)
        self.log("test_loss", loss, on_step=False, on_epoch=True, logger=True)
        return {"outputs": outputs,  "labels": labels}

    def test_step_end(self, outputs):
        return {"outputs": outputs["outputs"], "labels": outputs["labels"]}

    def test_epoch_end(self, outputs):
        preds = torch.as_tensor(torch.cat([output['outputs'] for output in outputs], dim=0))
        labels = torch.as_tensor(torch.cat([output['labels'] for output in outputs], dim=0))

        y_pred = preds.detach().cpu().numpy()
        # y_label = labels.detach().cpu().numpy()
        
        test_scaler = self.trainer.datamodule.test_scaler

        df_predData = copy.deepcopy(self.df_testColData)
        df_predData["Clint"] = y_pred

        df_pred = inverse_data(df_predData, test_scaler)
        df_label = inverse_data(self.df_testColData, test_scaler)

        pred, labels = df_pred['Clint'], df_label['Clint']

        MSE_score = mean_squared_error(labels, pred)
        MAE_score = mean_absolute_error(labels, pred)
        R2_score = r2_score(labels, pred)

        rm2_score = get_rm2(labels, pred)
        ci_score = get_cindex(labels, pred)

        self.test_result = {'preds': pred, 'labels': labels}
        self.test_log = {'rmse': MSE_score, 'MAE': MAE_score, 'r2': R2_score, 'rm2': rm2_score, 'ci':ci_score}

        self.log("test_MSE", MSE_score, on_step=False, on_epoch=True, logger=True)
        self.log("test_MAE", MAE_score, on_step=False, on_epoch=True, logger=True)
        self.log("test_R2", R2_score, on_step=False, on_epoch=True, logger=True)
        self.log("test_rm2", rm2_score, on_step=False, on_epoch=True, logger=True)
        self.log("test_CI", ci_score, on_step=False, on_epoch=True, logger=True)

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(filter(lambda p:p.requires_grad, self.parameters()), lr=self.lr)
        return optimizer
    


class clearanceDecoderModel(pl.LightningModule):
    def __init__(self, lr, dropout, model_name, feature_length,
                 df_trainColData, df_testColData, 
                 TransformerHeads = None, TransformerLayers = None, TransformerActFunc = None):
        super().__init__()
        ## -- learning rate setting -- ##
        self.lr = lr
        self.df_trainColData = df_trainColData
        self.df_testColData = df_testColData

        ## -- SMILES Attention model -- ##
        self.model = RobertaModel.from_pretrained(model_name)
        
        ## -- SMILES Feature Attention model -- ##
        if feature_length != 0:
            config = BertConfig.from_pretrained("bert-base-cased")
            config.max_length = feature_length
            config.hidden_size = self.model.config.hidden_size
            self.feature_model = BertModel(config)

        self.transformerdecoder = transformerDecoder(self.model.config.hidden_size)
        self.criterior = torch.nn.SmoothL1Loss()
        self.save_hyperparameters()

    def forward(self, smiles_input, features):
        outputs = self.model(smiles_input['input_ids'], smiles_input['attention_mask'])
        
        features[0] = features[0].unsqueeze(dim=2)
        features_output = self.feature_model(inputs_embeds=features[0])

        _, outs = self.transformerdecoder(outputs.last_hidden_state,  features_output.last_hidden_state)
        outs = outs.squeeze(dim=1)    

        return outs

    def training_step(self, batch):
        smiles_input, labels, *features = batch
        if len(features) != 0:
            features[0] = features[0].type(self.feature_model.dtype)

        outputs = self(smiles_input, features)

        loss = self.criterior(outputs, labels)
        # fup_loss = self.criterior(outputs[1], fup_labels)

        self.log("train_loss", loss)
        
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        smiles_input, labels, *features = batch
        if len(features) != 0:
            features[0] = features[0].type(self.feature_model.dtype)

        outputs = self(smiles_input, features)

        loss = self.criterior(outputs, labels)
        self.log("val_loss", loss, on_step=False, on_epoch=True, logger=True)

        return {"outputs": outputs, "labels": labels}

    def validation_step_end(self, outputs):
        return {"outputs": outputs["outputs"], "labels": outputs["labels"]}

    def validation_epoch_end(self, outputs):
        preds = torch.as_tensor(torch.cat([output['outputs'] for output in outputs], dim=0))
        labels = torch.as_tensor(torch.cat([output['labels'] for output in outputs], dim=0))

        y_pred = preds.detach().cpu().numpy()
        y_label = labels.detach().cpu().numpy()
        
        # valid_scaler = self.trainer.datamodule.train_scaler

        # df_predData = copy.deepcopy(self.df_trainColData[self.trainer.datamodule.train_pos:])
        # df_predData["Clint"] = y_pred

        # df_pred = inverse_data(df_predData, valid_scaler)
        # df_label = inverse_data(self.df_trainColData[self.trainer.datamodule.train_pos:], valid_scaler)
            
        # pred, labels = df_pred['Clint'], df_label['Clint']

        MSE_score = mean_squared_error(y_label, y_pred)
        MAE_score = mean_absolute_error(y_label, y_pred)
        R2_score = r2_score(y_label, y_pred)

        rm2_score = get_rm2(y_label, y_pred)
        ci_score = get_cindex(y_label, y_pred)

        self.valid_log = {'rmse': MSE_score, 'MAE': MAE_score, 'r2': R2_score, 'rm2': rm2_score, 'ci':ci_score}

        self.log("valid_MSE", MSE_score, on_step=False, on_epoch=True, logger=True)
        self.log("valid_MAE", MAE_score, on_step=False, on_epoch=True, logger=True)
        self.log("valid_R2", R2_score, on_step=False, on_epoch=True, logger=True)
        self.log("valid_rm2", rm2_score, on_step=False, on_epoch=True, logger=True)
        self.log("valid_CI", ci_score, on_step=False, on_epoch=True, logger=True)

    def test_step(self, batch, batch_idx):
        smiles_input, labels, *features = batch
        if len(features) != 0:
            features[0] = features[0].type(self.feature_model.dtype)

        outputs = self(smiles_input, features)

        loss = self.criterior(outputs, labels)

        self.log("test_loss", loss, on_step=False, on_epoch=True, logger=True)
        return {"outputs": outputs,  "labels": labels}

    def test_step_end(self, outputs):
        return {"outputs": outputs["outputs"], "labels": outputs["labels"]}

    def test_epoch_end(self, outputs):
        preds = torch.as_tensor(torch.cat([output['outputs'] for output in outputs], dim=0))
        labels = torch.as_tensor(torch.cat([output['labels'] for output in outputs], dim=0))

        y_pred = preds.detach().cpu().numpy()
        # y_label = labels.detach().cpu().numpy()
        
        test_scaler = self.trainer.datamodule.test_scaler

        df_predData = copy.deepcopy(self.df_testColData)
        df_predData["Clint"] = y_pred

        df_pred = inverse_data(df_predData, test_scaler)
        df_label = inverse_data(self.df_testColData, test_scaler)
            
        pred, labels = df_pred['Clint'], df_label['Clint']

        MSE_score = mean_squared_error(labels, pred)
        MAE_score = mean_absolute_error(labels, pred)
        R2_score = r2_score(labels, pred)
        rm2_score = get_rm2(labels, pred)
        ci_score = get_cindex(labels, pred)

        self.test_result = {'preds': pred, 'labels': labels}
        self.test_log = {'rmse': MSE_score, 'MAE': MAE_score, 'r2': R2_score, 'rm2': rm2_score, 'ci':ci_score}

        self.log("test_MSE", MSE_score, on_step=False, on_epoch=True, logger=True)
        self.log("test_MAE", MAE_score, on_step=False, on_epoch=True, logger=True)
        self.log("test_R2", R2_score, on_step=False, on_epoch=True, logger=True)
        self.log("test_rm2", rm2_score, on_step=False, on_epoch=True, logger=True)
        self.log("test_CI", ci_score, on_step=False, on_epoch=True, logger=True)


    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(filter(lambda p:p.requires_grad, self.parameters()), lr=self.lr)
        return optimizer