from transformers import BertPreTrainedModel, BertModel
from .toxicity_model_config import ToxicityModelConfig
import torch

class ToxicityModel(BertPreTrainedModel):
    """
    ToxicityModel class for PyTorch

    Args:
        config (transformers.configuration): model configuration

    Returns:    
        output (torch.tensor): tensor containing the output logits [-1,1]
    """
    config_class = ToxicityModelConfig

    def __init__(self, config):
        super().__init__(config)
        self.bert = BertModel(config)
        
        self.cls_layer1 = torch.nn.Linear(config.hidden_size,config.linear_layer)
        self.relu1 = torch.nn.ReLU()
        self.ff1 = torch.nn.Linear(config.linear_layer,config.linear_layer)
        self.tanh1 = torch.nn.Tanh()
        self.ff2 = torch.nn.Linear(config.linear_layer,config.linear_layer_output)

    def forward(self, input_ids, attention_mask, alpha=1, beta=1e-5):
        
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        
        logits = outputs.last_hidden_state[:,0,:]
        output = self.cls_layer1(logits)
        output = self.relu1(output)
        output = self.ff1(output)
        output = self.tanh1(output)
        output = self.ff2(output)

        # Apply alpha and beta to output (if not training)  
        if not self.training:
            
            # alpha multiplies the output by a scalar
            output = torch.mul(output, alpha)

            # beta clamps the output to a minimum value
            output = torch.clamp(output, min=beta)

        return output