src/mini_transformer/modules/lm_head.py

"""Lightweight language-model head projecting hidden states to vocab logits."""

from __future__ import annotations

import torch
import torch.nn as nn
from torch import Tensor

__all__ = ["LMHead"]


class LMHead(nn.Module):
    """
    Language modeling head: projects hidden states to vocabulary logits.

    Args:
        d_model (int): Model hidden dimension (>0).
        vocab_size (int): Vocabulary size (>0).

    Input:
        x (Tensor): shape (B, S, D) with D == d_model.

    Output:
        logits (Tensor): shape (B, S, V) with V == vocab_size.
    """

    def __init__(self, d_model: int, vocab_size: int):
        super().__init__()
        if not isinstance(vocab_size, int):
            raise TypeError(f"vocab_size must be an int, got {type(vocab_size)}")
        if not isinstance(d_model, int):
            raise TypeError(f"d_model must be an int, got {type(d_model)}")
        if vocab_size <= 0:
            raise ValueError(f"vocab_size must be strictly greater than 0, got {vocab_size}")
        if d_model <= 0:
            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")

        self.d_model = d_model
        self.vocab_size = vocab_size
        self.fc = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, x: Tensor) -> Tensor:
        if not isinstance(x, torch.Tensor):
            raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
        if x.dim() != 3:
            raise ValueError(
                f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
            )

        B, S, D = x.shape
        if D != self.d_model:
            raise ValueError(f"Last dim {D} must match d_model {self.d_model}")

        return self.fc(x)  # (B, S, V)