Config.

py

from pathlib import Path

# create an configuration

def get_config():

return {

"batch_size": 8,

"num_epochs": 20,

"lr": 10**-4,

"seq_len": 350,

"d_model": 512,

"lang_src": "en",

"lang_tgt": "it",

"model_folder": "weights",

"model_basename": "tmodel_",

"preload": None,

"tokenizer_file": "tokenizer_{0}.json",

"experiment_name": "runs/tmodel_",

# get weights

def get_weights_file_path(config, epoch: str):

model_folder = config["model_folder"]

model_basename = config["model_basename"]

model_filename = f"{model_basename}{epoch}.pt"

return str(Path('.') / model_folder / model_filename)

Corpus.py

import torch

import torch.nn as nn

from torch.utils.data import Dataset

from typing import Any

class BillingualDataset(Dataset):

def __init__(self, dataset, tokenizer_src, tokenizer_tgt, src_lang, tgt_lang, seq_len):

super().__init__()

self.dataset = dataset

self.tokenizer_src = tokenizer_src

self.tokenizer_tgt = tokenizer_tgt

self.src_lang = src_lang

self.tgt_lang = tgt_lang

self.seq_len = seq_len

self.sos_token = torch.tensor([

tokenizer_src.token_to_id('[SOS]')

], dtype=torch.int64)

self.eos_token = torch.tensor([

tokenizer_src.token_to_id('[EOS]')

], dtype=torch.int64)

self.pad_token = torch.tensor([

tokenizer_src.token_to_id('[PAD]')

], dtype=torch.int64)

def __len__(self):

return len(self.dataset)
def __getitem__(self, index) -> Any:

src_target_pair = self.dataset[index]

src_text = src_target_pair["translation"][self.src_lang]

tgt_text = src_target_pair["translation"][self.tgt_lang]

enc_input_tokens = self.tokenizer_src.encode(src_text).ids

dec_input_tokens = self.tokenizer_tgt.encode(tgt_text).ids

enc_num_padding_tokens = self.seq_len - len(enc_input_tokens) - 2

dec_num_padding_tokens = self.seq_len - len(dec_input_tokens) - 1

if enc_num_padding_tokens < 0 or dec_num_padding_tokens < 0:

raise ValueError("Someone is too long")

# add SOS and EOS to the source text

encoder_input = torch.cat(

self.sos_token,

torch.tensor(enc_input_tokens, dtype=torch.int64),

self.eos_token,

torch.tensor([self.pad_token] * enc_num_padding_tokens, dtype=torch.int64)

# add SOS to the decoder input

decoder_input = torch.cat(

self.sos_token,

torch.tensor(dec_input_tokens, dtype=torch.int64),

torch.tensor([self.pad_token] * dec_num_padding_tokens, dtype=torch.int64)

]
 )

# add EOS to the label

label = torch.cat(

torch.tensor(dec_input_tokens, dtype=torch.int64),

self.eos_token,

torch.tensor([self.pad_token] * dec_num_padding_tokens, dtype=torch.int64)

assert encoder_input.size(0) == self.seq_len

assert decoder_input.size(0) == self.seq_len

assert label.size(0) == self.seq_len

return {

"encoder_input": encoder_input, # seq_len

"decoder_input": decoder_input, # seq_len

"encoder_mask": (encoder_input != self.pad_token).unsqueeze(0).unsqueeze(0).int(), # (1, 1,

seq_len)

"decoder_mask": (decoder_input != self.pad_token).unsqueeze(0).unsqueeze(0).int() &

casual_mask(decoder_input.size(0)), # (seq_len, 1) & (1, seq_len, seq_len)

"label": label, # (seq_len)

"src_text": src_text,

"tgt_text": tgt_text

def casual_mask(size):

mask = torch.triu(torch.ones(1, size, size), diagonal=1).type(torch.int)

return mask == 0
Model.py

# import libraries

import torch

import math

import torch.nn as nn

# d_model -> size of embbeded vector

# h -> number of heads

class InputEmbeddings(nn.Module):

def __init__(self, d_model: int, vocab_size: int):

super().__init__()

self.d_model = d_model

self.vocab_size = vocab_size

self.embedding = nn.Embedding(vocab_size, d_model)

def forward(self, x):

return self.embedding(x) * math.sqrt(self.d_model)

class PositionalEncoding(nn.Module):

def __init__(self, d_model: int, sen_len: int, dropout: float) -> None:

super().__init__()

self.d_model = d_model

self.sen_len = sen_len

self.dropout = nn.Dropout(dropout)

# create a matrix of shape (sen_len, d_model)

 pe = torch.zeros(sen_len, d_model)

# create a vector of shape (sen_len)

position = torch.arange(0, sen_len, dtype=torch.float).unsqueeze(1) # shape: (sen_len, 1)

# formula

div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

# apply sin to even positions

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

# apply cos to odd positions

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

pe = pe.unsqueeze(0) # shape: (1, sen_len, d_model)

self.register_buffer("pe", pe)

def forward(self, x):

x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)

return self.dropout(x)

class LayerNormalization(nn.Module):

def __init__(self, eps: float = 10**-6) -> None:

super().__init__()

self.eps = eps

# alpha -> Multiplicative

self.alpha = nn.Parameter(torch.ones(1))

# beta -> Addictive

self.beta = nn.Parameter(torch.zeros(1))
 def forward(self, x):

mean = x.mean(dim = -1, keepdim=True)

std = x.std(dim = -1, keepdim=True)

return self.alpha * (x - mean) / (std + self.eps) + self.beta

class FeedForwardBlock(nn.Module):

def __init__(self, d_model: int, d_ff: int, dropout: float) -> None:

super().__init__()

self.linear_01 = nn.Linear(d_model, d_ff) # w1 and b1

self.dropout = nn.Dropout(dropout)

self.linear_02 = nn.Linear(d_ff, d_model) # w2 and b2

def forward(self, x):

# (Batch, sen_len, d_model) --> (Batch, sen_len, d_ff) --> (Batch, sen_len, d_model)

return self.linear_02(self.dropout(torch.relu(self.linear_01(x))))

class MultiHeadAttention(nn.Module):

def __init__(self, d_model: int, heads: int, dropout: float) -> None:

super().__init__()

self.d_model = d_model

self.heads = heads

assert d_model % heads == 0, "d_model is not divisible by heads"

self.d_k = d_model // heads

 # set the query, key and value vector

self.w_q = nn.Linear(d_model, d_model) # w_q

self.w_k = nn.Linear(d_model, d_model) # w_k

self.w_v = nn.Linear(d_model, d_model) # w_v

# output

self.w_o = nn.Linear(d_model, d_model) # w_o

self.dropout = nn.Dropout(dropout)

@staticmethod

def Attention(query, key, value, mask, dropout: nn.Dropout):

d_k = query.shape[-1]

# (Batch, h, sen_len, d_k) --> (Batch, h, sen_len, sen_len)

attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)

