src/streamlit_app.py

import json
from pathlib import Path

import streamlit as st
from model_utils import (
    load_model_and_tokenizer,
    predict_top_p,
    predict_top_n,
    explain_prediction,
    render_token_heatmap,
)

st.set_page_config(
    page_title="ArXiv Topic Predictor",
    page_icon=":material/description:",
    layout="centered"
)

ARXIV_MAP_PATH = Path(__file__).with_name("arxiv_map.json")
with ARXIV_MAP_PATH.open("r", encoding="utf-8") as map_file:
    arxiv_map = json.load(map_file)


MODEL_PATH = "src/best" 

if "model_loaded_toast_shown" not in st.session_state:
    st.session_state["model_loaded_toast_shown"] = False

try:
    tokenizer, model, device = load_model_and_tokenizer(MODEL_PATH)
    if not st.session_state["model_loaded_toast_shown"]:
        st.toast("The model has been uploaded successfully!", icon=":material/check_circle:")
        st.session_state["model_loaded_toast_shown"] = True
except Exception as e:
    st.error(f"Error'{MODEL_PATH}'.\n\nError details: {e}")
    st.stop()

SAMPLES = {
    "Computer Vision": {
        "title": "Modulate-and-Map: Crossmodal Feature Mapping with Cross-View Modulation for 3D Anomaly Detection",
        "summary": "We present ModMap, a natively multiview and multimodal framework for 3D anomaly detection and segmentation. Unlike existing methods that process views independently, our method draws inspiration from the crossmodal feature mapping paradigm to learn to map features across both modalities and views, while explicitly modelling view-dependent relationships through feature-wise modulation. We introduce a cross-view training strategy that leverages all possible view combinations, enabling effective anomaly scoring through multiview ensembling and aggregation. To process high-resolution 3D data, we train and publicly release a foundational depth encoder tailored to industrial datasets. Experiments on SiM3D, a recent benchmark that introduces the first multiview and multimodal setup for 3D anomaly detection and segmentation, demonstrate that ModMap attains state-of-the-art performance by surpassing previous methods by wide margins."
    },
    "Medical Imaging": {
        "title": "High-resolution ultra-low-field MRI with SNRAware denoising",
        "summary": "Ultra-low-field (ULF, <0.1 T) magnetic resonance imaging (MRI) systems offer advantages in cost, portability, and accessibility, but their current utility is still limited by low signal-to-noise ratio (SNR). Deep learning (DL)-based denoising has emerged as a potential strategy to mitigate this limitation. In this work, we present a systematic evaluation of a high-performance DL denoising model trained using the SNRAware framework and applied to 88 mT and 72 mT data. Using a series of controlled experiments, we assessed model performance as a function of spatial resolution, coil impedance matching, readout bandwidth, input noise level, k-space undersampling, anatomy, image contrast, and scanner platform, and compared against analytical denoising algorithms. The model consistently increased the effective SNR of ULF acquisitions, enabling images with nominal spatial resolutions comparable to those commonly used in clinical 3 T protocols. Residual analyses indicated that the model predominantly removed stochastic noise while preserving underlying signal structure. At the same time, the results highlight some constraints: denoising performance remains dependent on the starting SNR of the acquisition, and training-domain mismatch influences behavior under certain artifact conditions. These findings suggest that DL-based denoising can significantly expand the practical capabilities of ULF MRI, while emphasizing potential benefits from hardware-software co-optimization and the need for rigorous clinical validation to determine the diagnostic value of denoised images."
    },
    "Physics": {
        "title": "AstroConcepts: A Large-Scale Multi-Label Classification Corpus for Astrophysics",
        "summary": "Scientific multi-label text classification suffers from extreme class imbalance, where specialized terminology exhibits severe power-law distributions that challenge standard classification approaches. Existing scientific corpora lack comprehensive controlled vocabularies, focusing instead on broad categories and limiting systematic study of extreme imbalance. We introduce AstroConcepts, a corpus of English abstracts from 21,702 published astrophysics papers, labeled with 2,367 concepts from the Unified Astronomy Thesaurus. The corpus exhibits severe label imbalance, with 76% of concepts having fewer than 50 training examples. By releasing this resource, we enable systematic study of extreme class imbalance in scientific domains and establish strong baselines across traditional, neural, and vocabulary-constrained LLM methods. Our evaluation reveals three key patterns that provide new insights into scientific text classification. First, vocabulary-constrained LLMs achieve competitive performance relative to domain-adapted models in astrophysics classification, suggesting a potential for parameter-efficient approaches. Second, domain adaptation yields relatively larger improvements for rare, specialized terminology, although absolute performance remains limited across all methods. Third, we propose frequency-stratified evaluation to reveal performance patterns that are hidden by aggregate scores, thereby making robustness assessment central to scientific multi-label evaluation. These results offer actionable insights for scientific NLP and establish benchmarks for research on extreme imbalance."
    },
    "Mathematics": {
        "title": "p-adic Hodge theory of de Rham local systems, I: Newton polygon and monodromy",
        "summary": "We prove that the relative p-adic monodromy theorem holds over a dense open subset. Moreover, we establish the equivalence of the following two statements: the local constancy of the Newton polygon function associated with a de Rham local system around rank-1 points, and the relative p-adic monodromy theorem near rank-1 points. We demonstrate how to extend the relative p-adic monodromy conjecture from the neighborhood of rank-1 points to the entire interiors of Newton partitions."
    },
    "Finance": {
        "title": "Concave Continuation: Linking Routing to Arbitrage",
        "summary": "We extend AMM trade functions to negative inputs via the concave continuation, derived from the invariance of the local conservation law under allocation direction flips. This unifies routing and arbitrage into a single problem. We extend the one-hop transfer algorithm proposed in \cite{jiang} to this setting."
    },
    
}

