Upload 10 files

app.py

@@ -0,0 +1,166 @@
+import streamlit as st
+from transformers import AutoTokenizer, AutoModel, utils
+from bertviz import model_view
+import streamlit.components.v1 as components
+from train import get_or_build_tokenizer, greedy_decode
+from config import get_config, latest_weights_file_path
+from model import build_transformer
+import torch
+from bertviz import model_view
+import torch
+import altair as alt
+import pandas as pd
+import warnings
+warnings.filterwarnings("ignore")
+
+utils.logging.set_verbosity_error()  # Suppress standard warnings
+
+st.set_page_config(page_title='Attention Visualizer', layout='wide')
+
+def mtx2df(m, max_row, max_col, row_tokens, col_tokens):
+    return pd.DataFrame(
+        [
+            (
+                r,
+                c,
+                float(m[r, c]),
+                "%.2d - %s" % (r, row_tokens[r] if len(row_tokens) > r else "<blank>"),
+                "%.2d - %s" % (c, col_tokens[c] if len(col_tokens) > c else "<blank>"),
+            )
+            for r in range(m.shape[0])
+            for c in range(m.shape[1])
+            if r < max_row and c < max_col
+        ],
+        columns=["row", "column", "value", "row_token", "col_token"],
+    )
+    
+    
+def get_attn_map(attn_type: str, layer: int, head: int, model):
+    if attn_type == "encoder":
+        attn = model.encoder.layers[layer].self_attention_block.attention_scores
+    elif attn_type == "decoder":
+        attn = model.decoder.layers[layer].self_attention_block.attention_scores
+    elif attn_type == "encoder-decoder":
+        attn = model.decoder.layers[layer].cross_attention_block.attention_scores
+    return attn[0, head].data
+
+def attn_map(attn_type, layer, head, row_tokens, col_tokens, max_sentence_len, model):
+    df = mtx2df(
+        get_attn_map(attn_type, layer, head, model),
+        max_sentence_len,
+        max_sentence_len,
+        row_tokens,
+        col_tokens,
+    )
+    return (
+        alt.Chart(data=df)
+        .mark_rect()
+        .encode(
+            x=alt.X("col_token", axis=alt.Axis(title="")),
+            y=alt.Y("row_token", axis=alt.Axis(title="")),
+            color=alt.Color("value", scale=alt.Scale(scheme="blues")),
+            tooltip=["row", "column", "value", "row_token", "col_token"],
+        )
+        #.title(f"Layer {layer} Head {head}")
+        .properties(height=200, width=200, title=f"Layer {layer} Head {head}")
+        .interactive()
+    )
+
+def get_all_attention_maps(attn_type: str, layers: list[int], heads: list[int], row_tokens: list, col_tokens, max_sentence_len: int, model):
+    charts = []
+    for layer in layers:
+        rowCharts = []
+        for head in heads:
+            rowCharts.append(attn_map(attn_type, layer, head, row_tokens, col_tokens, max_sentence_len, model))
+        charts.append(alt.hconcat(*rowCharts))
+    return alt.vconcat(*charts)
+
+def initiate_model(config, device):
+    tokenizer_src = get_or_build_tokenizer(config, None, config["lang_src"])
+    tokenizer_tgt = get_or_build_tokenizer(config, None, config["lang_tgt"])
+
+    model = build_transformer(tokenizer_src.get_vocab_size(), tokenizer_tgt.get_vocab_size(), config["seq_len"], config['seq_len'], d_model=config['d_model']).to(device)
+
+    model_filename = latest_weights_file_path(config)
+    state = torch.load(model_filename)
+    model.load_state_dict(state['model_state_dict'])
+    return model, tokenizer_src, tokenizer_tgt
+
+def process_input(input_text, tokenizer_src, tokenizer_tgt, model, config, device):
+    src = tokenizer_src.encode(input_text)
+    src = torch.cat([
+        torch.tensor([tokenizer_src.token_to_id('[SOS]')], dtype=torch.int64), 
+        torch.tensor(src.ids, dtype=torch.int64),
+        torch.tensor([tokenizer_src.token_to_id('[EOS]')], dtype=torch.int64),
+        torch.tensor([tokenizer_src.token_to_id('[PAD]')] * (config['seq_len'] - len(src.ids) - 2), dtype=torch.int64)
+    ], dim=0).to(device)
+    source_mask = (src != tokenizer_src.token_to_id('[PAD]')).unsqueeze(0).unsqueeze(0).int().to(device)
+
+    encoder_input_tokens = [tokenizer_src.id_to_token(i) for i in src.cpu().numpy()]
+    encoder_input_tokens = [i for i in encoder_input_tokens if i != '[PAD]']
+
+    model_out = greedy_decode(model, src, source_mask, tokenizer_src, tokenizer_tgt, config['seq_len'], device)
+
+    decoder_input_tokens = [tokenizer_tgt.id_to_token(i) for i in model_out.cpu().numpy()]
+    
+    output = tokenizer_tgt.decode(model_out.detach().cpu().numpy())
+    
+    return encoder_input_tokens, decoder_input_tokens, output
+        
+
+# def get_html_data(model_name, input_text):
+#     model_name ="microsoft/xtremedistil-l12-h384-uncased"
+#     model = AutoModel.from_pretrained(model_name, output_attentions=True, cache_dir='__pycache__')  # Configure model to return attention values
+#     tokenizer = AutoTokenizer.from_pretrained(model_name)
+#     inputs = tokenizer.encode(input_text, return_tensors='pt')  # Tokenize input text
+#     outputs = model(inputs)  # Run model
+#     attention = outputs[-1]  # Retrieve attention from model outputs
+#     tokens = tokenizer.convert_ids_to_tokens(inputs[0])  # Convert input ids to token strings
+#     model_html = model_view(attention, tokens, html_action="return")  # Display model view
+#     with open("static/model_view.html", 'w') as file:
+#         file.write(model_html.data)
+
+def main():
+    st.title('Transformer Visualizer')
+    # st.info('Enter a sentence to visualize the attention of the model')
+    st.write('This app visualizes the attention of a transformer model on a given sentence.')
+    # add a side bar with model options and a prompt
+    config = get_config()
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model, tokenizer_src, tokenizer_tgt = initiate_model(config, device)
+    with st.sidebar:
+        input_text = st.text_input('Enter a sentence')
+        # put two buttons side by side in the sidebar
+        # translate_button = st.button('Translate', key='translate_button')
+        viz_button = st.button('Visualize Attention', key='viz_button')
+        attn_type = st.selectbox('Select attention type', ['encoder', 'decoder', 'encoder-decoder'])
+        layers = st.multiselect('Select layers', list(range(3)))
+        heads = st.multiselect('Select heads', list(range(7)))
+        # allow the user to select the all the layers and heads at once to visualize
+        if st.checkbox('Select all layers'):
+            layers = list(range(3))
+        if st.checkbox('Select all heads'):
+            heads = list(range(7))
+    
+    if viz_button and input_text != '':
+        encoder_input_tokens, decoder_input_tokens, output = process_input(input_text, tokenizer_src, tokenizer_tgt, model, config, device)
+        max_sentence_len = len(encoder_input_tokens)
+        row_tokens = encoder_input_tokens
+        col_tokens = decoder_input_tokens
+        st.write('Input:', ' '.join(encoder_input_tokens))
+        st.write('Output:', ' '.join(decoder_input_tokens))
+        st.write('Translated:', output)
+        st.write('Attention Visualization')
+        st.write(get_all_attention_maps(attn_type, layers, heads, row_tokens, col_tokens, max_sentence_len, model))
+    else:
+        st.write('Enter a sentence to visualize the attention of the model')
+    
+    # add a footer with the github repo link and dataset link 
+    st.markdown('---')
+    st.write('Made by [Pratik Dwivedi](https://github.com/Dekode1859)')
+    st.write('Check out the Scratch Implementation and Visualizer Code on [GitHub](https://github.com/Dekode1859/transformer-visualizer)')
+    st.write('Dataset: [Opus-books: english-Italian](https://huggingface.co/datasets/Helsinki-NLP/opus_books)')
+    # st.write('This app is a Streamlit implementation of the [BERTViz](
+
+if __name__ == '__main__':
+    main()