if mask is not None:

attention_scores.masked_fill_(mask == 0, -1e9)

attention_scores = attention_scores.softmax(dim = -1) # (Batch, h, sen_len, sen_len)

if dropout is not None:

attention_scores = dropout(attention_scores)

return (attention_scores @ value), attention_scores

def forward(self, q, k, v, mask):

# (Batch, sen_len, d_model) --> (Batch, sen_len, d_model)

query = self.w_q(q)
 key = self.w_k(k)

value = self.w_v(v)

# (Batch, sen_len, d_model) --> (Batch, sen_len, heads, d_k) --> (Batch, heads, sen_len, d_k)

query = query.view(query.shape[0], query.shape[1], self.heads, self.d_k).transpose(1, 2)

key = key.view(key.shape[0], key.shape[1], self.heads, self.d_k).transpose(1, 2)

value = value.view(value.shape[0], value.shape[1], self.heads, self.d_k).transpose(1, 2)

# call the attention mechanism

x, self.attention_scores = MultiHeadAttention.Attention(query, key, value, mask, self.dropout)

# (Batch, heads, sen_len, d_k) --> (Batch, sen_len, heads, d_k) --> (Batch, sen_len, d_model)

x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.heads * self.d_k)

# (Batch, sen_len, d_model) --> (Batch, sen_len, d_model)

