module.py

import torch
import torch.nn as nn
import math

class TransformerEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, max_len, dropout=0.1):
        super(TransformerEmbedding, self).__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_encoding = self._generate_positional_encoding(d_model, max_len)
        self.dropout = nn.Dropout(dropout)
        self.max_len = max_len  # Store max_len as instance variable

    def _generate_positional_encoding(self, d_model, max_len):
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        return pe

    def forward(self, x):
        seq_len = x.size(1)
        if seq_len > self.max_len:  # Truncate if sequence length exceeds max_len
            x = x[:, :self.max_len]
            seq_len = self.max_len
        token_embeddings = self.token_embedding(x)
        position_encodings = self.position_encoding[:, :seq_len, :].to(x.device)
        embeddings = token_embeddings + position_encodings
        return self.dropout(embeddings)
    
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        self.q_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)

        self.out_linear = nn.Linear(d_model, d_model)

        self.dropout = nn.Dropout(dropout)

        self.attention_weights = None

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        Q = self.q_linear(query)  # (batch_size, seq_len, d_model)
        K = self.k_linear(key)    # (batch_size, seq_len, d_model)
        V = self.v_linear(value)  # (batch_size, seq_len, d_model)
        
        # 멀티 헤드로 나누기: (batch_size, seq_len, num_heads, head_dim)
        Q = self._split_heads(Q, batch_size)
        K = self._split_heads(K, batch_size)
        V = self._split_heads(V, batch_size)

        attention_output, attention_weights = self._scaled_dot_product_attention(Q, K, V, mask)

        concat_attention = self._combine_heads(attention_output, batch_size)

        output = self.out_linear(concat_attention)  # (batch_size, seq_len, d_model)
        self.attention_weights = attention_weights
        return output, attention_weights

    def _scaled_dot_product_attention(self, Q, K, V, mask):
        d_k = Q.size(-1)  # head_dim

        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))  # (batch_size, num_heads, seq_len, seq_len)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attention_weights = torch.softmax(scores, dim=-1)  # (batch_size, num_heads, seq_len, seq_len)
        attention_weights = self.dropout(attention_weights)

        output = torch.matmul(attention_weights, V)  # (batch_size, num_heads, seq_len, head_dim)
        return output, attention_weights

    def _split_heads(self, x, batch_size):
        return x.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # (batch_size, num_heads, seq_len, head_dim)

    def _combine_heads(self, x, batch_size):
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)  # (batch_size, seq_len, d_model)
        return x
    
class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super(TransformerEncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.ReLU(),
            nn.Linear(d_model * 4, d_model)
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout)
        
    def forward(self, x, mask=None):
        x = x + self.dropout1(self.self_attn(x, x, x, mask)[0])
        x = self.norm1(x)
        x = x + self.dropout2(self.feed_forward(x))
        x = self.norm2(x)
        return x

class TransformerDecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super(TransformerDecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.cross_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.ReLU(),
            nn.Linear(d_model * 4, d_model)
        )
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, x, encoder_output, src_mask, tgt_mask):
        x = x + self.dropout1(self.self_attn(x, x, x, tgt_mask)[0])
        x = self.norm1(x)
        x = x + self.dropout2(self.cross_attn(x, encoder_output, encoder_output, src_mask)[0])
        x = self.norm2(x)
        x = x + self.dropout3(self.feed_forward(x))
        x = self.norm3(x)
        return x