Added files

pinns_practice/basic_pinn.py

@@ -2,28 +2,32 @@ import torch
 import torch.nn as nn
 import matplotlib.pyplot as plt
 
+
 class BasicPINN(nn.Module):
     def __init__(self):
         super().__init__()
         self.net = nn.Sequential(
             nn.Linear(1, 20),
-            nn.Tanh(), # for RELU 2 derivative is always 0, so use Tanh
+            nn.Tanh(),  # for RELU 2 derivative is always 0, so use Tanh
             nn.Linear(20, 20),
             nn.Tanh(),
-            nn.Linear(20, 1)
+            nn.Linear(20, 1),
         )
 
     def forward(self, x):
-        x =self.net(x)
+        x = self.net(x)
         return x
 
+
 model = BasicPINN()
 optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
 
 for epoch in range(2000):
     optimizer.zero_grad()
 
-    t_physics = torch.rand(100, 1).requires_grad_(True)  # requires_grad=True important for derivatives
+    t_physics = torch.rand(100, 1).requires_grad_(
+        True
+    )  # requires_grad=True important for derivatives
 
     y_pred = model(t_physics)
 
@@ -33,10 +37,12 @@ for epoch in range(2000):
     # Calculating derivative dy/dt
     # We use PyTorch auto-differentiation to find the rate of change of y_pred with respect to t_physics.
     dy_dt = torch.autograd.grad(
-        outputs=y_pred, # What we differentiate, y
-        inputs=t_physics, # What we differentiate with respect to, (time, t)
-        grad_outputs=torch.ones_like(y_pred), # vector from 1, for 100 examples, calculates gradients independently
-        create_graph=True, # history of calculations, critical for PINNs
+        outputs=y_pred,  # What we differentiate, y
+        inputs=t_physics,  # What we differentiate with respect to, (time, t)
+        grad_outputs=torch.ones_like(
+            y_pred
+        ),  # vector from 1, for 100 examples, calculates gradients independently
+        create_graph=True,  # history of calculations, critical for PINNs
     )[0]
 
     # Physical Loss dy/dt + 2y = 0
@@ -60,7 +66,9 @@ with torch.no_grad():
 
 
 y_exact = torch.exp(-2 * t_test)
-plt.plot(t_test.numpy(), y_test_pred.numpy(), label="PINN model", color="red", linestyle="--")
+plt.plot(
+    t_test.numpy(), y_test_pred.numpy(), label="PINN model", color="red", linestyle="--"
+)
 plt.plot(t_test.numpy(), y_exact.numpy(), label="Exact solution (Math)", alpha=0.5)
 plt.legend()
 plt.title("Solving the differential equation!!")

transformer_from_scratch/dataset.py

@@ -0,0 +1,74 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset
+
+
+class BilingualDataset(Dataset):
+    def __init__(self, ds, tokenizer_src, tokenizer_tgt, src_lang, tgt_lang, seq_len):
+        super().__init__()
+
+        self.ds = ds
+        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.ds)
+
+    def __getitem__(self, index):
+        src_target_pair = self.ds[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  # for SOS and EOS
+        dec_num_padding_tokens = self.seq_len - len(dec_input_tokens) - 1 # for SOS
+
+        if enc_num_padding_tokens < 0 or dec_num_padding_tokens < 0:
+            raise ValueError('Sentence is too long')
+
+        # Add SOS and EOS tokens to 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 token 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 token to the label (what we want )
+        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)), # (1, Seq_len) & (1, Seq_len, Seq_len)
+            "label": label,                 # (Seq_len)
+        }

transformer_from_scratch/model.py

@@ -174,20 +174,36 @@ class Encoder(nn.Module):
 
 
 class DecoderBlock(nn.Module):