config.py

@@ -0,0 +1,33 @@
+from pathlib import Path
+
+def get_config():
+    return {
+        "batch_size": 8,
+        "num_epochs": 20,
+        "lr": 10**-4,
+        "seq_len": 350,
+        "d_model": 512,
+        "datasource": 'opus_books',
+        "lang_src": "en",
+        "lang_tgt": "it",
+        "model_folder": "weights",
+        "model_basename": "tmodel_",
+        "preload": "latest",
+        "tokenizer_file": "tokenizer_{0}.json",
+        "experiment_name": "runs/tmodel"
+    }
+
+def get_weights_file_path(config, epoch: str):
+    model_folder = f"{config['datasource']}_{config['model_folder']}"
+    model_filename = f"{config['model_basename']}{epoch}.pt"
+    return str(Path('.') / model_folder / model_filename)
+
+# Find the latest weights file in the weights folder
+def latest_weights_file_path(config):
+    model_folder = f"{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])

dataset.py

@@ -0,0 +1,90 @@
+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.seq_len = seq_len
+
+        self.ds = ds
+        self.tokenizer_src = tokenizer_src
+        self.tokenizer_tgt = tokenizer_tgt
+        self.src_lang = src_lang
+        self.tgt_lang = tgt_lang
+
+        self.sos_token = torch.tensor([tokenizer_tgt.token_to_id("[SOS]")], dtype=torch.int64)
+        self.eos_token = torch.tensor([tokenizer_tgt.token_to_id("[EOS]")], dtype=torch.int64)
+        self.pad_token = torch.tensor([tokenizer_tgt.token_to_id("[PAD]")], dtype=torch.int64)
+
+    def __len__(self):
+        return len(self.ds)
+
+    def __getitem__(self, idx):
+        src_target_pair = self.ds[idx]
+        src_text = src_target_pair['translation'][self.src_lang]
+        tgt_text = src_target_pair['translation'][self.tgt_lang]
+
+        # Transform the text into tokens
+        enc_input_tokens = self.tokenizer_src.encode(src_text).ids
+        dec_input_tokens = self.tokenizer_tgt.encode(tgt_text).ids
+
+        # Add sos, eos and padding to each sentence
+        enc_num_padding_tokens = self.seq_len - len(enc_input_tokens) - 2  # We will add <s> and </s>
+        # We will only add <s>, and </s> only on the label
+        dec_num_padding_tokens = self.seq_len - len(dec_input_tokens) - 1
+
+        # Make sure the number of padding tokens is not negative. If it is, the sentence is too long
+        if enc_num_padding_tokens < 0 or dec_num_padding_tokens < 0:
+            raise ValueError("Sentence is too long")
+
+        # Add <s> and </s> token
+        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),
+            ],
+            dim=0,
+        )
+
+        # Add only <s> token
+        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),
+            ],
+            dim=0,
+        )
+
+        # Add only </s> token
+        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),
+            ],
+            dim=0,
+        )
+
+        # Double check the size of the tensors to make sure they are all seq_len long
+        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).int() & causal_mask(decoder_input.size(0)), # (1, seq_len) & (1, seq_len, seq_len),
+            "label": label,  # (seq_len)
+            "src_text": src_text,
+            "tgt_text": tgt_text,
+        }
+    
+def causal_mask(size):
+    mask = torch.triu(torch.ones((1, size, size)), diagonal=1).type(torch.int)
+    return mask == 0

