Upload 12 files

bleu.py

@@ -0,0 +1,49 @@
+import torch
+from torch.utils.data import DataLoader
+from torchtext.data.metrics import bleu_score
+
+from tqdm import tqdm
+
+from decode_method import beam_search_decode
+from transformer import Transformer
+
+from tokenizers import Tokenizer
+
+
+def calculate_bleu_score(
+    model: Transformer, 
+    bleu_dataloader: DataLoader, 
+    src_tokenizer: Tokenizer, 
+    tgt_tokenizer: Tokenizer, 
+    device = torch.device('cpu'),
+    num_samples: int = 9999999,
+):
+    """"""
+    model.eval()
+
+    # inferance
+    count = 0
+    expected = []
+    predicted = []
+
+    with torch.no_grad():
+        batch_iterator = tqdm(bleu_dataloader)
+        for batch in batch_iterator:
+            count += 1
+            encoder_input = batch['encoder_input'].to(device)
+            encoder_mask = batch['encoder_mask'].to(device)
+
+            assert encoder_input.size(0) == 1, "batch_size = 1 for bleu calculation"
+
+            model_out = beam_search_decode(model, encoder_input, encoder_mask, src_tokenizer, tgt_tokenizer, 300, device)
+
+            target_text = batch['tgt_text'][0]
+            model_out_text = tgt_tokenizer.decode(model_out.detach().cpu().numpy())
+
+            expected.append([target_text.split()])
+            predicted.append(model_out_text.split())
+
+            if count == num_samples:
+                break
+
+    return bleu_score(predicted, expected) * 100.0

bleu_test.py

@@ -0,0 +1,34 @@
+from torch.utils.data import DataLoader
+from bleu import calculate_bleu_score
+from load_dataset import load_local_bleu_dataset
+from dataset import BilingualDataset
+from config import load_config
+from load_and_save_model import load_model_tokenizer
+
+
+def get_bleu_of_model(config) -> float:
+    model, src_tokenizer, tgt_tokenizer = load_model_tokenizer(config)
+    bleu_ds_raw = load_local_bleu_dataset(
+        src_dataset_filename='datasets/'+config['dataset']['bleu_dataset']+'.'+config['dataset']['src_lang'],
+        tgt_dataset_filename='datasets/'+config['dataset']['bleu_dataset']+'.'+config['dataset']['tgt_lang'],
+        src_lang=config['dataset']['src_lang'],
+        tgt_lang=config['dataset']['tgt_lang'],
+    )
+    bleu_ds = BilingualDataset(
+        ds=bleu_ds_raw,
+        src_tokenizer=src_tokenizer,
+        tgt_tokenizer=tgt_tokenizer,
+        src_lang=config['dataset']['src_lang'],
+        tgt_lang=config['dataset']['tgt_lang'],
+        src_max_seq_len=config['dataset']['src_max_seq_len'],
+        tgt_max_seq_len=config['dataset']['tgt_max_seq_len'],
+    )
+    bleu_dataloader = DataLoader(bleu_ds, batch_size=1, shuffle=True)
+    return calculate_bleu_score(
+        model, bleu_dataloader, src_tokenizer, tgt_tokenizer,
+    )
+
+if __name__ == '__main__':
+    for file_name in {'config_final.yaml', 'config_huge.yaml', 'config_big.yaml', 'config_small.yaml'}:
+        config = load_config(file_name)
+        print(get_bleu_of_model(config), f" is the BLEU of {file_name}", sep='')

dataset.py

