File size: 2,793 Bytes
e3b9106 | 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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
from .configuration_approxdumb import ApproxDumbConfig
class ApproxDumbForCausalLM(PreTrainedModel):
config_class = ApproxDumbConfig
def __init__(self, config: ApproxDumbConfig):
super().__init__(config)
self.config = config
self.d_model = config.d_model
# Token & Position Embeddings
self.token_emb = nn.Embedding(config.vocab_size, config.d_model)
self.pos_emb = nn.Embedding(config.max_len, config.d_model)
# Causal Self-Attention
self.q = nn.Linear(config.d_model, config.d_model)
self.k = nn.Linear(config.d_model, config.d_model)
self.v = nn.Linear(config.d_model, config.d_model)
self.out_proj = nn.Linear(config.d_model, config.d_model)
# Feed Forward Network
self.ffn = nn.Sequential(
nn.Linear(config.d_model, 12),
nn.ReLU(),
nn.Linear(12, config.d_model)
)
# LM Head
self.lm_head = nn.Linear(config.d_model, config.vocab_size)
self.post_init()
def forward(
self,
input_ids=None,
labels=None,
output_attentions=None,
output_hidden_states=None,
return_dict=None,
**kwargs
):
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
B, T = input_ids.shape
pos = torch.arange(0, T, device=input_ids.device).unsqueeze(0)
h = self.token_emb(input_ids) + self.pos_emb(pos)
# Causal Attention
q, k, v = self.q(h), self.k(h), self.v(h)
weights = (q @ k.transpose(-2, -1)) * (self.d_model ** -0.5)
mask = torch.tril(torch.ones(T, T, device=input_ids.device))
weights = weights.masked_fill(mask == 0, float('-inf'))
weights = torch.softmax(weights, dim=-1)
attn_out = weights @ v
h = h + self.out_proj(attn_out)
h = h + self.ffn(h)
logits = self.lm_head(h)
loss = None
if labels is not None:
# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
loss_fct = nn.CrossEntropyLoss()
loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
if not return_dict:
output = (logits,)
return ((loss,) + output) if loss is not None else output
return CausalLMOutputWithPast(
loss=loss,
logits=logits,
)
def prepare_inputs_for_generation(self, input_ids, **kwargs):
return {"input_ids": input_ids}
|