redesign app as a three-step interface

.streamlit/config.toml

@@ -1,2 +1,15 @@
 [theme]
-base="light"
+base = "light"
+primaryColor = "#3B4BD8"
+backgroundColor = "#FFFFFF"
+secondaryBackgroundColor = "#F3F2EE"
+textColor = "#1A1C24"
+font = "sans serif"
+
+[client]
+# Navigation is rendered in the page header (see ui.py), so hide the sidebar page list.
+showSidebarNavigation = false
+toolbarMode = "minimal"
+
+[server]
+maxUploadSize = 5

ProtHGT_app.py

@@ -1,723 +1,582 @@
+"""ProtHGT — guided three-step interface (add proteins → choose what to predict → review results)."""
 import os
-import streamlit as st
-from rapidfuzz import process
 import random
+import re
 
-# with st.spinner("Initializing the environment... This may take up to 10 minutes at the start of each session."):
-#     # Create a temporary placeholder for the message
-#     loading_placeholder = st.empty()
-
-#     # Show the info message only while the spinner is active
-#     loading_placeholder.info("""
-#     **Note:** This initialization is required at the start of each session.  
-#     Once the app is ready, you can run multiple predictions without re-initializing by clicking the **Reset** button in the sidebar.
-#     """)
-
-#     # Run setup script if not already executed
-#     if not os.path.exists(".setup_done"):
-#         start_time = time.time()
-#         os.system("bash setup.sh")
-#         end_time = time.time()
-#         print(f"Environment prepared in {end_time - start_time:.2f} seconds")
-#         with open(".setup_done", "w") as f:
-#             f.write("done")
-
-# # Remove the info message after initialization is complete
-# loading_placeholder.empty()
-
-from run_prothgt_app import *
-from visualize_kg import *
-
-def convert_df(df):
-   return df.to_csv(index=False).encode('utf-8')
-
-# Initialize session state variables
-if 'predictions_df' not in st.session_state:
-    st.session_state.predictions_df = None
-if 'heterodata' not in st.session_state:
-    st.session_state.heterodata = None
-if 'submitted' not in st.session_state:
-    st.session_state.submitted = False
-if 'previous_inputs' not in st.session_state:
-    st.session_state.previous_inputs = None
-if 'generating_predictions' not in st.session_state:
-    st.session_state.generating_predictions = False
-if 'protein_visualizations' not in st.session_state:
-    st.session_state.protein_visualizations = {}
-
-
-def reset_prediction_state():
-    st.session_state.generating_predictions = False
-    st.session_state.submitted = False
-    st.session_state.predictions_df = None
-    st.session_state.previous_inputs = None
-    # Clean up visualization files
-    if 'protein_visualizations' in st.session_state:
-        for viz_info in st.session_state.protein_visualizations.values():
-            try:
-                os.unlink(viz_info['path'])
-            except:
-                pass
-        st.session_state.protein_visualizations = {}
-
-def set_generating_predictions():
-    st.session_state.generating_predictions = True
-    st.session_state.submitted = True
-
-with st.expander("🚀 Upcoming Features"):
-    st.info("""
-    We are actively working on enhancing ProtHGT application with new capabilities:
-    
-    - **Real-time data retrieval for new proteins**: Currently, ProtHGT can only generate predictions for proteins that already exist in our knowledge graph. We are developing a new feature that will allow users to **predict functions for entirely new proteins starting from their sequences**. This will work by **retrieving relevant relationship data in real time from external source databases** (e.g., UniProt, STRING, and other biological repositories). The system will dynamically construct a knowledge graph for the query protein, incorporating its interactions, domains, pathways, and other biological associations before running function prediction. This approach will enable ProtHGT to analyze newly discovered or less-studied proteins even if they are not pre-annotated in our dataset.  
-    - **Expanded embedding options**: Currently, this application represents proteins using **TAPE embeddings**, which serve as the initial numerical representations of protein sequences before being processed in the heterogeneous graph model. We are working on integrating **ProtT5** and **ESM-2** as alternative initial embeddings, allowing users to choose different sequence representations that may enhance performance for specific tasks. A detailed comparison of how these embeddings influence function prediction accuracy will be included in our upcoming publication.
-
-    Stay tuned for updates and future publications!
-    """)
-    st.success("""
-    ✅ **Knowledge graph visualization for interpretability** – Now available! Explore the biological relationships behind ProtHGT's predictions interactively. For each query protein, you can inspect protein interactions, GO annotations, domains, pathways, and other key connections in a structured graphical format. Navigate to the **"View Knowledge Graphs"** tab after generating predictions to try it out.
-    """)
-
-with st.sidebar:
-
-    disabled = st.session_state.generating_predictions
-
-    st.markdown("""
-        <style>
-        .title {
-            font-size: 35px;
-            font-weight: bold;
-            color: #424242;
-            margin-bottom: 0px;
-        }
-        .subtitle {
-            font-size: 20px;
-            color: #424242;
-            margin-bottom: 20px;
-            line-height: 1.5;
-        }
-        .badges {
-            margin-top: 10px;
-            margin-bottom: 20px;
-        }
-        </style>
-        
-        <div class="title">ProtHGT</div>
-        <div class="subtitle">Heterogeneous Graph Transformers for Automated Protein Function Prediction Using Knowledge Graphs and Language Models</div>
-        <div class="badges">
-            <a href="https://github.com/HUBioDataLab/ProtHGT">
-                <img src="https://img.shields.io/badge/GitHub-black?logo=github" alt="github-repository">
-            </a>
-        </div>
-    """, unsafe_allow_html=True)
-
-    available_proteins = get_available_proteins()
-    
-    if 'example_proteins' not in st.session_state:
-        st.session_state.example_proteins = random.sample(available_proteins, 5)
-
-    selected_proteins = []
-    
-    # Add protein selection methods
-    selection_method = st.radio(
-        "Choose input method:",
-        ["Use example query", "Search proteins", "Upload protein ID file"],
-        disabled=disabled
-    )
+import streamlit as st
 
-    if selection_method == "Use example query":
-        selected_proteins = st.session_state.example_proteins
-        st.write(f"Selected proteins:")
-        st.markdown(
-            f"""
-            <div style="
-                height: 150px; 
-                overflow-y: scroll;
-                border: 1px solid #ccc;
-                border-radius: 4px;
-                padding: 8px;
-                margin-bottom: 16px;
-                background-color: white;">
-                {'<br>'.join(selected_proteins)}
-            </div>
-            """, 
-            unsafe_allow_html=True
-        )
+st.set_page_config(
+    page_title="ProtHGT · Protein function prediction",
+    page_icon="🧬",
+    layout="wide",
+    initial_sidebar_state="collapsed",
+)
+
+from rapidfuzz import fuzz, process
+
+from run_prothgt_app import KG_PATH, generate_prediction_df, load_kg, load_model, load_name_info
+from visualize_kg import visualize_protein_subgraph
+from ui import GITHUB_URL, KG_VIEW_GUIDE, STRICTNESS, render_header, render_stepper
+
+# ---------------------------------------------------------------------------
+# Settings
+# ---------------------------------------------------------------------------
+MAX_PROTEINS = 100
+MAX_KG_VIZ_PROTEINS = 10
+SEARCH_LIMIT = 50
+KG_DOWNLOAD_URL = "https://huggingface.co/datasets/HUBioDataLab/ProtHGT/blob/main/knowledge_graphs/prothgt-kg.pt"
+
+ASPECTS = {
+    "GO_term_F": ("Molecular function", "What the protein does at the molecular level, e.g. binding or catalysis."),
+    "GO_term_P": ("Biological process", "Which larger pathways or programs it takes part in."),
+    "GO_term_C": ("Cellular component", "Where in the cell the protein is located."),
+}
+ASPECT_SHORT = {"GO_term_F": "MF", "GO_term_P": "BP", "GO_term_C": "CC"}
+CATEGORY_TO_KEY = {"Molecular Function": "GO_term_F", "Biological Process": "GO_term_P", "Cellular Component": "GO_term_C"}
+
+
+# Widget values that must survive moving between steps (Streamlit drops the state of
+# widgets that are not rendered; re-assigning keeps it).
+PERSISTENT_KEYS = ["proteins", "aspect_GO_term_F", "aspect_GO_term_P", "aspect_GO_term_C",
+                   "strictness", "use_custom_threshold", "custom_threshold"]
+for _k in PERSISTENT_KEYS:
+    if _k in st.session_state:
+        st.session_state[_k] = st.session_state[_k]
+
+DEFAULTS = {
+    "step": 1,
+    "proteins": [],
+    "aspect_GO_term_F": True,
+    "aspect_GO_term_P": True,
+    "aspect_GO_term_C": True,
+    "strictness": "Balanced",
+    "use_custom_threshold": False,
+    "custom_threshold": 0.50,
+    "notice": None,
+    "results": None,
+    "results_cache": {},
+    "viz": {},
+}
+for _k, _v in DEFAULTS.items():
+    st.session_state.setdefault(_k, _v)
+
+
+# ---------------------------------------------------------------------------
+# Cached heavy resources (shared by every visitor of the Space, loaded once)
+# ---------------------------------------------------------------------------
+@st.cache_resource(show_spinner=False)
+def get_protein_names():
+    """UniProt ID -> protein name, for every protein in the knowledge graph."""
+    return load_name_info()["Protein"]
+
+
+@st.cache_resource(show_spinner=False)
+def get_search_index():
+    """(id, lower-cased name) pairs plus fuzzy-match choices, built once per process."""
+    names = get_protein_names()
+    lowered = [(pid, name.lower()) for pid, name in names.items()]
+    choices = {pid: f"{pid} {name}" for pid, name in names.items()}
+    return lowered, choices
+
+
+@st.cache_resource(show_spinner=False)
+def get_kg():
+    return load_kg()
+
+
+@st.cache_resource(show_spinner=False)
+def get_model(go_cat):
+    return load_model(go_cat, get_kg())
+
+
+@st.cache_data(show_spinner=False, max_entries=200)
+def search_proteins(query):
+    """Exact ID → ID prefix → name substring (closest names first) → fuzzy fallback."""
+    q = query.strip()
+    qu, ql = q.upper(), q.lower()
+    names = get_protein_names()
+    lowered, choices = get_search_index()
+
+    hits = [qu] if qu in names else []
+    if re.fullmatch(r"[A-Za-z0-9_]+", q):
+        hits += sorted(pid for pid, _ in lowered if pid.startswith(qu))[:SEARCH_LIMIT]
+    substr = [(name != ql, len(name), pid) for pid, name in lowered if ql in name]
+    hits += [pid for _, _, pid in sorted(substr)[:SEARCH_LIMIT]]
+    if len(hits) < SEARCH_LIMIT:
+        hits += [m[2] for m in process.extract(q, choices, scorer=fuzz.WRatio, limit=SEARCH_LIMIT)]
+    return list(dict.fromkeys(hits))[:SEARCH_LIMIT]
+
+
+# ---------------------------------------------------------------------------
+# State helpers / callbacks
+# ---------------------------------------------------------------------------
+def go_to(step):
+    st.session_state.step = step
+
+
+def selected_aspects():
+    return [k for k in ASPECTS if st.session_state[f"aspect_{k}"]]
+
+
+def current_threshold():
+    if st.session_state.use_custom_threshold:
+        return float(st.session_state.custom_threshold)
+    return STRICTNESS[st.session_state.strictness][0]
+
+
+def add_proteins(candidates, source):
+    names = get_protein_names()
+    candidates = list(dict.fromkeys(c.strip().upper() for c in candidates if c and c.strip()))
+    known = [c for c in candidates if c in names]
+    unknown = [c for c in candidates if c not in names]
+
+    current = list(st.session_state.proteins)
+    new = [p for p in known if p not in current]
+    room = MAX_PROTEINS - len(current)
+    dropped = new[room:] if room < len(new) else []
+    current.extend(new[:max(room, 0)])
+    st.session_state.proteins = current
+
+    st.session_state.notice = {
+        "source": source,
+        "added": len(new) - len(dropped),
+        "duplicates": len(known) - len(new),
+        "unknown": unknown,
+        "over_limit": dropped,
+    }
 