return self.w_o(x)

class ResidualConnection(nn.Module):

def __init__(self, dropout: float):

super().__init__()

self.dropout = nn.Dropout(dropout)

self.norm = LayerNormalization()

def forward(self, x, sublayer):

return x + self.dropout(sublayer(self.norm(x)))

class EncoderBlock(nn.Module):
 def __init__(self, self_attention_block: MultiHeadAttention, feed_forward_network:
FeedForwardBlock, dropout: float):

super().__init__()

self.self_attention_block = self_attention_block

self.feed_forward_block = feed_forward_network

self.residual_connection = nn.ModuleList([ResidualConnection(dropout) for _ in range(2)])

def forward(self, x, src_mask):

x = self.residual_connection[0](x, lambda x: self.self_attention_block(x, x, x, src_mask))

x = self.residual_connection[1](x, self.feed_forward_block)

return x

class Encoder(nn.Module):

def __init__(self, layers: nn.ModuleList):

super().__init__()

self.layers = layers

self.norm = LayerNormalization()

def forward(self, x, mask):

for layer in self.layers:

x = layer(x, mask)

return self.norm(x)

class DecoderBlock(nn.Module):

def __init__(self, self_attention_block: MultiHeadAttention, cross_attention_block:

MultiHeadAttention, feed_forward_block: FeedForwardBlock, dropout: float):

super().__init__()
 self.self_attention_block = self_attention_block

self.cross_attention_block = cross_attention_block

self.feed_forward_block = feed_forward_block

self.residual_connections = nn.ModuleList([ResidualConnection(dropout) for _ in range(3)])

def forward(self, x, encoder_output, src_mask, tgt_mask):

x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, tgt_mask))

x = self.residual_connections[1](x, lambda x: self.cross_attention_block(x, encoder_output,

encoder_output, src_mask))

x = self.residual_connections[2](x, self.feed_forward_block)

return x

class Decoder(nn.Module):

def __init__(self, layers: nn.ModuleList):

super().__init__()

self.layers = layers

self.norm = LayerNormalization()

def forward(self, x,encoder_output, src_mask, tgt_mask):

for layer in self.layers:

x = layer(x, encoder_output, src_mask, tgt_mask)

return self.norm(x)

class ProjectionLayer(nn.Module):

def __init__(self, d_model: int, vocab_size: int):

super().__init__()
 self.linear = nn.Linear(d_model, vocab_size)

def forward(self, x):

# (Batch, sen_len, d_model) --> (Batch, sen_len, vocab_size)

return torch.log_softmax(self.linear(x), dim=-1)

class Transformer(nn.Module):

def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: InputEmbeddings, tgt_embed:

InputEmbeddings, src_pos: PositionalEncoding, tgt_pos: PositionalEncoding, proj: ProjectionLayer):

super().__init__()

self.encoder = encoder

self.decoder = decoder

self.src_embed = src_embed

self.tgt_embed = tgt_embed

self.src_pos = src_pos

self.tgt_pos = tgt_pos

self.projection = proj

def encode(self, src, src_mask):

src = self.src_embed(src)

src = self.src_pos(src)

return self.encoder(src, src_mask)

def decode(self, encoder_output, src_mask, tgt, tgt_mask):

tgt = self.tgt_embed(tgt)

tgt = self.tgt_pos(tgt)

return self.decoder(tgt, encoder_output, src_mask, tgt_mask)

def project(self, x):

 return self.projection(x)

# build the transformer

# N -> number of encoder and decoder blocks

def build_transformer(src_vocab_size: int, tgt_vocab_size: int, src_seq_len: int, tgt_seq_len: int,

d_model: int = 512, N: int = 6, heads: int = 8, dropout: float = 0.1, d_ff: int = 2048) -> Transformer:

# create the embedding layers

src_embed = InputEmbeddings(d_model, src_vocab_size)

