decoder only transformer learning

.gitignore

@@ -0,0 +1,169 @@
+# Initially taken from Github's Python gitignore file
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# tests and logs
+tests/fixtures/cached_*_text.txt
+logs/
+lightning_logs/
+lang_code_data/
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# vscode
+.vs
+.vscode
+
+# Pycharm
+.idea
+
+# TF code
+tensorflow_code
+
+# Models
+proc_data
+
+# examples
+runs
+/runs_old
+/wandb
+/examples/runs
+/examples/**/*.args
+/examples/rag/sweep
+
+# data
+/data
+serialization_dir
+
+# emacs
+*.*~
+debug.env
+
+# vim
+.*.swp
+
+#ctags
+tags
+
+# pre-commit
+.pre-commit*
+
+# .lock
+*.lock
+
+# DS_Store (MacOS)
+.DS_Store
+
+# ruff
+.ruff_cache

README.md

@@ -1,3 +1 @@
----
-license: unlicense
----
+# decoder only transformers

learning.py

@@ -0,0 +1,78 @@
+import torch # create tensors and provides helper functions
+import torch.nn as nn # for nn.Module(), nn.Embedding() and nn.Linear()
+import torch.nn.functional as F # gives us the softmax() and argmax()
+from torch.optim import Adam # Adam optimizer, stochastic gradient descent
+from torch.utils.data import TensorDataset, DataLoader # for storing data loader
+
+# first, create a dict that maps vocabulary tokens to id numbers 
+token_to_id = ({
+    'what': 0,
+    'is': 1,
+    'your': 2,
+    'name': 3,
+    'gpt': 4,
+    'my': 5,
+    '<EOS>': 10, # END OF SEQUENCE
+    '<PAD>': 11, # PADDING 
+})
+
+## create the dict that maps the ids to tokens, for interpretintg the model output.
+id_to_token = dict(map(reversed, token_to_id.items()))
+VOCAB_SIZE = len(token_to_id)
+SEQ_LEN = 6
+D_MODEL = 2
+# we use decoder only transformer, the inputs contain
+# the questions followed by <EOS> token followed by the response 'gpt'
+# this is because all of the tokens will be used as inputs to the decoder only
+# transformer during training. 
+# it's called teacher forcing 
+# teacher forcing helps us train the neural network faster 
+
+inputs = torch.tensor([
+    [
+        token_to_id['what'],
+        token_to_id['is'],
+        token_to_id['your'],
+        token_to_id['name'],
+    ],
+    [
+        token_to_id['gpt'],
+        token_to_id['is'],
+        token_to_id['my'],
+    ]
+])
+
+# we are using decoder only transformer the outputs, or
+# the predictions, are the input questions (minus the first word) followed by
+# <EOS> gpt <EOS>. the first <EOS> means we are dong processing the input question
+# and the second means we are done generating the output.
+labels = torch.tensor([
+    [
+        token_to_id['is'],
+        token_to_id['your'],
+        token_to_id['name'],
+        token_to_id['<EOS>'],
+        token_to_id['gpt'],
+        token_to_id['<EOS>'],
+    ],
+    [
+        token_to_id['is'],
+        token_to_id['my'],
+        token_to_id['<EOS>'],
+        token_to_id['name'],
+        token_to_id['<EOS>'],
+        token_to_id['<PAD>'],
+    ]
+])
+
+dataset = TensorDataset(inputs, labels)
+dataloader = DataLoader(dataset=dataset)
+
+print(f'Shape of the input: {inputs.shape}')
+print(f'Shape of the labels: {labels.shape}')
+
+x = inputs.unsqueeze(0)
+y = labels.unsqueeze(0)
+
+print(f'Batch input: {x.shape}')
+print(f'Batch labels: {y.shape}')

model.py