model.py

@@ -0,0 +1,267 @@
+import torch
+import torch.nn as nn
+import math
+
+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 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 InputEmbeddings(nn.Module):
+
+    def __init__(self, d_model: int, vocab_size: int) -> None:
+        super().__init__()
+        self.d_model = d_model
+        self.vocab_size = vocab_size
+        self.embedding = nn.Embedding(vocab_size, d_model)
+
+    def forward(self, x):
+        # (batch, seq_len) --> (batch, seq_len, d_model)
+        # Multiply by sqrt(d_model) to scale the embeddings according to the paper
+        return self.embedding(x) * math.sqrt(self.d_model)
+    
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_model: int, seq_len: int, dropout: float) -> None:
+        super().__init__()
+        self.d_model = d_model
+        self.seq_len = seq_len
+        self.dropout = nn.Dropout(dropout)
+        # Create a matrix of shape (seq_len, d_model)
+        pe = torch.zeros(seq_len, d_model)
+        # Create a vector of shape (seq_len)
+        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1) # (seq_len, 1)
+        # Create a vector of shape (d_model)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) # (d_model / 2)
+        # Apply sine to even indices
+        pe[:, 0::2] = torch.sin(position * div_term) # sin(position * (10000 ** (2i / d_model))
+        # Apply cosine to odd indices
+        pe[:, 1::2] = torch.cos(position * div_term) # cos(position * (10000 ** (2i / d_model))
+        # Add a batch dimension to the positional encoding
+        pe = pe.unsqueeze(0) # (1, seq_len, d_model)
+        # Register the positional encoding as a buffer
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False) # (batch, seq_len, d_model)
+        return self.dropout(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 MultiHeadAttentionBlock(nn.Module):
+
+    def __init__(self, d_model: int, h: int, dropout: float) -> None:
+        super().__init__()
+        self.d_model = d_model # Embedding vector size
+        self.h = h # Number of heads
+        # Make sure d_model is divisible by h
+        assert d_model % h == 0, "d_model is not divisible by h"
+
+        self.d_k = d_model // h # Dimension of vector seen by each head
+        self.w_q = nn.Linear(d_model, d_model, bias=False) # Wq
+        self.w_k = nn.Linear(d_model, d_model, bias=False) # Wk
+        self.w_v = nn.Linear(d_model, d_model, bias=False) # Wv
+        self.w_o = nn.Linear(d_model, 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]
+        # Just apply the formula from the paper
+        # (batch, h, seq_len, d_k) --> (batch, h, seq_len, seq_len)
+        attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)
+        if mask is not None:
+            # Write a very low value (indicating -inf) to the positions where mask == 0
+            attention_scores.masked_fill_(mask == 0, -1e9)
+        attention_scores = attention_scores.softmax(dim=-1) # (batch, h, seq_len, seq_len) # Apply softmax
+        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 EncoderBlock(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, src_mask):
+        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, src_mask))
+        x = self.residual_connections[1](x, self.feed_forward_block)
+        return x
+    
+class Encoder(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 DecoderBlock(nn.Module):
+
+    def __init__(self, features: int, self_attention_block: MultiHeadAttentionBlock, cross_attention_block: MultiHeadAttentionBlock, feed_forward_block: FeedForwardBlock, dropout: float) -> None:
+        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(features, 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, features: int, layers: nn.ModuleList) -> None:
+        super().__init__()
+        self.layers = layers
+        self.norm = LayerNormalization(features)
+
+    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, vocab_size) -> None:
+        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 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) -> None:
+        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_layer = projection_layer
+
+    def encode(self, src, src_mask):
+        # (batch, seq_len, d_model)
+        src = self.src_embed(src)
+        src = self.src_pos(src)
+        return self.encoder(src, src_mask)
+    
+    def decode(self, encoder_output: torch.Tensor, src_mask: torch.Tensor, tgt: torch.Tensor, tgt_mask: torch.Tensor):
+        # (batch, seq_len, d_model)
+        tgt = self.tgt_embed(tgt)
+        tgt = self.tgt_pos(tgt)
+        return self.decoder(tgt, encoder_output, src_mask, tgt_mask)
+    
+    def project(self, x):
+        # (batch, seq_len, vocab_size)
+        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: 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 encoding layers
+    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 = MultiHeadAttentionBlock(d_model, h, dropout)
+        feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
+        encoder_block = EncoderBlock(d_model, 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 = MultiHeadAttentionBlock(d_model, h, dropout)
+        decoder_cross_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)
+        feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
+        decoder_block = DecoderBlock(d_model, decoder_self_attention_block, decoder_cross_attention_block, feed_forward_block, dropout)
+        decoder_blocks.append(decoder_block)
+    
+    # Create the encoder and decoder
+    encoder = Encoder(d_model, nn.ModuleList(encoder_blocks))
+    decoder = Decoder(d_model, 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)
+    
+    # Initialize the parameters
+    for p in transformer.parameters():
+        if p.dim() > 1:
+            nn.init.xavier_uniform_(p)
+    
+    return transformer

