Transformer From Scratch (with PyTorch)

This repo contains a full implementation of the famous Transformer architecture built from scratch using PyTorch, from the reasearch paper "Attention Is All You Need", without relying on high level libraries such as HugginFace. This repo not only creates the transformer but also shows all its different usages, such as: Classification, Text Generation, and Translation.

This repo is mainly for study and research purposes. Feedback is welcomed :)

Repo Structure

Transformer_from_Scratch/
│
├── data/                    # Use folder for any data you want to try the code with!
│
├── models/                  # Save the models (tokenizer.pkl and best_model.pt) here!
│
├── notebooks/               # Each notebook contains training phase + trying the saved model
│   ├── classification_train.ipynb       
│   ├── generator_train.ipynb           
│   └── translation_train.ipynb         
│
├── transformer/             # Transformer components
│   ├── decoder.py
│   ├── encoder.py
│   ├── feed_foward.py
│   ├── multi_head_attention.py
│   └── transformer.py       # Full Transformer / Encoder-Decoder model
│
├── sentiment_classifier.py  # Classification / Encoder model
├── text_generator.py        # Text Generation / Decoder model
├── tokenizer.py             # Tokenizer model
├── requirements.txt
├── LICENSE
└── README.md

Index

Tokenizer
Transformer Architecture
Sentiment Classification (Encoder-only)
Text Generation (Decoder-only)
Translation (Encoder-Decoder)
References & Citations
License

Clone Repo

git clone https://github.com/JimWid/Transformer-From-Scratch.git
cd Transformer-From-Scratch

Set Up Virtual Environment

python -m venv venv
On Windows: venv\Scripts\activate
On Mac: source env/bin/activate

Install Dependencies

pip install -r requirements.txt

Tokenizer

I am using a custom tokenizer from Scratch too, but you can use any.

A tokenizer is used to build a vocabulary from the dataset, it maps the words to indices (it sets a word to a number, e.g. "the" → 5) and it adds special tokens such as:

<PAD> # Padding for short sentences to match max_len.
<UNK> # Used for Unknown words (words that are not inside vocabulary)
<SOS> # Used to declare when a setences begins
<EOS> # Used to declare when a sentence ends

We only use <SOS> and <EOS> in Encoder-Decoder models, such as translation.

NOTE: Tokenizers turns sentences -> tokens, by breaking this down into pieces, my tokenizer split them into whole words, I belive this is easier to understand and visualize.

Transformer Architecture

Let's first define all the paramerets we are gonna use:

vocab_size # Size of vocabulary
embedding_size # Vector size of each token
num_layers # number of layers
num_heads # number of attention-heads
d_ff # hidden layer of feed foward network
max_len # max length of sentences
dropout # dropout

There is 4 main components in a Transformer: Postional Encoding, Self-Attention, FeedFoward, Masking, LayerNorm. Let's go over them one by one.

Positional Encoding

The simpliest way to describe this is adding position information to the embeddings vector, therefore during training, the model doesn't just learn meanings of words, but also their positions. Smart rigth?

We inject position information with the Sinusoidal method:

