Torchformers: A Lightweight and Highly Customizable Transformers Library
import torch
from torchformers import Transformer


# embedding arity
d_model = 512
# number of encoder/decoder blocks
N = 6
# vocabulary length
vocab_len = 300

model = Transformer(d_model, N, vocab_len).eval()

# single input
encoder_input = torch.rand((vocab_len, d_model))
decoder_input = torch.rand((vocab_len, d_model))

output = model(encoder_input, decoder_input)
assert output.size() == (
    vocab_len,
    vocab_len,
), f"Output shape doesn't match the expected ({vocab_len}, {vocab_len})"

batch_size = 30

# batched input
encoder_input_batched = torch.rand((batch_size, vocab_len, d_model))
decoder_input_batched = torch.rand((batch_size, vocab_len, d_model))

output = model(encoder_input_batched, decoder_input_batched)
assert output.size() == (
    batch_size,
    vocab_len,
    vocab_len,
), f"Output shape doesn't match the expected ({batch_size}, {vocab_len}, {vocab_len})"

# if all checks pass print OK
print("OK")

demo.py
max-fofanov/torchformers
from setuptools import setup, find_packages


setup(
    name="torchformers",
    version="1.0.1",
    package_dir={"": "src"},
    packages=find_packages("src"),
    install_requires=["torch"],
    author="Max Fofanov",
    author_email="max.fofanov@gmail.com",
    description="Torchformers is a sleek and efficient Python library, providing essential and customizable transformer\
     models and components for rapid development and experimentation in Natural Language Processing, all with the\
      simplicity and power of PyTorch.",
    long_description=open("README.md").read(),
    long_description_content_type="text/markdown",
    url="https://github.com/max-fofanov/torchformers",
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
    ],
    python_requires=">=3.8",
)

setup.py
max-fofanov/torchformers
from .transformer import Transformer, TransformerDecoder, TransformerDecoderBlock
from .bert import BERT
from .gpt import GPT, GPTDecoder, GPTDecoderBlock
from .utils import (
    SelfAttentionBlock,
    MaskedSelfAttentionBlock,
    CrossAttentionBlock,
    FeedForwardBlock,
    Encoder,
    EncoderBlock,
)

src/torchformers/__init__.py
max-fofanov/torchformers
import torch
from torch import nn

from ..utils import Encoder


class BERT(nn.Module):
    def __init__(self, d_model: int, N: int, num_classes: int, num_heads: int = 1):
        super().__init__()

        self.encoder = Encoder(d_model, N, num_heads=num_heads)

        self.head = nn.Sequential(nn.Linear(d_model, num_classes), nn.Softmax(dim=-1))

    def forward(self, x) -> torch.Tensor:
        encoder_output = self.encoder(x)
        return self.head(encoder_output)

src/torchformers/bert/bert.py
max-fofanov/torchformers
from .gpt import GPT
from .decoder import GPTDecoder, GPTDecoderBlock

src/torchformers/gpt/__init__.py
max-fofanov/torchformers
import torch
from torch import nn

from ..utils import MaskedSelfAttentionBlock, FeedForwardBlock