-    def __init__(self, self_attention_block: MultiHeadAttention, cross_attention_block: MultiHeadAttention, feed_forward_block: FeedForwardBlock, dropout: float):
+    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)])
+        self.residual_connections = nn.ModuleList(
+            [ResidualConnection(dropout) for _ in range(3)]
+        )
 
     # x - input of the decoder, src_mask - mask for encoder, tgt_mask - mask applied to the decoder
     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[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__()
@@ -210,9 +226,17 @@ class ProjectionLayer(nn.Module):
         return torch.log_softmax(self.proj(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, projection_layer: ProjectionLayer):
+    def __init__(
+        self,
+        encoder: Encoder,
+        decoder: Decoder,
+        src_embed: InputEmbeddings,
+        tgt_embed: InputEmbeddings,
+        src_pos: PositionalEncoding,
+        tgt_pos: PositionalEncoding,
+        projection_layer: ProjectionLayer,
+    ):
         super().__init__()
         self.encoder = encoder
         self.decoder = decoder
@@ -236,7 +260,17 @@ class Transformer(nn.Module):
         return self.projection_layer(x)
 
 
-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, h: int = 8, dropout: int = 0.1, d_ff: int = 2048):
+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,
+    h: int = 8,
+    dropout: int = 0.1,
+    d_ff: int = 2048,
+):
     # Create the embedding layers
     src_embed = InputEmbeddings(d_model, src_vocab_size)
     tgt_embed = InputEmbeddings(d_model, tgt_vocab_size)
@@ -250,7 +284,9 @@ def build_transformer(src_vocab_size: int, tgt_vocab_size: int, src_seq_len: int
     for _ in range(N):
         encoder_self_attention_block = MultiHeadAttention(d_model, h, dropout)
         feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
-        encoder_block = EncoderBlock(encoder_self_attention_block, feed_forward_block, dropout)
+        encoder_block = EncoderBlock(
+            encoder_self_attention_block, feed_forward_block, dropout
+        )
         encoder_blocks.append(encoder_block)
 
     # Create the decoder blocks
@@ -259,7 +295,12 @@ def build_transformer(src_vocab_size: int, tgt_vocab_size: int, src_seq_len: int
         decoder_self_attention_block = MultiHeadAttention(d_model, h, dropout)
         decoder_cross_attention_block = MultiHeadAttention(d_model, h, 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_block = DecoderBlock(
+            decoder_self_attention_block,
+            decoder_cross_attention_block,
+            feed_forward_block,
+            dropout,
+        )
         decoder_blocks.append(decoder_block)
 
     # Create the encoder and decoder
@@ -270,7 +311,9 @@ def build_transformer(src_vocab_size: int, tgt_vocab_size: int, src_seq_len: int
     projection_layer = ProjectionLayer(d_model, tgt_vocab_size)
 
     # Build the transformer
-    transformer = Transformer(encoder, decoder, src_embed, tgt_embed, src_pos, tgt_pos, projection_layer)
+    transformer = Transformer(
+        encoder, decoder, src_embed, tgt_embed, src_pos, tgt_pos, projection_layer
+    )
 
     # Initialize the parameters
     for p in transformer.parameters():
@@ -278,4 +321,3 @@ def build_transformer(src_vocab_size: int, tgt_vocab_size: int, src_seq_len: int
             nn.init.xavier_uniform_(p)
 
     return transformer
-

transformer_from_scratch/train.py

@@ -0,0 +1,47 @@
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+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 torch.utils.data import Dataset, DataLoader, random_split
+
+
+def get_all_sentences(ds, lang):
+    for item in ds:
+        yield item["translation"][lang]
+
+
+def get_or_build_tokenizer(config, ds, lang):
+    # config['tokenizer_file'] = '../tokenizers/tokenizer_{0}.json'
+    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(ds, lang), trainer=trainer)
+        tokenizer.save(str(tokenizer_path))
+    else:
+        tokenizer = Tokenizer.from_file(str(tokenizer_path))
+    return tokenizer
+
+
+def get_ds(config):
+    ds_raw = load_dataset(
+        "opus_books", f"{config['lang_src']}-{config['lang_tgt']}", split="train"
+    )
+
+    # Build tokenizers
+    tokenizer_src = get_or_build_tokenizer(config, ds_raw, config["lang_src"])
+    tokenizer_tgt = get_or_build_tokenizer(config, ds_raw, config["lang_tgt"])
+
+    # Keep 90% for training, 10% for validation
+    train_ds_size = int(0.9 * len(ds_raw))
+    val_ds_size = len(ds_raw) - train_ds_size
+    train_ds_raw, val_ds_raw = random_split(ds_raw, [train_ds_size, val_ds_size])