@@ -0,0 +1,99 @@
+from typing import List, Dict, Any
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset
+from torch import Tensor
+
+from tokenizers import Tokenizer
+
+class BilingualDataset(Dataset):
+    """
+    A Bilingual Dataset that follows the structure of the 'opus_books' dataset.
+    """
+    def __init__(
+        self,
+        ds: List[Dict[str, Dict[str,str]]],
+        src_tokenizer: Tokenizer,
+        tgt_tokenizer: Tokenizer,
+        src_lang: str,
+        tgt_lang: str,
+        src_max_seq_len: int,
+        tgt_max_seq_len: int,
+    ) -> None:
+        super(BilingualDataset, self).__init__()
+
+        self.ds = ds
+        self.src_tokenizer = src_tokenizer
+        self.tgt_tokenizer = tgt_tokenizer
+        self.src_lang = src_lang
+        self.tgt_lang = tgt_lang
+
+        self.src_max_seq_len = src_max_seq_len
+        self.tgt_max_seq_len = tgt_max_seq_len
+
+        self.sos_token = torch.tensor([src_tokenizer.token_to_id('<sos>')], dtype=torch.int64)
+        self.eos_token = torch.tensor([src_tokenizer.token_to_id('<eos>')], dtype=torch.int64)
+        self.pad_token = torch.tensor([src_tokenizer.token_to_id('<pad>')], dtype=torch.int64)
+
+    def __len__(self):
+        return len(self.ds)
+    
+    def __getitem__(self, index: int) -> Dict[str, Any]:
+        src_tgt_pair = self.ds[index]
+        src_text = src_tgt_pair['translation'][self.src_lang]
+        tgt_text = src_tgt_pair['translation'][self.tgt_lang]
+
+        encoder_input_tokens = self.src_tokenizer.encode(src_text).ids
+        decoder_input_tokens = self.tgt_tokenizer.encode(tgt_text).ids
+
+        encoder_num_padding = self.src_max_seq_len - len(encoder_input_tokens) - 2 # <sos> + <eos>
+        decoder_num_padding = self.tgt_max_seq_len - len(decoder_input_tokens) - 1 # <sos>
+
+        # <sos> + source_text + <eos> + <pad> = encoder_input
+        encoder_input = torch.cat(
+            [
+                self.sos_token,
+                torch.tensor(encoder_input_tokens, dtype=torch.int64),
+                self.eos_token,
+                torch.tensor([self.pad_token] * encoder_num_padding, dtype=torch.int64)
+            ]
+        )
+
+        decoder_input_tokens = torch.tensor(decoder_input_tokens, dtype=torch.int64)
+        decoder_padding = torch.tensor([self.pad_token] * decoder_num_padding, dtype=torch.int64)
+        # <sos> + target_text + <pad> = decoder_input
+        decoder_input = torch.cat(
+            [
+                self.sos_token,
+                decoder_input_tokens,
+                decoder_padding
+            ]
+        )
+        # target_text + <eos> + <pad> = expected decoder_output (label)
+        label = torch.cat(
+            [
+                decoder_input_tokens,
+                self.eos_token,
+                decoder_padding
+            ]
+        )
+
+        assert encoder_input.size(0) == self.src_max_seq_len
+        assert decoder_input.size(0) == self.tgt_max_seq_len
+        assert label.size(0) == self.tgt_max_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).unsqueeze(0).int() & causal_mask(decoder_input.size(0)), # (1, seq_len, seq_len)
+            'label': label, # (seq_len)
+            'src_text': src_text,
+            'tgt_text': tgt_text,
+        }
+    
+def causal_mask(size: int) -> Tensor:
+    mask = torch.triu(torch.ones(1, size, size), diagonal=1).type(torch.int)
+    return mask == 0
+    

decode_method.py