tgt_embed = InputEmbeddings(d_model, tgt_vocab_size)

# create the positional encodings

src_pos = PositionalEncoding(d_model, src_seq_len, dropout)

tgt_pos = PositionalEncoding(d_model, tgt_seq_len, dropout)

# create the encoder blocks

encoder_blocks = []

for _ in range(N):

encoder_self_attention_block = MultiHeadAttention(d_model, heads, dropout)

feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)

encoder_block = EncoderBlock(encoder_self_attention_block, feed_forward_block, dropout)

encoder_blocks.append(encoder_block)

# create the decoder blocks

decoder_blocks = []

for _ in range(N):

decoder_self_attention_block = MultiHeadAttention(d_model, heads, dropout)

 decoder_cross_attention_block = MultiHeadAttention(d_model, heads, dropout)

feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)

decoder_block = DecoderBlock(decoder_self_attention_block, decoder_cross_attention_block,

feed_forward_block, dropout)

decoder_blocks.append(decoder_block)

# create encoder and decoder blocks

encoder = Encoder(nn.ModuleList(encoder_blocks))

decoder = Decoder(nn.ModuleList(decoder_blocks))

# create the projection layer

projection_layer = ProjectionLayer(d_model, tgt_vocab_size)

# create the transformer

transformer = Transformer(encoder, decoder, src_embed, tgt_embed, src_pos, tgt_pos,

projection_layer)

# Intialize the parameters

for p in transformer.parameters():

if p.dim() > 1:

nn.init.xavier_uniform_(p)

return transformer
Train.py

import torch

import warnings

import torch.nn as nn

from torch.utils.data import Dataset, DataLoader, random_split

from torch.utils.tensorboard import SummaryWriter

from tqdm import tqdm

from datasets import load_dataset

from tokenizers import Tokenizer

from tokenizers.models import WordLevel

from tokenizers.trainers import WordLevelTrainer

from tokenizers.pre_tokenizers import Whitespace

from pathlib import Path

from corpus import BillingualDataset, casual_mask

from model import build_transformer

from config import get_weights_file_path, get_config

def get_all_sentences(dataset, lang):

for item in dataset:

yield item['translation'][lang]

# build the tokenizer

def get_or_build_tokenizer(config, dataset, lang):

# tokenizer path

Tokenizer_path = Path(config["tokenizer_file"].format(lang))
 if not Path.exists(Tokenizer_path):

tokenizer = Tokenizer(WordLevel(unk_token='[UNK]'))

tokenizer.pre_tokenizer = Whitespace()

trainer = WordLevelTrainer(special_tokens=["[UNK]","[PAD]","[SOS]","[EOS]"], min_frequency=2)

tokenizer.train_from_iterator(get_all_sentences(dataset, lang), trainer=trainer)

tokenizer.save(str(Tokenizer_path))

else:

tokenizer = Tokenizer.from_file(str(Tokenizer_path))

return tokenizer

# get the dataset

def get_dataset(config):

dataset_raw = load_dataset("opus_books", f"{config['lang_src']}-{config['lang_tgt']}", split="train")

# build the tokenizer

tokenizer_src = get_or_build_tokenizer(config, dataset_raw, config["lang_src"])

tokenizer_tgt = get_or_build_tokenizer(config, dataset_raw, config["lang_tgt"])

# keep the 90% data for training and 10% for testing

train_ds_size = int(0.9 * len(dataset_raw))

val_ds_size = len(dataset_raw) - train_ds_size

train_ds_raw, val_ds_raw = random_split(dataset_raw, [train_ds_size, val_ds_size])

train_ds = BillingualDataset(train_ds_raw, tokenizer_src, tokenizer_tgt, config["lang_src"],

config["lang_tgt"], config["seq_len"])

val_ds = BillingualDataset(val_ds_raw, tokenizer_src, tokenizer_tgt, config["lang_src"],

config["lang_tgt"], config["seq_len"])

max_len_src = 0
 max_len_tgt = 0

for item in dataset_raw:

src_ids = tokenizer_src.encode(item["translation"][config["lang_src"]]).ids