-    elif selection_method == "Search proteins":
-        # User enters search term
-        search_query = st.text_input(
-            "1\\. Start typing a protein ID (at least 3 characters) and press Enter to see search results in the dropdown menu below (2)",
-            "", 
-            disabled=disabled
-        )
 
-        # Initialize selected_proteins in session state if not exists
-        if 'selected_proteins_search' not in st.session_state:
-            st.session_state.selected_proteins_search = []
-
-        # Apply fuzzy search only if query length is >= 3
-        filtered_proteins = []
-        if len(search_query) >= 3:
-            # Case-insensitive search by converting query and proteins to lowercase
-            matches = process.extract(
-                search_query.upper(), 
-                {p: p.upper() for p in available_proteins}, 
-                limit=50
-            )
-            filtered_proteins = [match[0] for match in matches]  # Show top 50 matches
-
-        with st.container():
-            # Include previously selected proteins in options
-            all_options = list(set(filtered_proteins + st.session_state.selected_proteins_search))
-            
-            selected_proteins = st.multiselect(
-                "2\\. Select proteins from search results",
-                options=all_options,
-                default=st.session_state.selected_proteins_search,
-                placeholder="Start typing a protein ID above (1) to see search results...",
-                max_selections=100,
-                disabled=disabled,
-                key="protein_selector"
-            )
-            
-            # Update session state with current selection
-            st.session_state.selected_proteins_search = selected_proteins
-            
-            # Apply custom CSS to make container scrollable
-            st.markdown("""
-                <style>
-                div[data-testid="stMultiSelect"] div:nth-child(2) {
-                    max-height: 250px;
-                    overflow-y: auto;
-                }
-                </style>
-                """, unsafe_allow_html=True)
-            
-    else:  # Upload file option
-        uploaded_file = st.file_uploader(
-            "Upload a text file with UniProt IDs (one per line, max 100)*",
-            type=['txt'],
-            disabled=disabled
-        )
+def add_from_paste():
+    text = st.session_state.get("paste_text", "")
+    add_proteins(re.split(r"[\s,;]+", text), "pasted list")
 
-        if uploaded_file:
-            protein_list = [line.strip() for line in uploaded_file.read().decode('utf-8').splitlines()]
-
-            # Remove empty lines and duplicates
-            protein_list = list(filter(None, protein_list))
-            protein_list = list(dict.fromkeys(protein_list))
-            
-            # Check for proteins not in available_proteins
-            proteins_not_found = [p for p in protein_list if p not in available_proteins]
-            # Filter to keep only available proteins
-            protein_list = [p for p in protein_list if p in available_proteins]
-
-            if len(protein_list) > 100:
-                st.error("Please upload a file with maximum 100 protein IDs.")
-                selected_proteins = []
-            else:
-                selected_proteins = protein_list
-                st.write(f"Loaded {len(selected_proteins)} proteins")
-
-                if proteins_not_found:
-                    st.warning(f"""
-                    The following proteins were not found in our input knowledge graph and have been discarded: 
-                    """)
-                    with st.expander("View Discarded Proteins"):
-                        # Create scrollable container with fixed height
-                        st.markdown(
-                            f"""
-                            <div style="
-                                height: 150px; 
-                                overflow-y: scroll;
-                                border: 1px solid #ccc;
-                                border-radius: 4px;
-                                padding: 8px;
-                                margin-bottom: 16px;
-                                background-color: white;">
-                                {'<br>'.join(proteins_not_found)}
-                            </div>
-                            """, 
-                            unsafe_allow_html=True
-                        )
-                        
-                    st.warning(f"""
-                    Currently, our system can only generate predictions for proteins that are already included in our knowledge graph. **Real-time retrieval of relationship data from external source databases is not yet supported.**  
-                    We are actively working on integrating this capability in future updates. Stay tuned!
-                    """)
-
-    if selected_proteins:
-        st.write(f"Total proteins selected: {len(selected_proteins)}")
-        # Add download button
-        proteins_text = '\n'.join(selected_proteins)
-        st.download_button(
-            label="Download Selected Proteins List",
-            data=proteins_text,
-            file_name="selected_proteins.txt",
-            mime="text/plain",
-            key="download_selected_proteins"
-        )
 
-        # Add GO category selection
-        go_category_options = {
-            'All Categories': None,
-            'Molecular Function': 'GO_term_F',
-            'Biological Process': 'GO_term_P',
-            'Cellular Component': 'GO_term_C'
-        }
-        selected_go_category = st.selectbox(
-            "Select GO Category for predictions",
-            options=list(go_category_options.keys()),
-            help="Choose which GO category to generate predictions for. Selecting 'All Categories' will generate predictions for all three categories.",
-            disabled=disabled
-        )
+def add_from_search():
+    add_proteins(st.session_state.get("search_pick", []), "search")
+    st.session_state.search_pick = []
 
-        generation_threshold = st.number_input(
-            "Generation threshold (optional)",
-            min_value=0.0,
-            max_value=1.0,
-            value=0.0,
-            step=0.05,
-            help="If > 0, only predictions with Probability >= threshold are generated. This reduces output size and speeds up the app.",
-            disabled=disabled,
-        )
 
-    if selected_proteins and selected_go_category:
-        
-        button_disabled = st.session_state.submitted
-
-        # Add custom CSS for red button
-        st.markdown("""
-            <style>
-            div.stButton > button:first-child {
-                background-color: #ff4b4b;
-                color: white;
-            }
-            div.stButton > button:hover {
-                background-color: #ff0000;
-                color: white;
-            }
-            div.stButton > button:focus {
-                background-color: #ff0000;
-                color: white;
-            }
-            </style>
-        """, unsafe_allow_html=True)
-
-        if st.button("Generate Predictions", 
-                    disabled=button_disabled, 
-                    key="generate_predictions",
-                    on_click=set_generating_predictions):
+def add_from_file():
+    uploaded = st.session_state.get("upload_file")
+    if uploaded is None:
+        return
+    text = uploaded.getvalue().decode("utf-8", errors="ignore")
+    add_proteins(re.split(r"[\s,;]+", text), uploaded.name)
+
+
+def add_example():
+    pool = list(get_protein_names().keys())
+    add_proteins(random.sample(pool, 5), "example set")
+
+
+def clear_proteins():
+    st.session_state.proteins = []
+    st.session_state.notice = None
+
+
+def clear_visualizations():
+    for path in st.session_state.viz.values():
+        try:
+            os.unlink(path["path"])
+        except OSError:
             pass
+    st.session_state.viz = {}
+
+
+def start_new_prediction():
+    st.session_state.results = None
+    st.session_state.results_cache = {}
+    st.session_state.notice = None
+    clear_visualizations()
+    go_to(1)
+
+
+# ---------------------------------------------------------------------------
+# Page
+# ---------------------------------------------------------------------------
+render_header()
+
+step = st.session_state.step
+hero_slot = st.empty()  # always created so the elements below keep a stable position across steps
+if step < 3:
+    hero_slot.markdown(
+        '<div class="phgt-hero"><h1>Protein function prediction with knowledge graphs</h1>'
+        "<p>Gene Ontology term predictions from a heterogeneous graph transformer that combines "
+        "protein language model embeddings with a large-scale biological knowledge graph.</p></div>",
+        unsafe_allow_html=True,
+    )
+else:
+    hero_slot.markdown('<div class="phgt-hero-spacer"></div>', unsafe_allow_html=True)
+render_stepper(step)
 