@@ -0,0 +1,448 @@
+'''
+A Decoder-Only Transformer components
+    -> Word Embedding
+    -> Position Encoding
+    -> Masked Self-Attention
+    -> Residual Connections
+    -> A fully connected layer
+    -> Classification Head
+'''
+import torch 
+import torch.nn as nn 
+import math 
+
+class WordEmbeddings(nn.Module):
+
+    def __init__(self, d_model: int, vocab_size: int):
+        ## d_model: The dimension of the transformer, which is also the number of embedding values per token.
+        ## vocab_size: Get the size of the underlying vocabulary
+       
+        super().__init__()
+        self.d_model = d_model
+        self.vocab_size = vocab_size
+        self.embedding = nn.Embedding(num_embeddings=vocab_size, 
+                                      embedding_dim=d_model)
+    
+    def forward(self, x):
+        # (batch, seq_len) --> (batch, seq_len, d_model)
+        # multiply by sqrt(d_model) to scale the embeddings 
+        return self.embedding(x) * math.sqrt(self.d_model)
+
+class PositionEncoding(nn.Module):
+    '''Ref: https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/decoder_transformers_with_pytorch_and_lightning_v2.ipynb
+    '''    
+ 
+    def __init__(self, d_model: int, seq_len: int, dropout: float):
+        ## d_model = The dimension of the transformer, which is also the number of embedding values per token.
+        ##           In the transformer I used in the StatQuest: Transformer Neural Networks Clearly Explained!!!
+        ##           d_model=2, so that's what we'll use as a default for now.
+        ##           However, in "Attention Is All You Need" d_model=512
+        ## seq_len = maximum number of tokens we allow as input.
+        ##           Since we are precomputing the position encoding values and storing them in a lookup table
+        ##           we can use d_model and seq_len to determine the number of rows and columns in that
+        ##           lookup table.
+        ##
+        ##           In this simple example, we are only using short phrases, so we are using
+        ##           seq_len=6 as the default setting.
+        ##           However, in The Annotated Transformer, they set the default value for seq_len to 5000
+        
+        ## We call the super's init because by creating our own __init__() method, we overwrite the one
+        ## we inherited from nn.Module. So we have to explicity call nn.Module's __init__(), otherwise it
+        ## won't get initialized. NOTE: If we didn't write our own __init__(), then we would not have
+        ## to call super().__init__(). Alternatively, if we didn't want to access any of nn.Module's methods, 
+        ## we wouldn't have to call it then either.        
+        super().__init__()
+
+        self.d_model = d_model
+        self.seq_len = seq_len
+        self.dropout = nn.Dropout(dropout)
+
+        ## Now we create a lookup table, pe, of position encoding values and initialize all of them to 0.
+        ## To do this, we will make a matrix of 0s that has seq_len rows and d_model columns.
+        ## for example...
+        ## torch.zeros(3, 2)
+        ## ...returns a matrix of 0s with 3 rows and 2 columns...
+        ## tensor([[0., 0.],
+        ##         [0., 0.],
+        ##         [0., 0.]])
+        pe = torch.zeros(seq_len, d_model)
+
+        ## Now we create a sequence of numbers for each position that a token can have in the input (or output).
+        ## For example, if the input tokens where "I'm happy today!", then "I'm" would get the first
+        ## position, 0, "happy" would get the second position, 1, and "today!" would get the third position, 2.
+        ## NOTE: Since we are going to be doing math with these position indices to create the 
+        ## positional encoding for each one, we need them to be floats rather than ints.
+        ## 
+        ## NOTE: Two ways to create floats are...
+        ##
+        ## torch.arange(start=0, end=3, step=1, dtype=torch.float)
+        ##
+        ## ...and...
+        ##
+        ## torch.arange(start=0, end=3, step=1).float()
+        ##
+        ## ...but the latter is just as clear and requires less typing.
+        ##
+        ## Lastly, .unsqueeze(1) converts the single list of numbers that torch.arange creates into a matrix with
+        ## one row for each index, and all of the indices in a single column. So if "seq_len" = 3, then we
+        ## would create a matrix with 3 rows and 1 column like this...
+        ##
+        ## torch.arange(start=0, end=3, step=1, dtype=torch.float).unsqueeze(1)
+        ##
+        ## ...returns...
+        ##
+        ## tensor([[0.],
+        ##         [1.],
+        ##         [2.]])        
+        position = torch.arange(start=0, end=seq_len, step=1).float().unsqueeze(1)
+
+
+        ## Here is where we start doing the math to determine the y-axis coordinates on the
+        ## sine and cosine curves.
+        ##
+        ## The positional encoding equations used in "Attention is all you need" are...
+        ##
+        ## PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
+        ## PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
+        ##
+        ## ...and we see, within the sin() and cos() functions, we divide "pos" by some number that depends
+        ## on the index (i) and total number of PE values we want per token (d_model). 
+        ##
+        ## NOTE: When the index, i, is 0 then we are calculating the y-axis coordinates on the **first pair** 
+        ##       of sine and cosine curves. When i=1, then we are calculating the y-axis coordiantes on the 
+        ##       **second pair** of sine and cosine curves. etc. etc.
+        ##
+        ## Now, pretty much everyone calculates the term we use to divide "pos" by first, and they do it with
+        ## code that looks like this...
+        ##
+        ## div_term = torch.exp(torch.arange(start=0, end=d_model, step=2).float() * -(math.log(10000.0) / d_model))
+        ##
+        ## Now, at least to me, it's not obvious that div_term = 1/(10000^(2i/d_model)) for a few reasons:
+        ##
+        ##    1) div_term wraps everything in a call to torch.exp() 
+        ##    2) It uses log()
+        ##    2) The order of the terms is different 
+        ##
+        ## The reason for these differences is, presumably, trying to prevent underflow (getting too close to 0).
+        ## So, to show that div_term = 1/(10000^(2i/d_model))...
+        ##
+        ## 1) Swap out math.log() for torch.log() (doing this requires converting 10000.0 to a tensor, which is my
+        ##    guess for why they used math.log() instead of torch.log())...
+        ##
+        ## torch.exp(torch.arange(start=0, end=d_model, step=2).float() * -(torch.log(torch.tensor(10000.0)) / d_model))
+        ##
+        ## 2) Rearrange the terms...
+        ##
+        ## torch.exp(-1 * (torch.log(torch.tensor(10000.0)) * torch.arange(start=0, end=d_model, step=2).float() / d_model))
+        ##
+        ## 3) Pull out the -1 with exp(-1 * x) = 1/exp(x)
+        ##
+        ## 1/torch.exp(torch.log(torch.tensor(10000.0)) * torch.arange(start=0, end=d_model, step=2).float() / d_model)
+        ##
+        ## 4) Use exp(a * b) = exp(a)^b to pull out the 2i/d_model term...
+        ##
+        ## 1/torch.exp(torch.log(torch.tensor(10000.0)))^(torch.arange(start=0, end=d_model, step=2).float() / d_model)
+        ##
+        ## 5) Use exp(log(x)) = x to get the original form of the denominator...
+        ##
+        ## 1/(torch.tensor(10000.0)^(torch.arange(start=0, end=d_model, step=2).float() / d_model))
+        ##
+        ## 6) Bam.
+        ## 
+        ## So, that being said, I don't think underflow is actually that big an issue. In fact, some coder at Hugging Face
+        ## also doesn't think so, and their code for positional encoding in DistilBERT (a streamlined version of BERT, which
+        ## is a transformer model)
+        ## calculates the values directly - using the form of the equation found in original Attention is all you need
+        ## manuscript. See...
+        ## https://github.com/huggingface/transformers/blob/455c6390938a5c737fa63e78396cedae41e4e87e/src/transformers/modeling_distilbert.py#L53
+        ## So I think we can simplify the code, but I'm also writing all these comments to show that it is equivalent to what
+        ## you'll see in the wild...
+        ##
+        ## Now let's create an index for the embedding positions to simplify the code a little more...
+        embedding_index = torch.arange(start=0, end=d_model, step=2).float()
+        ## NOTE: Setting step=2 results in the same sequence numbers that we would get if we multiplied i by 2.
+        ##       So we can save ourselves a little math by just setting step=2.
+
+        ## And now, finally, let's create div_term...
+        div_term = 1/torch.tensor(10000.0)**(embedding_index / d_model)
+        
+        ## Now we calculate the actual positional encoding values. Remember 'pe' was initialized as a matrix of 0s
+        ## with seq_len (max number of input tokens) rows and d_model (number of embedding values per token) columns.
+        pe[:, 0::2] = torch.sin(position * div_term) ## every other column, starting with the 1st, has sin() values
+        pe[:, 1::2] = torch.cos(position * div_term) ## every other column, starting with the 2nd, has cos() values
+        ## NOTE: If the notation for indexing 'pe[]' looks cryptic to you, read on...
+        ##
+        ## First, let's look at the general indexing notation:
+        ##
+        ## For each row or column in matrix we can select elements in that
+        ## row or column with the following indexs...
+        ##
+        ## i:j:k = select elements between i and j with stepsize = k.
+        ##
+        ## ...where...
+        ##
+        ## i defaults to 0
+        ## j defaults to the number of elements in the row, column or whatever.
+        ## k defaults to 1
+        ##
+        ## Now that we have looked at the general notation, let's look at specific
+        ## examples so that we can understand it.
+        ##
+        ## We'll start with: pe[:, 0::2]
+        ##
+        ## The stuff that comes before the comma (in this case ':') refers to the rows we want to select.
+        ## The ':' before the comma means "select all rows" because we are not providing specific 
+        ## values for i, j and k and, instead, just using the default values.
+        ##
+        ## The stuff after the comma refers to the columns we want to select.
+        ## In this case, we have '0::2', and that means we start with
+        ## the first column (column =  0) and go to the end (using the default value for j)
+        ## and we set the stepsize to 2, which means we skip every other column.
+        ##
+        ## Now to understand pe[:, 1::2]
+        ##
+        ## Again, the stuff before the comma refers to the rows, and, just like before
+        ## we use default values for i,j and k, so we select all rows.
+        ##
+        ## The stuff that comes after the comma refers to the columns.
+        ## In this case, we start with the 2nd column (column = 1), and go to the end
+        ## (using the default value for 'j') and we set the stepsize to 2, which
+        ## means we skip every other column.
+        ##
+        ## NOTE: using this ':' based notation is called "indexing" and also called "slicing"
+        ## Add a batch dimension to the positional encoding
+        pe = pe.unsqueeze(0) # (1, seq_len, d_model)
+        ## Now we "register 'pe'.
+        self.register_buffer('pe', pe) ## "register_buffer()" ensures that
+                                       ## 'pe' will be moved to wherever the model gets
+                                       ## moved to. So if the model is moved to a GPU, then,
+                                       ## even though we don't need to optimize 'pe', it will 
+                                       ## also be moved to that GPU. This, in turn, means
+                                       ## that accessing 'pe' will be relatively fast copared
+                                       ## to having a GPU have to get the data from a CPU.
+
+
+    def forward(self, word_embeddings):
+        ## Because this class, PositionEncoding, inherits from nn.Module, the forward() method 
+        ## is called by default when we use a PositionEncoding() object.
+        ## In other words, after we create a PositionEncoding() object, pe = PositionEncoding(),
+        ## then pe(word_embeddings) will call forward() and so this is where 
+        ## we will add the position encoding values to the word embedding values    
+        ## (batch, seq_len, d_model)
+        x =  word_embeddings + (self.pe[:,:word_embeddings.shape[1], :]).requires_grad_(False)
+
+        return self.dropout(x)
+
+class LayerNormalization(nn.Module):
+
+    def __init__(self, features: int, eps:float=10**-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.alpha = nn.Parameter(torch.ones(features)) # alpha is a learnable parameter
+        self.bias = nn.Parameter(torch.zeros(features)) # bias is a learnable parameter
+
+    def forward(self, x):
+        # x: (batch, seq_len, hidden_size)
+         # Keep the dimension for broadcasting
+        mean = x.mean(dim = -1, keepdim = True) # (batch, seq_len, 1)
+        # Keep the dimension for broadcasting
+        std = x.std(dim = -1, keepdim = True) # (batch, seq_len, 1)
+        # eps is to prevent dividing by zero or when std is very small
+        return self.alpha * (x - mean) / (std + self.eps) + self.bias
+
+class MultiHeadAttentionBlock(nn.Module):
+
+    def __init__(self, d_model: int, h: int, dropout: float) -> None:
+        super().__init__()
+        # Make sure d_model is divisible by h
+        assert d_model % h == 0, "d_model is not divisible by h"
+
+        self.d_model = d_model # Embedding vector size
+        self.h = h # Number of heads
+
+        self.d_k = d_model // h # Dimension of vector seen by each head
+        self.w_q = nn.Linear(in_features=d_model, out_features=d_model, bias=False) # Wq
+        self.w_k = nn.Linear(in_features=d_model, out_features=d_model, bias=False) # Wk
+        self.w_v = nn.Linear(in_features=d_model, out_features=d_model, bias=False) # Wv
+        self.w_o = nn.Linear(in_features=d_model, out_features=d_model, bias=False) # Wo
+        self.dropout = nn.Dropout(dropout)
+
+    @staticmethod
+    def attention(query, key, value, mask, dropout: nn.Dropout):
+        d_k = query.shape[-1]
+        ## (batch, h, seq_len, d_k) --> (batch, h, seq_len, seq_len)
+        ## Compute attention scores, the equation is (q * k^T)/sqrt(d_model)      
+        attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)
+        if mask is not None:
+            ## Here we are masking out things we don't want to pay attention to,
+            ## like tokens that come after the current token.
+            ## We can also use masking to block out the <PAD> token,
+            ## which is used when we have a batch of inputs sequences
+            ## and they are not all the exact same length. Because the batch is passed
+            ## in as a matrix, each input sequence has to have the same length, so we
+            ## add <PAD> to the shorter sequences so that they are all as long ast the
+            ## longest sequence.
+            ##
+            ## We replace <PAD>, or tokens that come after the current token
+            ## with a very large negative number so that the SoftMax() function
+            ## will give all masked elements an output value (or "probability") of 0.            
+            ## Write a very low value (indicating -inf) to the positions where mask == 0
+            attention_scores.masked_fill_(mask == 0, -1e9)
+
+        ## Apply softmax to determine what percent of each token's value to
+        ## use in the final attention values.
+        ## (batch, h, seq_len, seq_len)    
+        attention_scores = attention_scores.softmax(dim=-1) 
+
+        if dropout is not None:
+            attention_scores = dropout(attention_scores)
+
+        ## (batch, h, seq_len, seq_len) --> (batch, h, seq_len, d_k)
+        ## return attention scores which can be used for visualization
+        return (attention_scores @ value), attention_scores
+
+    def forward(self, q, k, v, mask):
+        query = self.w_q(q) # (batch, seq_len, d_model) --> (batch, seq_len, d_model)
+        key = self.w_k(k) # (batch, seq_len, d_model) --> (batch, seq_len, d_model)
+        value = self.w_v(v) # (batch, seq_len, d_model) --> (batch, seq_len, d_model)
+
+        # (batch, seq_len, d_model) --> (batch, seq_len, h, d_k) --> (batch, h, seq_len, d_k)
+        query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(1, 2)
+        key = key.view(key.shape[0], key.shape[1], self.h, self.d_k).transpose(1, 2)
+        value = value.view(value.shape[0], value.shape[1], self.h, self.d_k).transpose(1, 2)
+
+        # Calculate attention
+        x, self.attention_scores = MultiHeadAttentionBlock.attention(query, key, value, mask, self.dropout)
+        
+        # Combine all the heads together
+        # (batch, h, seq_len, d_k) --> (batch, seq_len, h, d_k) --> (batch, seq_len, d_model)
+        x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.h * self.d_k)
+
+        # Multiply by Wo
+        # (batch, seq_len, d_model) --> (batch, seq_len, d_model)  
+        return self.w_o(x)
+
+class ResidualConnection(nn.Module):
+    
+        def __init__(self, features: int, dropout: float) -> None:
+            super().__init__()
+            self.dropout = nn.Dropout(dropout)
+            self.norm = LayerNormalization(features)
+    
+        def forward(self, x, sublayer):
+            return x + self.dropout(sublayer(self.norm(x)))
+
+class FeedForwardBlock(nn.Module):
+
+    def __init__(self, d_model: int, d_ff: int, dropout: float) -> None:
+        super().__init__()
+        self.linear_1 = nn.Linear(d_model, d_ff) # w1 and b1
+        self.dropout = nn.Dropout(dropout)
+        self.linear_2 = nn.Linear(d_ff, d_model) # w2 and b2
+
+    def forward(self, x):
+        # (batch, seq_len, d_model) --> (batch, seq_len, d_ff) --> (batch, seq_len, d_model)
+        return self.linear_2(self.dropout(torch.relu(self.linear_1(x))))
+
+class DecoderBlock(nn.Module):
+
+    def __init__(self, 
+                 features: int, 
+                 self_attention_block: MultiHeadAttentionBlock, 
+                 feed_forward_block: FeedForwardBlock, 
+                 dropout: float) -> None:
+        super().__init__()
+        self.self_attention_block = self_attention_block
+        self.feed_forward_block = feed_forward_block
+        self.residual_connections = nn.ModuleList([ResidualConnection(features, dropout) for _ in range(2)])
+
+    def forward(self, x, mask):
+        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, mask))
+        x = self.residual_connections[1](x, self.feed_forward_block)
+        return x
+
+class Decoder(nn.Module):
+
+    def __init__(self, features: int, layers: nn.ModuleList) -> None:
+        super().__init__()
+        self.layers = layers
+        self.norm = LayerNormalization(features)
+
+    def forward(self, x, mask):
+        for layer in self.layers:
+            x = layer(x, mask)
+        return self.norm(x)
+    
+class ProjectionLayer(nn.Module):
+
+    def __init__(self, d_model, vocab_size):
+        super().__init__()
+        self.proj = nn.Linear(d_model, vocab_size)
+
+    def forward(self, x) -> None:
+        # (batch, seq_len, d_model) --> (batch, seq_len, vocab_size)
+        return self.proj(x)
+
+class DecoderOnlyTransformer(nn.Module):
+    def __init__(self, 
+                 word_embedding: WordEmbeddings,
+                 position_embedding: PositionEncoding,
+                 decoder: Decoder,
+                 projection_layer: ProjectionLayer):
+        super().__init__()
+        self.word_embedding = word_embedding
+        self.position_embedding = position_embedding
+        self.decoder = decoder
+        self.projection_layer = projection_layer
+
+    def decode(self, x: torch.Tensor, mask: torch.Tensor):
+        # x shape (batch, seq_len)
+        x = self.word_embedding(x)
+        x = self.position_embedding(x)
+        # x shape (batch, seq_len, d_model)
+        return self.decoder(x, mask)
+    
+    def project(self, x):
+        # (batch, seq_len, vocab_size)
+        return self.projection_layer(x)
+
+def build_transformer(vocab_size: int, 
+                      seq_len: int,
+                      d_model: int=512, 
+                      N: int=6, 
+                      h: int=8, 
+                      dropout: float=0.1, 
+                      d_ff: int=2048) -> DecoderOnlyTransformer:
+    # Create the embedding layers
+    word_embedding = WordEmbeddings(d_model, vocab_size)
+
+    # Create the positional encoding layers
+    position_encoding = PositionEncoding(d_model, seq_len, dropout)
+
+    # Create the decoder blocks
+    decoder_blocks = []
+    for _ in range(N):
+        multi_head_self_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)
+        feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
+        decoder_block = DecoderBlock(d_model, multi_head_self_attention_block, feed_forward_block, dropout)
+        decoder_blocks.append(decoder_block)
+    
+    # Create the encoder and decoder
+    decoder = Decoder(d_model, nn.ModuleList(decoder_blocks))
+    
+    # Create the projection layer
+    projection_layer = ProjectionLayer(d_model, vocab_size)
+    
+    # Create the transformer
+    transformer = DecoderOnlyTransformer(word_embedding, 
+                                         position_encoding,
+                                         decoder,
+                                         projection_layer)
+    
+    # Initialize the parameters
+    for p in transformer.parameters():
+        if p.dim() > 1:
+            nn.init.xavier_uniform_(p)
+    
+    return transformer
+    

