add buttons

src/streamlit_app.py

@@ -1,7 +1,8 @@
 """
-Streamlit app for interactive Jacobian and Temperature Scope visualizations.
+Streamlit app for interactive Semantic and Temperature Scope visualizations.
 """
 
+import base64
 import gc
 import os
 import sys
@@ -19,13 +20,33 @@ from torch import nn
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Add current directory to path for JCBScope_utils
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+_APP_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _APP_DIR)
 import JCBScope_utils
 
+DESIGN_DIR = os.path.join(_APP_DIR, "design")
+if not os.path.exists(DESIGN_DIR):
+    DESIGN_DIR = os.path.join(os.path.dirname(_APP_DIR), "design")
+
 # Device configuration: use CPU to match notebook and avoid device_map complexity
 device = torch.device("cpu")
 
 
+@st.cache_data
+def _load_svg(path: str) -> str | None:
+    """Load SVG file content; returns None if not found."""
+    if not os.path.exists(path):
+        return None
+    with open(path, encoding="utf-8") as f:
+        return f.read()
+
+
+def _render_svg_html(svg_content: str, max_width: int = 140) -> str:
+    """Return HTML to render SVG via base64 (reliable in Streamlit)."""
+    b64 = base64.b64encode(svg_content.encode("utf-8")).decode("utf-8")
+    return f'<img src="data:image/svg+xml;base64,{b64}" style="max-width:{max_width}px;height:auto;"/>'
+
+
 @st.cache_resource
 def load_model(model_name: str = "meta-llama/Llama-3.2-1B"):
     """Load and cache the tokenizer and model."""
@@ -347,13 +368,50 @@ def render_attribution_barplot(result, log_color: bool = False, cmap_name: str =
 
 
 def main():
-    st.set_page_config(page_title="Jacobian Scope Demo", page_icon="🔍", layout="centered")
-    st.title("🔍 Jacobian & Temperature Scopes")
+    st.set_page_config(page_title="Jacobian Scopes Demo", page_icon="🔍", layout="centered")
+    st.title("🔍 Jacobian Scopes Demo")
+    # Keep scope columns on one line (Streamlit stacks them below ~640px by default)
     st.markdown(
-        "**Semantic Scope** explains the predicted logit for a specific target token: enter your input "
-        "passage along with a target token.\n\n"
-        "**Temperature Scope** explains the overall predictive distribution and does not require a target."
+        '<style>div[data-testid="stHorizontalBlock"]{flex-wrap:nowrap!important}'
+        '[data-testid="column"]{min-width:120px!important}</style>',
+        unsafe_allow_html=True,
     )
+    scope_col1, scope_div1, scope_col2, scope_div2, scope_col3 = st.columns([1, 0.02, 1, 0.02, 1])
+    semantic_svg = _load_svg(os.path.join(DESIGN_DIR, "semantic_scope_button.svg"))
+    temp_svg = _load_svg(os.path.join(DESIGN_DIR, "temperature_scope_button.svg"))
+    fisher_svg = _load_svg(os.path.join(DESIGN_DIR, "fisher_scope_button.svg"))
+    with scope_col1:
+        if semantic_svg:
+            st.markdown(_render_svg_html(semantic_svg), unsafe_allow_html=True)
+        st.markdown(
+            "**Semantic Scope** — explains the predicted logit for a specific target token. "
+            "Enter your input passage along with a target token."
+        )
+    with scope_div1:
+        st.markdown(
+            '<div style="border-left: 5px solid #888; min-height: 200px; margin: 0;"></div>',
+            # '<div style="border-left: 5px solid steelblue; min-height: 160px; margin: 0;"></div>',
+            unsafe_allow_html=True,
+        )
+    with scope_col2:
+        if temp_svg:
+            st.markdown(_render_svg_html(temp_svg), unsafe_allow_html=True)
+        st.markdown(
+            "**Temperature Scope** — explains the confidence (effective inverse temperature) of the predictive distribution. "
+            "Target token not required."
+        )
+    with scope_div2:
+        st.markdown(
+            '<div style="border-left: 5px solid #888; min-height: 200px; margin: 0;"></div>',
+            unsafe_allow_html=True,
+        )
+    with scope_col3:
+        if fisher_svg:
+            st.markdown(_render_svg_html(fisher_svg), unsafe_allow_html=True)
+        st.markdown(
+            "**Fisher Scope** — A more refined attribution tool that explains the overall predictive distribution, motivated by information geometry. "
+            "Not shown in this demo due to limited compute."
+        )
 
     model_choice = st.selectbox(
         "Model",
@@ -374,6 +432,7 @@ def main():
         options=["Semantic Scope", "Temperature Scope"],
         index=0,
         horizontal=True,
+        key="attribution_type",
         help="Semantic Scope: attribute toward a target token. Temperature Scope: use hidden-state norm.",
     )
     mode = "Semantic" if attribution_type == "Semantic Scope" else "Temperature"