embedding_module.py

import torch
import torch.nn as nn
from dataclasses import dataclass

@dataclass
class EmbeddingConfig:
    vocab_size: int
    d_model: int
    block_size: int
    dropout: float = 0.1

class TokenPositionalEmbedding(nn.Module):
    def __init__(self, cfg: EmbeddingConfig):
        super().__init__()
        self.cfg = cfg

        # Token embedding table: [vocal_size, d_model]
        self.token_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)

        # Positional embedding table: [block_size, d_model]
        self.pos_emb = nn.Embedding(cfg.block_size, cfg.d_model)

        # Dropout for regularization
        self.drop = nn.Dropout(cfg.dropout)

        # Initialize weights
        nn.init.normal_(self.token_emb.weight, mean=0, std=0.02)
        nn.init.normal_(self.pos_emb.weight, mean=0, std=0.02)
    
    def forward(self, input_ids):
        """
        input_ids: (B, T)
        returns: (B, T, d_model)
        """
        B, T = input_ids.shape
        assert T<=self.cfg.block_size, (
            f"Sequence length {T} exceeds block_size {self.cfg.block_size}"
        )

        # Token embeddings: (B, T, d_model)
        tok_emb = self.token_emb(input_ids)

        # Positional indices [0..T-1]
        pos = torch.arange(T, device=input_ids.device).unsqueeze(0) # (1, T)
        pos_emb = self.pos_emb(pos) # (1, T, d_model)

        # Combine + dropout
        x = tok_emb + pos_emb

        return self.drop(x)
    

if __name__ == '__main__':
    cfg = EmbeddingConfig(vocab_size=5000, d_model=256, block_size=128)
    model = TokenPositionalEmbedding(cfg)

    input_ids = torch.randint(0, cfg.vocab_size, (2, 64))
    out = model(input_ids)
    print("Output shape:", out.shape)