# BERT: 정적 마스킹 (전처리 시 고정)
# 동일한 문장이 항상 같은 위치가 마스킹됨
masked_data = preprocess_with_static_masking(corpus)  # 1회 실행

# RoBERTa: 동적 마스킹 (에폭마다 새로 생성)
# 매 에폭마다 다른 위치가 마스킹됨
for epoch in range(num_epochs):
    masked_data = apply_dynamic_masking(corpus)  # 매번 다른 패턴
    train(model, masked_data)

import torch.nn as nn

# BERT: 임베딩과 Hidden이 같은 차원
class BertEmbedding(nn.Module):
    def __init__(self, vocab_size=30000, hidden_dim=768):
        super().__init__()
        # V x H = 30000 x 768 = 23,040,000 파라미터
        self.embedding = nn.Embedding(vocab_size, hidden_dim)

# ALBERT: 인수분해된 임베딩
class AlbertEmbedding(nn.Module):
    def __init__(self, vocab_size=30000, embed_dim=128, hidden_dim=768):
        super().__init__()
        # V x E = 30000 x 128 = 3,840,000
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        # E x H = 128 x 768 = 98,304
        self.projection = nn.Linear(embed_dim, hidden_dim)

    def forward(self, input_ids):
        """저차원 임베딩 후 Hidden 차원으로 투사합니다."""
        x = self.embedding(input_ids)  # (batch, seq_len, E)
        x = self.projection(x)         # (batch, seq_len, H)
        return x

# BERT: 각 레이어가 독립적인 파라미터
class BertEncoder(nn.Module):
    def __init__(self, num_layers=12):
        super().__init__()
        # 12개의 독립적인 레이어 → 12배의 파라미터
        self.layers = nn.ModuleList([
            TransformerLayer(hidden_dim=768) for _ in range(num_layers)
        ])

# ALBERT: 모든 레이어가 파라미터 공유
class AlbertEncoder(nn.Module):
    def __init__(self, num_layers=12):
        super().__init__()
        # 1개 레이어의 파라미터를 12번 재사용
        self.shared_layer = TransformerLayer(hidden_dim=768)
        self.num_layers = num_layers

    def forward(self, x):
        """동일한 레이어를 반복 적용합니다."""
        for _ in range(self.num_layers):
            x = self.shared_layer(x)
        return x

# DeBERTa의 개념적 아키텍처
class DeBERTaConceptual(nn.Module):
    """DeBERTa의 분리된 어텐션 개념을 보여주는 간소화 코드입니다."""
    def __init__(self, hidden_dim=768, max_position=512):
        super().__init__()
        # 내용 임베딩과 위치 임베딩을 분리
        self.content_embedding = nn.Embedding(30000, hidden_dim)
        self.position_embedding = nn.Embedding(max_position * 2, hidden_dim)

        # 분리된 어텐션을 위한 별도 프로젝션
        self.query_content = nn.Linear(hidden_dim, hidden_dim)
        self.key_content = nn.Linear(hidden_dim, hidden_dim)
        self.query_position = nn.Linear(hidden_dim, hidden_dim)
        self.key_position = nn.Linear(hidden_dim, hidden_dim)

    def disentangled_attention(self, content, position):
        """내용과 위치를 분리하여 어텐션을 계산합니다."""
        # 내용 기반 Query/Key
        q_c = self.query_content(content)
        k_c = self.key_content(content)

        # 위치 기반 Query/Key
        q_p = self.query_position(position)
        k_p = self.key_position(position)

        # 세 가지 어텐션 항 계산
        content_to_content = q_c @ k_c.transpose(-1, -2)
        content_to_position = q_c @ k_p.transpose(-1, -2)
        position_to_content = q_p @ k_c.transpose(-1, -2)

        # 최종 어텐션 점수
        attention = content_to_content + content_to_position + position_to_content
        return attention

from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 각 모델 로드 방법
models = {
    "RoBERTa": "roberta-base",
    "ALBERT": "albert-base-v2",
    "DeBERTa": "microsoft/deberta-v3-base",
}

for name, model_name in models.items():
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        num_labels=2,
    )
    num_params = sum(p.numel() for p in model.parameters()) / 1e6
    print(f"{name}: {num_params:.1f}M 파라미터")