encoder-decoder-transformer 0.1.0__py3-none-any.whl

encoder_decoder_transformer/__init__.py

@@ -0,0 +1,27 @@
+"""Encoder-decoder transformer implemented from scratch in PyTorch."""
+
+from .decoder import Decoder, DecoderLayer
+from .embeddings import Embeddings
+from .encoder import Encoder, EncoderLayer
+from .encoder_decoder import EncoderDecoder
+from .feed_forward import FeedForward
+from .make_model import make_model
+from .multi_head_attention import MultiHeadedAttention, attention
+from .utils import Generator, LayerNorm, SublayerConnection, clones
+
+__all__ = [
+    "Decoder",
+    "DecoderLayer",
+    "Embeddings",
+    "Encoder",
+    "EncoderDecoder",
+    "EncoderLayer",
+    "FeedForward",
+    "Generator",
+    "LayerNorm",
+    "MultiHeadedAttention",
+    "SublayerConnection",
+    "attention",
+    "clones",
+    "make_model",
+]

encoder_decoder_transformer/decoder.py

@@ -0,0 +1,129 @@
+"""Transformer decoder layers and stack."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+from torch import Tensor, nn
+
+from .utils import LayerNorm, SublayerConnection, clones
+
+if TYPE_CHECKING:
+    from .feed_forward import FeedForward
+    from .multi_head_attention import MultiHeadedAttention
+
+
+class DecoderLayer(nn.Module):
+    """A single decoder layer with self-attention, cross-attention and FFN.
+
+    Args:
+        d_model: Feature dimension of the input.
+        self_attn: Multi-head self-attention module.
+        cross_attn: Multi-head cross-attention module.
+        ff: Position-wise feed-forward network.
+        dropout: Dropout probability used inside the sublayer connections.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        self_attn: MultiHeadedAttention,
+        cross_attn: MultiHeadedAttention,
+        ff: FeedForward,
+        dropout: float = 0.1,
+    ) -> None:
+        """Initialize the attention and feed-forward sublayers.
+
+        Args:
+            d_model: Feature dimension of the input.
+            self_attn: Multi-head self-attention module.
+            cross_attn: Multi-head cross-attention module.
+            ff: Position-wise feed-forward network.
+            dropout: Dropout probability used inside the sublayer connections.
+        """
+        super().__init__()
+        self.self_attn: MultiHeadedAttention = self_attn
+        self.cross_attn: MultiHeadedAttention = cross_attn
+        self.ff: FeedForward = ff
+        self.sublayers: nn.ModuleList = nn.ModuleList(
+            [
+                SublayerConnection(d_model, dropout),
+                SublayerConnection(d_model, dropout),
+                SublayerConnection(d_model, dropout),
+            ]
+        )
+
+    def forward(
+        self,
+        x: Tensor,
+        encoder_output: Tensor,
+        src_mask: Tensor,
+        tgt_mask: Tensor,
+    ) -> Tensor:
+        """Apply self-attention, cross-attention and the feed-forward network.
+
+        Args:
+            x: Target input tensor of shape ``(batch, tgt_seq, d_model)``.
+            encoder_output: Encoder output of shape
+                ``(batch, src_seq, d_model)``.
+            src_mask: Source attention mask broadcastable to
+                ``(batch, 1, tgt_seq, src_seq)``.
+            tgt_mask: Target attention mask broadcastable to
+                ``(batch, 1, tgt_seq, tgt_seq)``.
+
+        Returns:
+            A tensor of shape ``(batch, tgt_seq, d_model)``.
+        """
+        x = self.sublayers[0](x, lambda t: self.self_attn(t, t, t, tgt_mask))
+        x = self.sublayers[1](
+            x, lambda t: self.cross_attn(t, encoder_output, encoder_output, src_mask)
+        )
+        x = self.sublayers[2](x, self.ff)
+        return x
+
+
+class Decoder(nn.Module):
+    """Stack of ``N`` decoder layers followed by layer normalization.
+
+    Args:
+        decoder_layer: The decoder layer to clone.
+        n: Number of decoder layers in the stack.
+        d_model: Feature dimension of the input.
+    """
+
+    def __init__(self, decoder_layer: DecoderLayer, n: int, d_model: int) -> None:
+        """Initialize the cloned layers and final layer norm.
+
+        Args:
+            decoder_layer: The decoder layer to clone.
+            n: Number of decoder layers in the stack.
+            d_model: Feature dimension of the input.
+        """
+        super().__init__()
+        self.layers: nn.ModuleList = clones(decoder_layer, n)
+        self.layer_norm: LayerNorm = LayerNorm(d_model)
+
+    def forward(
+        self,
+        x: Tensor,
+        encoder_output: Tensor,
+        src_mask: Tensor,
+        tgt_mask: Tensor,
+    ) -> Tensor:
+        """Run ``x`` through every decoder layer.
+
+        Args:
+            x: Target input tensor of shape ``(batch, tgt_seq, d_model)``.
+            encoder_output: Encoder output of shape
+                ``(batch, src_seq, d_model)``.
+            src_mask: Source attention mask broadcastable to
+                ``(batch, 1, tgt_seq, src_seq)``.
+            tgt_mask: Target attention mask broadcastable to
+                ``(batch, 1, tgt_seq, tgt_seq)``.
+
+        Returns:
+            A tensor of shape ``(batch, tgt_seq, d_model)``.
+        """
+        for layer in self.layers:
+            x = layer(x, encoder_output, src_mask, tgt_mask)
+        return cast("Tensor", self.layer_norm(x))