requirements.txt

@@ -0,0 +1,12 @@
+transformers>=4.34
+numpy
+tiktoken
+torch
+torchdata
+accelerate
+evaluate
+rouge_score
+loralib
+peft
+datasets
+torchmetrics

shakespeare/data/test.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff705d970c50caa6f8b25a9f35b2c168a8dae8d81ef0963347e63327402c9d60
+size 69612

shakespeare/data/train.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c3897b6613e06aa21e514d86d67214e3570fcfe20ccb8b983f4cb676e1d6c56
+size 584160

shakespeare/weights/tmodel_26000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9c59034983f6e1048d7d2d3699da8a6b4c25ead392c7c43c7eac6bdfe5a9b74
+size 846393125

shakespeare_config.py

@@ -0,0 +1,57 @@
+import os
+import torch
+
+from pathlib import Path
+from transformers import  GPT2Tokenizer
+
+
+def get_config():
+    return {
+        "batch_size": 8,
+        "num_epochs": 600000,
+        "lr": 10**-4,
+        "seq_len": 350,
+        "d_model": 512,
+        "vocab_size": 50304,
+        "datasource": 'shakespeare',
+        "model_folder": "weights",
+        "model_basename": "tmodel_",
+        "preload": "latest",
+        "tokenizer_file": "tokenizer.json",
+        "experiment_name": "runs/tmodel"
+    }
+
+current_directory = os.path.dirname(os.path.abspath(__file__))
+
+def get_weights_file_path(config, epoch: str):
+    model_folder = f"{current_directory}/{config['datasource']}/{config['model_folder']}"
+    # Create the folder and subfolders if they don't exist
+    Path(model_folder).mkdir(parents=True, exist_ok=True)
+    model_filename = f"{config['model_basename']}{epoch}.pt"
+    return model_folder + '/' + model_filename
+
+def get_data_folder_path(config):
+    model_folder = f"{current_directory}/{config['datasource']}/data"
+    Path(model_folder).mkdir(parents=True, exist_ok=True)
+    return model_folder 
+
+# Find the latest weights file in the weights folder
+def latest_weights_file_path(config):
+    model_folder = f"{current_directory}/{config['datasource']}/{config['model_folder']}"
+    model_filename = f"{config['model_basename']}*"
+    weights_files = list(Path(model_folder).glob(model_filename))
+    if len(weights_files) == 0:
+        return None
+    weights_files.sort()
+    return str(weights_files[-1])
+
+def get_gpt2_tokenizer(config):
+    tokenizer:GPT2Tokenizer = GPT2Tokenizer.from_pretrained(
+        pretrained_model_name_or_path="openai-community/gpt2",
+        model_max_length=config['seq_len'],
+        pad_token='[PAD]')
+    return tokenizer
+
+def causal_mask(size):
+    mask = torch.triu(torch.ones((1, size, size)), diagonal=1).type(torch.int)
+    return mask==0