-        # Create a tuple of current inputs to track changes
-        current_inputs = (tuple(selected_proteins), selected_go_category)
-        
-        # Check if inputs have changed
-        if st.session_state.previous_inputs != current_inputs:
-            st.session_state.predictions_df = None
-            st.session_state.submitted = False
-            st.session_state.previous_inputs = current_inputs
-            
-    st.warning("⚠️ Due to memory and computational constraints, the maximum number of proteins that can be processed at once is limited to 100 proteins. For larger datasets, please consider running the model locally using our [GitHub repository](https://github.com/HUBioDataLab/ProtHGT).")
-
-if st.session_state.submitted:
-    kg_path = "data/prothgt-kg.pt"
-    first_run = not os.path.exists(kg_path)
-
-    with st.spinner("Downloading dataset files for the first time, this may take a few minutes..." if first_run else "Generating predictions..."):
-        # Generate predictions only if not already in session state
-        if st.session_state.predictions_df is None:
-
-            # Load model config from JSON file
-            import json
-            import os
-
-            # Define data directory path
-            data_dir = "data"
-            models_dir = os.path.join(data_dir, "models")
-
-            # Load model configuration
-            model_config_paths = {
-                'GO_term_F': os.path.join(models_dir, "prothgt-config-molecular-function.yaml"),
-                'GO_term_P': os.path.join(models_dir, "prothgt-config-biological-process.yaml"),
-                'GO_term_C': os.path.join(models_dir, "prothgt-config-cellular-component.yaml")
-            }
-
-            # Paths for model and data
-            model_paths = {
-                'GO_term_F': os.path.join(models_dir, "prothgt-model-molecular-function.pt"),
-                'GO_term_P': os.path.join(models_dir, "prothgt-model-biological-process.pt"),
-                'GO_term_C': os.path.join(models_dir, "prothgt-model-cellular-component.pt")
-            }
-
-            # Get the selected GO category
-            go_category = go_category_options[selected_go_category]
-
-            # If a specific category is selected, use that model path
-            if go_category:
-                model_config_paths = [model_config_paths[go_category]]
-                model_paths = [model_paths[go_category]]
-                go_categories = [go_category]
-            else:
-                model_config_paths = [model_config_paths[cat] for cat in ['GO_term_F', 'GO_term_P', 'GO_term_C']]
-                model_paths = [model_paths[cat] for cat in ['GO_term_F', 'GO_term_P', 'GO_term_C']]
-                go_categories = ['GO_term_F', 'GO_term_P', 'GO_term_C']
-
-            # Generate predictions
-            heterodata, predictions_df = generate_prediction_df(
-                protein_ids=selected_proteins,
-                model_paths=model_paths,
-                model_config_paths=model_config_paths,
-                go_category=go_categories,
-                threshold=generation_threshold,
-            )
 
-            st.session_state.heterodata = heterodata
-            st.session_state.predictions_df = predictions_df
-
-            # Reset only the generating_predictions flag to release the sidebar
-            st.session_state.generating_predictions = False
-            st.rerun()
-
-        # Display and filter predictions
-        st.success("Predictions generated successfully!")
-
-        # tabs for predictions and visualizations
-        predictions_tab, kg_viz_tab = st.tabs(["View Predictions", "View Knowledge Graphs"])
-
-        with predictions_tab:
-            st.markdown("### Filter and View Predictions")
-            
-
-            # Create filters
-            col1, col2, col3, col4 = st.columns(4)
-            
-            with col1:
-                # Extract UniProt IDs from URLs for the selectbox
-                uniprot_ids = st.session_state.predictions_df['UniProt_ID'].unique().tolist()
-                
-                # Protein filter
-                selected_protein = st.selectbox(
-                    "Filter by Protein",
-                    options=['All'] + sorted(uniprot_ids)
-                )
+def protein_chips_html(proteins, limit=40):
+    shown = "".join(f'<span class="phgt-chip">{p}</span>' for p in proteins[:limit])
+    more = f'<span class="phgt-muted">+{len(proteins) - limit} more</span>' if len(proteins) > limit else ""
+    return f'<div class="phgt-chips">{shown}{more}</div>'
 
-            with col2:
-                # GO category filter
-                selected_category = st.selectbox(
-                    "Filter by GO Category",
-                    options=['All'] + sorted(st.session_state.predictions_df['GO_category'].unique().tolist())
-                )
 