predict.py

@@ -0,0 +1,16 @@
+from transformers import AutoTokenizer, AutoModel, utils
+from bertviz import model_view
+utils.logging.set_verbosity_error()  # Suppress standard warnings
+
+def get_predictions(input_text):
+    model_name = "microsoft/xtremedistil-l12-h384-uncased"
+    model = AutoModel.from_pretrained(model_name, output_attentions=True)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    inputs = tokenizer.encode(input_text, return_tensors='pt')
+    outputs = model(inputs)
+    attention = outputs[-1]
+    tokens = tokenizer.convert_ids_to_tokens(inputs[0])
+    model_html = model_view(attention, tokens, html_action="return")
+    with open("static/model_view.html", 'w') as file:
+        file.write(model_html.data)
+    

requirements.txt

@@ -0,0 +1,16 @@
+## Use python 3.9
+
+torch
+torchvision
+torchaudio
+torchtext
+datasets
+tokenizers
+torchmetrics
+tensorboard
+altair
+wandb
+transformers
+bertviz
+IPython
+streamlit

train.py

@@ -0,0 +1,283 @@
+from model import build_transformer
+from dataset import BilingualDataset, causal_mask
+from config import get_config, get_weights_file_path, latest_weights_file_path
+
+# import torchtext.datasets as datasets
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader, random_split
+from torch.optim.lr_scheduler import LambdaLR
+
+import warnings
+from tqdm import tqdm
+import os
+from pathlib import Path
+
+# Huggingface datasets and tokenizers
+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
+
+# import torchmetrics
+# from torch.utils.tensorboard import SummaryWriter
+
+def greedy_decode(model, source, source_mask, tokenizer_src, tokenizer_tgt, max_len, device):
+    sos_idx = tokenizer_tgt.token_to_id('[SOS]')
+    eos_idx = tokenizer_tgt.token_to_id('[EOS]')
+
+    # Precompute the encoder output and reuse it for every step
+    encoder_output = model.encode(source, source_mask)
+    # Initialize the decoder input with the sos token
+    decoder_input = torch.empty(1, 1).fill_(sos_idx).type_as(source).to(device)
+    while True:
+        if decoder_input.size(1) == max_len:
+            break
+
+        # build mask for target
+        decoder_mask = causal_mask(decoder_input.size(1)).type_as(source_mask).to(device)
+
+        # calculate output
+        out = model.decode(encoder_output, source_mask, decoder_input, decoder_mask)
+
+        # get next token
+        prob = model.project(out[:, -1])
+        _, next_word = torch.max(prob, dim=1)
+        decoder_input = torch.cat(
+            [decoder_input, torch.empty(1, 1).type_as(source).fill_(next_word.item()).to(device)], dim=1
+        )
+
+        if next_word == eos_idx:
+            break
+
+    return decoder_input.squeeze(0)
+
+
+def run_validation(model, validation_ds, tokenizer_src, tokenizer_tgt, max_len, device, print_msg, global_step, writer=None, num_examples=2):
+    model.eval()
+    count = 0
+
+    source_texts = []
+    expected = []
+    predicted = []
+
+    try:
+        # get the console window width
+        with os.popen('stty size', 'r') as console:
+            _, console_width = console.read().split()
+            console_width = int(console_width)
+    except:
+        # If we can't get the console width, use 80 as default
+        console_width = 80
+
+    with torch.no_grad():
+        for batch in validation_ds:
+            count += 1
+            encoder_input = batch["encoder_input"].to(device) # (b, seq_len)
+            encoder_mask = batch["encoder_mask"].to(device) # (b, 1, 1, seq_len)
+
+            # check that the batch size is 1
+            assert encoder_input.size(
+                0) == 1, "Batch size must be 1 for validation"
+
+            model_out = greedy_decode(model, encoder_input, encoder_mask, tokenizer_src, tokenizer_tgt, max_len, device)
+
+            source_text = batch["src_text"][0]
+            target_text = batch["tgt_text"][0]
+            model_out_text = tokenizer_tgt.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('-'*console_width)
+            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}")
+
+            if count == num_examples:
+                print_msg('-'*console_width)
+                break
+    
+    # if writer:
+    #     # Evaluate the character error rate
+    #     # Compute the char error rate 
+    #     metric = torchmetrics.CharErrorRate()
+    #     cer = metric(predicted, expected)
+    #     writer.add_scalar('validation cer', cer, global_step)
+    #     writer.flush()
+
+    #     # Compute the word error rate
+    #     metric = torchmetrics.WordErrorRate()
+    #     wer = metric(predicted, expected)
+    #     writer.add_scalar('validation wer', wer, global_step)
+    #     writer.flush()
+
+    #     # Compute the BLEU metric
+    #     metric = torchmetrics.BLEUScore()
+    #     bleu = metric(predicted, expected)
+    #     writer.add_scalar('validation BLEU', bleu, global_step)
+    #     writer.flush()
+
+def get_all_sentences(ds, lang):
+    for item in ds:
+        yield item['translation'][lang]
+
+def get_or_build_tokenizer(config, ds, lang):
+    print(f"Checking for existing tokenizer for {lang}")
+    tokenizer_path = Path(config['tokenizer_file'].format(lang))
+    if not Path.exists(tokenizer_path):
+        print(f"Building tokenizer for {lang}")
+        # Most code taken from: https://huggingface.co/docs/tokenizers/quicktour
+        tokenizer = Tokenizer(WordLevel(unk_token="[UNK]"))
+        tokenizer.pre_tokenizer = Whitespace()
+        trainer = WordLevelTrainer(special_tokens=["[UNK]", "[PAD]", "[SOS]", "[EOS]"], min_frequency=2)
+        print(f"Training tokenizer for {lang}")
+        tokenizer.train_from_iterator(get_all_sentences(ds, lang), trainer=trainer)
+        print(f"Saving tokenizer for {lang}")
+        tokenizer.save(str(tokenizer_path))
+    else:
+        print(f"Found existing tokenizer for {lang}")
+        tokenizer = Tokenizer.from_file(str(tokenizer_path))
+    return tokenizer
+
+def get_ds(config):
+    # It only has the train split, so we divide it overselves
+    print(f"Loading dataset {config['datasource']}")
+    ds_raw = load_dataset(f"{config['datasource']}", f"{config['lang_src']}-{config['lang_tgt']}", split='train')
+
+    # Build tokenizers
+    print(f"Building tokenizers for {config['lang_src']} and {config['lang_tgt']}")
+    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
+    print("Splitting dataset into training and 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])
+
+    train_ds = BilingualDataset(train_ds_raw, tokenizer_src, tokenizer_tgt, config['lang_src'], config['lang_tgt'], config['seq_len'])
+    val_ds = BilingualDataset(val_ds_raw, tokenizer_src, tokenizer_tgt, config['lang_src'], config['lang_tgt'], config['seq_len'])
+
+    # Find the maximum length of each sentence in the source and target sentence
+    print("Finding the maximum length of the source and target sentences")
+    max_len_src = 0
+    max_len_tgt = 0
+
+    for item in ds_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}')
+    
+
+    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):
+    model = build_transformer(vocab_src_len, vocab_tgt_len, config["seq_len"], config['seq_len'], d_model=config['d_model'])
+    return model
+
+def train_model(config):
+    # Define the device
+    device = "cuda" if torch.cuda.is_available() else "mps" if torch.has_mps 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.index)}")
+        print(f"Device memory: {torch.cuda.get_device_properties(device.index).total_memory / 1024 ** 3} GB")
+    elif (device == 'mps'):
+        print(f"Device name: <mps>")
+    else:
+        print("NOTE: If you have a GPU, consider using it for training.")
+        print("      On a Windows machine with NVidia GPU, check this video: https://www.youtube.com/watch?v=GMSjDTU8Zlc")
+        print("      On a Mac machine, run: pip3 install --pre torch torchvision torchaudio torchtext --index-url https://download.pytorch.org/whl/nightly/cpu")
+    device = torch.device(device)
+
+    # Make sure the weights folder exists
+    Path(f"{config['datasource']}_{config['model_folder']}").mkdir(parents=True, exist_ok=True)
+
+    train_dataloader, val_dataloader, tokenizer_src, tokenizer_tgt = get_ds(config)
+    model = get_model(config, tokenizer_src.get_vocab_size(), tokenizer_tgt.get_vocab_size()).to(device)
+    # Tensorboard
+    # writer = SummaryWriter(config['experiment_name'])
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'], eps=1e-9)
+
+    # 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_src.token_to_id('[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}")
+        for batch in batch_iterator:
+
+            encoder_input = batch['encoder_input'].to(device) # (b, seq_len)
+            decoder_input = batch['decoder_input'].to(device) # (B, seq_len)
+            encoder_mask = batch['encoder_mask'].to(device) # (B, 1, 1, seq_len)
+            decoder_mask = batch['decoder_mask'].to(device) # (B, 1, seq_len, seq_len)
+
+            # Run the tensors through the encoder, decoder and the projection layer
+            encoder_output = model.encode(encoder_input, encoder_mask) # (B, seq_len, d_model)
+            decoder_output = model.decode(encoder_output, encoder_mask, decoder_input, decoder_mask) # (B, seq_len, d_model)
+            proj_output = model.project(decoder_output) # (B, seq_len, vocab_size)
+
+            # Compare the output with the label
+            label = batch['label'].to(device) # (B, seq_len)
+
+            # Compute the loss using a simple cross entropy
+            loss = loss_fn(proj_output.view(-1, tokenizer_tgt.get_vocab_size()), label.view(-1))
+            batch_iterator.set_postfix({"loss": f"{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
+        run_validation(model, val_dataloader, tokenizer_src, tokenizer_tgt, config['seq_len'], device, lambda msg: batch_iterator.write(msg), global_step, writer=None)
+
+        # 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)

translate.py

@@ -0,0 +1,79 @@
+from pathlib import Path
+from config import get_config, latest_weights_file_path 
+from model import build_transformer
+from tokenizers import Tokenizer
+from datasets import load_dataset
+from dataset import BilingualDataset
+import torch
+import sys
+
+def translate(sentence: str):
+    # Define the device, tokenizers, and model
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+    config = get_config()
+    tokenizer_src = Tokenizer.from_file(str(Path(config['tokenizer_file'].format(config['lang_src']))))
+    tokenizer_tgt = Tokenizer.from_file(str(Path(config['tokenizer_file'].format(config['lang_tgt']))))
+    model = build_transformer(tokenizer_src.get_vocab_size(), tokenizer_tgt.get_vocab_size(), config["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'])
+
+    # if the sentence is a number use it as an index to the test set
+    label = ""
+    if type(sentence) == int or sentence.isdigit():
+        id = int(sentence)
+        ds = load_dataset(f"{config['datasource']}", f"{config['lang_src']}-{config['lang_tgt']}", split='all')
+        ds = BilingualDataset(ds, tokenizer_src, tokenizer_tgt, config['lang_src'], config['lang_tgt'], config['seq_len'])
+        sentence = ds[id]['src_text']
+        label = ds[id]["tgt_text"]
+    seq_len = config['seq_len']
+
+    # translate the sentence
+    model.eval()
+    with torch.no_grad():
+        # Precompute the encoder output and reuse it for every generation step
+        source = tokenizer_src.encode(sentence)
+        source = torch.cat([
+            torch.tensor([tokenizer_src.token_to_id('[SOS]')], dtype=torch.int64), 
+            torch.tensor(source.ids, dtype=torch.int64),
+            torch.tensor([tokenizer_src.token_to_id('[EOS]')], dtype=torch.int64),
+            torch.tensor([tokenizer_src.token_to_id('[PAD]')] * (seq_len - len(source.ids) - 2), dtype=torch.int64)
+        ], dim=0).to(device)
+        source_mask = (source != tokenizer_src.token_to_id('[PAD]')).unsqueeze(0).unsqueeze(0).int().to(device)
+        encoder_output = model.encode(source, source_mask)
+
+        # Initialize the decoder input with the sos token
+        decoder_input = torch.empty(1, 1).fill_(tokenizer_tgt.token_to_id('[SOS]')).type_as(source).to(device)
+
+        # Print the source sentence and target start prompt
+        if label != "": print(f"{f'ID: ':>12}{id}") 
+        print(f"{f'SOURCE: ':>12}{sentence}")
+        if label != "": print(f"{f'TARGET: ':>12}{label}") 
+        print(f"{f'PREDICTED: ':>12}", end='')
+
+        # Generate the translation word by word
+        while decoder_input.size(1) < seq_len:
+            # build mask for target and calculate output
+            decoder_mask = torch.triu(torch.ones((1, decoder_input.size(1), decoder_input.size(1))), diagonal=1).type(torch.int).type_as(source_mask).to(device)
+            out = model.decode(encoder_output, source_mask, decoder_input, decoder_mask)
+
+            # project next token
+            prob = model.project(out[:, -1])
+            _, next_word = torch.max(prob, dim=1)
+            decoder_input = torch.cat([decoder_input, torch.empty(1, 1).type_as(source).fill_(next_word.item()).to(device)], dim=1)
+
+            # print the translated word
+            print(f"{tokenizer_tgt.decode([next_word.item()])}", end=' ')
+
+            # break if we predict the end of sentence token
+            if next_word == tokenizer_tgt.token_to_id('[EOS]'):
+                break
+
+    # convert ids to tokens
+    return tokenizer_tgt.decode(decoder_input[0].tolist())
+    
+#read sentence from argument
+translate(sys.argv[1] if len(sys.argv) > 1 else "I am not a very good a student.")