change display_attention parameters' names

config.yml

@@ -23,7 +23,7 @@ checkpoint:
 
 # training hyperparameters
 batch_size: 32
-total_epoch: 5
+total_epoch: 10
 
 clip: 1.0
 

inference_streamlit.py

@@ -72,7 +72,6 @@ def add_bg_from_url():
          unsafe_allow_html=True
      )
 
-
 def main(config_fpath="config.yml"):
     config = data_utils.get_config(config_fpath)
     for key, value in config.items():
@@ -112,6 +111,7 @@ def main(config_fpath="config.yml"):
 
         st.header("Attention Matrix")
         src_tokens = [token.lower() for token in src_tok.tokenize(input_text)] 
+        src_tokens = [src_tok.vocab.bos_token] + src_tokens + [src_tok.vocab.eos_token]
         fig = model_utils.display_attention(src_tokens, pred_tokens[1:], 
                                             attention, n_heads=1, 
                                             n_rows=1, n_cols=1, fig_size=(5, 5))

preprocess.py

@@ -75,7 +75,12 @@ def main(config_fpath="config.yml"):
 
     print("Load tokenizers...")
     src_tok = EnTokenizer()
-    tgt_tok = ViTokenizer()
+
+    # Vietnamese Multi-word Tokenizer
+    # tgt_tok = ViTokenizer()
+
+    # Vietnamese Word Tokenizer
+    tgt_tok = EnTokenizer()
 
     print("Load DataLoaders")
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

utils/model_utils.py

@@ -66,8 +66,8 @@ def translate_sentence(sent, src_tok, tgt_tok, model, device, max_len=256):
 
     return pred_tokens, attention
 
-def display_attention(sentence, translation, attention, n_heads = 8, 
-                      n_rows = 4, n_cols = 2, fig_size=(15,25)):
+def display_attention(src_tokens, pred_tokens, attention, n_heads = 8, 
+                      n_rows=4, n_cols=2, fig_size=(15,25)):
 
     assert n_rows * n_cols == n_heads
 
@@ -82,9 +82,8 @@ def display_attention(sentence, translation, attention, n_heads = 8,
         cax = ax.matshow(_attention, cmap='bone')
 
         ax.tick_params(labelsize=12)
-        ax.set_xticklabels(['']+['<sos>']+[t.lower() for t in sentence]+['<eos>'],
-                           rotation=45)
-        ax.set_yticklabels(['']+translation)
+        ax.set_xticklabels([''] + src_tokens)
+        ax.set_yticklabels([''] + pred_tokens)
 
         ax.xaxis.set_major_locator(ticker.MultipleLocator(1))
         ax.yaxis.set_major_locator(ticker.MultipleLocator(1))