import torch import torch.nn as nn from transformers import PretrainedConfig class CustomDecoderConfig(PretrainedConfig): model_type = "custom-transformer-decoder" def __init__( self, vocab_size=30522, max_position_embeddings=128, hidden_size=256, num_attention_heads=8, intermediate_size=1024, num_hidden_layers=4, hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1, **kwargs ): super().__init__(**kwargs) # 어휘 크기(vocab size): 입력 token id의 전체 개수 self.vocab_size = vocab_size # 최대 위치 임베딩(max position embeddings): 입력 sequence의 최대 길이 self.max_position_embeddings = max_position_embeddings # 은닉 차원(hidden size): 각 token vector의 차원 self.hidden_size = hidden_size # attention head 개수(num attention heads): Multi-Head Attention의 head 수 self.num_attention_heads = num_attention_heads # FFN 내부 차원(intermediate size): Feed Forward Network의 중간 차원 self.intermediate_size = intermediate_size # decoder layer 개수(num hidden layers): DecoderBlock 반복 횟수 self.num_hidden_layers = num_hidden_layers # dropout 비율(hidden dropout probability): 과적합(overfitting) 방지 self.hidden_dropout_prob = hidden_dropout_prob # attention dropout 비율(attention probabilities dropout probability) self.attention_probs_dropout_prob = attention_probs_dropout_prob class PositionWiseFeedForward(nn.Module): def __init__(self, hidden_size, intermediate_size, dropout_prob): super().__init__() # 첫 번째 선형층(first linear layer): hidden size를 intermediate size로 확장 self.linear_1 = nn.Linear(hidden_size, intermediate_size) # 활성화 함수(activation function): 비선형성 추가 self.activation = nn.GELU() # 두 번째 선형층(second linear layer): intermediate size를 hidden size로 축소 self.linear_2 = nn.Linear(intermediate_size, hidden_size) # 드롭아웃(dropout): 과적합(overfitting) 방지 self.dropout = nn.Dropout(dropout_prob) def forward(self, x): # Position-wise FFN: 각 위치(position)의 token에 독립적으로 적용 x = self.linear_1(x) x = self.activation(x) x = self.dropout(x) x = self.linear_2(x) return x class DecoderBlock(nn.Module): def __init__(self, config): super().__init__() # Masked Self-Attention: 미래 token을 보지 못하게 하는 attention self.masked_self_attention = nn.MultiheadAttention( embed_dim=config.hidden_size, num_heads=config.num_attention_heads, dropout=config.attention_probs_dropout_prob, batch_first=True ) # 첫 번째 정규화(layer normalization): self-attention 출력 안정화 self.norm_1 = nn.LayerNorm(config.hidden_size) # Cross-Attention: encoder output과 decoder input 사이의 관계 학습 self.cross_attention = nn.MultiheadAttention( embed_dim=config.hidden_size, num_heads=config.num_attention_heads, dropout=config.attention_probs_dropout_prob, batch_first=True ) # 두 번째 정규화(layer normalization): cross-attention 출력 안정화 self.norm_2 = nn.LayerNorm(config.hidden_size) # Position-Wise Feed Forward Network: token별 비선형 변환 self.feed_forward = PositionWiseFeedForward( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, dropout_prob=config.hidden_dropout_prob ) # 세 번째 정규화(layer normalization): FFN 출력 안정화 self.norm_3 = nn.LayerNorm(config.hidden_size) # 드롭아웃(dropout): residual connection 전에 적용 self.dropout = nn.Dropout(config.hidden_dropout_prob) def forward(self, x, encoder_output=None, self_attention_mask=None, cross_attention_mask=None): # 잔차 연결(residual connection)을 위해 입력 저장 residual = x # Masked Self-Attention 계산 self_attention_output, _ = self.masked_self_attention( query=x, key=x, value=x, attn_mask=self_attention_mask ) # Add & Norm: residual connection + layer normalization x = self.norm_1(residual + self.dropout(self_attention_output)) if encoder_output is not None: # 잔차 연결(residual connection)을 위해 입력 저장 residual = x # Cross-Attention 계산 cross_attention_output, _ = self.cross_attention( query=x, key=encoder_output, value=encoder_output, key_padding_mask=cross_attention_mask ) # Add & Norm: residual connection + layer normalization x = self.norm_2(residual + self.dropout(cross_attention_output)) # 잔차 연결(residual connection)을 위해 FFN 입력 저장 residual = x # Position-Wise Feed Forward 계산 ff_output = self.feed_forward(x) # Add & Norm: residual connection + layer normalization x = self.norm_3(residual + self.dropout(ff_output)) return x class TransformerDecoder(nn.Module): def __init__(self, config): super().__init__() # 토큰 임베딩(token embedding): token id를 vector로 변환 self.token_embedding = nn.Embedding( config.vocab_size, config.hidden_size ) # 위치 임베딩(position embedding): token의 순서 정보 추가 self.position_embedding = nn.Embedding( config.max_position_embeddings, config.hidden_size ) # DecoderBlock stack: N개의 DecoderBlock 반복 self.layers = nn.ModuleList([ DecoderBlock(config) for _ in range(config.num_hidden_layers) ]) # 드롭아웃(dropout): embedding 이후 적용 self.dropout = nn.Dropout(config.hidden_dropout_prob) # 최종 정규화(final layer normalization) self.final_norm = nn.LayerNorm(config.hidden_size) # 출력층(lm head): hidden vector를 vocab size로 변환 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size) self.config = config def make_causal_mask(self, seq_length, device): # Causal mask: 현재 token이 미래 token을 보지 못하게 함 mask = torch.triu( torch.ones(seq_length, seq_length, device=device), diagonal=1 ).bool() return mask def forward(self, input_ids, encoder_output=None, cross_attention_mask=None): batch_size, seq_length = input_ids.size() # 위치 번호(position ids) 생성 position_ids = torch.arange( seq_length, dtype=torch.long, device=input_ids.device ).unsqueeze(0).expand(batch_size, seq_length) # token embedding과 position embedding을 더함 x = self.token_embedding(input_ids) + self.position_embedding(position_ids) x = self.dropout(x) # 미래 token 차단용 causal mask 생성 self_attention_mask = self.make_causal_mask(seq_length, input_ids.device) # DecoderBlock 반복 적용 for layer in self.layers: x = layer( x, encoder_output=encoder_output, self_attention_mask=self_attention_mask, cross_attention_mask=cross_attention_mask ) x = self.final_norm(x) # logits: 각 token 위치에서 vocab 전체에 대한 점수 logits = self.lm_head(x) return logits if __name__ == "__main__": config = CustomDecoderConfig() model = TransformerDecoder(config) input_ids = torch.randint(0, config.vocab_size, (2, 10)) encoder_output = torch.randn(2, 10, config.hidden_size) logits = model(input_ids, encoder_output=encoder_output) print("Input shape:", input_ids.shape) print("Logits shape:", logits.shape)