class GPTDecoder(nn.Module):
    def __init__(self, d_model: int, N: int, vocab_len: int, num_heads: int = 1):
        super().__init__()

        self.decoders = nn.Sequential(
            *[
                GPTDecoderBlock(d_model, vocab_len, num_heads=num_heads)
                for _ in range(N)
            ]
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.decoders(x)


class GPTDecoderBlock(nn.Module):
    def __init__(self, d_model: int, vocab_len: int, num_heads: int = 1):
        super().__init__()

        self.masked_self_att = MaskedSelfAttentionBlock(
            d_model, vocab_len, num_heads=num_heads
        )

        self.ff = FeedForwardBlock(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.masked_self_att(x)
        return self.ff(x)

src/torchformers/gpt/decoder.py
max-fofanov/torchformers
import torch
from torch import nn

from .decoder import GPTDecoder


class GPT(nn.Module):
    def __init__(self, d_model: int, N: int, vocab_len: int, num_heads: int = 1):
        super().__init__()

        self.decoder = GPTDecoder(d_model, N, vocab_len, num_heads=num_heads)

        self.head = nn.Sequential(nn.Linear(d_model, vocab_len), nn.Softmax(dim=-1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.decoder(x)
        return self.head(x)

src/torchformers/gpt/gpt.py
max-fofanov/torchformers
from .transformer import Transformer
from .decoder import TransformerDecoder, TransformerDecoderBlock

src/torchformers/transformer/__init__.py
max-fofanov/torchformers
import torch
from torch import nn

from ..utils import MaskedSelfAttentionBlock, CrossAttentionBlock, FeedForwardBlock


class TransformerDecoder(nn.Module):
    def __init__(self, d_model: int, N: int, vocab_len: int, num_heads: int = 1):
        super().__init__()

        self.decoders = [
            TransformerDecoderBlock(d_model, vocab_len, num_heads=num_heads)
            for _ in range(N)
        ]

    def forward(self, x: torch.Tensor, encoder_output: torch.Tensor) -> torch.Tensor:
        for decoder in self.decoders:
            x = decoder(x, encoder_output)
        return x


class TransformerDecoderBlock(nn.Module):
    def __init__(self, d_model: int, vocab_len: int, num_heads: int = 1):
        super().__init__()

        self.masked_self_att = MaskedSelfAttentionBlock(
            d_model, vocab_len, num_heads=num_heads
        )
        self.cross_att = CrossAttentionBlock(d_model, num_heads=num_heads)

        self.ff = FeedForwardBlock(d_model)

    def forward(self, x: torch.Tensor, encoder_output: torch.Tensor) -> torch.Tensor:
        x = self.masked_self_att(x)
        x = self.cross_att(x, encoder_output)
        return self.ff(x)

src/torchformers/transformer/decoder.py
max-fofanov/torchformers
import torch
from torch import nn

from ..utils import Encoder
from .decoder import TransformerDecoder


class Transformer(nn.Module):
    def __init__(self, d_model: int, N: int, vocab_len: int, num_heads: int = 1):
        super().__init__()

        self.encoder = Encoder(d_model, N, num_heads=num_heads)
        self.decoder = TransformerDecoder(d_model, N, vocab_len, num_heads=num_heads)

        self.head = nn.Sequential(nn.Linear(d_model, vocab_len), nn.Softmax(dim=-1))

    def forward(
        self, encoder_input: torch.Tensor, decoder_input: torch.Tensor
    ) -> torch.Tensor:
        encoder_output = self.encoder(encoder_input)
        decoder_output = self.decoder(decoder_input, encoder_output)
        return self.head(decoder_output)

src/torchformers/transformer/transformer.py
max-fofanov/torchformers
from .transformer_blocks import (
    SelfAttentionBlock,
    MaskedSelfAttentionBlock,
    CrossAttentionBlock,
    FeedForwardBlock,
)
from .encoder import Encoder, EncoderBlock

src/torchformers/utils/__init__.py
max-fofanov/torchformers
import torch
from torch import nn

from ..utils import SelfAttentionBlock, FeedForwardBlock


class Encoder(nn.Module):
    def __init__(self, d_model: int, N: int, num_heads: int = 1):
        super().__init__()

        self.encoders = nn.Sequential(
            *[EncoderBlock(d_model, num_heads=num_heads) for _ in range(N)]
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.encoders(x)


class EncoderBlock(nn.Module):
    def __init__(self, d_model: int, num_heads: int = 1):
        super().__init__()

        self.self_att = SelfAttentionBlock(d_model, num_heads=num_heads)
        self.ff = FeedForwardBlock(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.self_att(x)
        return self.ff(x)

src/torchformers/utils/encoder.py
max-fofanov/torchformers
from functools import cached_property
from typing import Optional

import torch
from torch import nn


class SelfAttentionBlock(nn.Module):
    def __init__(self, d_model: int, num_heads: int = 1):
        super().__init__()

        self.attn = nn.MultiheadAttention(
            embed_dim=d_model, num_heads=num_heads, batch_first=True
        )
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_att, _ = self.attn(x, x, x)
        return self.norm(x_att + x)


class CrossAttentionBlock(SelfAttentionBlock):
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        x_att, _ = self.attn(x, y, y)
        return self.norm(x_att + x)


class MaskedSelfAttentionBlock(SelfAttentionBlock):
    def __init__(self, d_model: int, vocab_len: int, num_heads: int = 1):
        super().__init__(d_model, num_heads=num_heads)

        self.vocab_len = vocab_len

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_att, _ = self.attn(x, x, x, attn_mask=self.attn_mask)
        return self.norm(x_att + x)

    @cached_property
    def attn_mask(self) -> torch.Tensor:
        mask = torch.triu(torch.ones(self.vocab_len, self.vocab_len), diagonal=1)
        return mask.masked_fill(mask == 1, float("-inf"))


class FeedForwardBlock(nn.Module):
    def __init__(self, d_model: int, hidden_size: Optional[int] = None):
        super().__init__()

        hidden_size = d_model if hidden_size is None else hidden_size

        self.ff = nn.Sequential(
            nn.Linear(in_features=d_model, out_features=hidden_size),
            nn.ReLU(),
            nn.Linear(in_features=hidden_size, out_features=d_model),
        )

        self.norm = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_ff = self.ff(x)
        return self.norm(x_ff + x)

src/torchformers/utils/transformer_blocks.py
max-fofanov/torchformers