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

class OutputProjection(nn.Module):
    def __init__(self, d_model: int, vocab_size: int):
        super().__init__()
        self.proj = nn.Linear(d_model, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        logits = self.proj(x)
        probabilities = F.softmax(logits, dim=-1)
        return probabilities

if __name__ == '__main__':
    D_MODEL = 512
    VOCAB_SIZE = 37000
    BATCH_SIZE = 4
    SEQ_LEN = 50

    prediction_head = OutputProjection(d_model=D_MODEL, vocab_size=VOCAB_SIZE)
    dummy_decoder_output = torch.randn(BATCH_SIZE, SEQ_LEN, D_MODEL)
    
    probs = prediction_head(dummy_decoder_output)
    print(f"1. 入力形状 (Decoder出力): {dummy_decoder_output.shape}")
    print(f"2. 出力形状 (確率分布): {probs.shape}")

    sum_probs = torch.sum(probs, dim=-1)
    print(f"3. 確率の合計値 (期待値 1.0): {sum_probs[0, 0].item():.4f}")

    predicted_tokens = torch.argmax(probs, dim=-1)
    print(f"4. 予測トークンIDの形状: {predicted_tokens.shape}")
    print(f"   (例) 最初の位置の予測ID: {predicted_tokens[0, 0].item()}")

# Attention Is All You Need
# 3.4 Embeddings and Softmax
#     PredictionHead
# PredictionHead.py
# 第2部 第14章