import math from typing import Optional, Tuple, Dict, Any import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM try: from .configuration_gpt2 import GPT2CustomConfig except ImportError: from configuration_gpt2 import GPT2CustomConfig class CausalSelfAttention(nn.Module): def __init__(self, config): super().__init__() assert config.n_embd % config.n_head == 0 self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=True) self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=True) self.attn_dropout = nn.Dropout(config.attn_pdrop) self.resid_dropout = nn.Dropout(config.resid_pdrop) self.n_head = config.n_head self.n_embd = config.n_embd # Causal mask buffer self.register_buffer("bias", torch.tril(torch.ones(config.n_positions, config.n_positions)) .view(1, 1, config.n_positions, config.n_positions)) def forward(self, x): B, T, C = x.size() q, k, v = self.c_attn(x).split(self.n_embd, dim=2) k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf')) att = F.softmax(att, dim=-1) att = self.attn_dropout(att) y = att @ v # (B, nh, T, hs) y = y.transpose(1, 2).contiguous().view(B, T, C) y = self.resid_dropout(self.c_proj(y)) return y class MLP(nn.Module): def __init__(self, config): super().__init__() self.c_fc = nn.Linear(config.n_embd, config.n_inner, bias=True) self.c_proj = nn.Linear(config.n_inner, config.n_embd, bias=True) self.act = nn.GELU(approximate="tanh") self.dropout = nn.Dropout(config.resid_pdrop) def forward(self, x): return self.dropout(self.c_proj(self.act(self.c_fc(x)))) class GPT2Block(nn.Module): def __init__(self, config): super().__init__() self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) self.attn = CausalSelfAttention(config) self.ln_2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) 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 GPT2CustomModel(PreTrainedModel): config_class = GPT2CustomConfig base_model_prefix = "transformer" def __init__(self, config): super().__init__(config) self.wte = nn.Embedding(config.vocab_size, config.n_embd) self.wpe = nn.Embedding(config.n_positions, config.n_embd) self.drop = nn.Dropout(config.embd_pdrop) self.h = nn.ModuleList([GPT2Block(config) for _ in range(config.n_layer)]) self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) self.post_init() def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) 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=self.config.initializer_range) def forward(self, input_ids, **kwargs): device = input_ids.device B, T = input_ids.size() pos = torch.arange(0, T, dtype=torch.long, device=device).unsqueeze(0) x = self.wte(input_ids) + self.wpe(pos) x = self.drop(x) for block in self.h: x = block(x) x = self.ln_f(x) from transformers.modeling_outputs import BaseModelOutputWithPast return BaseModelOutputWithPast(last_hidden_state=x) class GPT2CustomLMHeadModel(PreTrainedModel, GenerationMixin): config_class = GPT2CustomConfig base_model_prefix = "transformer" _tied_weights_keys = ["lm_head.weight"] def __init__(self, config): super().__init__(config) self.transformer = GPT2CustomModel(config) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) self.transformer.wte.weight = self.lm_head.weight # Weight tying self.post_init() def get_input_embeddings(self): return self.transformer.wte def set_input_embeddings(self, new_embeddings): self.transformer.wte = new_embeddings def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def tie_weights(self, **kwargs): self.lm_head.weight = self.transformer.wte.weight def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) 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=self.config.initializer_range) def forward(self, input_ids, labels=None, **kwargs): transformer_outputs = self.transformer(input_ids) hidden_states = transformer_outputs.last_hidden_state logits = self.lm_head(hidden_states) loss = None if labels is not None: shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100) from transformers.modeling_outputs import CausalLMOutputWithPast return CausalLMOutputWithPast( loss=loss, logits=logits ) def prepare_inputs_for_generation(self, input_ids, **kwargs): return {"input_ids": input_ids} # Register configuration and model for auto mapping AutoConfig.register("gpt2_custom", GPT2CustomConfig) AutoModel.register(GPT2CustomConfig, GPT2CustomModel) AutoModelForCausalLM.register(GPT2CustomConfig, GPT2CustomLMHeadModel)