def load_sample(sample_name):
    st.session_state["input_title"] = SAMPLES[sample_name]["title"]
    st.session_state["input_summary"] = SAMPLES[sample_name]["summary"]
    st.session_state.pop("explain_tag", None)
    st.session_state["has_predicted"] = False

with st.container():
    st.title(":material/find_in_page: ArXiv Classifier")
    
    st.markdown("""
        #### *Intelligent Academic Categorization*
        This assistant leverages state-of-the-art **Transformer architectures** to analyze scientific publications. 
        By performing semantic analysis on titles and abstracts, it maps complex research into 
        primary domains with high precision.
    """)

    st.divider() 

    col1, col2 = st.columns(2)

    with col1:
        st.markdown("##### :material/database: Data Source")
        st.info("[ArXiv Dataset on Kaggle](https://www.kaggle.com/datasets/Cornell-University/arxiv)")

    with col2:
        st.markdown("##### :material/memory: Architecture")
        st.success("[SciBERT (allenai)](https://huggingface.co/allenai/scibert_scivocab_uncased)")

    with st.expander("View Technical Implementation Details"):
        st.write("""
            The system utilizes **SciBERT**, a BERT model specifically trained on 1.14M scientific 
            papers. This ensures superior handling of domain-specific terminology and 
            mathematical contexts compared to general-purpose NLP models.
        """)


st.subheader(":material/edit_note: Enter the data")

st.write("Try a sample:")
s_col1, s_col2, s_col3, s_col4, s_col5 = st.columns(5)
with s_col1:
    if st.button("CV", use_container_width=True):
        load_sample("Computer Vision")
with s_col2:
    if st.button("Medical", use_container_width=True):
        load_sample("Medical Imaging")
with s_col3:
    if st.button("Physics", use_container_width=True):
        load_sample("Physics")
with s_col4:
    if st.button("Math", use_container_width=True):
        load_sample("Mathematics")
with s_col5:
    if st.button("Finance", use_container_width=True):
        load_sample("Finance")

input_title = st.text_input(
    "Title",
    placeholder="Enter paper title...",
    key="input_title"
)
input_summary = st.text_area(
    "Abstract",
    placeholder="Enter paper abstract...",
    height=300,
    key="input_summary"
)

predict_btn = st.button("Define the topic", type="primary", use_container_width=True)

with st.sidebar:
    st.header(":material/settings: Prediction settings")
    pred_mode = st.radio("Mode", ["Top-p", "Top-N"], horizontal=True)

    if pred_mode == "Top-p":
        top_p_value = st.slider("Top-p", 0.50, 1.00, 0.95)
        top_n_value = None
    else:
        top_n_value = st.slider("Top-N", 1, 20, 5)
        top_p_value = None
    st.divider()

results_container = st.container()

with results_container:
    st.subheader(":material/analytics: Results")

    if predict_btn:
        if input_title or input_summary:
            st.session_state["has_predicted"] = True
            st.session_state.pop("explain_tag", None)
        else:
            st.warning("First, fill in at least one of the fields on the left")
            st.session_state["has_predicted"] = False

    if st.session_state.get("has_predicted", False):
        with st.spinner("Analyzing text..."):
            if pred_mode == "Top-p":
                tags, probs = predict_top_p(
                    input_title,
                    input_summary,
                    tokenizer,
                    model,
                    device,
                    p=top_p_value,
                )
            else:
                tags, probs = predict_top_n(
                    input_title,
                    input_summary,
                    tokenizer,
                    model,
                    device,
                    n=top_n_value,
                )

        st.write("**Recommended Tags:**")
        for tag, prob in tags:
            prob = round(prob, 2)
            row1, row2 = st.columns([6, 1])

            with row1:
                if tag in arxiv_map:
                    st.progress(prob, text=f"{tag} : {arxiv_map[tag]} {int(prob * 100)}%")
                else:
                    st.progress(prob, text=f"{tag} {int(prob * 100)}%")

            with row2:
                if st.button("Explain", key=f"explain_{tag}"):
                    st.session_state["explain_tag"] = tag

        st.divider()
        st.write("**Explanation:**")

        selected_tag = st.session_state.get("explain_tag")
        if selected_tag:
            text = f"Title: {input_title} [SEP] Abstract: {input_summary}"

            with st.spinner("Building heatmap..."):
                try:
                    tokens, scores, class_prob = explain_prediction(
                        text, selected_tag, tokenizer, model, device
                    )
                    st.caption(f"Model score for this class: {class_prob:.2%}")
                    st.markdown(
                        render_token_heatmap(tokens, scores),
                        unsafe_allow_html=True
                    )
                    st.caption("Green words push the selected class up; red words push it down.")
                except Exception as e:
                    st.error(f"Explanation failed: {e}")
        else:
            st.info("Press **Explain** next to a prediction.")

st.divider()
st.caption("© 2026 ArXiv AI Classifier | Vladimir Nedugov")