| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import PreTrainedModel, PretrainedConfig |
|
|
| class NeoraConfig(PretrainedConfig): |
| model_type = "neora" |
| def __init__( |
| self, |
| vocab_size=24000, |
| d_model=1280, |
| n_layers=24, |
| n_heads=20, |
| ffn_dim=5120, |
| max_seq_len=2048, |
| **kwargs |
| ): |
| super().__init__(**kwargs) |
| self.vocab_size = vocab_size |
| self.d_model = d_model |
| self.n_layers = n_layers |
| self.n_heads = n_heads |
| self.ffn_dim = ffn_dim |
| self.max_seq_len = max_seq_len |
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.norm1 = nn.LayerNorm(config.d_model) |
| self.q_proj = nn.Linear(config.d_model, config.d_model) |
| self.k_proj = nn.Linear(config.d_model, config.d_model) |
| self.v_proj = nn.Linear(config.d_model, config.d_model) |
| self.o_proj = nn.Linear(config.d_model, config.d_model) |
| self.norm2 = nn.LayerNorm(config.d_model) |
| self.fc1 = nn.Linear(config.d_model, config.ffn_dim) |
| self.fc2 = nn.Linear(config.ffn_dim, config.d_model) |
| self.n_heads = config.n_heads |
|
|
| def forward(self, x): |
| residual = x |
| x = self.norm1(x) |
| q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x) |
| batch, seq, dim = q.shape |
| head_dim = dim // self.n_heads |
| q = q.view(batch, seq, self.n_heads, head_dim).transpose(1, 2) |
| k = k.view(batch, seq, self.n_heads, head_dim).transpose(1, 2) |
| v = v.view(batch, seq, self.n_heads, head_dim).transpose(1, 2) |
| x = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| x = x.transpose(1, 2).contiguous().view(batch, seq, dim) |
| x = self.o_proj(x) |
| x = residual + x |
| residual = x |
| x = self.norm2(x) |
| x = self.fc2(F.gelu(self.fc1(x))) |
| return residual + x |
|
|
| class NeoraModel(PreTrainedModel): |
| config_class = NeoraConfig |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.token_embedding = nn.Embedding(config.vocab_size, config.d_model) |
| self.position_embedding = nn.Embedding(config.max_seq_len, config.d_model) |
| self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)]) |
| self.final_norm = nn.LayerNorm(config.d_model) |
| self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) |
| self.lm_head.weight = self.token_embedding.weight |
|
|
| def forward(self, input_ids): |
| batch, seq = input_ids.shape |
| positions = torch.arange(seq, device=input_ids.device) |
| x = self.token_embedding(input_ids) + self.position_embedding(positions) |
| for block in self.blocks: |
| x = block(x) |
| return self.lm_head(self.final_norm(x)) |
|
|
