Create modelling_custom.py

modelling_custom.py

@@ -0,0 +1,18 @@
+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()