@@ -47,4 +47,69 @@ def greedy_decode(
             break
     if give_attn:
         return (decoder_input.squeeze(0), attn)
-    return decoder_input.squeeze(0)
+    return decoder_input.squeeze(0)
+
+def beam_search_decode(
+    model: Transformer, 
+    src: Tensor, 
+    src_mask: Tensor, 
+    src_tokenizer: Tokenizer, 
+    tgt_tokenizer: Tokenizer, 
+    tgt_max_seq_len: int, 
+    device,
+    beam_size: int = 3,
+):
+    sos_idx = src_tokenizer.token_to_id('<sos>')
+    eos_idx = src_tokenizer.token_to_id('<eos>')
+
+    # Precompute the encoder output and reuse it for every step
+    encoder_output = model.encode(src, src_mask)
+    # Initialize the decoder input with the sos token
+    decoder_initial_input = torch.empty(1,1).fill_(sos_idx).type_as(src).to(device)
+
+    # Create a candidate list
+    candidates = [(decoder_initial_input, 1)]
+
+    while True:
+
+        # If a candidate has reached the maximum length, it means we have run the decoding for at least max_len iterations, so stop the search
+        if any([cand.size(1) == tgt_max_seq_len for cand, _ in candidates]):
+            break
+
+        # Create a new list of candidates
+        new_candidates = []
+
+        for candidate, score in candidates:
+
+            # Do not expand candidates that have reached the eos token
+            if candidate[0][-1].item() == eos_idx:
+                continue
+
+            # Build the candidate's mask
+            candidate_mask = causal_mask(candidate.size(1)).type_as(src_mask).to(device)
+            # calculate output
+            out, attn = model.decode(encoder_output, src_mask, candidate, candidate_mask)
+            # get next token probabilities
+            prob = model.project(out[:, -1])
+            # get the top k candidates
+            topk_prob, topk_idx = torch.topk(prob, beam_size, dim=1)
+            for i in range(beam_size):
+                # for each of the top k candidates, get the token and its probability
+                token = topk_idx[0][i].unsqueeze(0).unsqueeze(0)
+                token_prob = topk_prob[0][i].item()
+                # create a new candidate by appending the token to the current candidate
+                new_candidate = torch.cat([candidate, token], dim=1)
+                # We sum the log probabilities because the probabilities are in log space
+                new_candidates.append((new_candidate, score + token_prob))
+
+        # Sort the new candidates by their score
+        candidates = sorted(new_candidates, key=lambda x: x[1], reverse=True)
+        # Keep only the top k candidates
+        candidates = candidates[:beam_size]
+
+        # If all the candidates have reached the eos token, stop
+        if all([cand[0][-1].item() == eos_idx for cand, _ in candidates]):
+            break
+
+    # Return the best candidate
+    return candidates[0][0].squeeze()

load_and_save_model.py

@@ -0,0 +1,73 @@
+from typing import Tuple
+
+import torch
+from transformer import get_model, Transformer
+from config import load_config, get_weights_file_path
+from train import get_local_dataset_tokenizer
+from tokenizer import get_or_build_local_tokenizer
+
+from tokenizers import Tokenizer
+
+
+def load_train_data_and_save_model(config, model_name):
+    """
+    loads training data (model, optim, scheduler,...) and saves ONLY the model.
+    """
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f'Using device {device}')    
+
+    train_dataloader, val_dataloader, src_tokenizer, tgt_tokenizer = get_local_dataset_tokenizer(config)
+    model = get_model(config, src_tokenizer.get_vocab_size(), tgt_tokenizer.get_vocab_size()).to(device)
+
+    assert config['model']['preload'], 'where to preload model.'
+
+    model_load_filename = get_weights_file_path(config, config['model']['preload'])
+    print(f'Preloading model from train data in {model_load_filename}')
+    state = torch.load(model_load_filename, map_location=device)
+
+    model.load_state_dict(state['model_state_dict'])
+
+    model_save_filename = get_weights_file_path(config, model_name)
+    torch.save(model.state_dict(), model_save_filename)
+    print(f'Model saved at {model_save_filename}')
+
+def load_model_tokenizer(
+    config,
+    device = torch.device('cpu'),
+    logs: bool = True,
+) -> Tuple[Transformer, Tokenizer, Tokenizer]:
+    """
+    Loads a local model and tokenizer from a given config
+    """
+    if config['model']['preload'] is None:
+        raise ValueError('Unspecified preload model')
+    
+    src_tokenizer = get_or_build_local_tokenizer(
+        config=config, 
+        ds=None, 
+        lang=config['dataset']['src_lang'], 
+        tokenizer_type=config['dataset']['src_tokenizer']
+    )
+    tgt_tokenizer = get_or_build_local_tokenizer(
+        config=config, 
+        ds=None, 
+        lang=config['dataset']['tgt_lang'], 
+        tokenizer_type=config['dataset']['tgt_tokenizer']
+    )
+
+    model = get_model(
+        config,
+        src_tokenizer.get_vocab_size(), 
+        tgt_tokenizer.get_vocab_size(),
+    ).to(device)
+
+    model_filename = get_weights_file_path(config, config['model']['preload'])
+    model.load_state_dict(
+        torch.load(model_filename, map_location=device)
+    )
+    print(f'Finish loading model and tokenizers')
+    return (model, src_tokenizer, tgt_tokenizer)
+
+if __name__ == '__main__':
+    config = load_config(file_name='config_huge.yaml')
+    load_train_data_and_save_model(config, 'huge')

load_dataset.py