encoder_decoder_transformer/embeddings.py

@@ -0,0 +1,39 @@
+"""Token embedding layer."""
+
+from __future__ import annotations
+
+import math
+from typing import cast
+
+from torch import Tensor, nn
+
+
+class Embeddings(nn.Module):
+    """Embed token indices and scale them by ``sqrt(d_model)``.
+
+    Args:
+        d_model: Dimension of the embedding vectors.
+        vocab_size: Number of tokens in the vocabulary.
+    """
+
+    def __init__(self, d_model: int, vocab_size: int) -> None:
+        """Initialize the embedding table.
+
+        Args:
+            d_model: Dimension of the embedding vectors.
+            vocab_size: Number of tokens in the vocabulary.
+        """
+        super().__init__()
+        self.embeddings: nn.Embedding = nn.Embedding(vocab_size, d_model)
+        self.d_model: int = d_model
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Embed ``x`` and scale by ``sqrt(d_model)``.
+
+        Args:
+            x: Token indices of shape ``(batch, seq)``.
+
+        Returns:
+            Embeddings of shape ``(batch, seq, d_model)``.
+        """
+        return cast("Tensor", self.embeddings(x) * math.sqrt(self.d_model))

encoder_decoder_transformer/encoder.py

@@ -0,0 +1,101 @@
+"""Transformer encoder layers and stack."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+from torch import Tensor, nn
+
+from .utils import LayerNorm, SublayerConnection, clones
+
+if TYPE_CHECKING:
+    from .feed_forward import FeedForward
+    from .multi_head_attention import MultiHeadedAttention
+
+
+class EncoderLayer(nn.Module):
+    """A single encoder layer with self-attention and a feed-forward network.
+
+    Args:
+        d_model: Feature dimension of the input.
+        self_attn: Multi-head self-attention module.
+        ff: Position-wise feed-forward network.
+        dropout: Dropout probability used inside the sublayer connections.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        self_attn: MultiHeadedAttention,
+        ff: FeedForward,
+        dropout: float = 0.1,
+    ) -> None:
+        """Initialize the attention and feed-forward sublayers.
+
+        Args:
+            d_model: Feature dimension of the input.
+            self_attn: Multi-head self-attention module.
+            ff: Position-wise feed-forward network.
+            dropout: Dropout probability used inside the sublayer connections.
+        """
+        super().__init__()
+        self.self_attn: MultiHeadedAttention = self_attn
+        self.ff: FeedForward = ff
+        self.sublayers: nn.ModuleList = nn.ModuleList(
+            [
+                SublayerConnection(d_model, dropout),
+                SublayerConnection(d_model, dropout),
+            ]
+        )
+
+    def forward(self, x: Tensor, mask: Tensor) -> Tensor:
+        """Apply self-attention then the feed-forward network.
+
+        Args:
+            x: Input tensor of shape ``(batch, seq, d_model)``.
+            mask: Attention mask broadcastable to
+                ``(batch, 1, seq, seq)``.
+
+        Returns:
+            A tensor of shape ``(batch, seq, d_model)``.
+        """
+        x = self.sublayers[0](x, lambda t: self.self_attn(t, t, t, mask))
+        x = self.sublayers[1](x, self.ff)
+        return x
+
+
+class Encoder(nn.Module):
+    """Stack of ``N`` encoder layers followed by layer normalization.
+
+    Args:
+        encoder_layer: The encoder layer to clone.
+        n: Number of encoder layers in the stack.
+        d_model: Feature dimension of the input.
+    """
+
+    def __init__(self, encoder_layer: EncoderLayer, n: int, d_model: int) -> None:
+        """Initialize the cloned layers and final layer norm.
+
+        Args:
+            encoder_layer: The encoder layer to clone.
+            n: Number of encoder layers in the stack.
+            d_model: Feature dimension of the input.
+        """
+        super().__init__()
+        self.layers: nn.ModuleList = clones(encoder_layer, n)
+        self.layer_norm: LayerNorm = LayerNorm(d_model)
+
+    def forward(self, x: Tensor, mask: Tensor) -> Tensor:
+        """Run ``x`` through every encoder layer.
+
+        Args:
+            x: Input tensor of shape ``(batch, seq, d_model)``.
+            mask: Attention mask broadcastable to
+                ``(batch, 1, seq, seq)``.
+
+        Returns:
+            A tensor of shape ``(batch, seq, d_model)``.
+        """
+        for layer in self.layers:
+            x = layer(x, mask)
+        return cast("Tensor", self.layer_norm(x))

encoder_decoder_transformer/encoder_decoder.py

@@ -0,0 +1,104 @@
+"""Full encoder-decoder transformer model."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+from torch import Tensor, nn
+
+if TYPE_CHECKING:
+    from .decoder import Decoder
+    from .embeddings import Embeddings
+    from .encoder import Encoder
+    from .utils import Generator
+
+
+class EncoderDecoder(nn.Module):
+    """Combine an encoder, decoder, generator and token embeddings.
+
+    Args:
+        encoder: The encoder stack.
+        decoder: The decoder stack.
+        generator: Output projection producing vocab log-probabilities.
+        src_emb: Source token embeddings.
+        tgt_emb: Target token embeddings.
+    """
+
+    def __init__(
+        self,
+        encoder: Encoder,
+        decoder: Decoder,
+        generator: Generator,
+        src_emb: Embeddings,
+        tgt_emb: Embeddings,
+    ) -> None:
+        """Store the encoder, decoder, generator and embeddings.
+
+        Args:
+            encoder: The encoder stack.
+            decoder: The decoder stack.
+            generator: Output projection producing vocab log-probabilities.
+            src_emb: Source token embeddings.
+            tgt_emb: Target token embeddings.
+        """
+        super().__init__()
+        self.encoder: Encoder = encoder
+        self.decoder: Decoder = decoder
+        self.generator: Generator = generator
+        self.src_emb: Embeddings = src_emb
+        self.tgt_emb: Embeddings = tgt_emb
+
+    def forward(
+        self,
+        src: Tensor,
+        target: Tensor,
+        src_mask: Tensor,
+        tgt_mask: Tensor,
+    ) -> Tensor:
+        """Encode ``src`` and decode ``target`` into vocab log-probabilities.
+
+        Args:
+            src: Source token indices of shape ``(batch, src_seq)``.
+            target: Target token indices of shape ``(batch, tgt_seq)``.
+            src_mask: Source attention mask.
+            tgt_mask: Target attention mask.
+
+        Returns:
+            Log-probabilities of shape ``(batch, tgt_seq, tgt_vocab_size)``.
+        """
+        memory = self.encode(src, src_mask)
+        decoder_output = self.decode(target, memory, src_mask, tgt_mask)
+        return cast("Tensor", self.generator(decoder_output))
+
+    def encode(self, src: Tensor, src_mask: Tensor) -> Tensor:
+        """Embed and encode the source tokens.
+
+        Args:
+            src: Source token indices of shape ``(batch, src_seq)``.
+            src_mask: Source attention mask.
+
+        Returns:
+            Encoder memory of shape ``(batch, src_seq, d_model)``.
+        """
+        return cast("Tensor", self.encoder(self.src_emb(src), src_mask))
+
+    def decode(
+        self,
+        tgt: Tensor,
+        encoder_output: Tensor,
+        src_mask: Tensor,
+        tgt_mask: Tensor,
+    ) -> Tensor:
+        """Embed and decode the target tokens given the encoder output.
+
+        Args:
+            tgt: Target token indices of shape ``(batch, tgt_seq)``.
+            encoder_output: Encoder memory of shape
+                ``(batch, src_seq, d_model)``.
+            src_mask: Source attention mask.
+            tgt_mask: Target attention mask.
+
+        Returns:
+            Decoder features of shape ``(batch, tgt_seq, d_model)``.
+        """
+        return cast("Tensor", self.decoder(self.tgt_emb(tgt), encoder_output, src_mask, tgt_mask))

