Upload folder using huggingface_hub

single_headed_transformer_v4_weights.pth

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

transformer.py

@@ -0,0 +1,365 @@
+"""
+Author: Eshan Jayasundara
+Last Updated: 2nd of March 2025
+Created: 28th of February 2025
+___
+
+About:
+    └── Single head transformer (Transformer with self-attention training with teacher-forcing)
+___
+
+Training:
+    └── Teacher Forcing (Baseline)
+            ├── During training, the actual ground-truth tokens (from the dataset) are fed as input to the decoder instead of using the model’s own predictions.
+            ├── This makes training faster and ensures the model learns accurate token-to-token mappings.
+            └── Drawback: At inference time, the model doesn't see ground-truth inputs, so errors can accumulate (called exposure bias).
+___
+
+vocabulary dataset (from huggingface):
+    └── "yukiarimo/english-vocabulary"
+___
+
+Architecture:
+
+    Encoder
+        ├── Input text
+        │        └── Eg: "Hello, how are you?"
+        ├── Remove punctuation from input text
+        ├── Input tokenization
+        ├── Embedding lookup with torch.nn.Embedding
+        ├── Positional encoding (sin, cosine)
+        ├── Self-attention
+        │        ├── single-head
+        │        ├── Q = Wq @ Embedding
+        │        ├── K = Wk @ Embedding
+        │        └── V = Wv @ Embedding
+        ├── Add and norm
+        ├── Feed forward layer
+        │        ├── 2 hidden layers
+        │        ├── ReLU as the activation in hidden layer
+        │        ├── No activation at the output layer
+        │        └── nn.Linear(in_features=embedding_dim, out_features=d_ff), nn.ReLU(), nn.Linear(in_features=d_ff, out_features=embedding_dim)
+        ├── Add and norm (again)
+        └── Save encoder out to be used in cross attention
+
+    Decoder
+        ├── Decoder teacher text (same as the target text but shifted right)
+        │        ├── Eg: Decoder teacher text - "<SOS> hello, I'm fine."
+        │        └── Eg: target text - "hello, I'm fine. <EOS>"
+        ├── Remove punctuation from input text
+        ├── Input tokenization
+        ├── Embedding lookup with torch.nn.Embedding
+        ├── Positional encoding (sin, cosine)
+        ├── Masked-self-attention (single-head, new class signature for masked self attention introduced)
+        │        ├── single-head
+        │        ├── causal mask with triangular matrix
+        │        ├── Q = Wq @ Embedding
+        │        ├── K = Wk @ Embedding
+        │        └── V = Wv @ Embedding
+        ├── Add and norm
+        ├── Cross attention (same class signature used in the encoder self-attention can be used)
+        │        ├── single-head
+        │        ├── Q = Wq @ Add and normalized output from masked-self-attention
+        │        ├── K = Wk @ Encoder output
+        │        └── V = Wv @ Encoder output
+        ├── Add and norm
+        ├── Feed forward layer
+        │        ├── 2 hidden layers
+        │        ├── ReLU as the activation in hidden layer
+        │        ├── No activation at the output layer
+        │        └── nn.Linear(in_features=embedding_dim, out_features=d_ff), nn.ReLU(), nn.Linear(in_features=d_ff, out_features=embedding_dim)
+        ├── Add and norm (again)
+        └── Linear layer (No activation or softmax as in 'Attention is all you need' is used here)
+        
+    Optimization
+        ├── Initialize the Adam optimizer with the model’s parameters and a specified learning rate.
+        │        └── self.optimizer = torch.optim.Adam(params=self.parameters, lr=learning_rate)
+        ├── Before computing gradients for the current batch, we reset any existing gradients from the previous iteration.
+        │        └── self.optimizer.zero_grad()
+        ├── The model takes in `input_tokens` and `decoder_teacher_tokens` and performs a forward pass to compute `logits`
+        │        └── logits = self.forward(input_tokens, decoder_teacher_tokens)
+        ├── The cross-entropy loss
+        │        ├── Measures the difference between the predicted token distribution (logits) and the actual target tokens (decoder_target_tokens).
+        │        ├── It expects logits to have raw scores (not probabilities), and it applies softmax internally.
+        │        └── loss = F.cross_entropy(logits, decoder_target_tokens)
+        ├── Compute the gradients of the loss with respect to all trainable parameters in the model using automatic differentiation (backpropagation).
+        │        └── loss.backward()
+        └── Optimizer updates the model's weights using the computed gradients.
+                 └── self.optimizer.step()
+    
+    After training, to calculate the output tokens -> text, 'Autoregressive text generation' is used (one word at a time)
+        ├── Start with <SOS>. (Initial input to the decoder) but input to the encoder is the `prompt`.
+        ├── Model predicts the next token.
+        ├── Append the predicted token to the sequence.
+        ├── Repeat until an <EOS> token or max length is reached.
+        └── For illustration let's use words instead of tokens(numerical representation)
+                <SOS>
+                <SOS> hello
+                <SOS> hello I'm
+                <SOS> hello I'm good
+                <SOS> hello I'm good <EOS>
+___
+
+Feauter Improvements:
+    ├── Multi-head attention instead of single-head attention.
+    ├── Layer normalization instead of simple mean-variance normalization.
+    └── Dropout layers for better generalization.
+"""
+
+
+from datasets import load_dataset
+import torch
+import torch.nn as nn
+import string
+import torch.nn.functional as F
+
+# SELECT DEVICE
+if torch.cuda.is_available():
+    device = torch.device('cuda:1')
+    print(f"Using Device: {device} | Name: {torch.cuda.get_device_name(0)}")
+else:
+    device = torch.device('cpu')
+    print(f"Using Device: {device}")
+
+# SINGLE HEAD ATTENTION
+class SingleHeadAttention(torch.nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.query_layer = torch.nn.Linear(in_features=embedding_dim, out_features=embedding_dim)
+        self.key_layer = torch.nn.Linear(in_features=embedding_dim, out_features=embedding_dim)
+        self.value_layer = torch.nn.Linear(in_features=embedding_dim, out_features=embedding_dim)
+
+    def forward(self, q_embedding, k_embedding, v_embedding, attention_mask):
+        Q = self.query_layer.forward(q_embedding)
+        K = self.key_layer.forward(k_embedding)
+        V = self.value_layer.forward(v_embedding)
+
+        # softmax over last dimension
+        attention_scores = (torch.matmul(Q, K.transpose(-2, -1)) / self.embedding_dim ** 0.5).float()
+        
+        # Apply attention mask (if provided)
+        if attention_mask is not None:
+            attention_scores = attention_scores.masked_fill(attention_mask == 0, torch.finfo(attention_scores.dtype).min)
+
+        # Compute attention weights using softmax
+        attention_weights = F.softmax(attention_scores, dim=-1)  # (batch_size, seq_len, seq_len)
+
+        # Compute attention output
+        attention_output = torch.matmul(attention_weights, V)  # (batch_size, seq_len, embedding_dim)
+
+        return attention_output, attention_weights
+
+# FEED FORWARD NN
+class FeedForwardLayer(torch.nn.Module):
+    def __init__(self, embedding_dim=64, d_ff=256):
+        super().__init__()
+        self.fc1 = torch.nn.Linear(in_features=embedding_dim, out_features=d_ff)
+        self.fc2 = torch.nn.Linear(in_features=d_ff, out_features=embedding_dim)
+        self.activation = torch.nn.ReLU()
+
+    def forward(self, x):
+        return self.fc2.forward(
+                   self.activation(
+                       self.fc1.forward(x)
+                   )
+                )
+    
+# MASKED ATTENTION
+class DecoderMaskedAttention(nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.query_layer = nn.Linear(in_features=embedding_dim, out_features=embedding_dim)
+        self.key_layer = nn.Linear(in_features=embedding_dim, out_features=embedding_dim)
+        self.value_layer = nn.Linear(in_features=embedding_dim, out_features=embedding_dim)
+
+    def forward(self, q_embedding, k_embedding, v_embedding, attention_mask=None):
+        # Linear transformations
+        Q = self.query_layer(q_embedding)  # (seq_len, embedding_dim)
+        K = self.key_layer(k_embedding)    # (seq_len, embedding_dim)
+        V = self.value_layer(v_embedding)  # (seq_len, embedding_dim)
+
+        # Scaled dot-product attention scores
+        attention_scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.embedding_dim ** 0.5)  # (batch_size, seq_len, seq_len)
+
+        # Create causal mask
+        seq_len = q_embedding.shape[0]
+        causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1).bool()  # Upper triangular matrix
+
+        # Apply causal mask to attention scores
+        attention_scores = attention_scores.masked_fill(causal_mask, torch.finfo(attention_scores.dtype).min)
+
+        # Apply additional attention mask (if provided)
+        if attention_mask is not None:
+            attention_scores = attention_scores.masked_fill(attention_mask == 0, torch.finfo(attention_scores.dtype).min)
+
+        # Compute attention weights using softmax
+        attention_weights = F.softmax(attention_scores, dim=-1)  # (seq_len, seq_len)
+
+        # Compute attention output
+        attention_output = torch.matmul(attention_weights, V)  # (seq_len, embedding_dim)
+
+        return attention_output, attention_weights
+    
+
+class Transformer(torch.nn.Module):
+    def __init__(self, embedding_dim, learning_rate=1e-3, vocab_dataset="yukiarimo/english-vocabulary", split="train"):
+        super().__init__()
+        
+        # SETUP VOCABULARY
+        self.vocab_df = load_dataset(vocab_dataset, split=split).to_pandas()
+
+        remove_indices = self.vocab_df[(self.vocab_df["text"]=='PAD') | (self.vocab_df["text"]=='SOS') | (self.vocab_df["text"]=='EOS')].index
+        self.vocab_df = self.vocab_df.drop(remove_indices, axis=0)
+
+        self.vocab_df.loc[0, "text"] = '<PAD>'
+        self.vocab_df.loc[1, "text"] = '<UNK>'
+        self.vocab_df.loc[2, "text"] = '<SOS>'
+        self.vocab_df.loc[3, "text"] = '<EOS>'
+
+        self.vocab_size = self.vocab_df.shape[0]
+
+        self.vocab_df['idx'] = range(0, self.vocab_size)
+        self.vocab_df = self.vocab_df.set_index("text")
+        self.vocab = self.vocab_df["idx"].to_dict()
+
+        # INITIALIZE ALL TRAINABLE MODELS
+        self.embedding_fn = nn.Embedding(num_embeddings=self.vocab_size, embedding_dim=embedding_dim)
+        self.encoder_self_attention = SingleHeadAttention(embedding_dim=embedding_dim)
+        self.encoder_ff = FeedForwardLayer(embedding_dim=embedding_dim, d_ff=embedding_dim * 4)
+        self.cross_attention = SingleHeadAttention(embedding_dim=embedding_dim)
+        self.decoder_masked_attention = DecoderMaskedAttention(embedding_dim=embedding_dim)
+        self.decoder_ff = FeedForwardLayer(embedding_dim=embedding_dim, d_ff=embedding_dim * 4)
+        self.linear = nn.Linear(in_features=embedding_dim, out_features=self.vocab_size)
+
+        # PARAMETERS OF LEARNABLE MODELS
+        self.parameters = list(self.embedding_fn.parameters()) + \
+                          list(self.encoder_self_attention.parameters()) + \
+                          list(self.encoder_ff.parameters()) + \
+                          list(self.cross_attention.parameters()) + \
+                          list(self.decoder_masked_attention.parameters()) + \
+                          list(self.decoder_ff.parameters()) + \
+                          list(self.linear.parameters())
+        
+        # OPTIMIZER
+        self.optimizer = torch.optim.Adam(params=self.parameters, lr=learning_rate)
+
+    # INPUT TEXT HANDLING
+    def remove_punctuation(self, text):
+        return text.translate(str.maketrans("", "", string.punctuation))
+
+    def tokenize(self, text, unk_token="<UNK>"):
+        tokens = text.strip().split()
+        return torch.tensor([self.vocab.get(token, self.vocab.get(unk_token)) for token in tokens], device=device)
+
+    def positional_encoding(self, embedding, max_len, embedding_dim=64):
+        pe = torch.zeros(max_len, embedding_dim, device=device)
+        
+        # Create a tensor of positions (0, 1, 2, ..., max_len - 1)
+        position = torch.arange(0, max_len, dtype=torch.float, device=device).unsqueeze(1)
+
+        # Compute the division term for the frequency
+        div_term = torch.exp(torch.arange(0, embedding_dim, 2, device=device).float() * (torch.log(torch.tensor(10000.0, device=device))) / embedding_dim)
+
+        # Apply sine to even indices and cosine to odd indices
+        pe[:, 0::2] = torch.sin(position / div_term)  # Even dimensions
+        pe[:, 1::2] = torch.cos(position / div_term)  # Odd dimensions
+        
+        return embedding + pe
+    
+    # ADD AND NORM
+    def add_norm(self, old_tensor, new_tensor):
+        addition = old_tensor + new_tensor
+        norm = (addition - addition.mean(dim=-1, keepdim=True)) / addition.std(dim=-1, keepdim=True)
+        return norm
+
+    # ENCODER
+    def encoder(self, encoder_input_tokens):
+        encoder_input_embeddings = self.embedding_fn(encoder_input_tokens).to(device=device)
+        encoder_input_pos_embeddings = self.positional_encoding(encoder_input_embeddings, max_len=encoder_input_embeddings.shape[0], embedding_dim=64).to(device=device)
+        encoder_self_attention_out, _ = self.encoder_self_attention.forward(
+                                        q_embedding=encoder_input_pos_embeddings, 
+                                        k_embedding=encoder_input_pos_embeddings, 
+                                        v_embedding=encoder_input_pos_embeddings, 
+                                        attention_mask=None
+                                        )
+        add_norm_encoder_self_attention_out = self.add_norm(old_tensor=encoder_input_pos_embeddings, new_tensor=encoder_self_attention_out.to(device=device)).to(device=device)
+        encoder_ff_out = self.encoder_ff.forward(add_norm_encoder_self_attention_out).to(device=device)
+        add_norm_encoder_ff_out = self.add_norm(old_tensor=add_norm_encoder_self_attention_out, new_tensor=encoder_ff_out).to(device=device)
+        return add_norm_encoder_ff_out
+        
+    # DECODER
+    def decoder(self, decoder_teacher_tokens, encoder_out):
+        decoder_teacher_embeddings = self.embedding_fn(decoder_teacher_tokens).to(device=device)
+        decoder_teacher_pos_embeddings = self.positional_encoding(decoder_teacher_embeddings, max_len=decoder_teacher_embeddings.shape[0], embedding_dim=64).to(device=device)
+        decoder_masked_attention_out, _ = self.decoder_masked_attention.forward(
+                                            q_embedding=decoder_teacher_pos_embeddings,
+                                            k_embedding=decoder_teacher_pos_embeddings,
+                                            v_embedding=decoder_teacher_pos_embeddings,
+                                            attention_mask=None
+                                            )
+        add_norm_decoder_masked_attention_out = self.add_norm(old_tensor=decoder_teacher_pos_embeddings, new_tensor=decoder_masked_attention_out.to(device=device)).to(device=device)
+        cross_attention_out, _ = self.cross_attention.forward(
+                                q_embedding=add_norm_decoder_masked_attention_out, 
+                                k_embedding=encoder_out, 
+                                v_embedding=encoder_out, 
+                                attention_mask=None
+                                )
+        add_norm_cross_attention_out = self.add_norm(old_tensor=add_norm_decoder_masked_attention_out, new_tensor=cross_attention_out.to(device=device)).to(device=device)
+        decoder_ff_out = self.decoder_ff.forward(add_norm_cross_attention_out).to(device=device)
+        add_norm_decoder_ff_out = self.add_norm(old_tensor=add_norm_cross_attention_out, new_tensor=decoder_ff_out).to(device=device)
+        logits = self.linear.forward(add_norm_decoder_ff_out).to(device=device)
+        return logits
+    
+    # FORWARD PASS THROUGH ENCODER and DECODER
+    def forward(self, encoder_input_tokens, decoder_teacher_tokens):
+        encoder_out = self.encoder(encoder_input_tokens)
+        decoder_out = self.decoder(decoder_teacher_tokens, encoder_out=encoder_out)
+        return decoder_out
+    
+    # TRAIN the TRANSFORMER
+    def train(self, dataset, epochs=100):
+        for epoch in range(epochs):
+            total_loss = 0
+            for input_text, output_text in dataset:
+                encoder_input_text = self.remove_punctuation(input_text)
+                target_text = self.remove_punctuation(output_text)
+                decoder_teacher_text = "<SOS> " + target_text
+                decoder_target_text = target_text + " <EOS>"
+
+                encoder_input_tokens = self.tokenize(encoder_input_text)
+                decoder_teacher_tokens = self.tokenize(decoder_teacher_text)
+                decoder_target_tokens = self.tokenize(decoder_target_text)
+
+                self.optimizer.zero_grad()
+                logits = self.forward(encoder_input_tokens=encoder_input_tokens, decoder_teacher_tokens=decoder_teacher_tokens).to(device=device)
+                loss = F.cross_entropy(logits, decoder_target_tokens)
+                loss.backward()
+                self.optimizer.step()
+
+                total_loss += loss.item()
+                
+            if (epoch+1) % 10 == 0:
+                print(f"Epoch {epoch+1:04d} - Loss: {total_loss:.4f}")
+
+        print("*** END ***\n")
+
+    # GET PREDICTED TOKENS
+    def predict_tokens(self, encoder_input_tokens, max_output_len=20):
+        encoder_out = self.encoder(encoder_input_tokens).to(device=device)
+        decoder_input = [self.vocab["<SOS>"]]
+        for _ in range(max_output_len):
+            current_decoder_tokens = torch.tensor(decoder_input).to(device=device)
+            pred_index = torch.argmax(self.decoder(current_decoder_tokens, encoder_out).to(device=device)[-1, :]).item()
+            decoder_input.append(pred_index)
+            if pred_index == self.vocab["<EOS>"]:
+                break
+        return decoder_input
+    
+    # GET PREDICTED TEXT
+    def predict_text(self, encoder_input_tokens):
+        return ' '.join(
+            [self.vocab_df[self.vocab_df['idx'] == token].index.values[0] \
+             for token in self.predict_tokens(encoder_input_tokens=encoder_input_tokens)]
+            )