@@ -0,0 +1,104 @@
+from typing import List, Dict, Any
+
+from pathlib import Path
+
+from utils import get_full_file_path
+
+# SENTENCE_STOPPERS = {'!', '"', '#', '$', '%', '&', "'", '(', ')', '*', '+', ',', '-', '.', '/', ':', ';', '<', '=', '>', '?', '@', '[', '\\', ']', '^', '_', '`', '{', '|', '}', '~'}
+# VIETNAMESE_SPECIAL_CHARACTERS = {'à', 'á', 'ả', 'ã', 'ạ', 'â', 'ầ', 'ấ', 'ẩ', 'ẫ', 'ậ', 'ă', 'ằ', 'ắ', 'ẳ', 'ẵ', 'ặ', 'è', 'é', 'ẻ', 'ẽ', 'ẹ', 'ê', 'ề', 'ế', 'ể', 'ễ', 'ệ', 'ì', 'í', 'ỉ', 'ĩ', 'ị', 'ò', 'ó', 'ỏ', 'õ', 'ọ', 'ô', 'ồ', 'ố', 'ổ', 'ỗ', 'ộ', 'ơ', 'ờ', 'ớ', 'ở', 'ỡ', 'ợ', 'ù', 'ú', 'ủ', 'ũ', 'ụ', 'ư', 'ừ', 'ứ', 'ử', 'ữ', 'ự', 'ỳ', 'ý', 'ỷ', 'ỹ', 'ỵ'}
+
+# def is_Vietnamese_character(char):
+#     return char.isalpha() or char in VIETNAMESE_SPECIAL_CHARACTERS
+
+# def categorize_word(word: str) -> str:
+#     """
+#     Categoize word into 3 types:
+#     - "vi": likely Vietnamese.
+#     - "lo": likely Laos.
+#     - "num": a number
+#     """
+#     if any(char.isdigit() for char in word):
+#         return "num"
+    
+#     for stopper in SENTENCE_STOPPERS:
+#         if word.endswith(stopper):
+#             word = word[:-1]
+#         if len(word) == 0:
+#             break
+    
+#     if len(word) > 0 and any(not is_Vietnamese_character(char) for char in word):
+#         return "lo"
+#     else:
+#         return "vi"
+# 
+# def open_dataset(
+#     dataset_filename: str, 
+#     src_lang: str = "lo", 
+#     tgt_lang: str = "vi"
+# ) -> List[Dict[str, Dict[str,str]]]:
+#     ds = []
+#     file_path = get_full_file_path(dataset_filename)
+#     with open(file_path, 'r', encoding='utf-8') as file:
+#         lines = file.readlines()
+
+#     for index, line in enumerate(lines):
+#         line = line.split(sep=None)
+
+#         lo_positions = [i for i, word in enumerate(line) if categorize_word(word) == "lo"]
+#         if len(lo_positions) == 0:
+#             # print(line)
+#             continue
+
+#         split_index = max(lo_positions)
+#         assert split_index is not None, f"Dataset error on line {index+1}."
+
+#         src_text = ' '.join(line[:split_index+1])
+#         tgt_text = line[split_index+1:]
+        
+#         if index <= 5:
+#             print(src_text, tgt_text, sep="\n", end="\n-------")
+
+#         # TODO: post process the tgt_text to split all numbers in to single digits.
+#         ds.append({'translation':{src_lang:src_text, tgt_lang:tgt_text}})
+#     return ds
+
+# open_dataset('datasets/dev_clean.dat')
+
+def load_local_dataset(
+    dataset_filename: str, 
+    src_lang: str = "lo", 
+    tgt_lang: str = "vi"
+) -> List[Dict[str, Dict[str,str]]]:
+    ds = []
+    file_path = get_full_file_path(dataset_filename)
+    with open(file_path, 'r', encoding='utf-8') as file:
+        lines = file.readlines()
+
+    for index, line in enumerate(lines):
+        src_text, tgt_text = line.split(sep="\t", maxsplit=1)
+        ds.append({'translation':{src_lang:src_text, tgt_lang:tgt_text}})
+    return ds
+
+def load_local_bleu_dataset(
+    src_dataset_filename: str, 
+    tgt_dataset_filename: str,
+    src_lang: str = "lo", 
+    tgt_lang: str = "vi"
+) -> List[Dict[str, Dict[str,str]]]:
+    def load_local_monolanguage_dataset(dataset_filename: str):
+        mono_ds = []
+        file_path = get_full_file_path(dataset_filename)
+        with open(file_path, 'r', encoding='utf-8') as file:
+            lines = file.readlines()
+        for line in lines:
+            mono_ds.append(line)
+        return mono_ds
+    
+    src_texts = load_local_monolanguage_dataset(src_dataset_filename)
+    tgt_texts = load_local_monolanguage_dataset(tgt_dataset_filename)
+
+    assert len(src_texts) == len(tgt_texts)
+    ds = []
+    for i in range(len(src_texts)):
+        ds.append({'translation':{src_lang:src_texts[i], tgt_lang:tgt_texts[i]}})
+    return ds

