embeddings.py

"""
embeddings.py

Handles converting token IDs into embeddings, and adding positional
information to those embeddings so the model knows token order.
"""

import torch
import torch.nn as nn
import math


class PositionalEncoding(nn.Module):
    """
    Adds sinusoidal positional information to token embeddings.

    Attention has no built-in sense of sequence order (unlike RNNs),
    so we inject position information directly into the embeddings.
    """

    def __init__(self, embed_dim: int, max_len: int = 5000):
        super().__init__()

        pe = torch.zeros(max_len, embed_dim)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, embed_dim, 2).float() * (-math.log(10000.0) / embed_dim)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)  # shape: (1, max_len, embed_dim)
        self.register_buffer("pe", pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Add positional encoding to the input embeddings."""
        seq_len = x.size(1)
        return x + self.pe[:, :seq_len, :]


class TokenEmbedding(nn.Module):
    """
    Wraps a standard embedding layer plus positional encoding
    into a single reusable component.
    """

    def __init__(self, vocab_size: int, embed_dim: int, max_len: int = 5000):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
        self.positional_encoding = PositionalEncoding(embed_dim, max_len)

    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        x = self.token_embedding(token_ids)
        x = self.positional_encoding(x)
        return x