tgt_ids = tokenizer_tgt.encode(item["translation"][config["lang_tgt"]]).ids

max_len_src = max(max_len_src, len(src_ids))

max_len_tgt = max(max_len_tgt, len(tgt_ids))

print(f"Max length of source sentence: {max_len_src}")

print(f"Max length of target sentence: {max_len_tgt}")

# create a data loader

train_dataloader = DataLoader(

train_ds,

batch_size=config["batch_size"],

shuffle=True

val_dataloader = DataLoader(

val_ds,

batch_size=1,

shuffle=True

return train_dataloader, val_dataloader, tokenizer_src, tokenizer_tgt

def get_model(config, vocab_src_len, vocab_tgt_len):

# build the model

 model = build_transformer(vocab_src_len, vocab_tgt_len, config["seq_len"], config["seq_len"],
config["d_model"])

return model

# create an model training loop

def train_model(config):

# define the device

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print(f"Using device: {device}")

# create an model folder

Path(config["model_folder"]).mkdir(parents=True, exist_ok=True)

# get the data loaders

train_dataloader, val_dataloader, tokenizer_src, tokenizer_tgt = get_dataset(config)

# get the tokenizer model

model = get_model(config, tokenizer_src.get_vocab_size(), tokenizer_tgt.get_vocab_size())

# change the model into device

model = model.to(device)

# Tensorboard

writer = SummaryWriter(config["experiment_name"])

# set the optimizer

optimizer = torch.optim.Adam(model.parameters(), lr=config["lr"], eps=1e-9)

intial_epoch = 0
 global_step = 0

if config["preload"]:

model_filename = get_weights_file_path(config, config["preload"])

print(f"Preloading model: {model_filename}")

state = torch.load(model_filename)

intial_epoch = state["epoch"] + 1

optimizer.load_state_dict(state["optimizer_state_dict"])

global_step = state["global_step"]

# set the loss function

loss_fn = nn.CrossEntropyLoss(ignore_index=tokenizer_src.token_to_id("[PAD]"),
label_smoothing=0.1).to(device)

for epoch in range(intial_epoch, config["num_epochs"]):

model.train()

batch_iterator = tqdm(train_dataloader, desc=f"Processing epoch {epoch:02d}")

for batch in batch_iterator:

encoder_input = batch["encoder_input"].to(device) # (Batch_size, seq_len)

decoder_input = batch["decoder_input"].to(device) # (Batch_size, seq_len)

encoder_mask = batch["encoder_mask"].to(device) # (Batch_size, 1, 1, seq_len)

decoder_mask = batch["decoder_mask"].to(device) # (Batch_size, 1, seq_len, seq_len)

# run the tensors through transformers

encoder_output = model.encode(encoder_input, encoder_mask) # (Batch_size, seq_len,

d_model)

decoder_output = model.decode(encoder_output, encoder_mask, decoder_input,

decoder_mask) # (Batch_size, seq_len, d_model)

projec_output = model.project(decoder_output) # (Batch_size, seq_len, tgt_vocab_size)

 label = batch['label'].to(device) # (Batch, seq_len)

# (Batch_size, seq_len, tgt_vocab_size) --> (Batch_size, seq_len, tgt_vocab_size)

loss = loss_fn(projec_output.view(-1, tokenizer_tgt.get_vocab_size()), label.view(-1))

batch_iterator.set_postfix({f"loss": f"{loss.item():6.3f}"})

# log the loss

writer.add_scalar("train loss", loss.item(), global_step)

writer.flush()

# backpropagation the loss

loss.backward()

# update the weights

optimizer.step()

optimizer.zero_grad()

global_step += 1

# save the model

model_filename = get_weights_file_path(config, f"{epoch:02d}")

torch.save(

"epoch": epoch,

"model_state_dict": model.state_dict(),

"optimizer_state_dict": optimizer.state_dict(),

"global_step": global_step

},

model_filename

)
if __name__ == "__main__":

warnings.filterwarnings("ignore")

config = get_config()

train_model(config)

Project Structure