Update app.py

app.py

@@ -132,31 +132,33 @@ def main():
         input_text = st.text_input('Enter a sentence')
         # put two buttons side by side in the sidebar
         # translate_button = st.button('Translate', key='translate_button')
-        viz_button = st.button('Visualize Attention', key='viz_button')
+        # viz_button = st.button('Visualize Attention', key='viz_button')
         attn_type = st.selectbox('Select attention type', ['encoder', 'decoder', 'encoder-decoder'])
-        layers = st.multiselect('Select layers', list(range(3)))
-        heads = st.multiselect('Select heads', list(range(7)))
+        layers = st.multiselect('Select layers', list(range(6)))
+        heads = st.multiselect('Select heads', list(range(8)))
         # allow the user to select the all the layers and heads at once to visualize
         if st.checkbox('Select all layers'):
-            layers = list(range(3))
+            layers = list(range(6))
         if st.checkbox('Select all heads'):
-            heads = list(range(7))
+            heads = list(range(8))
     
-    if viz_button and input_text != '':
-        encoder_input_tokens, decoder_input_tokens, output = process_input(input_text, tokenizer_src, tokenizer_tgt, model, config, device)
-        max_sentence_len = len(encoder_input_tokens)
-        row_tokens = encoder_input_tokens
-        col_tokens = decoder_input_tokens
-        st.write('Input:', ' '.join(encoder_input_tokens))
-        st.write('Output:', ' '.join(decoder_input_tokens))
-        st.write('Translated:', output)
-        st.write('Attention Visualization')
-        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))
+    if input_text != '':
+        with st.spinner("Translating..."):
+            encoder_input_tokens, decoder_input_tokens, output = process_input(input_text, tokenizer_src, tokenizer_tgt, model, config, device)
+            max_sentence_len = len(encoder_input_tokens)
+            row_tokens = encoder_input_tokens
+            col_tokens = decoder_input_tokens
+            st.write('Input:', ' '.join(encoder_input_tokens))
+            st.write('Output:', ' '.join(decoder_input_tokens))
+            st.write('Translated:', output)
+            st.write('Attention Visualization')
+        with st.spinner("Visualizing Attention..."):
+            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')