Upload 5 files

config.py

@@ -0,0 +1,26 @@
+from pathlib import Path
+
+def get_config():
+    return {
+        "batch_size":1,
+        "num_epochs": 20,
+        "lr": 1e-5,
+        "seq_len": 261,
+        "d_model": 768,
+        "lang_src": "en",
+        "lang_tgt": "it",
+        "model_folder": "weights",
+        "model_basename": "tmodel_",
+        "preload": None,
+        "tokenizer_file": "tokenizer.json",
+        "experiment_name": "runs/tmodel",
+        'project_name': 'proj1'
+    }
+
+def get_weights_file_path(config, epoch: str):
+    model_folder = config["model_folder"]
+    model_basename = config["model_basename"]
+    model_filename = f"{model_basename}"
+    return str(Path('.') / model_folder / model_filename)
+
+

dataset.py

@@ -0,0 +1,131 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import IterableDataset, Dataset
+from transformers import ViTFeatureExtractor
+from transformers import ViTImageProcessor
+from io import BytesIO
+from base64 import b64decode
+from PIL import Image,ImageFile
+import base64
+import itertools
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+import io
+import urllib
+import random
+import PIL.Image
+from datasets import load_dataset
+from datasets.utils.file_utils import get_datasets_user_agent
+
+
+USER_AGENT = get_datasets_user_agent()
+
+# import model
+model_id = 'google/vit-base-patch16-224-in21k'
+feature_extractor = ViTFeatureExtractor.from_pretrained(
+    model_id
+)
+class BilingualDataset(Dataset):
+
+    def __init__(self, ds,tokenizer_tgt, seq_len):
+        super().__init__()
+        self.seq_len = seq_len
+        ImageFile.LOAD_TRUNCATED_IMAGES = True
+        self.ds = ds
+        self.tokenizer_tgt = tokenizer_tgt
+        # self.tgt_lang = tgt_lang
+        self.processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k')
+        self.sos_token = torch.tensor([tokenizer_tgt.convert_tokens_to_ids("[SOS]")], dtype=torch.int64)
+        self.eos_token = torch.tensor([tokenizer_tgt.convert_tokens_to_ids("[EOS]")], dtype=torch.int64)
+        self.pad_token = torch.tensor([tokenizer_tgt.convert_tokens_to_ids("[PAD]")], dtype=torch.int64)
+
+    def __len__(self):
+        return len(self.ds)
+    # def __getitem__(self):
+    #     pass
+  
+    def __getitem__(self, idx):
+        data_pair = self.ds[idx]
+
+           
+        src_image = data_pair['image_base64_str']
+        tgt_text = data_pair['outputs']
+       
+       
+
+
+            
+
+        src_image  = Image.open(BytesIO(b64decode(''.join(src_image))))
+        if src_image.mode != 'RGB':
+            src_image = src_image.convert('RGB')
+        src_image = self.processor(src_image, return_tensors='pt')
+
+   
+   
+
+                # Transform the text into tokens
+        dec_input_tokens = self.tokenizer_tgt.encode(tgt_text)
+
+                # # 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_input_tokens = dec_input_tokens[:self.seq_len-1]
+        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 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,
+                )
+
+
+        assert decoder_input.size(0) == self.seq_len
+        assert label.size(0) == self.seq_len
+           
+        return {     
+        'encoder_input' : src_image['pixel_values'].squeeze(0).squeeze(0).squeeze(0).squeeze(0).squeeze(0),  # (seq_len)
+        'decoder_input' : decoder_input,  # (seq_len)
+        ## encoder mask not used :)
+        "encoder_mask" : (torch.cat((torch.ones(197,),torch.zeros(63),),)).unsqueeze(0).unsqueeze(0), # (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,  
+                    # "src_text": src_text,
+        "tgt_text" : tgt_text
+
+        }
+        # yield encoder_input, decoder_input, encoder_mask, decoder_mask, label
+
+
+def causal_mask(size):
+    mask = torch.triu(torch.ones((1, size, size)), diagonal=1).type(torch.int)
+    return mask == 0

