| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| def last_hidden(backbone, input_ids, attention_mask): |
| output = backbone(input_ids=input_ids, attention_mask=attention_mask, use_cache=False) |
| positions = attention_mask.sum(dim=1) - 1 |
| if torch.any(positions < 0): |
| raise ValueError("Each prompt requires at least one token") |
| return output.last_hidden_state[ |
| torch.arange(input_ids.shape[0], device=input_ids.device), positions |
| ] |
|
|
|
|
| class TokenDecisionModel(nn.Module): |
| def __init__(self, causal_lm, answer_ids): |
| super().__init__() |
| self.backbone = causal_lm.model |
| self.vocabulary_head = causal_lm.lm_head |
| self.register_buffer("answer_ids", torch.tensor(answer_ids, dtype=torch.long)) |
|
|
| def forward(self, input_ids, attention_mask): |
| hidden = last_hidden(self.backbone, input_ids, attention_mask) |
| return self.vocabulary_head(hidden) |
|
|
| def loss(self, logits, labels): |
| return F.cross_entropy(logits.float(), self.answer_ids[labels]) |
|
|
| def class_logits(self, logits): |
| return logits[:, self.answer_ids] |
|
|
| def outside_probability(self, logits): |
| return 1 - logits.float().softmax(-1)[:, self.answer_ids].sum(-1) |
|
|