tokenizer.py

@@ -0,0 +1,36 @@
+from tokenizers import Tokenizer
+from tokenizers.models import WordLevel, BPE
+from tokenizers.trainers import WordLevelTrainer, BpeTrainer
+from tokenizers.pre_tokenizers import Whitespace, ByteLevel
+
+from pathlib import Path
+
+
+def get_all_sentences(ds, lang: str):
+    for item in ds:
+        yield item['translation'][lang]
+
+def get_or_build_local_tokenizer(config, ds, lang: str, tokenizer_type: str, force_build: bool = False) -> Tokenizer:
+    tokenizer_path = Path(config['dataset']['tokenizer_file'].format(lang))
+    if not Path.exists(tokenizer_path) or force_build:
+        if ds is None:
+            raise ValueError("Cannot find local tokenizer, dataset given is None")
+        
+        if tokenizer_type == "WordLevel":
+            tokenizer = Tokenizer(WordLevel(unk_token='<unk>'))
+            tokenizer.pre_tokenizer = Whitespace()
+            trainer = WordLevelTrainer(special_tokens=['<unk>', '<pad>', '<sos>', '<eos>'], min_frequency=2)
+        elif tokenizer_type == "BPE":
+            tokenizer = Tokenizer(BPE(unk_token='<unk>'))
+            tokenizer.pre_tokenizer = Whitespace()
+            trainer = BpeTrainer(special_tokens=['<unk>', '<pad>', '<sos>', '<eos>'], min_frequency=2)
+        else:
+            raise ValueError("Unsupported Tokenizer type")
+        
+        tokenizer.train_from_iterator(
+            get_all_sentences(ds, lang), trainer=trainer
+        )
+        tokenizer.save(str(tokenizer_path))
+    else:
+        tokenizer = Tokenizer.from_file(str(tokenizer_path))
+    return tokenizer

train.py