inference.py

@@ -0,0 +1,155 @@
+import torch
+import torch.nn as nn
+import transformers
+from torch.utils.data import Dataset
+from transformers import ViTFeatureExtractor
+from io import BytesIO
+from base64 import b64decode
+from PIL import Image
+from accelerate import Accelerator
+import base64
+from config import get_config
+from pathlib import Path
+from tokenizers import Tokenizer
+from tokenizers.models import WordLevel
+from tokenizers.trainers import WordLevelTrainer
+from tokenizers.pre_tokenizers import Whitespace
+from model import build_transformer
+import torch.nn.functional as F
+from transformers import GPT2TokenizerFast
+
+def process(model,image, tokenizer, device):
+    image = get_image(image)
+    model.eval()
+    with torch.no_grad():
+        encoder_input = image.unsqueeze(0).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)
+
+        model_out = greedy_decode(model, encoder_input, None, tokenizer, 196,device)
+        model_text  = tokenizer.decode(model_out.detach().cpu().numpy())
+        print(model_text)
+
+
+
+
+
+       
+
+
+
+# get image prompt 
+def get_image(image):
+# import model
+    model_id = 'google/vit-base-patch16-224-in21k'
+    feature_extractor = ViTFeatureExtractor.from_pretrained(
+    model_id
+    )
+
+    
+    image  = Image.open(BytesIO(b64decode(''.join(image))))
+
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+        
+    enc_input = feature_extractor(
+    image,
+    return_tensors='pt'
+    )
+
+    return enc_input['pixel_values'].squeeze(0).squeeze(0).squeeze(0).squeeze(0).squeeze(0)
+
+    
+
+
+#get tokenizer
+def get_or_build_tokenizer(config):
+    tokenizer = GPT2TokenizerFast.from_pretrained("openai-community/gpt2", unk_token ='[UNK]', bos_token = '[SOS]', eos_token = '[EOS]' , 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
+
+
+# get model
+def get_model(config, vocab_tgt_len):
+    model = build_transformer(vocab_tgt_len, config['seq_len'], d_model=config['d_model'])
+    return model
+
+# greedy decode 
+def greedy_decode(model, source, source_mask, tokenizer_tgt, max_len, device):
+    sos_idx = tokenizer_tgt.convert_tokens_to_ids('[SOS]')
+    eos_idx = tokenizer_tgt.convert_tokens_to_ids('[EOS]')
+
+    # Precompute the encoder output and reuse it for every step
+    encoder_output = model.encode(source, None)
+   
+    # Initialize the decoder input with the sos token
+    decoder_input = torch.empty(1, 1).fill_(sos_idx).long().to(device)
+    while True:
+        if decoder_input.size(1) == max_len:
+            break
+
+        # build mask for target
+        decoder_mask = causal_mask(decoder_input.size(1)).long().to(device)
+      
+
+        # calculate output
+        out = model.decode(encoder_output, source_mask, decoder_input, decoder_mask)
+        # print(f'out: {out.shape}')
+
+        # Get next token probabilities with temperature applied
+        logits = model.project(out[:, -1]) 
+        probabilities = F.softmax(logits, dim=-1)
+
+        # Greedily select the next word
+        next_word = torch.argmax(probabilities, dim=1)
+        
+        # Append next word
+        decoder_input = torch.cat([decoder_input, next_word.unsqueeze(0)], dim=1)
+        # # get next token
+        # prob = model.project(out[:, -1])
+        # _, next_word = torch.max(prob, dim=1)
+        # # print(f'prob: {prob.shape}')
+        # decoder_input = torch.cat(
+        #     [decoder_input, torch.empty(1, 1).long().fill_(next_word.item()).to(device)], dim=1
+        # )
+
+        if next_word.item() == eos_idx:
+            break
+
+    return decoder_input.squeeze(0)
+
+def image_base64():
+    
+
+    with open('C:/AI/projects/vision_model_pretrained/validation/content/memory_image_23330.jpg', 'rb') as image_file:
+        base64_bytes = base64.b64encode(image_file.read())
+       
+
+        base64_string = base64_bytes.decode()
+        return base64_string
+
+
+def start():
+    print('start')
+    accelerator = Accelerator()
+    device = accelerator.device
+
+    config = get_config()
+    tokenizer = get_or_build_tokenizer(config)
+    model = get_model(config, len(tokenizer))
+    model = accelerator.prepare(model)
+    accelerator.load_state('C:/AI/projects/vision_model_pretrained/Vision_Model_pretrained/models/vision_model_04')
+
+    image = image_base64()
+
+    
+
+    process(model, image, tokenizer, device)
+
+start()
+

model.py

@@ -0,0 +1,411 @@
+import torch
+import torch.nn as nn
+import math
+import matplotlib.pyplot as plt
+import numpy as np
+from PIL import Image
+from torchvision.transforms.functional import to_pil_image, to_tensor
+import time
+import numpy as np
+from matplotlib.image import imread
+from transformers import ViTFeatureExtractor
+from io import BytesIO
+from base64 import b64decode
+import base64
+from transformers import ViTImageProcessor, ViTModel
+## code from @jankrepl on github
+
+class PretrainedVit():
+    def __init__(self):
+       
+        self.model = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k')
+    def forward(self, x):
+        
+        self.model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        self.model.config.output_hidden_states = True
+        outputs = self.model(x)
+        # print(outputs)
+        last_hidden_states = outputs.hidden_states
+        return list(last_hidden_states)
+
+class PatchEmbed(nn.Module):
+    """Split image into patches and then embed them.
+
+    Parameters
+    ----------
+    img_size : int
+        Size of the image (it is a square).
+
+    patch_size : int
+        Size of the patch (it is a square).
+
+    in_chans : int
+        Number of input channels.
+
+    embed_dim : int
+        The emmbedding dimension.
+
+    Attributes
+    ----------
+    n_patches : int
+        Number of patches inside of our image.
+
+    proj : nn.Conv2d
+        Convolutional layer that does both the splitting into patches
+        and their embedding.
+    """
+    def __init__(self, img_size, patch_size, in_chans=3, embed_dim=1024, num_registers = 6):
+        super().__init__()
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.norm = RMSNorm()
+        self.n_patches = (img_size // patch_size) ** 2
+        self.pos_embed = nn.Parameter(
+                torch.zeros(1, self.n_patches+1+num_registers, embed_dim)
+        )
+         # Adding CLS token as a learnable parameter
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.register_token = nn.Parameter(torch.zeros(num_registers, embed_dim))
+
+        self.proj = nn.Conv2d(
+                in_chans,
+                embed_dim,
+                kernel_size=patch_size,
+                stride=patch_size,
+        )
+
+    def forward(self, x):
+        """Run forward pass.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Shape `(n_samples, in_chans, img_size, img_size)`.
+
+        Returns
+        -------
+        torch.Tensor
+            Shape `(n_samples, n_patches, embed_dim)`.
+        """
+        x = self.proj(x)  # (n_samples, embed_dim, n_patches ** 0.5, n_patches ** 0.5)
+        x = x.flatten(2)  # (n_samples, embed_dim, n_patches)
+        x = x.transpose(1, 2) # (n_samples, n_patches, embed_dim)
+        batch_size = x.shape[0]
+
+
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # Expand CLS tokens for the batch
+        x = torch.cat([cls_tokens, x], dim=1)
+
+        # x: (n_samples, n_patches + 1 + num_registers, embed_dimension) add register tokens
+        register_tokens = self.register_token.unsqueeze(0).expand(batch_size, -1, -1) 
+        x = torch.cat([x, register_tokens], dim=1)
+        X = self.norm(x)
+        x = x + self.pos_embed  # Learnable pos embed -> (n_samples, n_patches_embed_dim) 
+    
+        return x
+
+
+## not used 
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int = 1024, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.dim = dim
+        # The gamma parameter
+        self.weight = nn.Parameter(torch.ones(self.dim))
+
+    def _norm(self, x: torch.Tensor):
+        # (B, Seq_Len, Dim) * (B, Seq_Len, 1) = (B, Seq_Len, Dim)
+        # rsqrt: 1 / sqrt(x)
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x: torch.Tensor):
+        # (Dim) * (B, Seq_Len, Dim) = (B, Seq_Len, Dim)
+        return self.weight * self._norm(x.float()).type_as(x)
+
+class LayerNormalization(nn.Module):
+
+    def __init__(self, eps:float=1e-12) -> None:
+        super().__init__()
+        self.eps = eps
+        self.alpha = nn.Parameter(torch.ones(1)) # alpha is a learnable parameter
+        self.bias = nn.Parameter(torch.zeros(1)) # 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
+        # print(f'mean shape {mean.squeeze(-1).shape}')
+        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, dropout: float) -> None:
+            super().__init__()
+            self.dropout = nn.Dropout(dropout)
+            self.norm = LayerNormalization()
+    
+        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) # Wq
+        self.w_k = nn.Linear(d_model, d_model) # Wk
+        self.w_v = nn.Linear(d_model, d_model) # Wv
+        self.w_o = nn.Linear(d_model, d_model) # 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
+
+        # attention_viz(attention_scores)
+        return (attention_scores @ value), attention_scores
+
+    def forward(self, q, k, v, mask, is_cross=False):
+        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)
+        
+        if is_cross:
+            attention_viz(self.attention_scores)
+        # 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, self_attention_block: MultiHeadAttentionBlock, feed_forward_block: FeedForwardBlock, dropout: float,  layer: int ) -> None:
+        super().__init__()
+        self.self_attention_block = self_attention_block
+        self.feed_forward_block = feed_forward_block
+        self.residual_connections = nn.ModuleList([ResidualConnection(dropout) for _ in range(2)])
+        self.layer = layer
+
+    def forward(self, x, src_mask, index):
+        # print(x.shape)
+        # print(self.layer)
+        
+        out = x[11]
+        # out = self.residual_connections[1](out, self.feed_forward_block)
+        return out
+    
+class Encoder(nn.Module):
+
+    def __init__(self, layers: nn.ModuleList) -> None:
+        super().__init__()
+        self.layers = layers
+        self.norm = LayerNormalization()
+
+    def forward(self, x, mask):
+        for index, layer in enumerate(self.layers):
+        #     print(index)
+            x = layer(x, mask, index)
+            break
+        return self.norm(x)
+
+class DecoderBlock(nn.Module):
+
+    def __init__(self, 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(dropout) for _ in range(3)])
+
+    def forward(self, x, encoder_output, src_mask, tgt_mask):
+        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, tgt_mask))
+        x = self.residual_connections[1](x, lambda x: self.cross_attention_block(x, encoder_output, encoder_output, src_mask))
+        x = self.residual_connections[2](x, self.feed_forward_block)
+        
+        return x
+    
+class Decoder(nn.Module):
+
+    def __init__(self, layers: nn.ModuleList) -> None:
+        super().__init__()
+        self.layers = layers
+        self.norm = LayerNormalization()
+
+    def forward(self, x, encoder_output, src_mask, tgt_mask):
+        for layer in self.layers:
+            x = layer(x, encoder_output, src_mask, tgt_mask)
+        return self.norm(x)
+
+class ProjectionLayer(nn.Module):
+
+    def __init__(self, d_model, 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 torch.log_softmax(self.proj(x), dim = -1)
+    
+class Transformer(nn.Module):
+
+    def __init__(self, encoder: Encoder, decoder: Decoder, tgt_embed: InputEmbeddings, tgt_pos: PositionalEncoding, projection_layer: ProjectionLayer, att: PretrainedVit) -> 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
+        self.patch_embed = PatchEmbed(img_size=224, patch_size=14)
+        self.att = att
+
+    def encode(self, src, src_mask):
+        # (batch, seq_len, d_model)
+        attention_list = self.att.forward(src)
+        # src = self.src_pos(src)
+        return self.encoder(attention_list[1:], 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(tgt_vocab_size: int, tgt_seq_len: int, d_model: int=768, N: int=10, h: int=12, dropout: float=0.1, d_ff: int=3072) -> Transformer:
+    # Create the embedding layers
+  
+    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)
+
+    #attention from pretrained vit
+    att = PretrainedVit()
+    
+    
+    # Create the encoder blocks
+    encoder_blocks = []
+    for _ in range(N):
+        print()
+        encoder_self_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)
+        feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
+        encoder_block = EncoderBlock(encoder_self_attention_block, feed_forward_block, dropout, _)
+        encoder_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(decoder_self_attention_block, decoder_cross_attention_block, feed_forward_block, dropout)
+        decoder_blocks.append(decoder_block)
+    
+    # Create the encoder and decoder
+    encoder = Encoder(nn.ModuleList(encoder_blocks))
+    decoder = Decoder(nn.ModuleList(decoder_blocks))
+    
+    # Create the projection layer
+    projection_layer = ProjectionLayer(d_model, tgt_vocab_size)
+    
+    # Create the transformer
+    transformer = Transformer(encoder, decoder,  tgt_embed, tgt_pos, projection_layer, att)
+    
+    # Initialize the parameters
+    for p in transformer.parameters():
+        if p.dim() > 1:
+            nn.init.xavier_uniform_(p)
+    
+    return transformer

