File size: 4,081 Bytes
ffeff6b | 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """
SOVYN-85M 모델 아키텍처
https://huggingface.co/SOVYN/SOVYN-85M
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class ModelConfig:
vocab_size: int = 16384
context_length: int = 512
embed_dim: int = 768
num_heads: int = 12
num_layers: int = 12
dropout: float = 0.1
bias: bool = False
class CausalSelfAttention(nn.Module):
def __init__(self, cfg):
super().__init__()
self.num_heads = cfg.num_heads
self.head_dim = cfg.embed_dim // cfg.num_heads
self.embed_dim = cfg.embed_dim
self.qkv = nn.Linear(cfg.embed_dim, 3 * cfg.embed_dim, bias=cfg.bias)
self.proj = nn.Linear(cfg.embed_dim, cfg.embed_dim, bias=cfg.bias)
self.resid_drop = nn.Dropout(cfg.dropout)
self.dropout_p = cfg.dropout
def forward(self, x):
B, T, C = x.shape
qkv = self.qkv(x).reshape(B, T, 3, self.num_heads, self.head_dim)
qkv = qkv.permute(2, 0, 3, 1, 4)
q, k, v = qkv.unbind(0)
out = F.scaled_dot_product_attention(
q, k, v, is_causal=True,
dropout_p=self.dropout_p if self.training else 0.0,
)
out = out.transpose(1, 2).reshape(B, T, C)
return self.resid_drop(self.proj(out))
class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
hidden = 4 * cfg.embed_dim
self.fc1 = nn.Linear(cfg.embed_dim, hidden, bias=cfg.bias)
self.fc2 = nn.Linear(hidden, cfg.embed_dim, bias=cfg.bias)
self.drop = nn.Dropout(cfg.dropout)
def forward(self, x):
return self.drop(self.fc2(F.gelu(self.fc1(x))))
class Block(nn.Module):
def __init__(self, cfg):
super().__init__()
self.ln1 = nn.LayerNorm(cfg.embed_dim)
self.attn = CausalSelfAttention(cfg)
self.ln2 = nn.LayerNorm(cfg.embed_dim)
self.ffn = FeedForward(cfg)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.ffn(self.ln2(x))
return x
class SOVYN85M(nn.Module):
def __init__(self, cfg=None):
super().__init__()
if cfg is None:
cfg = ModelConfig()
self.cfg = cfg
self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.embed_dim)
self.pos_emb = nn.Embedding(cfg.context_length, cfg.embed_dim)
self.drop = nn.Dropout(cfg.dropout)
self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.num_layers)])
self.ln_f = nn.LayerNorm(cfg.embed_dim)
self.head = nn.Linear(cfg.embed_dim, cfg.vocab_size, bias=False)
self.head.weight = self.tok_emb.weight
@property
def num_params(self):
return sum(p.numel() for p in self.parameters()) - self.tok_emb.weight.numel()
def forward(self, idx, targets=None):
B, T = idx.shape
pos = torch.arange(T, device=idx.device)
x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
for block in self.blocks:
x = block(x)
x = self.ln_f(x)
logits = self.head(x)
loss = None
if targets is not None:
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1),
ignore_index=0)
return logits, loss
@torch.no_grad()
def generate(self, idx, max_new_tokens=200, temperature=0.8, top_k=50):
self.eval()
for _ in range(max_new_tokens):
ctx = idx[:, -self.cfg.context_length:]
logits, _ = self(ctx)
logits = logits[:, -1, :] / max(temperature, 1e-8)
if top_k > 0:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, -1:]] = float('-inf')
probs = F.softmax(logits, dim=-1)
nxt = torch.multinomial(probs, 1)
idx = torch.cat([idx, nxt], dim=1)
if nxt.item() == 2: # EOS
break
return idx
|