@@ -0,0 +1,199 @@
+import os
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.utils.data import DataLoader, Dataset
+import torchmetrics
+from torch.utils.tensorboard import SummaryWriter
+
+from tqdm import tqdm
+
+# from datasets import load_dataset
+from load_dataset import load_local_dataset
+from transformer import get_model
+from config import load_config, get_weights_file_path
+from validate import run_validation
+from tokenizer import get_or_build_local_tokenizer
+
+from pathlib import Path
+
+from dataset import BilingualDataset
+from bleu import calculate_bleu_score
+from decode_method import greedy_decode
+
+def get_local_dataset_tokenizer(config):
+    train_ds_raw = load_local_dataset(
+        dataset_filename='datasets/'+config['dataset']['train_dataset'],
+        src_lang=config['dataset']['src_lang'],
+        tgt_lang=config['dataset']['tgt_lang']
+    )
+    val_ds_raw = load_local_dataset(
+        dataset_filename='datasets/'+config['dataset']['validate_dataset'],
+        src_lang=config['dataset']['src_lang'],
+        tgt_lang=config['dataset']['tgt_lang']
+    )
+
+    src_tokenizer = get_or_build_local_tokenizer(
+        config=config, 
+        ds=train_ds_raw + val_ds_raw, 
+        lang=config['dataset']['src_lang'], 
+        tokenizer_type=config['dataset']['src_tokenizer']
+    )
+    tgt_tokenizer = get_or_build_local_tokenizer(
+        config=config, 
+        ds=train_ds_raw + val_ds_raw, 
+        lang=config['dataset']['tgt_lang'], 
+        tokenizer_type=config['dataset']['tgt_tokenizer']
+    )
+
+    train_ds = BilingualDataset(
+        ds=train_ds_raw,
+        src_tokenizer=src_tokenizer,
+        tgt_tokenizer=tgt_tokenizer,
+        src_lang=config['dataset']['src_lang'],
+        tgt_lang=config['dataset']['tgt_lang'],
+        src_max_seq_len=config['dataset']['src_max_seq_len'],
+        tgt_max_seq_len=config['dataset']['tgt_max_seq_len'],
+    )
+    val_ds = BilingualDataset(
+        ds=val_ds_raw,
+        src_tokenizer=src_tokenizer,
+        tgt_tokenizer=tgt_tokenizer,
+        src_lang=config['dataset']['src_lang'],
+        tgt_lang=config['dataset']['tgt_lang'],
+        src_max_seq_len=config['dataset']['src_max_seq_len'],
+        tgt_max_seq_len=config['dataset']['tgt_max_seq_len'],
+    )
+
+    src_max_seq_len = 0
+    tgt_max_seq_len = 0    
+    for item in (train_ds_raw + val_ds_raw):
+        src_ids = src_tokenizer.encode(item['translation'][config['dataset']['src_lang']]).ids
+        tgt_ids = tgt_tokenizer.encode(item['translation'][config['dataset']['tgt_lang']]).ids
+        src_max_seq_len = max(src_max_seq_len, len(src_ids))
+        tgt_max_seq_len = max(tgt_max_seq_len, len(tgt_ids))
+    print(f'Max length of source sequence: {src_max_seq_len}')
+    print(f'Max length of target sequence: {tgt_max_seq_len}')
+
+    train_dataloader = DataLoader(train_ds, batch_size=config['train']['batch_size'], shuffle=True)
+    val_dataloader = DataLoader(val_ds, batch_size=1, shuffle=True)
+
+    return train_dataloader, val_dataloader, src_tokenizer, tgt_tokenizer
+
+def train_model(config):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f'Using device {device}')    
+
+    Path(config['model']['model_folder']).mkdir(parents=True, exist_ok=True)
+
+    train_dataloader, val_dataloader, src_tokenizer, tgt_tokenizer = get_local_dataset_tokenizer(config)
+    model = get_model(config, src_tokenizer.get_vocab_size(), tgt_tokenizer.get_vocab_size()).to(device)
+
+    print(f'{src_tokenizer.get_vocab_size()}, {tgt_tokenizer.get_vocab_size()}')
+
+    #Tensorboard
+    writer = SummaryWriter(config['experiment_name'])
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=config['train']['lr'], eps=1e-9)
+    
+    from transformers import get_linear_schedule_with_warmup
+    scheduler = get_linear_schedule_with_warmup(
+        optimizer,
+        num_warmup_steps=config['train']['warm_up_steps'],
+        num_training_steps=len(train_dataloader) * config['train']['num_epochs']+1
+    )
+
+    initial_epoch = 0 
+    global_step = 0
+    if config['model']['preload']:
+        model_filename = get_weights_file_path(config, config['model']['preload'])
+        print(f'Preloading model from {model_filename}')
+        state = torch.load(model_filename, map_location=device)
+
+        initial_epoch = state['epoch']+1
+        model.load_state_dict(state['model_state_dict'])
+        optimizer.load_state_dict(state['optimizer_state_dict'])
+        scheduler.load_state_dict(state['scheduler_state_dict'])
+        global_step = state['global_step']
+    
+    loss_fn = nn.CrossEntropyLoss(
+        ignore_index=src_tokenizer.token_to_id('<pad>'), 
+        label_smoothing=config['train']['label_smoothing'],
+    ).to(device)
+
+    print(f"Training model with {model.count_parameters()} params.")
+
+    patience = config['train']['patience']
+    best_state = {
+        'model_state_dict': model.state_dict(),
+        'scheduler_state_dict': scheduler.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'loss': 9999999.99
+    }
+
+    for epoch in range(initial_epoch, config['train']['num_epochs']):
+        batch_iterator = tqdm(train_dataloader, desc=f'Proceesing epoch {epoch:02d}')
+        for batch in batch_iterator:
+            model.train()
+
+            encoder_input = batch['encoder_input'].to(device) # (batch, seq_len)
+            decoder_input = batch['decoder_input'].to(device) # (batch. seq_len)
+            encoder_mask = batch['encoder_mask'].to(device) # (batch, 1, 1, seq_len)
+            decoder_mask = batch['decoder_mask'].to(device) # (batch, 1, seq_len, seq_len)
+            
+            encoder_output = model.encode(encoder_input, encoder_mask) # (batch, seq_len, d_model)
+            decoder_output, attn = model.decode(encoder_output, encoder_mask, decoder_input, decoder_mask) # (batch, seq_len, d_model)
+            proj_output = model.project(decoder_output) # (batch, seq_len, tgt_vocab_size)
+
+            label = batch['label'].to(device) # (batch, seq_len)
+
+            loss = loss_fn(proj_output.view(-1, tgt_tokenizer.get_vocab_size()), label.view(-1))
+            batch_iterator.set_postfix({f"loss":f"{loss.item():6.3f}"})
+
+            writer.add_scalar('train_loss', loss.item(), global_step)
+            writer.flush()
+
+            global_step += 1
+            if global_step % patience == 0:
+                if loss > best_state['loss']:
+                    model.load_state_dict(best_state['model_state_dict'])
+                    optimizer.load_state_dict(best_state['optimizer_state_dict'])
+                    scheduler.load_state_dict(best_state['scheduler_state_dict'])
+                    continue
+                else:
+                    best_state = {
+                        'model_state_dict': model.state_dict(),
+                        'scheduler_state_dict': scheduler.state_dict(),
+                        'optimizer_state_dict': optimizer.state_dict(),
+                        'loss': 9999999.99
+                    }
+            loss.backward()
+
+            optimizer.step()
+            scheduler.step()
+            optimizer.zero_grad()
+
+        run_validation(model, val_dataloader, src_tokenizer, tgt_tokenizer, device, lambda msg: batch_iterator.write(msg), global_step, writer)
+
+        model_filename = get_weights_file_path(config, f'{epoch:02d}')
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': best_state['model_state_dict'],
+            'scheduler_state_dict': best_state['scheduler_state_dict'],
+            'optimizer_state_dict': best_state['optimizer_state_dict'],
+            'global_step': global_step,
+        }, model_filename)
+
+        # print(f"Bleu score: {calculate_bleu_score(model, val_dataloader, src_tokenizer, tgt_tokenizer, device)}")
+
+        if config['train']['on_colab']:
+            # if (epoch % 5) == 0:
+            #     model_zip_filename = f'model_epoch_{epoch}.zip'
+            #     os.system(f'zip -r {model_zip_filename} /content/silver-spoon/weights')
+            runs_zip_filename = f'runs_epoch_{epoch}.zip'
+            os.system(f"zip -r {runs_zip_filename} /content/silver-spoon/{config['experiment_name']}")
+
+
+if __name__ == '__main__':
+    config = load_config(file_name='config.yaml')
+    train_model(config)