Postional Encoding Code:

    def get_positional_encoding(self, max_len, embedding_size):
        pos = torch.arange(max_len, dtype=torch.float32).unsqueeze(1)
        i = torch.arange(embedding_size, dtype=torch.float32).unsqueeze(0)

        angle_rates = 1 / torch.pow(10000, (2 * (i // 2) / embedding_size))
        angle_rads = pos * angle_rates

        positional_encoding = torch.zeros(max_len, embedding_size)

        positional_encoding[:, 0::2] = torch.sin(angle_rads[:, 0::2]) # Sin for even values/dimensions
        positional_encoding[:, 1::2] = torch.cos(angle_rads[:, 1::2]) # Cos for odd values/dimensions
        
        return positional_encoding.unsqueeze(0)

Self-Attention

Self-Attention is basically a formula:

Where Q, K and V are:

Query (Q)
Key (K)
Value (V)

Each one of these are a matrix full of parameters (weights) that are trainable, with a size of embedding_size.

Inside this equation is it perform what is called Scaled Dot Product:

Multi-Head Attention

Basically concatenating all the outputs of all the attention-heads.

Multi-Head Attention Code:

class MultiHeadAttention(nn.Module):
    def __init__(self, embedding_size, num_heads):
        super(MultiHeadAttention, self).__init__()

        self.embedding_size = embedding_size
        self.num_heads = num_heads
        self.head_dim = embedding_size // num_heads

        # Linear Projections for Q, K, V
        self.Q = nn.Linear(embedding_size, embedding_size)
        self.K = nn.Linear(embedding_size, embedding_size)
        self.V = nn.Linear(embedding_size, embedding_size)

        # Final output projection
        self.output = nn.Linear(embedding_size, embedding_size)

    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        # Q = Query matrix, Shape = [batch, seq_len_q, d_k]
        # K = Key matrix, Shape = [batch, seq_len_k, d_k]
        # V = Value matrix, Shape = [batch, seq_len_k, d_k]

        head_dim = Q.size(-1)
        scores = torch.matmul(Q, K.transpose(-2, -1)) 

        # Scaling scores
        scores = scores / math.sqrt(head_dim)

        # Optional mask
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))

        # Softmax
        attn_weights = F.softmax(scores, dim=-1)

        # Multiplying attention weights with V
        output = torch.matmul(attn_weights, V)
        return output

    def forward(self, Q, K, V, mask=None):

        # Getting Batch Size and Sequence Lengths
        N = Q.shape[0]
        value_len, key_len, query_len = V.shape[1], K.shape[1], Q.shape[1]

        # Linear Projection
        Q = self.Q(Q)
        K = self.K(K)
        V = self.V(V)

        # Spliting into heads
        Q = Q.view(N, query_len, self.num_heads, self.head_dim).transpose(1, 2)
        K = K.view(N, key_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = V.view(N, value_len, self.num_heads, self.head_dim).transpose(1, 2)

        # Using the scaled dot product attention
        attn_output = self.scaled_dot_product_attention(Q, K, V, mask)
        attn_output = attn_output.transpose(1, 2).contiguous().view(N, -1, self.embedding_size)

        # Final Linear Projection
        output = self.output(attn_output)

        return output

Masking

Feed Foward Network

Feed Foward Network provides a non-linear and position-wise transformation, it increases model capacity and it mixes the information across features.

Feed Foward Network Code:

class FeedForward(nn.Module):
    def __init__(self, embedding_size, d_ff, dropout):
        super(FeedForward, self).__init__()

        self.linear1 = nn.Linear(embedding_size, d_ff)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(d_ff, embedding_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.linear2(self.dropout(self.relu(self.linear1(x))))

Add & LayerNorm

I do not use a custom layer norm, I instead use the PyTorch one, LayerNorm is basically another equation, that computes mean and the variance of the features of each token, then it normalizes, turning mean close to 0 and variance close to 1. It has its own parameters scale and shift. During training, it changes scale and shift to determine how much to normalize.

We also have Residual Connections (Add), an easy way to see it is that the Residual Connection simply adds the input to the new changed input. We are basically adding a correction to the token representation, not a full re-write.

I use LayerNorm and Residual in both Encoder and Decoder:

Decoder and Encoder are basically wrappers of all these components.

Encoder:

Encoder Code:

class EncoderBlock(nn.Module):
    def __init__(self, embedding_size, num_heads, d_ff, dropout):
        super(EncoderBlock, self).__init__()

        # Initializing MHA and FFN
        self.attention = MultiHeadAttention(embedding_size, num_heads)
        self.feed_forward = FeedForward(embedding_size, d_ff, dropout)

        # Norm Layers
        self.norm1 = nn.LayerNorm(embedding_size)
        self.norm2 = nn.LayerNorm(embedding_size)

        # Dropouts
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Multi-Head Attention
        attn_output = self.attention(x, x, x, mask)

        # Add & Norm
        x = self.norm1(x + self.dropout1(attn_output))

        # Feed Forward
        ffn_output = self.feed_forward(x)
        output = self.norm2(x + self.dropout2(ffn_output))

        return output

class Encoder(nn.Module):
    def __init__(self, vocab_size, embedding_size, num_layers, num_heads, d_ff, max_len, dropout, device):
        super(Encoder, self).__init__()

        self.embedding = Embedding(vocab_size, embedding_size)
        self.pos_encoding = self.get_positional_encoding(max_len, embedding_size).to(device)
        self.dropout = nn.Dropout(dropout)

        self.layers = nn.ModuleList([
            EncoderBlock(embedding_size, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])

    def forward(self, x, mask=None):

        seq_len = x.size(1)
        x = self.embedding(x)
        x = x + self.pos_encoding[:, :seq_len, :]
        x = self.dropout(x)

        for layer in self.layers:
            x = layer(x, mask)

        return x
    
    def get_positional_encoding(self, max_len, embedding_size):
        pos = torch.arange(max_len, dtype=torch.float32).unsqueeze(1)
        i = torch.arange(embedding_size, dtype=torch.float32).unsqueeze(0)

        angle_rates = 1 / torch.pow(10000, (2 * (i // 2) / embedding_size))
        angle_rads = pos * angle_rates

        positional_encoding = torch.zeros(max_len, embedding_size)

        positional_encoding[:, 0::2] = torch.sin(angle_rads[:, 0::2]) # Sin for even values/dimensions
        positional_encoding[:, 1::2] = torch.cos(angle_rads[:, 1::2]) # Cos for odd values/dimensions
        
        return positional_encoding.unsqueeze(0)

Notice how we have an Encoder/Decoder Block and then an Encoder/Decoder, this way we can actually determine how many times we want to loop throughout our Encoder/Decoder.

Decoder:

Decoder Code:

class DecoderBlock(nn.Module):
    def __init__(self, embedding_size, num_heads, d_ff, dropout):
        super(DecoderBlock, self).__init__()

       # Initializing MHA and FFN
        self.masked_attention = MultiHeadAttention(embedding_size, num_heads)
        self.cross_attention = MultiHeadAttention(embedding_size, num_heads)
        self.feed_forward = FeedForward(embedding_size, d_ff, dropout)
        
        # Norm Layers
        self.norm1 = nn.LayerNorm(embedding_size)
        self.norm2 = nn.LayerNorm(embedding_size)
        self.norm3 = nn.LayerNorm(embedding_size)

        # Dropout
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, enc_out=None, src_mask=None, trg_mask=None):

        # Masked Multi-Head Attention
        _x = self.masked_attention(x, x, x, trg_mask)
        x = self.norm1(x + self.dropout(_x))

        # Cross-Attention
        if enc_out != None:
            _x = self.cross_attention(x, enc_out, enc_out, src_mask)
            x = self.norm2(x + self.dropout(_x))

        # Feed Forward
        _x = self.feed_forward(x)
        output = self.norm3(x + self.dropout(_x))
        return output

class Decoder(nn.Module):
    def __init__(self, vocab_size, embedding_size, num_layers, num_heads, d_ff, max_len, dropout, device):
        super(Decoder, self).__init__()

        self.device = device
        
        self.word_embedding = nn.Embedding(vocab_size, embedding_size)
        self.position_embedding = nn.Embedding(max_len, embedding_size)
        self.layers = nn.ModuleList(
            [DecoderBlock(embedding_size, num_heads, d_ff, dropout) for _ in range(num_layers)]
        )

        # Dropout
        self.dropout = nn.Dropout(dropout)

        # Causal Mask
        self.register_buffer("trg_mask", torch.tril(
            torch.ones(max_len, max_len, device=self.device)
            ).bool())
        
        # Final Linear Projection
        self.fc_out = nn.Linear(embedding_size, vocab_size)

    def forward(self, x, enc_out=None, src_mask=None):

        # x: [batch_size, target_seq_len]
        N, seq_len = x.shape
        mask = self.trg_mask[:seq_len, :seq_len]

        position = torch.arange(0, seq_len, device=self.device).unsqueeze(0).expand(N, seq_len)

        # x: [batch_size, target_seq_len, embedding_size]
        x = self.dropout((self.word_embedding(x) +  self.position_embedding(position)))

        # Pass through each DecoderBlock
        for layer in self.layers:
            x = layer(x, enc_out, src_mask, mask)

        logits = self.fc_out(x)

        return logits

enc_out is used for Translation model, and it means the encoded output from the Encoder. I skip it if there is no enc_out passed to the model. Therefore no Cross Attention.

Sentiment Classification

Once we have our Encoder, we can assign it to do Classification. We do not need the Decoder, since we are not doing anything sequence-to-sequence operation, simply attention. We create a custom dataset to simply transform the text into tokens and add padding (<PAD>) to short sentences to match max_len.

ClassificationDataset Class

class ClassificationDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, seq_len):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.seq_len = seq_len

    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
            text = self.texts[idx]
            label = self.labels[idx]

            # Tokenize single example
            tokens = self.tokenizer.transform(text)

            # Pad / truncate
            tokens = self.tokenizer.pad_sequence(tokens, max_len=self.seq_len)

            return torch.tensor(tokens, dtype=torch.long), torch.tensor(label, dtype=torch.long)

Classification Model

Sentiment Classifier Code:

class SentimentClassifier(nn.Module):
    def __init__(self, vocab_size, embedding_dim, num_layers, num_heads, d_ff, max_len, num_classes, dropout, device):
        super().__init__()
        
        self.encoder = Encoder(
            vocab_size=vocab_size,
            embedding_size=embedding_dim,
            num_layers=num_layers,
            num_heads=num_heads,
            d_ff=d_ff,
            max_len=max_len,
            dropout=dropout,
            device=device
        )

        self.classifier = nn.Sequential(
            nn.Linear(embedding_dim, embedding_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(embedding_dim, num_classes)
        )

    def forward(self, x, src_mask=None):
        # x: [batch_size, seq_len]
        x = self.encoder(x, src_mask) # [batch_size, seq_len, embed_dim]

        # Mean pooling of the sequence_length (dim=1)
        x = x.mean(dim=1)
        #x = nn.AdaptiveAvgPool1d(1)

        logits = self.classifier(x)
        return logits

Training Details:

Optimizer: Adam
Learning rate: 3e-4
Loss Function: CrossEntropyLoss
Epochs: 100
Datasets: Amazon, IMDB, Yelp

Model Specifications:

Total Parameters: 1 M
Model Size: ~3.77MB

Configuration:

model = SentimentClassifier(
    vocab_size=vocab_size, (5209)
    embedding_dim=128,
    num_layers=2,
    num_heads=4,
    d_ff=256,
    max_len=100,
    num_classes=2,
    dropout=0.1,
    device=device).to(device)

Final Metrics

Epoch	Train Loss	Train Accuracy	Validation Loss	Validation Accuracy
100	0.2728	88.33%	0.5007	79.00%

Text Generation

With the Decoder we can do Text Generation, and this is actually what most LLM models are, Decoder-only model. We need thought to create a custom dataset class, that will make the data sutiable for our model.

TextDataset

class TextDataset(Dataset):
    def __init__(self, tokens, seq_len):
        self.tokens = tokens
        self.seq_len = seq_len
        self.stride = seq_len // 2

    def __len__(self):
        return (len(self.tokens) - self.seq_len) // self.stride
    
    def __getitem__(self, idx):

        start = idx * self.stride

        x = self.tokens[start:start + self.seq_len]
        y = self.tokens[start + 1:start + self.seq_len + 1]
       
        return torch.tensor(x), torch.tensor(y)

Here we are basically introducing stride, which will determine how much we slide through tokens, in this case I am sliding/overlapping by half of the length of the sentence, this is good for better data usage!

Text Generation Model

Text Generator Code:

class TextGenerator(nn.Module):
    def __init__(self, vocab_size, embedding_dim, num_layers, num_heads, d_ff, device, max_len, dropout=0.1):
        super().__init__()
        
        self.decoder = Decoder(
            vocab_size=vocab_size,
            embedding_size=embedding_dim,
            num_layers=num_layers,
            num_heads=num_heads,
            d_ff=d_ff,
            device=device,
            max_len=max_len,
            dropout=dropout
        )

        self.generator = nn.Linear(embedding_dim, vocab_size)

    def forward(self, x):
        # x: [batch_size, seq_len]
        decoded = self.decoder(x)
        # logits = self.generator(decoded)
        return decoded

Training Details:

Optimizer: Adam
Learning rate: 3e-4
Loss Function: CrossEntropyLoss
Epochs: 400
Dataset: Personal Data

Model Specifications:

Total Parameters: 1.58 M
Model Size: ~6.3MB

Configuration:

model = TextGenerator(
    vocab_size=vocab_size, (3,818)
    embedding_dim=128,
    num_layers=3,
    num_heads=4,
    d_ff=512,
    max_len=128,
    dropout=0.1,
    device=device).to(device)

Final Metrics

Epoch	Train Loss	Validation Loss	Perplexity
400	0.3832	2.5705	13.07

Translation

In case with translation, we use both encoder and decoder, aka the Full Transformer.

Transformer Cross-Attention

Translation Model

Translation Model Code

class Transformer(nn.Module):
    def __init__(self, src_vocab_size, trg_vocab_size, src_pad_idx, trg_pad_idx, embedding_size, num_layers, d_ff, num_heads, max_len, dropout, device):
        super(Transformer, self).__init__()

        self.encoder = Encoder(src_vocab_size, embedding_size, num_layers, num_heads, d_ff, max_len,dropout, device)
        self.decoder = Decoder(trg_vocab_size, embedding_size, num_layers, num_heads, d_ff, max_len, dropout, device)
        self.src_pad_idx = src_pad_idx
        self.trg_pad_idx = trg_pad_idx
        self.device = device

    def make_src_mask(self, src):
        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)
        return src_mask

    def make_trg_mask(self, trg):
        N, trg_len = trg.shape

        trg_pad_mask = (trg != self.trg_pad_idx).unsqueeze(1).unsqueeze(2)

        trg_sub_mask = torch.tril(torch.ones(trg_len, trg_len)).to(self.device).bool()
        trg_mask = trg_pad_mask & trg_sub_mask
        return trg_mask
    
    def forward(self, src, trg):
        src_mask = self.make_src_mask(src)
        trg_mask = self.make_trg_mask(trg)
        enc_out = self.encoder(src, src_mask)

        output = self.decoder(trg, enc_out, src_mask, trg_mask)
        return output

License

MIT License

Copyright (c) 2026 Widmark La Rosa

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

Contributing

This code is free to use, modify and distribute!

Everyone is welcome to open pull requests for:

Bug Fixes
Code Improvements
Documentation enhancements

Thank You :)

References & Citations

Research Paper: Attention Is All You Need - Vaswani et al., 2017
Github Repo: transformer – by hyunwoongko, 2019

Name		Name	Last commit message	Last commit date
Latest commit History 43 Commits
data		data
models		models
notebooks		notebooks
transformer		transformer
.gitignore		.gitignore
LICENSE		LICENSE
README.md		README.md
requirements.txt		requirements.txt
sentiment_classifier.py		sentiment_classifier.py
text_generator.py		text_generator.py
tokenizer.py		tokenizer.py

Folders and files

Latest commit

History

Repository files navigation

Transformer From Scratch (with PyTorch)

Repo Structure

Index

Clone Repo

Set Up Virtual Environment

Install Dependencies

Tokenizer

Transformer Architecture

Positional Encoding

Postional Encoding Code:

Self-Attention

Multi-Head Attention

Multi-Head Attention Code:

Masking

Feed Foward Network

Feed Foward Network Code:

Add & LayerNorm

Encoder:

Encoder Code:

Decoder:

Decoder Code:

Sentiment Classification

ClassificationDataset Class

Classification Model

Sentiment Classifier Code:

Training Details:

Model Specifications:

Configuration:

Final Metrics

Text Generation

TextDataset

Text Generation Model

Text Generator Code:

Training Details:

Model Specifications:

Configuration:

Final Metrics

Translation

Transformer Cross-Attention

Translation Model

Translation Model Code

License

Contributing

References & Citations

About

Topics

Resources

License

Uh oh!

Stars

Watchers

Forks

Releases

Packages 0

Uh oh!

Contributors

Uh oh!

Languages

Packages