File size: 624 Bytes
9452cf2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 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() |