translate.py

@@ -6,7 +6,7 @@ from torch import Tensor
 from tokenizers import Tokenizer
 
 from transformer import Transformer
-from decode_method import greedy_decode
+from decode_method import greedy_decode, beam_search_decode
 
 def translate(
     model: Transformer,
@@ -19,51 +19,61 @@ def translate(
     """
     Translation function.
 
+    Supported `decode_method`: 'greedy' or 'beam-search'
+    
+    'beam-search' doesn't give attn scores.
+
     Output:
         - translation (str): the translated string.
         - attn (Tensor): The decoder's attention (for visualization)
     """
-    sos_token = torch.tensor([src_tokenizer.token_to_id('<sos>')], dtype=torch.int64)
-    eos_token = torch.tensor([src_tokenizer.token_to_id('<eos>')], dtype=torch.int64)
-    pad_token = torch.tensor([src_tokenizer.token_to_id('<pad>')], dtype=torch.int64)
-    
-    encoder_input_tokens = src_tokenizer.encode(text).ids
-    # <sos> + source_text + <eos> = encoder_input
-    encoder_input = torch.cat(
-        [
-            sos_token,
-            torch.tensor(encoder_input_tokens, dtype=torch.int64),
-            eos_token,
-        ]
-    )
-    encoder_mask = (encoder_input != pad_token).unsqueeze(0).unsqueeze(0).unsqueeze(0).int() # (1, 1, seq_len)
+    model.eval()
+    with torch.no_grad():
+        sos_token = torch.tensor([src_tokenizer.token_to_id('<sos>')], dtype=torch.int64)
+        eos_token = torch.tensor([src_tokenizer.token_to_id('<eos>')], dtype=torch.int64)
+        pad_token = torch.tensor([src_tokenizer.token_to_id('<pad>')], dtype=torch.int64)
+        
+        encoder_input_tokens = src_tokenizer.encode(text).ids
+        # <sos> + source_text + <eos> = encoder_input
+        encoder_input = torch.cat(
+            [
+                sos_token,
+                torch.tensor(encoder_input_tokens, dtype=torch.int64),
+                eos_token,
+            ]
+        )
+        encoder_mask = (encoder_input != pad_token).unsqueeze(0).unsqueeze(0).unsqueeze(0).int() # (1, 1, seq_len)
 
-    encoder_input = encoder_input.unsqueeze(0)
-    # encoder_mask = torch.tensor(encoder_mask)
-    
-    assert encoder_input.size(0) == 1
+        encoder_input = encoder_input.unsqueeze(0)
+        # encoder_mask = torch.tensor(encoder_mask)
+        
+        assert encoder_input.size(0) == 1
 
-    if decode_method == 'greedy':
-        model_out, attn = greedy_decode(
-            model, encoder_input, encoder_mask, src_tokenizer, tgt_tokenizer, 400, device, 
-            give_attn=True,
-        )
-    elif decode_method == 'beam-search':
-        raise NotImplementedError
-    else:
-        raise ValueError("Unsuppored decode method")
+        if decode_method == 'greedy':
+            model_out, attn = greedy_decode(
+                model, encoder_input, encoder_mask, src_tokenizer, tgt_tokenizer, 349, device, 
+                give_attn=True,
+            )
+        elif decode_method == 'beam-search':
+            model_out = beam_search_decode(
+                model, encoder_input, encoder_mask, src_tokenizer, tgt_tokenizer, 349, device, 
+            )
+            attn = None # Beam search doesn't give attention score
+        else:
+            raise ValueError("Unsuppored decode method")
 
-    model_out_text = tgt_tokenizer.decode(model_out.detach().cpu().numpy())
-    return model_out_text, attn
+        model_out_text = tgt_tokenizer.decode(model_out.detach().cpu().numpy())
+        return model_out_text, attn
     
 
 from config import load_config
 from load_and_save_model import load_model_tokenizer
 if __name__ == '__main__':
-    config = load_config(file_name='config_small.yaml')
+    config = load_config(file_name='config_huge.yaml')
     model, src_tokenizer, tgt_tokenizer = load_model_tokenizer(config)
     text = "ສະບາຍດີ" # Hello.
     translation, attn = translate(
-        model, src_tokenizer, tgt_tokenizer, text
+        model, src_tokenizer, tgt_tokenizer, text,
+        decode_method='beam-search',
     )
     print(translation)

utils.py

@@ -0,0 +1,6 @@
+from pathlib import Path
+
+def get_full_file_path(file_name: str) -> str:
+    script_dir = Path(__file__).resolve().parent
+    file_path = script_dir / file_name
+    return file_path

validate.py

@@ -0,0 +1,72 @@
+import os
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.utils.data import DataLoader, Dataset
+import torchmetrics
+from torch.utils.tensorboard import SummaryWriter
+
+from tqdm import tqdm
+
+# from datasets import load_dataset
+from load_dataset import load_local_dataset
+from transformer import get_model, Transformer
+from config import load_config, get_weights_file_path
+
+from tokenizers import Tokenizer
+from tokenizers.models import WordLevel, BPE
+from tokenizers.trainers import WordLevelTrainer, BpeTrainer
+from tokenizers.pre_tokenizers import Whitespace
+
+from pathlib import Path
+
+from dataset import BilingualDataset
+from bleu import calculate_bleu_score
+from decode_method import greedy_decode
+
+
+def run_validation(
+    model: Transformer, 
+    validation_ds: DataLoader,
+    src_tokenizer: Tokenizer,
+    tgt_tokenizer: Tokenizer,
+    device, 
+    print_msg,
+    global_state,
+    writer,
+    num_examples:int = 2
+):
+    model.eval()
+
+    # inferance
+    count = 0
+    source_texts = []
+    expected = []
+    predicted = []
+
+    console_width = 50
+    with torch.no_grad():
+        for batch in validation_ds:
+            count += 1
+            encoder_input = batch['encoder_input'].to(device)
+            encoder_mask = batch['encoder_mask'].to(device)
+
+            assert encoder_input.size(0) == 1, "batch_size = 1 for validation"
+
+            model_out = greedy_decode(model, encoder_input, encoder_mask, src_tokenizer, tgt_tokenizer, 300, device)
+
+            source_text = batch['src_text'][0]
+            target_text = batch['tgt_text'][0]
+            model_out_text = tgt_tokenizer.decode(model_out.detach().cpu().numpy())
+
+            source_texts.append(source_text)
+            expected.append(target_text)
+            predicted.append(model_out_text)
+
+            print_msg("-"*console_width)
+            print_msg(f"SOURCE: {source_text}")
+            print_msg(f"TARGET: {target_text}")
+            print_msg(f"PREDICTED: {model_out_text}")
+
+            if count == num_examples:
+                break