shakespeare_data.py

@@ -0,0 +1,49 @@
+'''Ref: https://github.com/karpathy/nanoGPT/blob/master/data/shakespeare/prepare.py
+'''
+import os
+import requests
+import numpy as np
+from transformers import  GPT2Tokenizer
+from shakespeare_config import get_data_folder_path, get_config, get_gpt2_tokenizer
+from pathlib import Path
+
+if __name__=='__main__':
+    config=get_config()
+    data_folder_path = get_data_folder_path(config=config)
+    # download the tiny shakespeare dataset
+    input_file_path = os.path.join(data_folder_path, 'input.txt')
+    tokenizer:GPT2Tokenizer = get_gpt2_tokenizer(config=config)
+
+    print(tokenizer.model_max_length)
+
+    if not Path(input_file_path).exists():
+        data_url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
+        with open(input_file_path, 'w', encoding='utf-8') as f:
+            f.write(requests.get(data_url).text)
+
+        data=''
+        with open(input_file_path, 'r', encoding='utf-8') as f:
+            for line in f.readlines():
+                if len(line.rstrip())>0:
+                    data += ' ' + line    
+        
+        print(data)
+        n = len(data)
+        train_split = int(n*0.9)
+        train_data = data[:train_split]
+        test_data = data[train_split:]    
+
+        train_ids = tokenizer.encode(train_data)
+        test_ids = tokenizer.encode(test_data)
+        print(f"train has {len(train_ids):,} tokens")
+        print(f"test has {len(test_ids):,} tokens")
+
+        # export to bin files
+        train_ids = np.array(train_ids, dtype=np.uint16)
+        test_ids = np.array(test_ids, dtype=np.uint16)
+        train_ids.tofile(os.path.join(data_folder_path, 'train.bin'))
+        test_ids.tofile(os.path.join(data_folder_path, 'test.bin'))
+        # train has 292,080 tokens
+        # test has 34,806 tokens
+        print(tokenizer.convert_ids_to_tokens(tokenizer.eos_token_id))
+        print(tokenizer.convert_ids_to_tokens(tokenizer.pad_token_id))

