import torch
from torch.nn import functional as F
from transformers.models.mistral.modeling_mistral import MistralForSequenceClassification

class NormalizedLinear(torch.nn.Linear):
    def forward(self, x):
        x = F.normalize(x, p=2, dim=-1)
        return super().forward(x)
    

class MistralForAttributePrediction(MistralForSequenceClassification):
    def __init__(self, config):
        super().__init__(config)
        
        del self.score
        self.score = NormalizedLinear(config.hidden_size, config.num_labels, bias=True)
        # Initialize weights and apply final processing
        self.post_init()