| from torch import nn | |
| from torch.nn import functional as F | |
| from decision_lab.models.token_model import last_hidden | |
| class HeadDecisionModel(nn.Module): | |
| def __init__(self, backbone, hidden_size, classes): | |
| super().__init__() | |
| self.backbone = backbone | |
| self.head = nn.Linear(hidden_size, classes) | |
| def forward(self, input_ids, attention_mask): | |
| return self.head(last_hidden(self.backbone, input_ids, attention_mask)) | |
| def loss(self, logits, labels): | |
| return F.cross_entropy(logits.float(), labels) | |
| def class_logits(self, logits): | |
| return logits | |