| import torch |
| import torch.nn.functional as F |
|
|
| from transformers import PreTrainedModel |
| from .original import TransformerModel, LMHead |
| ''' |
| Code for HuggingFace Hub Compatability |
| ''' |
|
|
| class HF_LMModel(PreTrainedModel): |
| """ Transformer with language model head only """ |
| def __init__(self, config): |
| super().__init__(config) |
| self.transformer = TransformerModel(config, vocab=config.n_vocab, n_ctx=config.n_ctx) |
| self.lm_head = LMHead(self.transformer, config, trunc_and_reshape=False) |
| self.return_probs = config.return_probs |
| self.return_acts = config.return_acts |
| if self.return_probs or self.return_acts: |
| pos_emb_mask = torch.zeros(1, 1, config.n_vocab) |
| pos_emb_mask[:, :, -config.n_ctx:] = -1e12 |
| self.register_buffer('pos_emb_mask', pos_emb_mask) |
|
|
| def forward(self, x, sequence_mask=None): |
| h = self.transformer(x, sequence_mask) |
| lm_logits = self.lm_head(h) |
| if self.return_probs: |
| lm_logits = F.softmax(lm_logits + self.pos_emb_mask, dim=-1) |
| elif self.return_acts: |
| lm_logits = lm_logits + self.pos_emb_mask |
| return { "logits": lm_logits } |
