Update app.py

app.py

@@ -6,9 +6,11 @@ import logging
 from dataclasses import dataclass
 from typing import Optional, Dict, List, Tuple
 
-# --- HIDE STREAMLIT MENU ---
+# --- HIDE STREAMLIT MENU / PAGE CONFIG ---
 st.set_page_config(
-    initial_sidebar_state="collapsed"
+    initial_sidebar_state="collapsed",
+    layout="wide",
+    page_title="LinkBERT by DEJAN AI"
 )
 
 hide_streamlit_style = """
@@ -86,8 +88,8 @@ def windowize_inference(
     total_tokens = len(full_encoding["input_ids"])
 
     if total_tokens == 0 and len(plain_text) > 0:
-         logger.warning("Tokenizer produced 0 tokens for a non-empty string.")
-         return []
+        logger.warning("Tokenizer produced 0 tokens for a non-empty string.")
+        return []
 
     while start_token_idx < total_tokens:
         end_token_idx = min(start_token_idx + cap, total_tokens)
@@ -168,8 +170,8 @@ def classify_text(
 
     for i, word_id in enumerate(word_ids):
         if word_id is not None and i < len(offsets):
-             start_char, end_char = offsets[i]
-             if start_char < end_char:
+            start_char, end_char = offsets[i]
+            if start_char < end_char:
                 current_token_max_prob = np.max(char_link_probabilities[start_char:end_char]) if start_char < len(char_link_probabilities) else 0.0
 
                 if word_id not in word_max_prob_map:
@@ -209,7 +211,8 @@ def classify_text(
             base_text_color = "#155724"
 
             html_parts.append(f"<span style='background-color: {base_bg_color}; color: {base_text_color}; "
-                              f"padding: 0.1em 0.2em; border-radius: 0.2em; opacity: {normalized_opacity:.2f};'>"
+                              f"padding: 0.1em 0.2em; border-radius: 0.2em; opacity: {normalized_opacity:.2f};' "
+                              f"title='Link Probability: {word_prob:.1%}'>"
                               f"{word_text}</span>")
         else:
             html_parts.append(word_text)
@@ -223,12 +226,11 @@ def classify_text(
 # ----------------------------------
 # Streamlit UI
 # ----------------------------------
-st.set_page_config(layout="wide", page_title="LinkBERT by DEJAN AI")
 st.title("LinkBERT")
 
 DEFAULT_THRESHOLD = 70.0
 THRESHOLD_STEP = 10.0
-THRESHOLD_BOUNDARY_PERCENT = 10.0
+THRESHOLD_BOUNDARY_PERCENT = 10.0  # Top/Bottom 10% for finer control
 
 if 'current_threshold' not in st.session_state:
     st.session_state.current_threshold = DEFAULT_THRESHOLD
@@ -258,24 +260,34 @@ def run_classification(new_threshold: float):
         st.warning("Please enter some text to classify.")
         st.session_state.output_html = ""
     else:
-        with st.spinner("Processing..."):
+        with st.spinner("Analyzing text..."):
             html, warning = classify_text(st.session_state.user_input, st.session_state.current_threshold)
             if warning: st.warning(warning)
             st.session_state.output_html = html
     st.rerun()
 
-if st.button("Classify Text", type="primary"):
+if st.button("Classify Text", type="primary", use_container_width=True):
     run_classification(slider_threshold)
 
 if st.session_state.output_html:
     st.markdown("---")
-    st.subheader(f"Results (Threshold: {st.session_state.current_threshold:.1f}%)")
     st.markdown(st.session_state.output_html, unsafe_allow_html=True)
+    st.markdown("---")
+
+    st.markdown(
+        f"<p style='text-align: center;'>Confidence Threshold: {st.session_state.current_threshold:.1f}%</p>",
+        unsafe_allow_html=True
+    )
 
     col1, col2, col3 = st.columns(3)
 
     with col1:
-        if st.button("Less", icon=":material/playlist_remove:", use_container_width=True, disabled=not st.session_state.output_html):
+        if st.button(
+            "Less",
+            icon=":material/playlist_remove:",
+            use_container_width=True,
+            help="Show fewer, more probable links"
+        ):
             current_thr = st.session_state.current_threshold
             if current_thr >= (100.0 - THRESHOLD_BOUNDARY_PERCENT):
                 new_threshold = current_thr + (100.0 - current_thr) / 2.0
@@ -284,14 +296,24 @@ if st.session_state.output_html:
             run_classification(min(100.0, new_threshold))
 
     with col2:
-        if st.button("Default", icon=":material/notes:", use_container_width=True, disabled=not st.session_state.output_html):
+        if st.button(
+            "Default",
+            icon=":material/notes:",
+            use_container_width=True,
+            help="Reset to default threshold (70%)"
+        ):
             run_classification(DEFAULT_THRESHOLD)
 
     with col3:
-        if st.button("More", icon=":material/docs_add_on:", use_container_width=True, disabled=not st.session_state.output_html):
+        if st.button(
+            "More",
+            icon=":material/docs_add_on:",
+            use_container_width=True,
+            help="Show more potential links"
+        ):
             current_thr = st.session_state.current_threshold
             if current_thr <= THRESHOLD_BOUNDARY_PERCENT:
                 new_threshold = current_thr / 2.0
             else:
                 new_threshold = current_thr - THRESHOLD_STEP
-            run_classification(max(0.0, new_threshold))
+            run_classification(max(0.0, new_threshold))