import torch
import torch.nn as nn
import torch.nn.functional as F


class TransformerOutputLayer(nn.Module):
    def __init__(self, hidden_size, vocab_size, task_type='generation'):
        """
        Transformer 輸出層的實現
        參數:
        - hidden_size: Transformer 的隱藏層維度 (hidden_size)
        - vocab_size: 目標詞匯表大小（生成任務）或類別數（分類任務）
        - task_type: 任務類型，可選 'generation'（序列生成）或 'classification'（分類任務）
        """
        
        super(TransformerOutputLayer, self).__init__()
        self.hidden_size = hidden_size
        self.vocab_size = vocab_size
        self.task_type = task_type
        
        # 輸出層的線性映射：hidden_size -> vocab_size
        self.linear = nn.Linear(hidden_size, vocab_size)
    
    def forward(self, hidden_states, labels=None):
        """
        前向傳播
        參數:
        - hidden_states: 解碼器或編碼器的輸出，形狀為 (batch_size, seq_len, hidden_size)
        - labels: 目標標簽（可選），訓練時用于計算損失
                  - 對于生成任務: 形狀為 (batch_size, seq_len)
                  - 對于分類任務: 形狀為 (batch_size,)
        返回:
        - 如果是訓練階段且提供了 labels，則返回 (loss, logits)
        - 否則返回 logits
        """
        
        # 通過線性層映射到詞匯表或類別空間，形狀為 (batch_size, seq_len, vocab_size)
        logits = self.linear(hidden_states)
        
        # 如果是生成任務（語言建模或翻譯）
        if self.task_type == 'generation':
            # 計算 Softmax 概率分布（用于推理階段）
            probs = F.softmax(logits, dim=-1)
            if labels is not None:
                # 將 logits 轉換為 (batch_size * seq_len, vocab_size)，以適配 CrossEntropyLoss
                loss_fn = nn.CrossEntropyLoss()
                loss = loss_fn(logits.view(-1, self.vocab_size), labels.view(-1))
                return loss, logits
            return logits  # 推理階段返回 logits
        
        # 如果是分類任務（如文本分類）
        elif self.task_type == 'classification':
            # 通常只使用序列的第一個時間步 [CLS] 的輸出進行分類
            # 假設 hidden_states 的形狀為 (batch_size, seq_len, hidden_size)
            # 取第一個時間步的輸出 (batch_size, hidden_size)
            cls_hidden_state = hidden_states[:, 0, :]  # 提取 [CLS] token 的隱藏狀態
            
            # 通過線性層映射到類別空間，形狀為 (batch_size, vocab_size)
            logits = self.linear(cls_hidden_state)
            if labels is not None:
                # 計算分類任務的交叉熵損失
                loss_fn = nn.CrossEntropyLoss()
                loss = loss_fn(logits, labels)  # labels 形狀為 (batch_size,)
                return loss, logits
            return logits  # 推理階段返回 logits
        else:
            raise ValueError("Unsupported task type: {}".format(self.task_type))


# 測試代碼
if __name__ == "__main__":
    batch_size = 2
    seq_len = 5
    hidden_size = 768
    vocab_size = 10000
    
    # 創建輸出層
    output_layer = TransformerOutputLayer(hidden_size, vocab_size, task_type='generation')
    
    # 模擬解碼器的輸出 (batch_size, seq_len, hidden_size)
    hidden_states = torch.randn(batch_size, seq_len, hidden_size)
    
    # 模擬目標標簽 (batch_size, seq_len)
    labels = torch.randint(0, vocab_size, (batch_size, seq_len))
    
    # 前向傳播
    loss, logits = output_layer(hidden_states, labels)
    print("Loss:", loss.item())
    print("Logits shape:", logits.shape)  # 應為 (batch_size, seq_len, vocab_size)

Loss: 9.21034049987793
Logits shape: torch.Size([2, 5, 10000])