Add heatmap button to Streamlit UI

app.py

@@ -1,13 +1,17 @@
 import streamlit as st
 from assist.chat import chat as chat_plugin
+from assist.heatmap import get_heatmap_figure
 
 st.title("RepoSage Chatbot Demo")
 
-# 1) Change the label to make it obvious we're asking a question
-question = st.text_input("Ask RepoSage a question:", "")
+# Tab or expander for heatmap
+with st.expander("🔢 View Token Similarity Heatmap"):
+    if st.button("Show Heatmap"):
+        fig = get_heatmap_figure("tensor.pt")
+        st.pyplot(fig)
 
-# 2) Only run when clicked
+# Chat UI
+question = st.text_input("Ask RepoSage a question:", "")
 if st.button("Ask RepoSage"):
-    # 3) Pass that question into your stub
     response = chat_plugin(question)
     st.write(response)

assist/heatmap.py

@@ -2,13 +2,14 @@ import torch
 import matplotlib.pyplot as plt
 from sklearn.metrics.pairwise import cosine_similarity
 
-def show_heatmap(tensor_path="tensor.pt"):
+def get_heatmap_figure(tensor_path="tensor.pt"):
     # Load embeddings
     weights = torch.load(tensor_path).detach().numpy()
     # Compute similarity
     sim = cosine_similarity(weights)
-    # Plot
-    plt.imshow(sim, cmap="viridis")
-    plt.colorbar()
-    plt.title("Token Similarity Heatmap")
-    plt.show()
+    # Build figure
+    fig, ax = plt.subplots()
+    cax = ax.imshow(sim, cmap="viridis")
+    fig.colorbar(cax, ax=ax)
+    ax.set_title("Token Similarity Heatmap")
+    return fig