shakespeare_inference.py

@@ -0,0 +1,115 @@
+from shakespeare_config import (get_config, 
+                                latest_weights_file_path, 
+                                get_gpt2_tokenizer,
+                                causal_mask,
+                                current_directory)
+import torch 
+import warnings
+import heapq
+from train import build_transformer
+
+def predict_with_greedy_search(start_str:str)-> None:
+    config:dict=get_config() 
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    tokenizer = get_gpt2_tokenizer(config=config)
+    model = build_transformer(vocab_size=config['vocab_size'],
+                              seq_len=config['seq_len'],
+                              d_model=config['d_model']).to(device)
+    # load the pretrained weights
+    model_filename = latest_weights_file_path(config)
+    state = torch.load(model_filename)
+    model.load_state_dict(state['model_state_dict'])
+    model.eval()
+
+    output = start_str
+    with torch.no_grad():
+        start_tokens = tokenizer.encode(start_str)
+        print(start_tokens)
+        input = torch.tensor(data=start_tokens, dtype=torch.int64).unsqueeze(dim=0).to(device)
+        # print(input)
+        while input.size(1) <= config['seq_len']:
+            # use mask otheriwse model may generate repetitive words in prediction
+            mask = causal_mask(input.size(1)).to(device)
+            out = model.decode(input,mask)
+            prob = model.project(out[:, -1])
+            _, next_word = torch.max(prob, dim=1)
+            input = torch.cat(
+                    [
+                        input, 
+                        torch.empty(1,1).type_as(input).fill_(next_word.item()).to(device)
+                    ],
+                    dim=1
+                )
+            output += tokenizer.decode(next_word.item())
+    
+    print(f'Model output: {output}')
+
+
+def predict_with_beam_search(start_str: str, 
+                             beam_width: int = 3) -> None:
+    config: dict = get_config() 
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    tokenizer = get_gpt2_tokenizer(config=config)
+    model = build_transformer(vocab_size=config['vocab_size'],
+                              seq_len=config['seq_len'],
+                              d_model=config['d_model']).to(device)
+
+    # Load the pretrained weights
+    model_filename = latest_weights_file_path(config)
+    state = torch.load(model_filename)
+    model.load_state_dict(state['model_state_dict'])
+    model.eval()
+
+    # Initial input
+    start_tokens = tokenizer.encode(start_str)
+    input = torch.tensor(data=start_tokens, dtype=torch.int64).unsqueeze(dim=0).to(device) # (1, seq_len)
+
+    # Beam search variables
+    beams = [(0, input, [])]  # Each beam is a tuple of (score, sequence, tokens_generated)
+
+    for _ in range(config['seq_len']):
+        all_candidates = []
+
+        # Process each beam
+        for score, seq, tokens in beams:
+            # use mask otheriwse model may generate repetitive words in prediction
+            mask = causal_mask(seq.size(1)).to(device)
+            out = model.decode(seq, mask)
+            prob = model.project(out[:, -1])
+
+            # Get the top k predictions
+            top_k_probabilities, top_k_indices = torch.topk(prob, beam_width, dim=1)
+
+            # Generate new beams for each of the top k tokens
+            for i in range(beam_width):
+                new_token = top_k_indices[0, i].item()
+                new_score = score - torch.log(top_k_probabilities[0, i]).item()  # We negate because we want to maximize
+                new_seq = torch.cat([seq, torch.tensor([[new_token]], device=device)], dim=1)
+                new_tokens = tokens + [new_token]
+                all_candidates.append((new_score, new_seq, new_tokens))
+
+        # Sort all candidates based on their score and keep the top `beam_width` beams
+        beams = heapq.nsmallest(beam_width, all_candidates, key=lambda x: x[0])
+
+        # Optionally, stop early if all beams end with an EOS token
+        if all(beam[1].shape[1] >= config['seq_len'] for beam in beams):
+            break
+
+    # Retrieve the best beam (with the highest score)
+    best_beam = beams[0]
+    best_tokens = best_beam[2]
+
+    # Decode the final sequence
+    output = tokenizer.decode(best_tokens, skip_special_tokens=True)
+    print(f'Model output: {output}')
+
+if __name__ == '__main__':
+    warnings.filterwarnings("ignore")
+    start_str = 'Now sadder, that you come so'
+    predict_with_greedy_search(start_str=start_str)
+    print('--'*100)
+    predict_with_beam_search(start_str=start_str)