-            with col3:
-                # GO term filter
-                go_term_filter = st.text_input(
-                    "Filter by GO Term ID",
-                    placeholder="e.g., GO:0003674",
-                    help="Enter a GO term ID to filter results"
-                ).strip()
-                
-            with col4:
-                # Probability threshold range slider
-                probability_range = st.slider(
-                    "Probability Range",
-                    min_value=0.0,
-                    max_value=1.0,
-                    value=(0.5, 1.0),  # (min, max) default values
-                    step=0.05
+def render_limit_aside():
+    n = len(st.session_state.proteins)
+    with st.container(border=True):
+        st.markdown("**Your proteins**")
+        if n:
+            st.markdown(protein_chips_html(st.session_state.proteins), unsafe_allow_html=True)
+        else:
+            st.markdown('<span class="phgt-muted">None added yet.</span>', unsafe_allow_html=True)
+        st.markdown(
+            f'<div class="phgt-meter"><div style="width:{n / MAX_PROTEINS * 100:.0f}%"></div></div>'
+            f'<span class="phgt-muted phgt-count">{n} of {MAX_PROTEINS} proteins per run</span>',
+            unsafe_allow_html=True,
+        )
+        if step == 2:
+            st.button("Edit proteins", on_click=go_to, args=(1,), use_container_width=True)
+    with st.container(border=True):
+        st.markdown("**Need more than 100?**")
+        st.markdown(
+            f'<span class="phgt-muted">Run ProtHGT locally with the code and pretrained models on '
+            f'<a href="{GITHUB_URL}" target="_blank">GitHub</a>.</span>',
+            unsafe_allow_html=True,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Step 1 — add proteins
+# ---------------------------------------------------------------------------
+def render_step_1():
+    main, aside = st.columns([2.3, 1], gap="large")
+    with main:
+        with st.container(border=True):
+            st.subheader("Which proteins do you want to annotate?")
+            st.caption("ProtHGT works on the ~326K UniProtKB/Swiss-Prot proteins in its knowledge graph. "
+                       "Up to 100 proteins per run.")
+
+            paste_tab, search_tab, upload_tab, example_tab = st.tabs(
+                ["Paste IDs", "Search", "Upload file", "Try an example"])
+
+            with paste_tab:
+                with st.form("paste_form", clear_on_submit=True, border=False):
+                    st.text_area(
+                        "UniProt accessions",
+                        key="paste_text",
+                        placeholder="P69905, P68871\nQ9Y6K9",
+                        help="Separate IDs with commas, spaces or new lines.",
+                        height=110,
+                    )
+                    st.form_submit_button("Add proteins", on_click=add_from_paste)
+
+            with search_tab:
+                query = st.text_input(
+                    "Search by UniProt ID or protein name",
+                    key="search_q",
+                    placeholder="e.g. hemoglobin, kinase, P699…",
                 )
-                min_probability_threshold, max_probability_threshold = probability_range
-
-            # Filter the dataframe using session state data
-            filtered_df = st.session_state.predictions_df.copy()
-
-            if selected_protein != 'All':
-                filtered_df = filtered_df[filtered_df['UniProt_ID'].str.contains(selected_protein)]
-                            
-            if selected_category != 'All':
-                filtered_df = filtered_df[filtered_df['GO_category'] == selected_category]
-
-            if go_term_filter:
-                filtered_df = filtered_df[filtered_df['GO_ID'] == go_term_filter]
-                
-            filtered_df = filtered_df[(filtered_df['Probability'] >= min_probability_threshold) & 
-                                    (filtered_df['Probability'] <= max_probability_threshold)]
-
-            filtered_df['UniProt_ID'] = [f"https://www.uniprot.org/uniprotkb/{pid}/entry" for pid in filtered_df['UniProt_ID']]
-            filtered_df['GO_ID'] = [f"https://www.ebi.ac.uk/QuickGO/term/{go_id}" for go_id in filtered_df['GO_ID']]
-
-            # Custom CSS to increase table width and improve layout
-            st.markdown("""
-                <style>
-                    .stDataFrame {
-                        width: 100%;
-                    }
-                    .stDataFrame > div {
-                        width: 100%;
-                    }
-                    .stDataFrame [data-testid="stDataFrameResizable"] {
-                        width: 100%;
-                        min-width: 100%;
-                    }
-                    .pagination-info {
-                        font-size: 14px;
-                        color: #666;
-                        padding: 10px 0;
-                    }
-                    .page-controls {
-                        display: flex;
-                        align-items: center;
-                        justify-content: center;
-                        gap: 20px;
-                        padding: 10px 0;
-                    }
-                </style>
-            """, unsafe_allow_html=True)
-
-            # Add pagination controls
-            col1, col2, col3 = st.columns([2, 1, 2])
-            with col2:
-                rows_per_page = st.selectbox("Rows per page", [50, 100, 200, 500], index=1)
-
-            total_rows = len(filtered_df)
-            total_pages = (total_rows + rows_per_page - 1) // rows_per_page
-
-            # Initialize page number in session state
-            if "page_number" not in st.session_state:
-                st.session_state.page_number = 0
-
-            # Calculate start and end indices for current page
-            start_idx = st.session_state.page_number * rows_per_page
-            end_idx = min(start_idx + rows_per_page, total_rows)
-
-            st.dataframe(
-                filtered_df.iloc[start_idx:end_idx],
-                hide_index=True,
-                use_container_width=True,
-                column_config={
-                    "UniProt_ID": st.column_config.LinkColumn(
-                        "UniProt ID",
-                        help="Click to view protein in UniProt",
-                        validate="^https://www\\.uniprot\\.org/uniprotkb/[A-Z0-9]+/entry$",
-                        display_text="^https://www\\.uniprot\\.org/uniprotkb/([A-Z0-9]+)/entry$"
-                    ),
-                    "GO_ID": st.column_config.LinkColumn(
-                        "GO ID",
-                        help="Click to view GO term in QuickGO",
-                        validate="^https://www\\.ebi\\.ac\\.uk/QuickGO/term/GO:[0-9]+$",
-                        display_text="^https://www\\.ebi\\.ac\\.uk/QuickGO/term/(GO:[0-9]+)$"
-                    ),
-                    "Probability": st.column_config.ProgressColumn(
-                        "Probability",
-                        format="%.2f",
-                        min_value=0,
-                        max_value=1,
-                    ),
-                    "Protein": st.column_config.TextColumn(
-                        "Protein",
-                        help="Protein Name",
-                    ),
-                    "GO_category": st.column_config.TextColumn(
-                        "GO Category",
-                        help="Gene Ontology Category",
-                    ),
-                    "GO_term": st.column_config.TextColumn(
-                        "GO Term",
-                        help="Gene Ontology Term Name",
-                    ),
-                }
-            )
-            # Pagination controls with better layout
-            col1, col2, col3 = st.columns([1, 3, 1])
-            with col1:
-                if st.button("Previous", disabled=st.session_state.page_number == 0):
-                    st.session_state.page_number -= 1
-                    st.rerun()
-
-            with col2:
-                st.markdown(f"""
-                    <div class="pagination-info" style="text-align: center">
-                        Page {st.session_state.page_number + 1} of {total_pages}<br>
-                        Showing rows {start_idx + 1} to {end_idx} of {total_rows}
-                    </div>
-                """, unsafe_allow_html=True)
-
-            with col3:
-                if st.button("Next", disabled=st.session_state.page_number >= total_pages - 1):
-                    st.session_state.page_number += 1
-                    st.rerun()
-
-            downloadable_df = filtered_df.copy()
-            downloadable_df['UniProt_ID'] = downloadable_df['UniProt_ID'].apply(
-                lambda x: x.split('/')[-2]  # Gets the ID part from the URL
-            )
-            downloadable_df['GO_ID'] = downloadable_df['GO_ID'].apply(
-                lambda x: x.split('/')[-1]  # Gets the ID part from the URL
+                if len(query.strip()) >= 3:
+                    names = get_protein_names()
+                    with st.spinner("Searching…"):
+                        hits = search_proteins(query)
+                    st.multiselect(
+                        f"Top {len(hits)} matches" if len(hits) == SEARCH_LIMIT else f"Matches ({len(hits)})",
+                        options=hits,
+                        key="search_pick",
+                        format_func=lambda p: f"{p} — {names.get(p, '')}",
+                        placeholder="Pick one or more matches",
+                    )
+                    st.button("Add selected", on_click=add_from_search, key="add_search",
+                              disabled=not st.session_state.get("search_pick"))
+                else:
+                    st.caption("Type at least 3 characters.")
+
+            with upload_tab:
+                st.file_uploader("Text file with one UniProt ID per line", type=["txt", "csv", "tsv"],
+                                 key="upload_file")
+                st.button("Add from file", on_click=add_from_file, key="add_file",
+                          disabled=st.session_state.get("upload_file") is None)
+
+            with example_tab:
+                st.markdown('<span class="phgt-muted">Load five random proteins from the knowledge graph '
+                            "to see how ProtHGT works.</span>", unsafe_allow_html=True)
+                st.button("Load example proteins", on_click=add_example, key="add_example")
+
+            notice = st.session_state.notice
+            if notice:
+                parts = []
+                if notice["added"]:
+                    parts.append(f"Added {notice['added']} protein{'s' if notice['added'] != 1 else ''} "
+                                 f"from {notice['source']}.")
+                if notice["duplicates"]:
+                    parts.append(f"{notice['duplicates']} already in your list.")
+                if parts:
+                    st.success(" ".join(parts))
+                elif not notice["unknown"]:
+                    st.info("No UniProt IDs found in the input.")
+                if notice["unknown"]:
+                    st.warning(
+                        f"{len(notice['unknown'])} ID(s) are not in the ProtHGT knowledge graph and were skipped. "
+                        "Predictions for proteins outside the graph are not supported yet."
+                    )
+                    with st.expander("Show skipped IDs"):
+                        st.code("\n".join(notice["unknown"]), language=None)
+                if notice["over_limit"]:
+                    st.error(f"{len(notice['over_limit'])} protein(s) were not added: the limit is "
+                             f"{MAX_PROTEINS} per run.")
+
+            st.divider()
+            st.multiselect(
+                f"Selected proteins ({len(st.session_state.proteins)}/{MAX_PROTEINS})",
+                options=st.session_state.proteins,
+                key="proteins",
+                placeholder="Add proteins using one of the tabs above",
+                help="Click × on a protein to remove it.",
             )
-            # Download filtered results
-            st.download_button(
-                label="Download Filtered Results",
-                data=convert_df(downloadable_df),
-                file_name="filtered_predictions.csv",
-                mime="text/csv",
-                key="download_filtered_predictions"
+            if st.session_state.proteins:
+                c1, c2, _ = st.columns([1, 1, 2])
+                c1.download_button("Download list", "\n".join(st.session_state.proteins),
+                                   file_name="selected_proteins.txt", mime="text/plain")
+                c2.button("Clear all", on_click=clear_proteins)
+
+        _, nav = st.columns([3, 1.3])
+        nav.button("Continue →", type="primary", use_container_width=True,
+                   disabled=not st.session_state.proteins, on_click=go_to, args=(2,))
+    with aside:
+        render_limit_aside()
+
+
+# ---------------------------------------------------------------------------
+# Step 2 — choose what to predict
+# ---------------------------------------------------------------------------
+def render_step_2():
+    if not st.session_state.proteins:
+        go_to(1)
+        st.rerun()
+
+    main, aside = st.columns([2.3, 1], gap="large")
+    with main:
+        with st.container(border=True):
+            st.subheader("Which functions are you interested in?")
+            cols = st.columns(3)
+            for col, (key, (label, desc)) in zip(cols, ASPECTS.items()):
+                with col:
+                    with st.container(border=True):
+                        st.checkbox(f"**{label}**", key=f"aspect_{key}")
+                        st.caption(desc)
+
+            st.subheader("How strict should predictions be?")
+            st.radio(
+                "Prediction strictness",
+                options=list(STRICTNESS),
+                key="strictness",
+                horizontal=True,
+                label_visibility="collapsed",
+                captions=[f"{d} · score ≥ {t:.2f}" for t, d in STRICTNESS.values()],
+                disabled=st.session_state.use_custom_threshold,
             )
-        
-        with kg_viz_tab:
-            st.markdown("### Knowledge Graph Visualization")
-
-            if not selected_proteins:
-                st.info("Please select proteins from the sidebar to visualize their knowledge graphs.")
-            elif len(selected_proteins) <= 10:
-                st.markdown("Visualize the knowledge graph for query proteins to explore the biological relationships that contributed to the predictions. The current tool supports up to second-degree connections. For deeper exploration beyond these, you can download the complete knowledge graph from [here](https://drive.google.com/file/d/18u1o2sm8YjMo9joFw4Ilwvg0-rUU0PXK/view?usp=share_link).")
-
-                protein_tabs = st.tabs([f"{protein_id}" for protein_id in selected_proteins])
-
-                # Create visualizations in each tab
-                for idx, protein_id in enumerate(selected_proteins):
-                    with protein_tabs[idx]:
-                        col1, col2 = st.columns([1, 1])
-                        with col1:
-                            max_node_count = st.slider(
-                                "Maximum neighbors per edge type (first-degree)",
-                                min_value=5,
-                                max_value=50,
-                                value=10,
-                                step=5,
-                                help="Control the maximum number of direct neighbors of the query protein shown for each relationship type",
-                                key=f"slider_{protein_id}"
-                            )
-                        with col2:
-                            second_degree_limit = st.slider(
-                                "Maximum neighbors per edge type (second-degree)",
-                                min_value=2,
-                                max_value=10,
-                                value=3,
-                                step=1,
-                                help="Control the maximum number of second-degree neighbors of the query protein shown for each relationship type. Second-degree edge limit is intentionally kept low to maintain visual clarity. Higher values may make the graph cluttered and difficult to interpret.",
-                                key=f"second_degree_slider_{protein_id}"
-                            )
-
-                        # Check if both visualizations exist for this protein
-                        viz_exists = (protein_id in st.session_state.protein_visualizations and 
-                                     'first_degree' in st.session_state.protein_visualizations[protein_id] and 
-                                     'second_degree' in st.session_state.protein_visualizations[protein_id])
-
-                        if not viz_exists:
-                            if st.button(f"Generate Visualization", key=f"viz_{protein_id}"):
-                                # Initialize the protein's visualizations if not exists
-                                if protein_id not in st.session_state.protein_visualizations:
-                                    st.session_state.protein_visualizations[protein_id] = {}
-
-                                # Generate both visualizations upfront
-                                # First degree only
-                                html_path_1st, edges_1st = visualize_protein_subgraph(
-                                    st.session_state.heterodata, 
-                                    protein_id, 
-                                    st.session_state.predictions_df, 
-                                    limit=max_node_count,
-                                    second_degree_limit=second_degree_limit,
-                                    include_second_degree=False
-                                )
-                                
-                                # With second degree
-                                html_path_2nd, edges_2nd = visualize_protein_subgraph(
-                                    st.session_state.heterodata, 
-                                    protein_id, 
-                                    st.session_state.predictions_df, 
-                                    limit=max_node_count,
-                                    second_degree_limit=second_degree_limit,
-                                    include_second_degree=True
-                                )
-                                
-                                # Store both visualizations in session state
-                                st.session_state.protein_visualizations[protein_id]['first_degree'] = {
-                                    'path': html_path_1st,
-                                    'edges': edges_1st,
-                                    'settings': {
-                                        'max_node_count': max_node_count,
-                                        'second_degree_limit': second_degree_limit
-                                    }
-                                }
-                                st.session_state.protein_visualizations[protein_id]['second_degree'] = {
-                                    'path': html_path_2nd,
-                                    'edges': edges_2nd,
-                                    'settings': {
-                                        'max_node_count': max_node_count,
-                                        'second_degree_limit': second_degree_limit
-                                    }
-                                }
-                                st.rerun()
-
-                        # If visualization exists, show the toggle and display appropriate version
-                        if viz_exists:
-                            # Check if settings have changed, but handle cases where settings don't exist
-                            current_settings = {
-                                'max_node_count': max_node_count,
-                                'second_degree_limit': second_degree_limit
-                            }
-                            
-                            # Safely get stored settings or use None if they don't exist
-                            stored_settings = (st.session_state.protein_visualizations[protein_id]['first_degree'].get('settings')
-                                             if 'first_degree' in st.session_state.protein_visualizations[protein_id] 
-                                             else None)
-                            
-                            include_second_degree = st.checkbox(
-                                "Include second-degree edges",
-                                value=False,
-                                key=f"second_degree_{protein_id}",
-                                help="Show connections between neighbor nodes"
-                            )
-
-                            # Get the appropriate visualization based on checkbox
-                            viz_type = 'second_degree' if include_second_degree else 'first_degree'
-                            viz_info = st.session_state.protein_visualizations[protein_id][viz_type]
-
-                            kg_viz_button_columns = st.columns([1, 1, 1])
-
-                            with kg_viz_button_columns[0]:
-                                # Format edges for download
-                                formatted_edges = {}
-                                for edge_type, edges in viz_info['edges'].items():
-                                    edge_type_str = f"{edge_type[0]}_{edge_type[1]}_{edge_type[2]}"
-                                    formatted_edges[edge_type_str] = [
-                                        {"source": edge[0][0], "target": edge[0][1], "probability": edge[1]}
-                                        for edge in edges
-                                    ]
-
-                                st.download_button(
-                                    label='Download Visualized Edges',
-                                    data=json.dumps(formatted_edges, indent=2),
-                                    file_name=f'{protein_id}_visualized_edges{"_with_2nd_degree" if include_second_degree else ""}.json',
-                                    mime='application/json'
-                                )
-
-                            with kg_viz_button_columns[1]:
-                                if st.button("Regenerate Visualization", key=f"regenerate_{protein_id}"):
-                                    # Clean up old files
-                                    if protein_id in st.session_state.protein_visualizations:
-                                        for viz_type in ['first_degree', 'second_degree']:
-                                            if viz_type in st.session_state.protein_visualizations[protein_id]:
-                                                try:
-                                                    old_path = st.session_state.protein_visualizations[protein_id][viz_type]['path']
-                                                    os.unlink(old_path)
-                                                except:
-                                                    pass
-                                        # Remove from session state
-                                        del st.session_state.protein_visualizations[protein_id]
-                                    st.rerun()
-
-                            if stored_settings is not None and current_settings != stored_settings:
-                                st.warning("⚠️ Settings have changed. Click 'Regenerate Visualization' to apply new settings.")
-                            elif stored_settings is None:
-                                st.warning("⚠️ Visualization was generated with default settings. Consider regenerating to apply custom settings.")
-
-                            # Display the appropriate visualization
-                            with open(viz_info['path'], 'r', encoding='utf-8') as f:
-                                html_content = f.read()
-
-                            st.components.v1.html(html_content, height=1200)
-
-            else:
-                st.warning("Knowledge graph visualization is only available when 10 or fewer proteins are selected.")
+            with st.expander("Advanced: set a custom score threshold",
+                             expanded=st.session_state.use_custom_threshold):
+                st.checkbox("Use a custom threshold", key="use_custom_threshold")
+                st.slider("Minimum score", 0.0, 1.0, step=0.05, key="custom_threshold",
+                          disabled=not st.session_state.use_custom_threshold,
+                          help="Only GO terms scoring at or above this value are returned. "
+                               "Lower values return many more terms and take longer.")
+                if st.session_state.use_custom_threshold and st.session_state.custom_threshold < 0.1:
+                    st.caption("A very low threshold returns almost every GO term for each protein; "
+                               "expect a large table.")
+
+        aspects = selected_aspects()
+        back, _, run = st.columns([1, 1.7, 1.3])
+        back.button("← Back", use_container_width=True, on_click=go_to, args=(1,))
+        run.button("Run prediction →", type="primary", use_container_width=True,
+                   disabled=not aspects, on_click=go_to, args=(3,))
+        if not aspects:
+            st.caption("Select at least one function category.")
+    with aside:
+        render_limit_aside()
+
+
+# ---------------------------------------------------------------------------
+# Step 3 — run & review
+# ---------------------------------------------------------------------------
+def run_prediction(proteins, aspects, threshold):
+    with st.status("Running ProtHGT…", expanded=True) as status:
+        if os.path.exists(KG_PATH):
+            st.write("Loading the knowledge graph…")
+        else:
+            st.write("Downloading dataset files for the first time, this may take a few minutes…")
+        heterodata = get_kg()
+        progress = st.progress(0.0)
+
+        def on_category(i, n, go_cat):
+            progress.progress(i / n, text=f"Scoring {ASPECTS[go_cat][0].lower()} terms ({i + 1}/{n})")
+
+        models = {cat: get_model(cat) for cat in aspects}
+        _, df = generate_prediction_df(
+            protein_ids=proteins,
+            go_category=aspects,
+            threshold=threshold,
+            heterodata=heterodata,
+            models=models,
+            progress_callback=on_category,
+        )
+        progress.progress(1.0, text="Done")
+        status.update(label=f"Predicted {len(df):,} protein–GO term associations", state="complete",
+                      expanded=False)
+    return df.sort_values("Probability", ascending=False, ignore_index=True)
+
+
+@st.fragment
+def render_predictions(df, threshold):
+    if df.empty:
+        st.info("No GO terms passed the score threshold. Try a lower strictness in step 2.")
+        return
+
+    f1, f2, f3 = st.columns([1.2, 1.2, 1.6])
+    proteins = ["All proteins"] + sorted(df["UniProt_ID"].unique())
+    protein = f1.selectbox("Protein", proteins, format_func=lambda p: p if p == "All proteins"
+                           else f"{p} — {df.loc[df.UniProt_ID == p, 'Protein'].iat[0]}")
+    available = [c for c in CATEGORY_TO_KEY if c in set(df["GO_category"])]
+    aspects = f2.pills("Aspect", available, selection_mode="multi", default=available,
+                       format_func=lambda c: ASPECT_SHORT[CATEGORY_TO_KEY[c]])
+    text = f3.text_input("Search GO ID or term", placeholder="e.g. GO:0005524 or binding")
+    min_score = st.slider("Minimum score", float(threshold), 1.0, float(threshold), step=0.01)
+
+    view = df
+    if protein != "All proteins":
+        view = view[view["UniProt_ID"] == protein]
+    view = view[view["GO_category"].isin(aspects or [])]
+    if text.strip():
+        t = text.strip()
+        view = view[view["GO_ID"].str.contains(t, case=False, regex=False)
+                    | view["GO_term"].str.contains(t, case=False, regex=False)]
+    view = view[view["Probability"] >= min_score]
+
+    display = view.assign(
+        UniProt_ID="https://www.uniprot.org/uniprotkb/" + view["UniProt_ID"] + "/entry",
+        GO_ID="https://www.ebi.ac.uk/QuickGO/term/" + view["GO_ID"],
+        GO_category=view["GO_category"].map(lambda c: ASPECT_SHORT[CATEGORY_TO_KEY[c]]),  # MF/BP/CC; CSV keeps full names
+    )
+    # Fixed widths for the long text columns so the table fits the page instead of scrolling sideways;
+    # truncated cells show their full text on hover. Short columns keep their automatic width.
+    st.dataframe(
+        display,
+        hide_index=True,
+        use_container_width=True,
+        height=520,
+        column_config={
+            "UniProt_ID": st.column_config.LinkColumn("UniProt ID", display_text=r"uniprotkb/([A-Z0-9]+)/entry"),
+            "Protein": st.column_config.TextColumn("Protein", width="large"),
+            "GO_ID": st.column_config.LinkColumn("GO ID", display_text=r"term/(GO:[0-9]+)"),
+            "GO_term": st.column_config.TextColumn("GO term", width="medium"),
+            "GO_category": st.column_config.TextColumn(
+                "Aspect", help="MF: molecular function · BP: biological process · CC: cellular component"),
+            # Plain number: scores cluster near 1, so a bar adds little; 3 decimals keep close scores apart
+            "Probability": st.column_config.NumberColumn("Score", format="%.3f", width="small"),
+        },
+    )
+    c1, c2 = st.columns([3, 1.3], vertical_alignment="center")
+    c1.caption(f"Showing {len(view):,} of {len(df):,} predictions")
+    c2.download_button("Download CSV", view.to_csv(index=False).encode("utf-8"),
+                       file_name="prothgt_predictions.csv", mime="text/csv",
+                       use_container_width=True)
+
+
+@st.fragment
+def render_knowledge_graph(df, proteins, aspects):
+    if len(proteins) > MAX_KG_VIZ_PROTEINS:
+        st.info(f"Knowledge graph views are available for runs of up to {MAX_KG_VIZ_PROTEINS} proteins.")
+        return
+
+    st.markdown('<span class="phgt-muted">Explore the biological relationships around a protein that '
+                "the model draws on: GO annotations, domains, interactions, pathways and more. "
+                "Views go up to second-degree connections; for deeper exploration, "
+                f'<a href="{KG_DOWNLOAD_URL}" target="_blank">download the complete knowledge graph</a>.</span>',
+                unsafe_allow_html=True)
+    with st.expander("How to read this graph"):
+        st.markdown(KG_VIEW_GUIDE)
+    c1, c2, c3, c4 = st.columns([1.3, 1, 1, 1])
+    protein = c1.selectbox("Protein", proteins, key="viz_protein")
+    pred_limit = c2.slider(
+        "Predicted GO terms", 1, 20, 5,
+        help="Highest-scoring predictions shown per predicted GO category. Purple: also annotated "
+             "in the knowledge graph; red: novel. Known annotations that were predicted at a lower rank "
+             "are also shown in purple.")
+    limit = c3.slider(
+        "Neighbours per edge type", 5, 50, 10, step=5,
+        help="Maximum direct neighbours shown for each relationship type, including known GO "
+             "annotations outside the top predictions (purple if predicted, blue if not).")
+    second_limit = c4.slider(
+        "Second-degree neighbours", 0, 10, 0,
+        help="Second-degree neighbours are the neighbours of the protein's direct neighbours, "
+             "e.g. other proteins sharing a domain or pathway with it. 0 shows direct neighbours only; "
+             "any higher value adds up to that many per edge type. Large values quickly make the "
+             "graph dense and hard to read.")
+    include_second = second_limit > 0
+
+    viz_key = f"{protein}|{pred_limit}|{limit}|{second_limit}"
+    viz = st.session_state.viz.get(viz_key)
+    if viz is None:
+        if st.button("Show graph", type="primary", key="viz_go"):
+            with st.spinner("Building subgraph…"):
+                path, edges = visualize_protein_subgraph(
+                    get_kg(), protein, df, limit=limit, prediction_limit=pred_limit, predicted_categories=aspects,
+                    second_degree_limit=second_limit, include_second_degree=include_second)
+            st.session_state.viz[viz_key] = {"path": path, "edges": edges}
+            st.rerun(scope="fragment")
+        return
+
+    with open(viz["path"], "r", encoding="utf-8") as f:
+        # 760px graph with both legends inside it; PNG and edges-JSON downloads are in its toolbox
+        # (kg_graph_template.html)
+        st.components.v1.html(f.read(), height=775, scrolling=True)
+
+
+def render_step_3():
+    proteins = list(st.session_state.proteins)
+    aspects = selected_aspects()
+    if not proteins or not aspects:
+        go_to(1 if not proteins else 2)
+        st.rerun()
+
+    threshold = current_threshold()
+    run_key = (tuple(proteins), tuple(aspects), threshold)
+    results = st.session_state.results
+    if results is None or results["key"] != run_key:
+        clear_visualizations()
+        cache = st.session_state.results_cache
+        if run_key in cache:  # switching back to settings already run in this session is instant
+            df = cache[run_key]
+        else:
+            try:
+                df = run_prediction(proteins, aspects, threshold)
+            except Exception as exc:  # surface failures instead of a blank page
+                st.error(f"Prediction failed: {exc}")
+                st.button("← Back to settings", on_click=go_to, args=(2,))
+                return
+            cache[run_key] = df
+            while len(cache) > 3:  # keep memory per visitor bounded
+                cache.pop(next(iter(cache)))
+        st.session_state.results = {"key": run_key, "df": df}
+        results = st.session_state.results
+    df = results["df"]
+
+    head, b1, b2 = st.columns([3, 1.1, 1.1], vertical_alignment="bottom")
+    head.subheader("Results")
+    b1.button("Change settings", on_click=go_to, args=(2,), use_container_width=True)
+    b2.button("New prediction", on_click=start_new_prediction, use_container_width=True)
+
+    m1, m2, m3, m4 = st.columns(4)
+    m1.metric("Proteins", len(proteins))
+    m2.metric("Predicted GO terms", f"{len(df):,}")
+    m3.metric("Aspects", " · ".join(ASPECT_SHORT[a] for a in aspects))
+    m4.metric("Score threshold", f"{threshold:.2f}")
+
+    pred_tab, kg_tab = st.tabs(["Predictions", "Knowledge graph"])
+    with pred_tab:
+        render_predictions(df, threshold)
+    with kg_tab:
+        render_knowledge_graph(df, proteins, aspects)
+
+
+# Render the step inside one st.empty() slot so the previous step's widgets are replaced
+# immediately instead of lingering (greyed out) while a long prediction runs.
+with st.empty().container():
+    {1: render_step_1, 2: render_step_2, 3: render_step_3}[step]()

pages/User_Guide.py

@@ -1,257 +1,68 @@
 import streamlit as st
 
-st.sidebar.markdown('''
-# Sections
-- [How to use](#how-to-use)
-- [Knowledge Graph Visualization](#knowledge-graph-visualization)
-- [Running Locally](#running-locally)
-''', unsafe_allow_html=True)
+st.set_page_config(page_title="User guide · ProtHGT", page_icon="🧬", layout="wide", initial_sidebar_state="collapsed")
 
-st.markdown('''
-# ProtHGT User Guide
-''')
+from ui import KG_VIEW_GUIDE, STRICTNESS, render_header
+render_header()
 
+st.title("User guide")
 st.markdown("""
-ProtHGT is a web-based tool for **automated protein function prediction** using heterogeneous graph transformers and knowledge graphs. Follow the steps below to generate predictions for your proteins.
+ProtHGT predicts **Gene Ontology (GO) terms** for proteins using a heterogeneous graph transformer
+trained on a biological knowledge graph. A prediction takes three steps.
 """)
 
-st.subheader("1. Select Proteins")
+st.subheader("1. Add proteins")
 st.markdown("""
-In the **sidebar**, choose how to input your proteins:
-- **Use example query**: Loads a random set of 5 proteins from the dataset to quickly explore the tool.
-- **Search Proteins**: Select or search UniProt IDs from the available dataset.
-- **Upload a File**: Upload a text file (.txt) containing UniProt IDs (one per line, max 100).    
-""")
-st.warning("⚠️ Currently, our system can only generate predictions for proteins that are already included in our knowledge graph. Real-time retrieval of relationship data from external source databases is not yet supported. We are actively working on integrating this capability in future updates. Stay tuned!")
-st.info("📥 Selected proteins can be downloaded as a txt file.")
-
-
-st.subheader("2. Choose Gene Ontology (GO) Category")
-st.markdown("""
-Select which **Gene Ontology (GO) sub-ontology** to use for function prediction:
-- **Molecular Function (MF)** – Biochemical activity of the protein
-- **Biological Process (BP)** – Biological roles and pathways
-- **Cellular Component (CC)** – Location within the cell
-- **All Categories** – Runs predictions for all three categories
-""")
-
-st.subheader("3. Set Generation Threshold (Optional)")
-st.markdown("""
-Use the **Generation threshold** slider to filter predictions at the point of generation. Only predictions with a probability equal to or above the threshold will be included in the output.
-
-Setting a threshold above 0 can significantly **reduce output size** and **speed up the app**, especially when processing many proteins or all GO categories.
-""")
-
-st.subheader("4. Generate Predictions")
-st.markdown("""
-Click **"Generate Predictions"** to start the analysis. The model will process the selected proteins and return predicted functional annotations.
-
-🔄 **Processing time**: A few minutes (depending on input size). On the **first run of a new session**, the knowledge graph file (~1 GB) will be downloaded automatically — this may take a few additional minutes.
-""")
-
-st.subheader("5. View and Filter Results")
-st.markdown("""
-Once predictions are generated, use the filter options to refine the output:
-- **Filter by Protein** (UniProt ID)
-- **Filter by GO Category**
-- **Filter by GO Term ID** – Enter a specific GO term (e.g., GO:0003674) to search for it directly
-- **Set Probability Range** – Adjust prediction confidence thresholds
-
-Results are displayed in a sortable table, with **probabilities** indicating prediction confidence. UniProt IDs and GO IDs are clickable links that open the corresponding entry in UniProt and QuickGO, respectively.
-""")
-
-st.info("📥 Filtered predictions can be downloaded as a CSV file.")
+Add up to **100 UniProt accessions** per run, in any combination of:
+- **Paste IDs** – type or paste accessions separated by commas, spaces or new lines.
+- **Search** – find proteins by UniProt ID or protein name and add the matches you want.
+- **Upload file** – a .txt/.csv file with one accession per line.
+- **Try an example** – five random proteins from the knowledge graph.
 
-st.subheader("6. Start a New Query")
-st.markdown("""
-After generating predictions, you can start a new query by selecting different options from the sidebar. Changing the protein selection or GO category will automatically reset the results.
+Remove a protein by clicking × on its chip. The list can be downloaded as a text file.
 """)
+st.info("ProtHGT can currently only predict for proteins that are part of its knowledge graph "
+        "(~326K UniProtKB/Swiss-Prot entries). IDs outside the graph are listed and skipped.")
 
-st.divider()
-
-st.markdown("""## Knowledge Graph Visualization <span style="background-color: #28a745; color: white; font-size: 14px; font-weight: bold; padding: 3px 10px; border-radius: 12px; vertical-align: middle;">NEW</span>""", unsafe_allow_html=True)
-
+st.subheader("2. Choose what to predict")
 st.markdown("""
-After generating predictions, switch to the **"View Knowledge Graphs"** tab to explore the biological context behind the model's predictions. This feature is available when **10 or fewer proteins** are selected.
+- **Function categories** – any of *Molecular function* (biochemical activity), *Biological process*
+  (pathways and programs) and *Cellular component* (location in the cell).
+- **Strictness** – sets the minimum score a GO term needs to be returned. Each preset keeps the
+  terms whose predicted score falls in the range below:
 """)
 
-st.subheader("Generating a Visualization")
+_presets = list(STRICTNESS.items())
+_lines = ["| Preset | Score range returned | Added on top of the stricter preset | Use it for |",
+          "|---|---|---|---|"]
+for i, (name, (low, desc)) in enumerate(_presets):
+    added = "—" if i == 0 else f"{low:.2f} – {_presets[i - 1][1][0]:.2f}"
+    _lines.append(f"| **{name}** | {low:.2f} – 1.00 | {added} | {desc} |")
+st.markdown("\n".join(_lines))
 st.markdown("""
-Each query protein gets its own subtab. For each protein:
-
-1. Use the **"Maximum neighbors per edge type (first-degree)"** slider to control how many direct neighbors are shown per relationship type. Higher values show a denser graph.
-2. Use the **"Maximum neighbors per edge type (second-degree)"** slider to control neighbors-of-neighbors when second-degree edges are enabled. This is intentionally kept low (2–10) to avoid cluttered graphs.
-3. Click **"Generate Visualization"** to render the graph.
+Presets are cumulative: *Balanced* includes everything *Precise* returns plus terms scoring between
+the two cut-offs, and *Exploratory* adds the next band down. Under *Advanced* you can set any
+threshold between 0 and 1 instead; on the results page you can still raise the minimum score further,
+but not below the threshold the run used.
 """)
 
-st.subheader("What the Graph Shows")
-st.markdown("""
-The visualization renders a **subgraph of the full ProtHGT knowledge graph**, centered on the query protein. Edges and nodes are selected as follows:
-
-- **First-degree edges**: All relationship types that directly connect the query protein to other nodes in the knowledge graph are included (e.g., protein-domain, PPI, GO term annotations). For GO term edges, nodes are ranked by predicted probability and the top-N are shown based on the slider setting.
-- **Second-degree edges**: When enabled via the **"Include second-degree edges"** checkbox, neighbor nodes are also expanded — their own connections (excluding edges back to the query protein) are added to the graph, again limited by the second-degree slider.
-""")
-
-st.subheader("Reading the Graph")
-
-st.components.v1.html("""
-<style>
-    .kg-legend {
-        margin-top: 10px;
-        margin-bottom: 10px;
-        padding: 20px;
-        border: 1px solid #ddd;
-        border-radius: 5px;
-        font-family: Arial, sans-serif;
-        display: flex;
-        gap: 20px;
-    }
-    .legend-section-nodes {
-        flex: 2;
-    }
-    .legend-section-edges {
-        flex: 1;
-    }
-    .legend-title {
-        margin-bottom: 15px;
-        color: #333;
-        font-size: 16px;
-        font-weight: bold;
-    }
-    .nodes-grid {
-        display: grid;
-        grid-template-columns: repeat(2, 1fr);
-        gap: 12px;
-    }
-    .edges-grid {
-        display: grid;
-        grid-template-columns: 1fr;
-        gap: 12px;
-    }
-    .legend-item {
-        display: flex;
-        align-items: center;
-        padding: 4px;
-    }
-    .node-indicator {
-        width: 15px;
-        height: 15px;
-        border-radius: 50%;
-        margin-right: 10px;
-        flex-shrink: 0;
-    }
-    .edge-indicator {
-        width: 40px;
-        height: 3px;
-        margin-right: 10px;
-        flex-shrink: 0;
-    }
-    .legend-label {
-        font-size: 14px;
-    }
-</style>
-<div class="kg-legend">
-    <div class="legend-section-nodes">
-        <div class="legend-title">Node Types</div>
-        <div class="nodes-grid">
-            <div class="legend-item">
-                <div class="node-indicator" style="background-color: #079dbb;"></div>
-                <span class="legend-label">Disease</span>
-            </div>
-            <div class="legend-item">
-                <div class="node-indicator" style="background-color: #58d0e8;"></div>
-                <span class="legend-label">Phenotype</span>
-            </div>
-            <div class="legend-item">
-                <div class="node-indicator" style="background-color: #815ac0;"></div>
-                <span class="legend-label">Drug</span>
-            </div>
-            <div class="legend-item">
-                <div class="node-indicator" style="background-color: #d2b7e5;"></div>
-                <span class="legend-label">Compound</span>
-            </div>
-            <div class="legend-item">
-                <div class="node-indicator" style="background-color: #6bbf59;"></div>
-                <span class="legend-label">Domain</span>
-            </div>
-            <div class="legend-item">
-                <div class="node-indicator" style="background-color: #ff8800;"></div>
-                <span class="legend-label">Biological Process</span>
-            </div>
-            <div class="legend-item">
-                <div class="node-indicator" style="background-color: #ffaa00;"></div>
-                <span class="legend-label">Molecular Function</span>
-            </div>
-            <div class="legend-item">
-                <div class="node-indicator" style="background-color: #ffc300;"></div>
-                <span class="legend-label">Cellular Component</span>
-            </div>
-            <div class="legend-item">
-                <div class="node-indicator" style="background-color: #720026;"></div>
-                <span class="legend-label">Pathway</span>
-            </div>
-            <div class="legend-item">
-                <div class="node-indicator" style="background-color: #ce4257;"></div>
-                <span class="legend-label">EC Number</span>
-            </div>
-            <div class="legend-item">
-                <div class="node-indicator" style="background-color: #3aa6a4;"></div>
-                <span class="legend-label">Protein</span>
-            </div>
-        </div>
-    </div>
-    <div class="legend-section-edges">
-        <div class="legend-title">Edge Colors</div>
-        <div class="edges-grid">
-            <div class="legend-item">
-                <div class="edge-indicator" style="background-color: #8338ec;"></div>
-                <span class="legend-label">Confirmed Prediction (Found in Ground Truth)</span>
-            </div>
-            <div class="legend-item">
-                <div class="edge-indicator" style="background-color: #c1121f;"></div>
-                <span class="legend-label">Novel Prediction (Not in Ground Truth)</span>
-            </div>
-            <div class="legend-item">
-                <div class="edge-indicator" style="background-color: #219ebc;"></div>
-                <span class="legend-label">Existing GO Term Annotation</span>
-            </div>
-            <div class="legend-item">
-                <div class="edge-indicator" style="background-color: #666666;"></div>
-                <span class="legend-label">Other Relationships</span>
-            </div>
-        </div>
-    </div>
-</div>
-""", height=320)
-
+st.subheader("3. Review results")
 st.markdown("""
-**Node colors** indicate the biological entity type of each node in the graph. The **query protein** is always shown as a larger node with a white background and red border, making it easy to identify at a glance.
+- **Predictions** – a sortable table with links to UniProt and QuickGO. Filter by protein, aspect,
+  GO ID / term text and minimum score, then download the filtered table as CSV.
+- **Knowledge graph** – for runs of up to 10 proteins, view the subgraph around a protein (with optional
+  second-degree edges) and download the shown edges as JSON.
 
-**Edge colors** indicate the nature of each connection:
-- **Purple** – The predicted GO term was also found in the ground truth annotations, confirming the model's prediction.
-- **Red** – The predicted GO term was not in the ground truth — this is a novel prediction made by the model.
-- **Blue** – An existing GO term annotation already present in the knowledge graph, included for context but not generated as a new prediction.
-- **Gray** – Other biological relationships such as protein-protein interactions, protein-domain associations, drug-target links, and pathway memberships.
+Use **Change settings** to rerun with different categories or strictness, or **New prediction** to start over.
 
-Hovering over any node or edge reveals a **tooltip** with the entity name, type, and prediction probability where applicable. Nodes and GO IDs are clickable links to their respective database entries (UniProt, QuickGO, InterPro, etc.).
+The first run after the Space starts downloads the knowledge graph (~1 GB), which can take a few extra minutes.
 """)
 
-st.subheader("Controls and Downloads")
-st.markdown("""
-- **Include second-degree edges** checkbox – Toggles between first-degree-only and expanded graph views.
-- **Regenerate Visualization** – Re-renders the graph with updated slider settings. Use this after changing the neighbor limits.
-- **Download Visualized Edges** – Downloads a JSON file containing all edges shown in the current visualization, including source/target node IDs and prediction probabilities.
-""")
-
-st.info("📥 For deeper exploration beyond second-degree connections, the complete knowledge graph can be downloaded from the link provided in the visualization tab.")
-
-st.divider()
+st.subheader("Knowledge graph view")
+st.markdown(KG_VIEW_GUIDE)
 
-st.markdown("## Running Locally")
+st.subheader("Running locally")
 st.markdown("""
-For **larger datasets** or **custom analyses**, you can run ProtHGT locally using our [**GitHub repository**](https://github.com/HUBioDataLab/ProtHGT).
-""")
-
-st.markdown("""
-
+For more than 100 proteins or custom analyses, run ProtHGT locally using the
+[GitHub repository](https://github.com/HUBioDataLab/ProtHGT).
 """)

ui.py

@@ -0,0 +1,164 @@
+"""Shared look-and-feel for the ProtHGT Streamlit pages (header, CSS, stepper)."""
+import streamlit as st
+
+ACCENT = "#3B4BD8"
+GITHUB_URL = "https://github.com/HUBioDataLab/ProtHGT"
+PAPER_URL = "https://www.biorxiv.org/content/10.1101/2025.04.19.649272"
+
+# Minimum score behind each strictness preset (used by the app and the user guide).
+# Replace "Balanced" with your validation-set Fmax threshold(s) to match the paper exactly.
+STRICTNESS = {
+    "Precise": (0.70, "Fewer, high-confidence terms"),
+    "Balanced": (0.50, "Recommended default"),
+    "Exploratory": (0.30, "More candidates to screen"),
+}
+
+# How to read the knowledge graph view. Shown in the User guide and in the "How to read this graph"
+# expander of the Knowledge graph tab, so the two stay in sync.
+KG_VIEW_GUIDE = """
+The view shows the part of the ProtHGT knowledge graph around one query protein, drawn at the centre
+with a dark ring. Direct neighbours sit on an inner ring, grouped by kind (confirmed predictions, novel
+predictions, known annotations, then other node types); second-degree neighbours sit on an outer ring
+next to the neighbour they connect to. Choose the protein, then click **Show graph**.
+
+- **Predicted GO terms** – how many of the highest-scoring predictions are shown for each predicted
+  GO category.
+- **Neighbours per edge type** – how many knowledge-graph neighbours are shown for each relationship
+  type (domains, interactions, pathways, …), including known GO annotations outside the top predictions.
+- **Second-degree neighbours** – set above 0 to also expand the neighbours' own connections, excluding
+  edges back to the query protein, up to that many per edge type. Keep it low: the graph quickly
+  becomes cluttered. Second-degree edges always come from the knowledge graph, never from predictions.
+
+Node colours mark the entity type (see the legend at the top of the graph). Predictions are drawn for the
+selected protein only; select another protein in **Protein** to see its predictions.
+
+Edge colours mark the kind of link:
+- **Purple** – predicted GO term that is also annotated in the knowledge graph. Known annotations that
+  the model predicted are purple even when they rank below the top terms set by **Predicted GO terms**.
+- **Red** – novel prediction (one of the top predicted terms) not annotated in the knowledge graph.
+- **Blue** – GO annotation in the knowledge graph that the model did not predict in this run.
+- **Grey** – other relationships, such as protein–protein interactions, domains, drug targets, pathways,
+  domain–GO mappings and GO hierarchy links.
+
+Interacting with the graph:
+- **Hover** a node to highlight its connections and show their relationship names, with a card giving the
+  node's name, type and a link to the source database (UniProt, QuickGO, InterPro, …). Hover an edge to see
+  the relationship and, for GO terms, the score.
+- **Drag** nodes to rearrange them; **scroll** to zoom and drag the background to pan.
+- **Click a legend item** to hide that node type or edge kind; nodes left without visible edges are hidden
+  too. Click it again to bring it back.
+- Use the icons at the top right to **reset** the view, save it as a **PNG**, or download the edges currently
+  shown as **JSON** (with their scores and edge kind; hidden edge kinds and node types are left out).
+
+For exploration beyond second-degree connections, download the complete knowledge graph from the link at
+the top of the Knowledge graph tab.
+"""
+
+_CSS = """
+<style>
+@import url('https://fonts.googleapis.com/css2?family=Manrope:wght@400;500;600;700;800&family=DM+Mono:wght@400;500&display=swap');
+
+/* Material icons (incl. st.status's stExpanderIcon*) must keep their icon font */
+.stApp :is(p, h1, h2, h3, h4, h5, label, li, a, button, input, textarea, select, div, span):not([data-testid^="stIcon"]):not([data-testid^="stExpanderIcon"]) {
+    font-family: 'Manrope', sans-serif;
+}
+code, .mono { font-family: 'DM Mono', monospace !important; }
+
+/* Hide the sidebar entirely: navigation lives in the header */
+[data-testid="stSidebar"], [data-testid="stSidebarCollapsedControl"] { display: none; }
+
+/* Streamlit's fixed top bar (3.75rem) overlaps the page: make it transparent and start content below it */
+header[data-testid="stHeader"] { background: transparent; }
+.block-container { max-width: 1100px; padding-top: 4.25rem; padding-bottom: 4rem; }
+
+/* Header */
+.phgt-brand { font-weight: 800; font-size: 1.35rem; letter-spacing: -0.01em; margin: 0 !important; line-height: 2.4rem; }
+div[data-testid="stPageLink"] a { justify-content: flex-end; }
+div[data-testid="stPageLink"] p { font-weight: 500; color: #4A4E5C; }
+
+/* Hero */
+.phgt-hero h1 { font-size: 2.1rem; font-weight: 800; letter-spacing: -0.02em; line-height: 1.15; margin: 1.6rem 0 0.4rem; padding: 0; }
+.phgt-hero p { font-size: 1.05rem; color: #4A4E5C; margin: 0 0 1.4rem; }  /* full width, like the stepper */
+.phgt-hero-spacer { height: 1.5rem; }  /* step 3 has no hero: keep the stepper clear of the header */
+
+/* Stepper */
+.phgt-steps { display: grid; grid-template-columns: repeat(3, minmax(0, 1fr)); gap: 12px; list-style: none; margin: 0 0 1.5rem; padding: 0; }
+.phgt-step { display: flex; align-items: center; gap: 12px; padding: 12px 14px; border-radius: 12px; background: #FFFFFF; border: 1px solid #E4E2DC; }
+.phgt-step.done { background: #EEF0FC; border-color: #EEF0FC; }
+.phgt-step.current { border: 2px solid ACCENT; }
+.phgt-dot { flex: none; width: 30px; height: 30px; border-radius: 999px; display: flex; align-items: center; justify-content: center; font-weight: 700; font-size: 14px; background: #EDECE7; color: #4A4E5C; box-sizing: border-box; }
+.phgt-step.done .phgt-dot { background: ACCENT; color: #FFFFFF; }
+.phgt-step.current .phgt-dot { background: #1E288A; color: #FFFFFF; }
+.phgt-step small { display: block; font-size: 12px; color: #4A4E5C; }
+.phgt-step span.t { font-weight: 700; font-size: 15px; }
+@media (max-width: 640px) {
+  /* keep the header links on one wrapped row instead of stacking */
+  div[data-testid="stHorizontalBlock"]:has(.phgt-brand) { flex-direction: row !important; flex-wrap: wrap !important; gap: 0 14px !important; }
+  div[data-testid="stHorizontalBlock"]:has(.phgt-brand) > div[data-testid="stColumn"] { width: max-content !important; flex: 0 0 max-content !important; min-width: 0 !important; }
+  div[data-testid="stHorizontalBlock"]:has(.phgt-brand) > div[data-testid="stColumn"]:first-child { flex: 1 0 100% !important; }
+  div[data-testid="stHorizontalBlock"]:has(.phgt-brand) div[data-testid="stPageLink"] { width: auto !important; }
+  .phgt-step small { display: none; }
+  .phgt-step span.t { font-size: 13px; }
+  .phgt-hero h1 { font-size: 1.8rem; }
+}
+
+/* Cards & chips */
+div[data-testid="stVerticalBlockBorderWrapper"] { border-radius: 16px; }
+/* Dataframes use the same frame as the cards (Streamlit sets the border inline, hence !important) */
+div[data-testid="stDataFrameResizable"] { border: 1px solid #E4E2DC !important; border-radius: 16px !important; overflow: hidden; }
+div[data-testid="stDataFrame"] .stDataFrameGlideDataEditor { border-radius: 16px; }
+.phgt-chips { display: flex; flex-wrap: wrap; gap: 6px; margin: 4px 0 10px; }
+.phgt-chip { padding: 4px 9px; border-radius: 8px; background: #F3F2EE; font: 500 13px 'DM Mono', monospace; }
+.phgt-muted { color: #4A4E5C; font-size: 0.9rem; }
+.phgt-meter { height: 6px; border-radius: 3px; background: #EDECE7; margin: 6px 0 4px; }
+.phgt-meter > div { height: 6px; border-radius: 3px; background: ACCENT; }
+.phgt-count { display: block; padding-bottom: 0.6rem; }  /* breathing room below the protein count */
+
+/* Multiselect chips use the accent tint instead of Streamlit red */
+span[data-baseweb="tag"] { background-color: #E2E5FB !important; color: #1E288A !important; font-family: 'DM Mono', monospace; }
+span[data-baseweb="tag"] svg { fill: #1E288A !important; }
+
+/* Primary buttons */
+.stButton > button, .stDownloadButton > button { border-radius: 12px; min-height: 44px; font-weight: 600; }
+.stButton > button[kind="primary"] { font-weight: 700; }
+
+/* Metrics */
+div[data-testid="stMetric"] { background: #FFFFFF; border: 1px solid #E4E2DC; border-radius: 14px; padding: 14px 16px; }
+</style>
+""".replace("ACCENT", ACCENT)
+
+
+def apply_style():
+    st.markdown(_CSS, unsafe_allow_html=True)
+
+
+def render_header(current="predict"):
+    apply_style()
+    cols = st.columns([3, 1.1, 1.1, 0.9, 1.1], vertical_alignment="center")
+    with cols[0]:
+        st.markdown('<p class="phgt-brand">ProtHGT</p>', unsafe_allow_html=True)
+    with cols[1]:
+        st.page_link("ProtHGT_app.py", label="Predict")
+    with cols[2]:
+        st.page_link("pages/User_Guide.py", label="User guide")
+    with cols[3]:
+        st.page_link("pages/About.py", label="About")
+    with cols[4]:
+        st.page_link(GITHUB_URL, label="GitHub")
+
+
+def render_stepper(step):
+    labels = ["Add proteins", "Choose what to predict", "Review results"]
+    items = []
+    for i, label in enumerate(labels, start=1):
+        if i < step:
+            cls, mark, state = "done", "&#10003;", ""
+        elif i == step:
+            cls, mark, state = "current", str(i), ' aria-current="step"'
+        else:
+            cls, mark, state = "", str(i), ""
+        items.append(
+            f'<li class="phgt-step {cls}"{state}><span class="phgt-dot">{mark}</span>'
+            f'<span><small>Step {i}</small><span class="t">{label}</span></span></li>'
+        )
+    st.markdown(f'<ol class="phgt-steps" aria-label="Progress">{"".join(items)}</ol>', unsafe_allow_html=True)