Upload 2 files

app.py

@@ -148,7 +148,7 @@ if btn:
                     html_action='return'
                 )._repr_html_(), height=775, width=1000, scrolling=True)
             
-        with st.expander('Attention'):
+        with st.expander('Attention at selected layer and head'):
             attention_matrix = attentions[layer][0, head].detach().numpy()
             separator_token = roberta_tokenizer.sep_token
             sep_token_index = tokens.index(separator_token) if separator_token in tokens else len(tokens) - 1
@@ -181,7 +181,7 @@ if btn:
                         html_action='return'
                     )._repr_html_(), height=375, width=1000, scrolling=True)
                 
-        with st.expander('Attention'):
+        with st.expander('Attention at selected layer and head'):
             attention_matrix = attentions[layer][0, head].detach().numpy()
             separator_token = distilbert_tokenizer.sep_token
             sep_token_index = tokens.index(separator_token) if separator_token in tokens else len(tokens) - 1

requirements.txt

@@ -3,4 +3,5 @@ safetensors
 transformers
 streamlit
 plotly
-bertviz
+bertviz
+ipython