train.py

@@ -0,0 +1,374 @@
+from model import build_transformer
+from dataset import BilingualDataset, causal_mask
+from config import get_config, get_weights_file_path
+
+
+import datasets
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import IterableDataset, 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
+import wandb
+import accelerate
+from torch.utils.tensorboard import SummaryWriter
+from safetensors.torch import load_model, save_model
+from accelerate import Accelerator
+from transformers import GPT2TokenizerFast
+import threading
+
+
+def greedy_decode(model, source, source_mask, tokenizer_tgt, max_len, device):
+    sos_idx = tokenizer_tgt.convert_tokens_to_ids('[SOS]')
+    eos_idx = tokenizer_tgt.convert_tokens_to_ids('[EOS]')
+
+    # Precompute the encoder output and reuse it for every step
+    encoder_output = model.module.encode(source, None)
+   
+    # Initialize the decoder input with the sos token
+    decoder_input = torch.empty(1, 1).fill_(sos_idx).long().to(device)
+    while True:
+        if decoder_input.size(1) == max_len:
+            break
+
+        # build mask for target
+        decoder_mask = causal_mask(decoder_input.size(1)).long().to(device)
+      
+
+        # calculate output
+        out = model.module.decode(encoder_output, source_mask, decoder_input, decoder_mask)
+        # print(f'out: {out.shape}')
+
+        # Get next token probabilities with temperature applied
+        logits = model.module.project(out[:, -1]) 
+        probabilities = F.softmax(logits, dim=-1)
+
+        # Greedily select the next word
+        next_word = torch.argmax(probabilities, dim=1)
+        
+        # Append next word
+        decoder_input = torch.cat([decoder_input, next_word.unsqueeze(0)], dim=1)
+        # # get next token
+        # prob = model.project(out[:, -1])
+        # _, next_word = torch.max(prob, dim=1)
+        # # print(f'prob: {prob.shape}')
+        # decoder_input = torch.cat(
+        #     [decoder_input, torch.empty(1, 1).long().fill_(next_word.item()).to(device)], dim=1
+        # )
+
+        if next_word.item() == eos_idx:
+            break
+
+    return decoder_input.squeeze(0)
+
+
+def run_validation(model, validation_ds,tokenizer_tgt, max_len, device, print_msg, global_step, num_examples=3):
+    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, None, 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):
+    for item in ds:
+        yield item['text']
+def batch_iterator(data):
+    for i in range(0, len(data)):
+        yield data[i]['text'] 
+
+# Assuming batch_iterator is a function that yields batches
+def tqdm_batch_iterator(data, *args, **kwargs):
+    for batch in tqdm(batch_iterator(data, *args, **kwargs), total=len(data)):
+        yield batch               
+
+def get_or_build_tokenizer(config, ds):
+    tokenizer = GPT2TokenizerFast.from_pretrained("openai-community/gpt2", unk_token ='[UNK]', bos_token = '[SOS]', eos_token = '[EOS]' , pad_token = '[PAD]')
+    return tokenizer
+    # tokenizer_path = Path(config['tokenizer_file'])
+    # if not Path.exists(tokenizer_path):
+    #     # 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)
+    #     tokenizer.train_from_iterator(get_all_sentences(ds), trainer=trainer)
+    #     tokenizer.save(str(tokenizer_path))
+    # else:
+    #     tokenizer = Tokenizer.from_file(str(tokenizer_path))
+    # return tokenizer
+
+def get_ds(config):
+    # It only has the train split, so we divide it overselves
+    # ds_raw = load_dataset("HausaNLP/HausaVG", split='train+validation+test+challenge_test')
+    train_ds_raw = load_dataset("MMInstruction/M3IT", 'coco', split ='train')
+    
+    val_ds_raw =   load_dataset("MMInstruction/M3IT", 'coco', split ='validation[:2%]')
+    
+    # ds_raw = load_dataset('opus_books', f"{config['lang_src']}-{config['lang_tgt']}", split='train')
+
+    # Build tokenizers
+    
+    tokenizer_tgt = get_or_build_tokenizer(config, train_ds_raw,)
+    seed = 20  # You can choose any integer as your seed
+    torch.manual_seed(seed)
+    # # Keep 90% for training, 10% for validation
+    # train_ds_size = int(0.9 * len(ds_raw))
+    # val_ds_size = len(ds_raw) - train_ds_size
+    # train_ds_raw, val_ds_raw = random_split(ds_raw, [train_ds_size, val_ds_size])
+
+    train_ds = BilingualDataset(train_ds_raw, tokenizer_tgt,  config['seq_len'])
+    val_ds = BilingualDataset(val_ds_raw, tokenizer_tgt, config['seq_len'])
+    
+
+    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_tgt
+
+def get_model(config, vocab_tgt_len):
+    model = build_transformer(vocab_tgt_len, config['seq_len'], d_model=config['d_model'])
+    return model
+
+def train_model(config):
+
+    accelerator = Accelerator()
+  
+
+    print()
+    wandb.login(key = 'c20a1022142595d7d1324fdc53b3ccb34c0ded22')
+    wandb.init(project="Vision", name=config['project_name'])
+
+    # Initialize WandB configuration
+    wandb.config.epochs = config['num_epochs']
+    wandb.config.batch_size = config['batch_size']
+    wandb.config.learning_rate = config['lr'] 
+    # Define the devic
+    # Define the device
+    device = accelerator.device
+    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    # Make sure the weights folder exists
+    Path(config['model_folder']).mkdir(parents=True, exist_ok=True)
+
+    train_dataloader, val_dataloader, tokenizer_tgt = get_ds(config)
+    model = get_model(config, len(tokenizer_tgt)).to(device)
+   
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'], betas=(0.9, 0.98),eps=1e-9)
+
+    model, optimizer, train_dataloader, val_dataloader = accelerator.prepare(
+    model, optimizer, train_dataloader, val_dataloader
+)
+
+    # If the user specified a model to preload before training, load it
+    initial_epoch = 0
+    global_step = 0
+
+    def save_models():
+        accelerator.save_state(output_dir=f'/kaggle/working/weights/tmodel_00')
+        print(f'saving global step {global_step}')
+    
+    if config['preload']:
+        model_filename = get_weights_file_path(config, config['preload'])
+        print(f'Preloading model {model_filename}')
+        accelerator.load_state(model_filename)
+        initial_epoch = 4
+   
+        # 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']
+
+    loss_fn = nn.CrossEntropyLoss(ignore_index=tokenizer_tgt.convert_tokens_to_ids('[PAD]'), label_smoothing=0.1).to(device)
+   
+    for epoch in range(initial_epoch, config['num_epochs']):
+
+        # timer = threading.Timer(5*60, save_models)
+        # timer.start()
+
+        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.module.encode(encoder_input, None) # (B, seq_len, d_model)
+            decoder_output = model.module.decode(encoder_output, None, decoder_input, decoder_mask) # (B, seq_len, d_model)
+            proj_output = model.module.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, len(tokenizer_tgt)), label.view(-1))
+            batch_iterator.set_postfix({"loss": f"{loss.item():6.3f}"})
+
+            # Log the loss
+            wandb.log({"Training Loss": loss.item(), "Global Step": global_step})
+
+            # # Backpropagate the loss
+            # loss.backward()
+            accelerator.backward(loss)
+
+            # Update the weights
+            optimizer.step()
+            optimizer.zero_grad(set_to_none=True)
+
+            global_step += 1
+            # if global_step == 20000 or global_step == 25000:
+            #     print(f'saved state at {global_step}')
+            #     accelerator.save_state(output_dir=f'/kaggle/working/weights/tmodel_{epoch:02d}')  
+            if global_step == 1000 or global_step == 5000  or global_step == 10000 or global_step == 15000 or global_step == 20000 or global_step == 30000:
+                run_validation(model, val_dataloader, tokenizer_tgt, config['seq_len'], device, lambda msg: batch_iterator.write(msg), global_step)
+                model.train()
+                
+
+        # # Run validation at the end of every epoch
+            # 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)
+        # accelerator.save_model(model, model_filename)
+    
+        accelerator.save_state(output_dir=f'/kaggle/working/weights/tmodel_{epoch:02d}')
+        # run_validation(model, val_dataloader, tokenizer_src, tokenizer_tgt, config['seq_len'], device, lambda msg: batch_iterator.write(msg), global_step, writer)
+        model.eval()
+        eval_loss = 0.0
+
+        #accelerate 
+        accurate = 0
+        num_elems = 0
+        # batch_iterator = tqdm(v_dataloader, desc=f"Processing Epoch {epoch:02d}")
+        with torch.no_grad():
+            batch_itere = tqdm(val_dataloader, desc=f"Processing loss")
+            for batch in batch_itere:
+            
+
+                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.module.encode(encoder_input, None) # (B, seq_len, d_model)
+                decoder_output = model.module.decode(encoder_output, None, decoder_input, decoder_mask)# (B, seq_len, d_model)
+                proj_output = model.module.project(decoder_output)
+            
+                # (B, seq_len, vocab_size)
+
+                # Compare the output with the label
+                # label = batch['label'].to(device) # (B, seq_len)
+                proj_output, label = accelerator.gather_for_metrics((
+                proj_output, batch["label"]
+            ))
+
+                # Compute the loss using a simple cross entropy
+                ls = loss_fn(proj_output.view(-1, len(tokenizer_tgt)), label.view(-1))
+                batch_itere.set_postfix({"loss": f"{ls.item():6.3f}"})
+                eval_loss += ls
+                # loss_fn(proj_output.view(-1, tokenizer_tgt.get_vocab_size()), label.view(-1))
+           
+                
+        avg_val_loss = eval_loss / len(val_dataloader)
+        accelerator.print(f"Epoch {epoch},Validation Loss: {avg_val_loss})Validation Loss: {avg_val_loss}")
+        # print(f'Epoch {epoch},Validation Loss: {avg_val_loss.item()}')
+        wandb.log({"Validation Loss": avg_val_loss.item(), "Global Step": global_step})
+
+    
+        run_validation(model, val_dataloader, tokenizer_tgt, config['seq_len'], device, lambda msg: batch_iterator.write(msg), global_step)
+
+
+if __name__ == '__main__':
+    warnings.filterwarnings("ignore")
+    config = get_config()
+    train_model(config)