Update src/streamlit_app.py

src/streamlit_app.py

@@ -144,7 +144,7 @@ def load_explanations(task, layer):
     return None
 
 def main():
-    st.title("Token Analysis (Lang Class-layer 11 1 file left")
+    st.title("Token Analysis")
     
     # Task and Layer Selection
     col1, col2 = st.columns(2)
@@ -213,57 +213,51 @@ def main():
                 
                 st.metric("Predicted Cluster", selected_row['Top 1'])
                 
-                # Load dev sentences once
+                                # Load dev sentences once
                 dev_sentences = load_dev_sentences(selected_task, selected_layer)
                 predicted_cluster = str(selected_row['Top 1'])
                 
-                if is_cls_token(token):
-                    # For CLS tokens, show the original context
-                    if dev_sentences and selected_row['line_idx'] < len(dev_sentences):
-                        st.subheader("Original Context")
-                        st.code(dev_sentences[selected_row['line_idx']].strip())
-                    else:
-                        st.info("No original context found for this token.")
+                # Show original context for all tokens
+                if dev_sentences and selected_row['line_idx'] < len(dev_sentences):
+                    st.subheader("Original Context")
+                    st.code(dev_sentences[selected_row['line_idx']].strip())
                 else:
-                    # For non-CLS tokens, show the wordcloud
-                    if clusters and predicted_cluster in clusters:
-                        # Create frequency dict for wordcloud
-                        token_frequencies = {}
-                        for token_info in clusters[predicted_cluster]:
-                            token = token_info['token']
-                            if not is_cls_token(token):  # Skip CLS tokens in wordcloud
-                                token_frequencies[token] = token_frequencies.get(token, 0) + token_info['occurrence']
-                        
-                        if token_frequencies:
-                            st.subheader("Cluster Word Cloud")
-                            wordcloud = create_wordcloud(token_frequencies)
-                            if wordcloud:
-                                plt.figure(figsize=(10, 5))
-                                plt.imshow(wordcloud, interpolation='bilinear')
-                                plt.axis('off')
-                                st.pyplot(plt)
+                    st.info("No original context found for this token.")
+                
+                # Show wordcloud for all tokens
+                if clusters and predicted_cluster in clusters:
+                    # Create frequency dict for wordcloud
+                    token_frequencies = {}
+                    for token_info in clusters[predicted_cluster]:
+                        token = token_info['token']
+                        token_frequencies[token] = token_frequencies.get(token, 0) + token_info['occurrence']
+                    
+                    if token_frequencies:
+                        st.subheader("Cluster Word Cloud")
+                        wordcloud = create_wordcloud(token_frequencies)
+                        if wordcloud:
+                            plt.figure(figsize=(10, 5))
+                            plt.imshow(wordcloud, interpolation='bilinear')
+                            plt.axis('off')
+                            st.pyplot(plt)
                 
                 # Show cluster contexts in expander
                 with st.expander(f"👀 View Contexts (Cluster {predicted_cluster})"):
                     if clusters and predicted_cluster in clusters:
-                        current_line = selected_row['line_idx']  # Get current line number
-                        shown_contexts = set()  # Keep track of shown contexts to avoid duplicates
+                        current_line = selected_row['line_idx']
+                        shown_contexts = set()
                         
                         for token_info in clusters[predicted_cluster]:
                             line_num = token_info['line_num']
-                            # Skip if it's the same line as current token or if we've shown this context before
                             if (line_num >= 0 and line_num < len(dev_sentences) and 
                                 line_num != current_line):
                                 context = dev_sentences[line_num].strip()
-                                # Only show if we haven't shown this exact context before
                                 if context not in shown_contexts:
                                     st.code(context)
                                     shown_contexts.add(context)
                         
                         if not shown_contexts:
                             st.info("No other similar contexts found in this cluster.")
-            else:
-                        st.info("No similar contexts found in this cluster.")
 
 if __name__ == "__main__":
     main()