Update app.py

app.py

@@ -82,7 +82,7 @@ def initiate_model(config, device):
     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, map_location=torch.device('cpu'))
+    state = torch.load(model_filename)
     model.load_state_dict(state['model_state_dict'])
     return model, tokenizer_src, tokenizer_tgt
 
@@ -151,7 +151,12 @@ def main():
         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))
+        if attn_type == 'encoder':
+            st.write(get_all_attention_maps(attn_type, layers, heads, row_tokens, row_tokens, max_sentence_len, model))
+        elif attn_type == 'decoder':
+            st.write(get_all_attention_maps(attn_type, layers, heads, col_tokens, col_tokens, max_sentence_len, model))
+        elif attn_type == 'encoder-decoder':
+            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')