|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import math
|
| import torch
|
| import torch.nn as nn
|
| from torch.nn import functional as F
|
| from dataclasses import dataclass, field
|
| import json
|
| from typing import List, Optional
|
|
|
| @dataclass
|
| class GPTConfig:
|
| vocab_size: int = 40960
|
| hidden_size: int = 768
|
| num_hidden_layers: int = 12
|
| num_attention_heads: int = 12
|
| max_position_embeddings: int = 1024
|
|
|
| intermediate_size: int = 3072
|
| hidden_act: str = "gelu"
|
| initializer_range: float = 0.02
|
| layer_norm_eps: float = 1e-05
|
| dropout: float = 0.1
|
| tie_word_embeddings: bool = False
|
|
|
| pad_token_id: int = 0
|
| bos_token_id: int = 8
|
| eos_token_id: int = 8
|
|
|
| model_type: str = "vexion_gpt"
|
| architectures: List[str] = field(default_factory=lambda: ["VexionGPTForCausalLM"])
|
| transformers_version: Optional[str] = None
|
|
|
| @classmethod
|
| def from_json(cls, json_path):
|
| with open(json_path, 'r', encoding='utf-8') as f:
|
| config_dict = json.load(f)
|
|
|
| valid_keys = {f.name for f in cls.__dataclass_fields__.values()}
|
| filtered_dict = {k: v for k, v in config_dict.items() if k in valid_keys}
|
|
|
| return cls(**filtered_dict)
|
|
|
| class MLP(nn.Module):
|
| def __init__(self, config):
|
| super().__init__()
|
| self.c_fc = nn.Linear(config.hidden_size, config.intermediate_size)
|
|
|
| if config.hidden_act == "gelu":
|
| self.act = nn.GELU()
|
| elif config.hidden_act == "silu":
|
| self.act = nn.SiLU()
|
| elif config.hidden_act == "relu":
|
| self.act = nn.ReLU()
|
| else:
|
| raise ValueError(f"Неизвестная активация: {config.hidden_act}")
|
|
|
| self.c_proj = nn.Linear(config.intermediate_size, config.hidden_size)
|
| self.dropout = nn.Dropout(config.dropout)
|
|
|
| def forward(self, x):
|
| x = self.c_fc(x)
|
| x = self.act(x)
|
| x = self.c_proj(x)
|
| x = self.dropout(x)
|
| return x
|
|
|
| class CausalSelfAttention(nn.Module):
|
| def __init__(self, config):
|
| super().__init__()
|
| assert config.hidden_size % config.num_attention_heads == 0
|
|
|
| self.c_attn = nn.Linear(config.hidden_size, 3 * config.hidden_size)
|
| self.c_proj = nn.Linear(config.hidden_size, config.hidden_size)
|
|
|
| self.resid_dropout = nn.Dropout(config.dropout)
|
|
|
| self.n_head = config.num_attention_heads
|
| self.embed_dim = config.hidden_size
|
|
|
| self.dropout_p = config.dropout
|
|
|
| def forward(self, x):
|
| B, T, C = x.size()
|
|
|
| qkv = self.c_attn(x)
|
| q, k, v = qkv.split(self.embed_dim, dim=2)
|
|
|
| k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
|
| q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
|
| v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
|
|
|
| y = F.scaled_dot_product_attention(
|
| q, k, v,
|
| attn_mask=None,
|
| dropout_p=self.dropout_p if self.training else 0.0,
|
| is_causal=True
|
| )
|
|
|
| y = y.transpose(1, 2).contiguous().view(B, T, C)
|
|
|
| y = self.resid_dropout(self.c_proj(y))
|
| return y
|
|
|
| class TransformerBlock(nn.Module):
|
| def __init__(self, config):
|
| super().__init__()
|
| self.ln_1 = nn.LayerNorm(config.hidden_size)
|
| self.attn = CausalSelfAttention(config)
|
| self.ln_2 = nn.LayerNorm(config.hidden_size)
|
| self.mlp = MLP(config)
|
|
|
| def forward(self, x):
|
| x = x + self.attn(self.ln_1(x))
|
| x = x + self.mlp(self.ln_2(x))
|
| return x
|
|
|
| class GPT(nn.Module):
|
| def __init__(self, config):
|
| super().__init__()
|
| self.config = config
|
|
|
| self.transformer = nn.ModuleDict(dict(
|
| wte = nn.Embedding(config.vocab_size, config.hidden_size),
|
| wpe = nn.Embedding(config.max_position_embeddings, config.hidden_size),
|
| drop = nn.Dropout(config.dropout),
|
| h = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_hidden_layers)]),
|
| ln_f = nn.LayerNorm(config.hidden_size),
|
| ))
|
|
|
| self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
|
|
| self.transformer.wte.weight = self.lm_head.weight
|
|
|
| self.apply(self._init_weights)
|
|
|
| def _init_weights(self, module):
|
| if isinstance(module, nn.Linear):
|
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
| if module.bias is not None:
|
| torch.nn.init.zeros_(module.bias)
|
| elif isinstance(module, nn.Embedding):
|
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
|
|
| for pn, p in self.named_parameters():
|
| if pn.endswith('c_proj.weight'):
|
| torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * self.config.num_hidden_layers))
|
|
|
| def forward(self, idx, targets=None):
|
| device = idx.device
|
| b, t = idx.size()
|
|
|
| assert t <= self.config.max_position_embeddings, f"Cannot forward sequence of length {t}, block size is only {self.config.max_seq_len}"
|
|
|
| pos = torch.arange(0, t, dtype=torch.long, device=device)
|
|
|
| tok_emb = self.transformer.wte(idx)
|
| pos_emb = self.transformer.wpe(pos)
|
|
|
| x = self.transformer.drop(tok_emb + pos_emb)
|
|
|
| for block in self.transformer.h:
|
| x = block(x)
|
|
|
| x = self.transformer.ln_f(x)
|
|
|
| if targets is None:
|
| logits = self.lm_head(x)
|
| loss = None
|
| else:
|
| logits = self.lm_head(x)
|
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
|
|
|
| return logits, loss
|
|
|