train.py

@@ -0,0 +1,234 @@
+import torchmetrics.classification
+from torchmetrics.text import ROUGEScore
+from model import build_transformer
+from shakespeare_config import (get_config, 
+                                get_data_folder_path, 
+                                get_weights_file_path, 
+                                latest_weights_file_path,
+                                current_directory,
+                                causal_mask,
+                                get_gpt2_tokenizer)
+
+import torch 
+import torch.nn as nn 
+from torch.utils.tensorboard import SummaryWriter
+import torchmetrics
+import numpy as np
+import warnings
+import os
+from pathlib import Path
+
+def get_model(config):
+    model = build_transformer(vocab_size=config['vocab_size'],
+                              seq_len=config['seq_len'],
+                              d_model=config['d_model'])
+    return model
+
+def get_batch(split, data_dir, block_size, batch_size, device='gpu', device_type='cuda'):
+    # We recreate np.memmap every batch to avoid a memory leak, as per
+    # https://stackoverflow.com/questions/45132940/numpy-memmap-memory-usage-want-to-iterate-once/61472122#61472122
+    if split == 'train':
+        data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
+    else:
+        data = np.memmap(os.path.join(data_dir, 'test.bin'), dtype=np.uint16, mode='r')
+    ix = torch.randint(len(data) - block_size, (batch_size,))
+    x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
+    y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
+    # if device_type == 'cuda':
+    #     # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
+    #     x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
+    # else:
+    #     x, y = x.to(device), y.to(device)
+    return x, y
+
+def greedy_decode(model, 
+                  input, 
+                  mask,
+                  tokenizer,
+                  max_len,
+                  device):
+    while True:
+        if input.size(1) == max_len:
+            break
+
+        out = model.decode(input, mask)
+        prob = model.project(out[:, -1])
+        _, next_word = torch.max(prob, dim=1)
+        input = torch.cat(
+            [input, torch.empty(1,1).type_as(input).fill_(next_word.item()).to(device)],
+            dim=1
+        )
+        if next_word == tokenizer.eos_token_id:
+            break 
+    return input.squeeze(0)
+
+def run_validation(model,
+                   x,
+                   y,
+                   tokenizer,
+                   max_len,
+                   device,
+                   print_msg,
+                   global_step,
+                   writer,
+                   rouge:ROUGEScore):
+    
+    model.eval()
+    source_texts = []
+    expected = []
+    predicted = []
+    
+    with torch.no_grad():
+        decoder_input = x.to(device) # (b, seq)
+        mask = causal_mask(x.size(1)).to(device) #(b,1,1,seq)
+
+        # check that batch size is 1
+        assert decoder_input.size(0)==1, "batch size  must be 1 for validation"
+
+        model_out = greedy_decode(model,
+                                  decoder_input, 
+                                  mask, 
+                                  tokenizer,
+                                  max_len, 
+                                  device)
+        
+        source_text = tokenizer.decode(x[0])
+        target_text = tokenizer.decode(y[0])
+        model_out_text = tokenizer.decode(model_out.detach().cpu().numpy())
+
+        source_texts.append(source_text)
+        expected.append(target_text)
+        predicted.append(model_out_text)
+        
+        # Print the source, target and model output
+        print_msg('-'*100)
+        print_msg(f"{f'SOURCE: ':>12}{source_text}")
+        print_msg(f"{f'TARGET: ':>12}{target_text}")
+        print_msg(f"{f'PREDICTED: ':>12}{model_out_text}")
+
+        rouge_score = rouge(predicted, expected)  
+        print_msg(f"{f'ROUGE-1 Score: ':>12}{rouge_score['rouge1_fmeasure'].item()}")
+        print_msg(f"{f'ROUGE-2 Score: ':>12}{rouge_score['rouge2_fmeasure'].item()}")
+        print_msg(f"{f'ROUGE-L Score: ':>12}{rouge_score['rougeL_fmeasure'].item()}")      
+        print_msg('-'*100)
+
+        if writer:
+            writer.add_scalar('validation ROUGE/ROUGE-1', rouge_score["rouge1_fmeasure"].item(), global_step)
+            writer.add_scalar('validation ROUGE/ROUGE-2', rouge_score["rouge2_fmeasure"].item(), global_step)
+            writer.add_scalar('validation ROUGE/ROUGE-L', rouge_score["rougeL_fmeasure"].item(), global_step)
+            writer.add_scalar('validation ROUGE/ROUGE-L', rouge_score["rougeLsum_fmeasure"].item(), global_step)
+            writer.flush() 
+
+def train_model(config):
+    # define the device 
+    device = "cuda" if torch.cuda.is_available() else "mps" if torch.has_ms or torch.backends.mps.is_available else "cpu"
+    print("Using device:", device)
+
+    if (device == 'cuda'):
+        print(f"Device name: {torch.cuda.get_device_name(device=device.index)}")
+        print(f"Device memory: {torch.cuda.get_device_properties(device.index).total_memory / 1024 ** 3} GB")
+    elif (device == 'cpu'):
+        print(f"device name: <mps>")
+    else:
+        print("It's cpu")
+
+    device = torch.device(device)
+
+    # make sure the weights folder exists 
+    Path(f"{current_directory}/{config['datasource']}_{config['model_folder']}").mkdir(parents=True, exist_ok=True)
+
+    tokenizer = get_gpt2_tokenizer(config=config)
+    model = get_model(config).to(device) 
+    # tensorboard
+    writer = SummaryWriter(f"{current_directory}/{config['experiment_name']}")
+
+    optimizer = torch.optim.Adam(model.parameters(), 
+                                 lr=config['lr'],
+                                 eps=1e-9)
+    rouge:ROUGEScore = ROUGEScore()
+
+    # if the user specified a model to preload before training, load it 
+    initial_epoch = 0
+    global_step = 0
+    preload = config['preload']
+    model_filename = (latest_weights_file_path(config) if preload == 'latest' else get_weights_file_path(config, preload) if preload else None)
+    if model_filename:
+        print(f'Preloading model {model_filename}')
+        state = torch.load(model_filename)
+        model.load_state_dict(state['model_state_dict'])
+        initial_epoch = state['epoch'] + 1 
+        optimizer.load_state_dict(state['optimizer_state_dict'])
+        global_step = state['global_step']
+    else:
+        print('No model to preload, starting from scratch')
+
+    loss_fn = nn.CrossEntropyLoss(ignore_index=tokenizer.convert_tokens_to_ids('[PAD]'), label_smoothing=0.1).to(device)
+    for epoch in range(initial_epoch, config['num_epochs']):
+        torch.cuda.empty_cache()
+        model.train()
+        #batch_iterator = tqdm(train_dataloader, desc=f"Processing Epoch {epoch:02d}")
+        X, y = get_batch(split='train',
+                        data_dir=get_data_folder_path(config=config),
+                        block_size=config['seq_len'],
+                        batch_size=config['batch_size'])
+        print(f'length of the batch: {len(X)}, type:{X.shape}')
+
+        decoder_input = X.to(device) # (b, seq_len)
+        decoder_mask = causal_mask(config['seq_len']).to(device) # (1, seq_len, seq_len)
+
+        # run the tensors through the encoder, decoder and the projection layer
+        decoder_output = model.decode(decoder_input, decoder_mask) # (b, seq, d_model)
+        proj_output = model.project(decoder_output) # (B, seq_len, vocab_size)
+
+        # compare the output with the label
+        label = y.to(device) #(b, seq_len)
+
+        # compute the loss using a simple cross entropy
+        loss = loss_fn(proj_output.view(-1, config['vocab_size']), 
+                        label.view(-1))
+        #batch_iterator.set_postfix({"loss": f"{loss.item():6.3f}"})
+        print(f"loss: {loss.item():6.3f}")
+
+        # log the loss 
+        writer.add_scalar('train loss', loss.item(), global_step)
+        writer.flush()
+
+        # backpropagate the loss 
+        loss.backward()
+
+        # update the weights
+        optimizer.step()
+        optimizer.zero_grad(set_to_none=True)
+
+        global_step += 1
+
+        # run validation at the end of every epoch
+        X_val, y_val = get_batch(split='val',
+                        data_dir=get_data_folder_path(config=config),
+                        block_size=config['seq_len'],
+                        batch_size=1)
+        run_validation(model, 
+                       X_val,
+                       y_val, 
+                       tokenizer, 
+                       config['seq_len'], 
+                       device, 
+                       lambda msg: print(msg), 
+                       global_step, 
+                       writer,
+                       rouge)
+
+        if epoch%1000==0 or epoch >= (config['num_epochs']-1):
+            # save the model at the end of every epoch
+            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)