File size: 1,605 Bytes
33050db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | from __future__ import annotations
import torch
from torch import nn
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput
class LayerfaultTinyConfig(PretrainedConfig):
model_type = "layerfault_tiny"
def __init__(self, vocab_size=10, hidden_size=8, **kwargs):
super().__init__(**kwargs)
self.vocab_size = vocab_size
self.hidden_size = hidden_size
class LayerfaultTinyForCausalLM(PreTrainedModel):
config_class = LayerfaultTinyConfig
main_input_name = "input_ids"
def __init__(self, config):
super().__init__(config)
self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
self.post_init()
def get_input_embeddings(self):
return self.embed
def set_input_embeddings(self, value):
self.embed = value
def get_output_embeddings(self):
return self.lm_head
def set_output_embeddings(self, value):
self.lm_head = value
def forward(self, input_ids=None, labels=None, **kwargs):
h = self.embed(input_ids)
logits = self.lm_head(h)
loss = None
if labels is not None:
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
loss = nn.functional.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
)
return CausalLMOutput(loss=loss, logits=logits)
|