lib/ui.py

"""Shared Streamlit UI helpers — the cross-page model toggle.

Keeping this in one place means every page shows the same picker and shares the
same selection (via st.session_state) and the same per-ref cache.
"""
import streamlit as st

import config
from lib.model import load_model

_CUSTOM_LABEL = "✏️ Custom HF model ID…"


@st.cache_resource(show_spinner="Loading NER model…")
def _load_cached(ref):
    """Cache one LoadedModel per distinct ref (None = demo fallback)."""
    return load_model(ref=ref)


def model_selector():
    """Render the sidebar model picker and return the selected LoadedModel.

    The choice persists across pages through st.session_state["model_label"].
    A "Custom HF model ID" option lets anyone load any Hub repo live, and the
    refresh button clears the cache to pick up a freshly-uploaded model.
    """
    options = config.available_models()           # {label: ref}
    labels = list(options.keys()) + [_CUSTOM_LABEL]

    st.sidebar.subheader("Model")
    current = st.session_state.get("model_label", labels[0])
    index = labels.index(current) if current in labels else 0
    choice = st.sidebar.selectbox(
        "Active NER model", labels, index=index, key="model_label",
        label_visibility="collapsed",
    )

    if choice == _CUSTOM_LABEL:
        ref = st.sidebar.text_input(
            "HF model repo id", placeholder="e.g. Zeqhx/cv-parser-ner",
            key="custom_model_id",
        ).strip() or None
    else:
        ref = options[choice]

    if st.sidebar.button("🔄 Reload model", use_container_width=True,
                         help="Clear the cache and re-pull (use after updating a model)."):
        _load_cached.clear()
        st.rerun()

    lm = _load_cached(ref)

    if lm.is_fallback:
        st.sidebar.warning("Demo mode — untrained head; predictions are not meaningful.",
                           icon="⚠️")
    else:
        st.sidebar.success("Model loaded", icon="✅")
    st.sidebar.caption(lm.source)

    return lm