add module script

app.py

@@ -0,0 +1,117 @@
+import torch
+import argparse
+import sentencepiece as spm
+from utils import utils_cls
+from model import BanglaTransformer
+from config import config as cfg
+torch.manual_seed(0)
+# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cpu')
+uobj = utils_cls(device=device)
+
+__MODULE__       = "Bangla Language Translation"
+__MAIL__         = "saifulbrur79@gmail.com"
+__MODIFICAIOTN__ = "28/03/2023"
+__LICENSE__      = "MIT"
+
+
+
+class Bn2EnTranslation:
+    def __init__(self):
+        self.bn_tokenizer='./model/bn_model.model'
+        self.en_tokenizer='./model/en_model.model'
+        self.bn_vocab='./model/bn_vocab.pkl'
+        self.en_vocab='./model/en_vocab.pkl'
+        self.model='./model/pytorch_model.pt'
+
+    def read_data(self, data_path):
+        with open(data_path, "r") as f:
+            data = f.readlines()
+        data = list(map(lambda x: [x.split("\t")[0], x.split("\t")[1].replace("\n", "")], data))
+        return data
+
+    def load_tokenizer(self, tokenizer_path:str = "")->object:
+        _tokenizer = spm.SentencePieceProcessor(model_file=tokenizer_path)
+        return _tokenizer
+
+    def get_vocab(self, BN_VOCAL_PATH:str="", EN_VOCAL_PATH:str=""):
+        bn_vocal, en_vocal = uobj.load_bn_vocal(BN_VOCAL_PATH), uobj.load_en_vocal(EN_VOCAL_PATH)
+        return bn_vocal, en_vocal
+
+    def load_model(self, model_path:str = "", SRC_VOCAB_SIZE:int=0, TGT_VOCAB_SIZE:int=0):
+        model = BanglaTransformer(
+            cfg.NUM_ENCODER_LAYERS, cfg.NUM_DECODER_LAYERS, cfg.EMB_SIZE,  SRC_VOCAB_SIZE,
+            TGT_VOCAB_SIZE, cfg.FFN_HID_DIM, nhead= cfg.NHEAD)
+        model.to(device)
+        checkpoint = torch.load(model_path)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        model.eval()
+        return model
+
+    def greedy_decode(self, model, src, src_mask, max_len, start_symbol, eos_index):
+        src = src.to(device)
+        src_mask = src_mask.to(device)
+        memory = model.encode(src, src_mask)
+        ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(device)
+        for i in range(max_len-1):
+            memory = memory.to(device)
+            memory_mask = torch.zeros(ys.shape[0], memory.shape[0]).to(device).type(torch.bool)
+            tgt_mask = (uobj.generate_square_subsequent_mask(ys.size(0))
+                                        .type(torch.bool)).to(device)
+            out = model.decode(ys, memory, tgt_mask)
+            out = out.transpose(0, 1)
+            prob = model.generator(out[:, -1])
+            _, next_word = torch.max(prob, dim = 1)
+            next_word = next_word.item()
+            ys = torch.cat([ys,torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)
+            if next_word == eos_index:
+                break
+        return ys
+    def get_bntoen_model(self):
+        print("Tokenizer Loading ...... : ", end="", flush=True)
+        bn_tokenizer = self.load_tokenizer(tokenizer_path=self.bn_tokenizer)
+        print("Done")
+        print("Vocab Loading ......     : ", end="", flush=True)
+        bn_vocab, en_vocab = self.get_vocab(BN_VOCAL_PATH=self.bn_vocab, EN_VOCAL_PATH=self.en_vocab)
+        print("Done")
+        print("Model Loading ......     : ", end="", flush=True)
+        model = self.load_model(model_path=self.model, SRC_VOCAB_SIZE=len(bn_vocab), TGT_VOCAB_SIZE=len(en_vocab))
+        print("Done") 
+
+        models = {
+            "bn_tokenizer" : bn_tokenizer, 
+            "bn_vocab" : bn_vocab,
+            "en_vocab" :  en_vocab, 
+            "model" : model
+        }
+        return models 
+    def translate(self, text, models):
+        model = models["model"]
+        src_vocab = models["bn_vocab"]
+        tgt_vocab = models["en_vocab"]
+        src_tokenizer = models["bn_tokenizer"]
+        src = text
+        PAD_IDX, BOS_IDX, EOS_IDX= src_vocab['<pad>'], src_vocab['<bos>'], src_vocab['<eos>']
+        tokens = [BOS_IDX] + [src_vocab.get_stoi()[tok] for tok in src_tokenizer.encode(src, out_type=str)]+ [EOS_IDX]
+        num_tokens = len(tokens)
+        src = (torch.LongTensor(tokens).reshape(num_tokens, 1) )
+        src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)
+        tgt_tokens = self.greedy_decode(model,  src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX, eos_index= EOS_IDX).flatten()
+        p_text = " ".join([tgt_vocab.get_itos()[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")
+        pts = " ".join(list(map(lambda x : x , p_text.replace(" ", "").split("▁"))))
+        return pts.strip()
+
+
+
+
+if __name__ == "__main__":
+    print(torch.cuda.get_device_name(0))
+    text = "এই উপজেলায় ১টি সরকারি কলেজ রয়েছে"
+
+    obj = Bn2EnTranslation()
+    models = obj.get_bntoen_model()
+    pre = obj.translate(text, models)
+    print("="*20)
+    print(f"input : {text}")
+    print(f"prediction: {pre}")
+

config.py

@@ -0,0 +1,12 @@
+class config:
+    EMB_SIZE = 512
+    NHEAD = 8
+    FFN_HID_DIM = 512
+    NUM_ENCODER_LAYERS = 6
+    NUM_DECODER_LAYERS = 6
+    NUM_EPOCHS = 300
+    BN_TOKENIZER_PATH = "./model/bn_model.model"
+    EN_TOKENIZER_PATH = "./model/en_model.model"
+    BN_VOCAL_PATH = "./model/bn_vocab.pkl"
+    EN_VOCAL_PATH = "./model/en_vocab.pkl"
+    MODEL_PATH = "./model/model_checkpoint.pt"

model.py

@@ -0,0 +1,84 @@
+import math
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import TransformerEncoder, TransformerDecoder, \
+    TransformerEncoderLayer, TransformerDecoderLayer
+torch.manual_seed(0)
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, emb_size: int, dropout, maxlen: int = 5000):
+        super(PositionalEncoding, self).__init__()
+        den = torch.exp(- torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)
+        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
+        pos_embedding = torch.zeros((maxlen, emb_size))
+        pos_embedding[:, 0::2] = torch.sin(pos * den)
+        pos_embedding[:, 1::2] = torch.cos(pos * den)
+        pos_embedding = pos_embedding.unsqueeze(-2)
+
+        self.dropout = nn.Dropout(dropout)
+        self.register_buffer('pos_embedding', pos_embedding)
+
+    def forward(self, token_embedding: Tensor):
+        return self.dropout(token_embedding +
+                            self.pos_embedding[:token_embedding.size(0),:])
+
+class TokenEmbedding(nn.Module):
+    def __init__(self, vocab_size: int, emb_size):
+        super(TokenEmbedding, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, emb_size)
+        self.emb_size = emb_size
+    def forward(self, tokens: Tensor):
+        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)
+    
+
+
+
+class BanglaTransformer(nn.Module):
+    def __init__(self, num_encoder_layers: int, num_decoder_layers: int,
+                 emb_size: int, src_vocab_size: int, tgt_vocab_size: int,
+                 dim_feedforward:int = 512, dropout:float = 0.1,  nhead:int=8):
+        super(BanglaTransformer, self).__init__()
+        encoder_layer = TransformerEncoderLayer(
+            d_model=emb_size, 
+            nhead=nhead,
+            dim_feedforward=dim_feedforward
+            )
+        self.transformer_encoder = TransformerEncoder(
+            encoder_layer, 
+            num_layers=num_encoder_layers
+            )
+        decoder_layer = TransformerDecoderLayer(
+            d_model=emb_size, 
+            nhead=nhead,
+            dim_feedforward=dim_feedforward
+            )
+        self.transformer_decoder = TransformerDecoder(
+            decoder_layer, 
+            num_layers=num_decoder_layers
+            )
+
+        self.generator = nn.Linear(emb_size, tgt_vocab_size)
+        self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)
+        self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)
+        self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout)
+
+    def forward(self, src: Tensor, trg: Tensor, src_mask: Tensor,
+                tgt_mask: Tensor, src_padding_mask: Tensor,
+                tgt_padding_mask: Tensor, memory_key_padding_mask: Tensor):
+        src_emb = self.positional_encoding(self.src_tok_emb(src))
+        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
+        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
+        outs = self.transformer_decoder(tgt_emb, memory, tgt_mask, None,
+                                        tgt_padding_mask, memory_key_padding_mask)
+        return self.generator(outs)
+
+    def encode(self, src: Tensor, src_mask: Tensor):
+        return self.transformer_encoder(self.positional_encoding(
+                            self.src_tok_emb(src)), src_mask)
+
+    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
+        return self.transformer_decoder(self.positional_encoding(
+                          self.tgt_tok_emb(tgt)), memory,
+                          tgt_mask)
+    

requirements.txt

@@ -0,0 +1,9 @@
+numpy==1.19.5
+sentencepiece==0.1.97
+tokenizers==0.8.1rc1
+torch==1.11.0+cu113
+torchsummary==1.5.1
+torchtext==0.12.0
+torchvision==0.12.0+cu113
+transformers==3.0.2
+nltk==3.7

utils.py

@@ -0,0 +1,25 @@
+import pickle
+import torch
+from config import config as cfg
+
+class utils_cls:
+    def __init__(self, device):
+        self.device = device
+
+    def load_bn_vocal(self, bn_vocal_path):
+        file = open(bn_vocal_path, 'rb')
+        bn_vocal = pickle.load(file)
+        file.close()
+        return bn_vocal
+    
+    def load_en_vocal(self, en_vocal_path):
+        file = open(en_vocal_path, 'rb')
+        en_vocal = pickle.load(file)
+        file.close()
+        return en_vocal
+
+    def generate_square_subsequent_mask(self, sz):
+        mask = (torch.triu(torch.ones((sz, sz), device=self.device)) == 1).transpose(0, 1)
+        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+        return mask
+