encoder_decoder_transformer/feed_forward.py

@@ -0,0 +1,42 @@
+"""Position-wise feed-forward network."""
+
+from __future__ import annotations
+
+from typing import cast
+
+from torch import Tensor, nn
+
+
+class FeedForward(nn.Module):
+    """Two-layer position-wise feed-forward network with GELU activation.
+
+    Args:
+        d_model: Input and output feature dimension.
+        d_expand: Hidden (expanded) feature dimension.
+        dropout: Dropout probability applied between the linear layers.
+    """
+
+    def __init__(self, d_model: int, d_expand: int, dropout: float = 0.1) -> None:
+        """Initialize the linear layers and activation.
+
+        Args:
+            d_model: Input and output feature dimension.
+            d_expand: Hidden (expanded) feature dimension.
+            dropout: Dropout probability applied between the linear layers.
+        """
+        super().__init__()
+        self.dropout: nn.Dropout = nn.Dropout(dropout)
+        self.gelu: nn.GELU = nn.GELU()
+        self.a_1: nn.Linear = nn.Linear(d_model, d_expand)
+        self.a_2: nn.Linear = nn.Linear(d_expand, d_model)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Apply the feed-forward network to ``x``.
+
+        Args:
+            x: Input tensor of shape ``(batch, seq, d_model)``.
+
+        Returns:
+            A tensor of shape ``(batch, seq, d_model)``.
+        """
+        return cast("Tensor", self.a_2(self.dropout(self.gelu(self.a_1(x)))))

encoder_decoder_transformer/make_model.py

@@ -0,0 +1,58 @@
+"""Factory for building a complete encoder-decoder transformer."""
+
+from __future__ import annotations
+
+import copy
+
+from torch import nn
+
+from .decoder import Decoder, DecoderLayer
+from .embeddings import Embeddings
+from .encoder import Encoder, EncoderLayer
+from .encoder_decoder import EncoderDecoder
+from .feed_forward import FeedForward
+from .multi_head_attention import MultiHeadedAttention
+from .utils import Generator
+
+
+def make_model(
+    src_vocab_size: int,
+    tgt_vocab_size: int,
+    max_context: int = 1024,
+    N: int = 6,
+    d_model: int = 512,
+    d_ff: int = 2048,
+    h: int = 8,
+    dropout: float = 0.1,
+) -> EncoderDecoder:
+    """Build an encoder-decoder transformer with Xavier initialization.
+
+    Args:
+        src_vocab_size: Size of the source vocabulary.
+        tgt_vocab_size: Size of the target vocabulary.
+        max_context: Maximum sequence length for the RoPE tables.
+        N: Number of encoder and decoder layers.
+        d_model: Feature dimension of the model.
+        d_ff: Hidden dimension of the feed-forward network.
+        h: Number of attention heads.
+        dropout: Dropout probability used throughout the model.
+
+    Returns:
+        An initialized :class:`EncoderDecoder` model.
+    """
+    attn = MultiHeadedAttention(h, d_model, max_context, dropout)
+    cross_attn = copy.deepcopy(attn)
+    ff = FeedForward(d_model, d_ff, dropout)
+    model = EncoderDecoder(
+        Encoder(EncoderLayer(d_model, attn, ff, dropout), N, d_model),
+        Decoder(DecoderLayer(d_model, attn, cross_attn, ff, dropout), N, d_model),
+        Generator(d_model, tgt_vocab_size),
+        Embeddings(d_model, src_vocab_size),
+        Embeddings(d_model, tgt_vocab_size),
+    )
+
+    for p in model.parameters():
+        if p.dim() > 1:
+            nn.init.xavier_uniform_(p)
+
+    return model

encoder_decoder_transformer/multi_head_attention.py

@@ -0,0 +1,146 @@
+"""Multi-head attention with rotary position embeddings (RoPE)."""
+
+from __future__ import annotations
+
+import math
+from typing import cast
+
+import torch
+from torch import Tensor, nn
+
+from .utils import clones
+
+
+def attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    mask: Tensor | None,
+    dropout: nn.Dropout | None,
+) -> tuple[Tensor, Tensor]:
+    """Compute scaled dot-product attention.
+
+    Args:
+        query: Query tensor of shape ``(batch, heads, seq_q, d_k)``.
+        key: Key tensor of shape ``(batch, heads, seq_k, d_k)``.
+        value: Value tensor of shape ``(batch, heads, seq_k, d_v)``.
+        mask: Optional mask broadcastable to ``(batch, heads, seq_q, seq_k)``.
+            Positions where the mask equals zero are ignored.
+        dropout: Optional dropout applied to the attention weights.
+
+    Returns:
+        A tuple ``(output, attn)`` where ``output`` has shape
+        ``(batch, heads, seq_q, d_v)`` and ``attn`` has shape
+        ``(batch, heads, seq_q, seq_k)``.
+    """
+    d_k = query.size(-1)
+    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
+    if mask is not None:
+        scores = scores.masked_fill(mask == 0, -1e9)
+    p_attn = scores.softmax(dim=-1)
+    if dropout is not None:
+        p_attn = dropout(p_attn)
+    return torch.matmul(p_attn, value), p_attn
+
+
+class MultiHeadedAttention(nn.Module):
+    """Multi-head self/cross attention using rotary position embeddings.
+
+    Args:
+        h: Number of attention heads.
+        d_model: Feature dimension of the input.
+        max_context: Maximum sequence length supported by the RoPE tables.
+        dropout: Dropout probability applied to the attention weights.
+    """
+
+    def __init__(
+        self,
+        h: int,
+        d_model: int,
+        max_context: int,
+        dropout: float = 0.1,
+    ) -> None:
+        """Initialize projections, dropout, and the RoPE buffers.
+
+        Args:
+            h: Number of attention heads.
+            d_model: Feature dimension of the input.
+            max_context: Maximum sequence length supported by the RoPE tables.
+            dropout: Dropout probability applied to the attention weights.
+
+        Raises:
+            AssertionError: If ``d_model`` is not divisible by ``h``.
+        """
+        super().__init__()
+        assert d_model % h == 0
+        self.d_k: int = d_model // h
+        self.h: int = h
+        self.linears: nn.ModuleList = clones(nn.Linear(d_model, d_model), 4)
+        self.attn: Tensor | None = None
+        self.dropout: nn.Dropout = nn.Dropout(p=dropout)
+
+        y_pos = torch.arange(max_context)
+        freqs = 10000 ** (-2 * torch.arange(self.d_k // 2, dtype=torch.float) / self.d_k)
+        angles = torch.outer(y_pos, freqs)
+        self.cos_table: Tensor
+        self.sin_table: Tensor
+        self.register_buffer("cos_table", torch.cos(angles))
+        self.register_buffer("sin_table", torch.sin(angles))
+
+    def rope_rotate(self, tensor: Tensor) -> Tensor:
+        """Apply rotary position embeddings to ``tensor``.
+
+        Args:
+            tensor: Tensor of shape ``(batch, heads, seq, d_k)``.
+
+        Returns:
+            The rotated tensor with the same shape as ``tensor``.
+        """
+        input_len = tensor.shape[2]
+
+        cos = self.cos_table[:input_len].unsqueeze(0).unsqueeze(0)
+        sin = self.sin_table[:input_len].unsqueeze(0).unsqueeze(0)
+
+        first_half = tensor[..., : self.d_k // 2]
+        second_half = tensor[..., self.d_k // 2 :]
+
+        a = first_half * cos - second_half * sin
+        b = first_half * sin + second_half * cos
+        return torch.concat((a, b), -1)
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        mask: Tensor | None,
+    ) -> Tensor:
+        """Apply multi-head attention to ``query``, ``key`` and ``value``.
+
+        Args:
+            query: Query tensor of shape ``(batch, seq_q, d_model)``.
+            key: Key tensor of shape ``(batch, seq_k, d_model)``.
+            value: Value tensor of shape ``(batch, seq_k, d_model)``.
+            mask: Optional mask broadcastable to
+                ``(batch, 1, seq_q, seq_k)``. Positions where the mask equals
+                zero are ignored.
+
+        Returns:
+            The attention output of shape ``(batch, seq_q, d_model)``.
+        """
+        if mask is not None:
+            mask = mask.unsqueeze(1)
+        nbatches = query.size(0)
+
+        query, key, value = [
+            lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
+            for lin, x in zip(self.linears, (query, key, value), strict=False)
+        ]
+
+        rotated_query = self.rope_rotate(query)
+        rotated_key = self.rope_rotate(key)
+
+        x, self.attn = attention(rotated_query, rotated_key, value, mask=mask, dropout=self.dropout)
+
+        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
+        return cast("Tensor", self.linears[-1](x))

encoder_decoder_transformer/utils.py

@@ -0,0 +1,126 @@
+"""Reusable building blocks for the encoder-decoder transformer."""
+
+from __future__ import annotations
+
+import copy
+from typing import TYPE_CHECKING, cast
+
+import torch
+from torch import Tensor, nn
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+def clones[ModuleT: nn.Module](module: ModuleT, n: int) -> nn.ModuleList:
+    """Create ``n`` independent deep copies of ``module``.
+
+    Args:
+        module: The module to duplicate.
+        n: Number of copies to create.
+
+    Returns:
+        A ``ModuleList`` holding the ``n`` deep copies of ``module``.
+    """
+    return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])
+
+
+class LayerNorm(nn.Module):
+    """Layer normalization over the last dimension of the input.
+
+    Args:
+        features: Size of the normalized (last) dimension.
+        eps: Small constant added to the standard deviation for numerical
+            stability.
+    """
+
+    def __init__(self, features: int, eps: float = 1e-6) -> None:
+        """Initialize the learnable affine parameters.
+
+        Args:
+            features: Size of the normalized (last) dimension.
+            eps: Small constant added to the standard deviation.
+        """
+        super().__init__()
+        self.a_2: nn.Parameter = nn.Parameter(torch.ones(features))
+        self.b_2: nn.Parameter = nn.Parameter(torch.zeros(features))
+        self.eps: float = eps
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Normalize ``x`` along its last dimension.
+
+        Args:
+            x: Input tensor of shape ``(..., features)``.
+
+        Returns:
+            A tensor with the same shape as ``x``, normalized along the last
+            dimension.
+        """
+        mean = x.mean(-1, keepdim=True)
+        std = x.std(-1, keepdim=True)
+        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
+
+
+class SublayerConnection(nn.Module):
+    """Residual connection wrapped around a pre-normalized sublayer.
+
+    Args:
+        d_model: Feature dimension of the input.
+        dropout: Dropout probability applied to the sublayer output.
+    """
+
+    def __init__(self, d_model: int, dropout: float = 0.1) -> None:
+        """Initialize the layer norm and dropout module.
+
+        Args:
+            d_model: Feature dimension of the input.
+            dropout: Dropout probability applied to the sublayer output.
+        """
+        super().__init__()
+        self.dropout: nn.Dropout = nn.Dropout(p=dropout)
+        self.layer_norm: LayerNorm = LayerNorm(d_model)
+
+    def forward(self, x: Tensor, layer: Callable[[Tensor], Tensor]) -> Tensor:
+        """Apply ``layer`` to the normalized input and add the residual.
+
+        Args:
+            x: Input tensor of shape ``(..., d_model)``.
+            layer: Callable applied to the normalized input.
+
+        Returns:
+            A tensor of shape ``(..., d_model)`` produced by
+            ``x + dropout(layer(layer_norm(x)))``.
+        """
+        x_norm = self.layer_norm(x)
+        output_x = layer(x_norm)
+        return cast("Tensor", x + self.dropout(output_x))
+
+
+class Generator(nn.Module):
+    """Map decoder features to log-probabilities over a vocabulary.
+
+    Args:
+        d_model: Feature dimension of the decoder output.
+        vocab_size: Number of tokens in the target vocabulary.
+    """
+
+    def __init__(self, d_model: int, vocab_size: int) -> None:
+        """Initialize the output projection.
+
+        Args:
+            d_model: Feature dimension of the decoder output.
+            vocab_size: Number of tokens in the target vocabulary.
+        """
+        super().__init__()
+        self.vocab_map: nn.Linear = nn.Linear(d_model, vocab_size)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Project ``x`` and apply a log-softmax over the vocabulary.
+
+        Args:
+            x: Decoder features of shape ``(batch, seq, d_model)``.
+
+        Returns:
+            Log-probabilities of shape ``(batch, seq, vocab_size)``.
+        """
+        return torch.log_softmax(self.vocab_map(x), -1)

encoder_decoder_transformer-0.1.0.dist-info/METADATA

@@ -0,0 +1,85 @@
+Metadata-Version: 2.4
+Name: encoder-decoder-transformer
+Version: 0.1.0
+Summary: A from-scratch PyTorch implementation of the Transformer encoder-decoder architecture.
+Keywords: transformer,pytorch,attention,encoder-decoder,deep-learning
+Author: Joe Rice
+Author-email: Joe Rice <2Joerice@gmail.com>
+License-Expression: MIT
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Dist: torch>=2.14.1
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+
+# encoder-decoder-transformer
+
+A from-scratch PyTorch implementation of the Transformer encoder-decoder
+architecture, including multi-head attention with rotary position embeddings
+(RoPE), pre-layer normalization, and a position-wise feed-forward network.
+
+## Installation
+
+Install from source with [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv add encoder-decoder-transformer
+```
+
+or with pip:
+
+```bash
+pip install encoder-decoder-transformer
+```
+
+The package targets Python 3.13 and depends on PyTorch.
+
+## Usage
+
+Build a model and run a forward pass. Token indices are integer tensors of
+shape `(batch, seq)`; the model returns log-probabilities of shape
+`(batch, tgt_seq, tgt_vocab_size)`. Attention masks are boolean/integer tensors
+of shape `(batch, seq_q, seq_k)` (the model adds the head dimension).
+
+```python
+import torch
+
+from encoder_decoder import make_model
+
+model = make_model(
+    src_vocab_size=1000,
+    tgt_vocab_size=1000,
+    max_context=128,
+    N=2,
+    d_model=64,
+    d_ff=256,
+    h=4,
+    dropout=0.1,
+)
+
+src = torch.randint(0, 1000, (2, 8))
+tgt = torch.randint(0, 1000, (2, 8))
+src_mask = torch.ones(2, 8, 8)
+tgt_mask = torch.ones(2, 8, 8)
+
+log_probs = model(src, tgt, src_mask, tgt_mask)
+print(log_probs.shape)  # torch.Size([2, 8, 1000])
+```
+
+## Development
+
+Run the linters and type checker with:
+
+```bash
+uv run ruff check .
+uv run ruff format --check .
+uv run mypy
+```
+
+## License
+
+MIT

encoder_decoder_transformer-0.1.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+encoder_decoder_transformer/__init__.py,sha256=PyP_wymKYM9ltnNf4uObAF0vQNJxipdolNQLIm4c--c,725
+encoder_decoder_transformer/decoder.py,sha256=-whNa1m6UyMc5R18IkHQqxMgfyfJNqU-SHo6zBS6VQ0,4360
+encoder_decoder_transformer/embeddings.py,sha256=JcgOwaolD51ewZ7qKiwTkKqQkc5t3KfJwWRjWYF1e-w,1079
+encoder_decoder_transformer/encoder.py,sha256=qPwr-dn2GiwJXfDNtNBWeBhRBzpxmuuKY13AEAYNl_o,3242
+encoder_decoder_transformer/encoder_decoder.py,sha256=H-h7fqhVrJRyIY1bDFctY6TJzrmNeIuuvTYSd86-zeY,3288
+encoder_decoder_transformer/feed_forward.py,sha256=7vztLtmg2Z6FMYGs_2eBVL6sfskmDfD6ONUp8XRv8eI,1377
+encoder_decoder_transformer/make_model.py,sha256=U-cVQDMjUxDJWHAZ30N-M2Dj5XQj58CfX_2-86NdjFo,1828
+encoder_decoder_transformer/multi_head_attention.py,sha256=kkDAexGpPKizRAsq3IUh9OQtYHulfHTzk85R0fklCU8,4958
+encoder_decoder_transformer/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+encoder_decoder_transformer/utils.py,sha256=_WGUivOegGokiE1JAflOowqlIzrtQGHHdABG31TYZ8A,4001
+encoder_decoder_transformer-0.1.0.dist-info/WHEEL,sha256=cmC5s21ojypbVslldL7IJq3hjZH-tINy4rziKePFsG0,81
+encoder_decoder_transformer-0.1.0.dist-info/METADATA,sha256=jePIVQ5u6KHdJfifuuwK-8Sm7CgcsONvOEVVDICLtFA,2143
+encoder_decoder_transformer-0.1.0.dist-info/RECORD,,

encoder_decoder_transformer-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.12.23
+Root-Is-Purelib: true
+Tag: py3-none-any