import torch
import torch.nn as nn

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)

    def forward(self, input_ids):
        return self.token_embedding(input_ids)

if __name__ == '__main__':
    VOCAB_SIZE = 10000
    D_MODEL = 512
    BATCH_SIZE = 4
    SEQ_LEN = 60

    token_embed = TokenEmbedding(vocab_size=VOCAB_SIZE, d_model=D_MODEL)
    dummy_input_ids = torch.randint(low=0, high=VOCAB_SIZE, size=(BATCH_SIZE, SEQ_LEN))
    output_embeddings = token_embed(dummy_input_ids)
    
    print(f"入力 (ID) 形状: {dummy_input_ids.shape}")
    print(f"出力 (埋め込み) 形状: {output_embeddings.shape}")

    expected_shape = (BATCH_SIZE, SEQ_LEN, D_MODEL)
    assert output_embeddings.shape == expected_shape

# Attention Is All You Need
# 3.4 Embeddings（Token Embedding 部分）
# TokenEmbedding.py
# 第2部 第1章