src/streamlit_app.py

import streamlit as st
from transformers import (
    AutoTokenizer,
    AutoModel,
)
import lightning as L
import matplotlib.pyplot as plt
from huggingface_hub import hf_hub_download
import time
import torch
import torch.nn as nn

repo_id = "Doub1e05/ML2_HW4"
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model_name = "oracat/bert-paper-classifier-arxiv"

model_path = hf_hub_download(
    repo_id=repo_id,
    filename="last-v1.ckpt", 
    repo_type="space" 
)



ARXIV_CATEGORIES = {
    "cs": [
        ("cs.AI", "Artificial Intelligence"),
        ("cs.AR", "Hardware Architecture"),
        ("cs.CC", "Computational Complexity"),
        ("cs.CE", "Computational Engineering, Finance, and Science"),
        ("cs.CG", "Computational Geometry"),
        ("cs.CL", "Computation and Language"),
        ("cs.CR", "Cryptography and Security"),
        ("cs.CV", "Computer Vision and Pattern Recognition"),
        ("cs.CY", "Computers and Society"),
        ("cs.DB", "Databases"),
        ("cs.DC", "Distributed, Parallel, and Cluster Computing"),
        ("cs.DL", "Digital Libraries"),
        ("cs.DM", "Discrete Mathematics"),
        ("cs.DS", "Data Structures and Algorithms"),
        ("cs.ET", "Emerging Technologies"),
        ("cs.FL", "Formal Languages and Automata Theory"),
        ("cs.GL", "General Literature"),
        ("cs.GR", "Graphics"),
        ("cs.GT", "Computer Science and Game Theory"),
        ("cs.HC", "Human-Computer Interaction"),
        ("cs.IR", "Information Retrieval"),
        ("cs.IT", "Information Theory"),
        ("cs.LG", "Machine Learning"),
        ("cs.LO", "Logic in Computer Science"),
        ("cs.MA", "Multiagent Systems"),
        ("cs.MM", "Multimedia"),
        ("cs.MS", "Mathematical Software"),
        ("cs.NA", "Numerical Analysis"),
        ("cs.NE", "Neural and Evolutionary Computing"),
        ("cs.NI", "Networking and Internet Architecture"),
        ("cs.OH", "Other Computer Science"),
        ("cs.OS", "Operating Systems"),
        ("cs.PF", "Performance"),
        ("cs.PL", "Programming Languages"),
        ("cs.RO", "Robotics"),
        ("cs.SC", "Symbolic Computation"),
        ("cs.SD", "Sound"),
        ("cs.SE", "Software Engineering"),
        ("cs.SI", "Social and Information Networks"),
        ("cs.SY", "Systems and Control"),
    ],
    "math": [
        ("math.AC", "Commutative Algebra"),
        ("math.AG", "Algebraic Geometry"),
        ("math.AP", "Analysis of PDEs"),
        ("math.AT", "Algebraic Topology"),
        ("math.CA", "Classical Analysis and ODEs"),
        ("math.CO", "Combinatorics"),
        ("math.CT", "Category Theory"),
        ("math.CV", "Complex Variables"),
        ("math.DG", "Differential Geometry"),
        ("math.DS", "Dynamical Systems"),
        ("math.FA", "Functional Analysis"),
        ("math.GM", "General Mathematics"),
        ("math.GN", "General Topology"),
        ("math.GR", "Group Theory"),
        ("math.GT", "Geometric Topology"),
        ("math.HO", "History and Overview"),
        ("math.IT", "Information Theory"),
        ("math.KT", "K-Theory and Homology"),
        ("math.LO", "Logic"),
        ("math.MG", "Metric Geometry"),
        ("math.MP", "Mathematical Physics"),
        ("math.NA", "Numerical Analysis"),
        ("math.NT", "Number Theory"),
        ("math.OA", "Operator Algebras"),
        ("math.OC", "Optimization and Control"),
        ("math.PR", "Probability"),
        ("math.QA", "Quantum Algebra"),
        ("math.RA", "Rings and Algebras"),
        ("math.RT", "Representation Theory"),
        ("math.SG", "Symplectic Geometry"),
        ("math.SP", "Spectral Theory"),
        ("math.ST", "Statistics Theory"),
    ],
    "astro-ph": [  # Astrophysics
        ("astro-ph.CO", "Cosmology and Nongalactic Astrophysics"),
        ("astro-ph.EP", "Earth and Planetary Astrophysics"),
        ("astro-ph.GA", "Astrophysics of Galaxies"),
        ("astro-ph.HE", "High Energy Astrophysical Phenomena"),
        ("astro-ph.IM", "Instrumentation and Methods for Astrophysics"),
        ("astro-ph.SR", "Solar and Stellar Astrophysics"),
    ],
    "cond-mat": [  # Condensed Matter
        ("cond-mat.dis-nn", "Disordered Systems and Neural Networks"),
        ("cond-mat.mes-hall", "Mesoscale and Nanoscale Physics"),
        ("cond-mat.mtrl-sci", "Materials Science"),
        ("cond-mat.other", "Other Condensed Matter"),
        ("cond-mat.quant-gas", "Quantum Gases"),
        ("cond-mat.soft", "Soft Condensed Matter"),
        ("cond-mat.stat-mech", "Statistical Mechanics"),
        ("cond-mat.str-el", "Strongly Correlated Electrons"),
        ("cond-mat.supr-con", "Superconductivity"),
    ],
    "gr-qc": [
        ("gr-qc", "General Relativity and Quantum Cosmology"),
    ],
    "hep-ex": [
        ("hep-ex", "High Energy Physics - Experiment"),
    ],
    "hep-lat": [
        ("hep-lat", "High Energy Physics - Lattice"),
    ],
    "hep-ph": [
        ("hep-ph", "High Energy Physics - Phenomenology"),
    ],
    "hep-th": [
        ("hep-th", "High Energy Physics - Theory"),
    ],
    "math-ph": [
        ("math-ph", "Mathematical Physics"),
    ],
    "nlin": [
        ("nlin.AO", "Adaptation and Self-Organizing Systems"),
        ("nlin.CD", "Chaotic Dynamics"),
        ("nlin.CG", "Cellular Automata and Lattice Gases"),
        ("nlin.PS", "Pattern Formation and Solitons"),
        ("nlin.SI", "Exactly Solvable and Integrable Systems"),
    ],
    "nucl-ex": [
        ("nucl-ex", "Nuclear Experiment"),
    ],
    "nucl-th": [
        ("nucl-th", "Nuclear Theory"),
    ],
    "physics": [
        ("physics.acc-ph", "Accelerator Physics"),
        ("physics.ao-ph", "Atmospheric and Oceanic Physics"),
        ("physics.app-ph", "Applied Physics"),
        ("physics.atm-clus", "Atomic and Molecular Clusters"),
        ("physics.atom-ph", "Atomic Physics"),
        ("physics.bio-ph", "Biological Physics"),
        ("physics.chem-ph", "Chemical Physics"),
        ("physics.class-ph", "Classical Physics"),
        ("physics.comp-ph", "Computational Physics"),
        ("physics.data-an", "Data Analysis, Statistics and Probability"),
        ("physics.ed-ph", "Physics Education"),
        ("physics.flu-dyn", "Fluid Dynamics"),
        ("physics.gen-ph", "General Physics"),
        ("physics.geo-ph", "Geophysics"),
        ("physics.hist-ph", "History and Philosophy of Physics"),
        ("physics.ins-det", "Instrumentation and Detectors"),
        ("physics.med-ph", "Medical Physics"),
        ("physics.optics", "Optics"),
        ("physics.plasm-ph", "Plasma Physics"),
        ("physics.pop-ph", "Popular Physics"),
        ("physics.soc-ph", "Physics and Society"),
        ("physics.space-ph", "Space Physics"),
    ],
    "q-bio": [
        ("q-bio.BM", "Biomolecules"),
        ("q-bio.CB", "Cell Behavior"),
        ("q-bio.GN", "Genomics"),
        ("q-bio.MN", "Molecular Networks"),
        ("q-bio.NC", "Neurons and Cognition"),
        ("q-bio.OT", "Other Quantitative Biology"),
        ("q-bio.PE", "Populations and Evolution"),
        ("q-bio.QM", "Quantitative Methods"),
        ("q-bio.SC", "Subcellular Processes"),
        ("q-bio.TO", "Tissues and Organs"),
    ],
    "q-fin": [
        ("q-fin.CP", "Computational Finance"),
        ("q-fin.EC", "Economics"),
        ("q-fin.GN", "General Finance"),
        ("q-fin.MF", "Mathematical Finance"),
        ("q-fin.PM", "Portfolio Management"),
        ("q-fin.PR", "Pricing of Securities"),
        ("q-fin.RM", "Risk Management"),
        ("q-fin.ST", "Statistical Finance"),
        ("q-fin.TR", "Trading and Market Microstructure"),
    ],
    "stat": [
        ("stat.AP", "Applications"),
        ("stat.CO", "Computation"),
        ("stat.ME", "Methodology"),
        ("stat.ML", "Machine Learning"),
        ("stat.OT", "Other Statistics"),
        ("stat.TH", "Theory"),
    ],
    "eess": [
        ("eess.AS", "Audio and Speech Processing"),
        ("eess.IV", "Image and Video Processing"),
        ("eess.SP", "Signal Processing"),
        ("eess.SY", "Systems and Control"),
    ],
    "econ": [
        ("econ.EM", "Econometrics"),
        ("econ.GN", "General Economics"),
        ("econ.TH", "Theoretical Economics"),
    ],
}


list_arxiv = [
    "Artificial Intelligence",
    "Hardware Architecture",
    "Computational Complexity",
    "Computational Engineering, Finance, and Science",
    "Computational Geometry",
    "Computation and Language",
    "Cryptography and Security",
    "Computer Vision and Pattern Recognition",
    "Computers and Society",
    "Databases",
    "Distributed, Parallel, and Cluster Computing",
    "Digital Libraries",
    "Discrete Mathematics",
    "Data Structures and Algorithms",
    "Emerging Technologies",
    "Formal Languages and Automata Theory",
    "General Literature",
    "Graphics",
    "Computer Science and Game Theory",
    "Human-Computer Interaction",
    "Information Retrieval",
    "Information Theory",
    "Machine Learning",
    "Logic in Computer Science",
    "Multiagent Systems",
    "Multimedia",
    "Mathematical Software",
    "Numerical Analysis",
    "Neural and Evolutionary Computing",
    "Networking and Internet Architecture",
    "Other Computer Science",
    "Operating Systems",
    "Performance",
    "Programming Languages",
    "Robotics",
    "Symbolic Computation",
    "Sound",
    "Software Engineering",
    "Social and Information Networks",
    "Systems and Control",
    "Commutative Algebra",
    "Algebraic Geometry",
    "Analysis of PDEs",
    "Algebraic Topology",
    "Classical Analysis and ODEs",
    "Combinatorics",
    "Category Theory",
    "Complex Variables",
    "Differential Geometry",
    "Dynamical Systems",
    "Functional Analysis",
    "General Mathematics",
    "General Topology",
    "Group Theory",
    "Geometric Topology",
    "History and Overview",
    "Information Theory",
    "K-Theory and Homology",
    "Logic",
    "Metric Geometry",
    "Mathematical Physics",
    "Numerical Analysis",
    "Number Theory",
    "Operator Algebras",
    "Optimization and Control",
    "Probability",
    "Quantum Algebra",
    "Rings and Algebras",
    "Representation Theory",
    "Symplectic Geometry",
    "Spectral Theory",
    "Statistics Theory",
    "Cosmology and Nongalactic Astrophysics",
    "Earth and Planetary Astrophysics",
    "Astrophysics of Galaxies",
    "High Energy Astrophysical Phenomena",
    "Instrumentation and Methods for Astrophysics",
    "Solar and Stellar Astrophysics",
    "Disordered Systems and Neural Networks",
    "Mesoscale and Nanoscale Physics",
    "Materials Science",
    "Other Condensed Matter",
    "Quantum Gases",
    "Soft Condensed Matter",
    "Statistical Mechanics",
    "Strongly Correlated Electrons",
    "Superconductivity",
    "General Relativity and Quantum Cosmology",
    "High Energy Physics - Experiment",
    "High Energy Physics - Lattice",
    "High Energy Physics - Phenomenology",
    "High Energy Physics - Theory",
    "Mathematical Physics",
    "Adaptation and Self-Organizing Systems",
    "Chaotic Dynamics",
    "Cellular Automata and Lattice Gases",
    "Pattern Formation and Solitons",
    "Exactly Solvable and Integrable Systems",
    "Nuclear Experiment",
    "Nuclear Theory",
    "Accelerator Physics",
    "Atmospheric and Oceanic Physics",
    "Applied Physics",
    "Atomic and Molecular Clusters",
    "Atomic Physics",
    "Biological Physics",
    "Chemical Physics",
    "Classical Physics",
    "Computational Physics",
    "Data Analysis, Statistics and Probability",
    "Physics Education",
    "Fluid Dynamics",
    "General Physics",
    "Geophysics",
    "History and Philosophy of Physics",
    "Instrumentation and Detectors",
    "Medical Physics",
    "Optics",
    "Plasma Physics",
    "Popular Physics",
    "Physics and Society",
    "Space Physics",
    "Biomolecules",
    "Cell Behavior",
    "Genomics",
    "Molecular Networks",
    "Neurons and Cognition",
    "Other Quantitative Biology",
    "Populations and Evolution",
    "Quantitative Methods",
    "Subcellular Processes",
    "Tissues and Organs",
    "Computational Finance",
    "Economics",
    "General Finance",
    "Mathematical Finance",
    "Portfolio Management",
    "Pricing of Securities",
    "Risk Management",
    "Statistical Finance",
    "Trading and Market Microstructure",
    "Applications",
    "Computation",
    "Methodology",
    "Machine Learning",
    "Other Statistics",
    "Theory",
    "Audio and Speech Processing",
    "Image and Video Processing",
    "Signal Processing",
    "Systems and Control",
    "Econometrics",
    "General Economics",
    "Theoretical Economics",
]


class MyBert(L.LightningModule):
    def __init__(
        self,
        pretrain_model,
        name_mlp,
        num_class=10,
        in_feature=768,
        criterion=None,
        optimizer=None,
        metrice=None,
        name_metrice=None,
        scheduler_config=None,
        learning_rate=1e-2,
    ):
        super(MyBert, self).__init__()
        self.save_hyperparameters(
            ignore=[
                "pretrain_model",
                "criterion",
                "metrice",
                "optimizer",
                "scheduler_config",
            ]
        )

        self.model = pretrain_model
        self.num_class = num_class
        self.learning_rate = learning_rate

        # Не сохраняем метрику как параметр модели
        self._name_metrice = name_metrice if name_metrice else "metrice"

        parts = name_mlp.split(".")
        module = self.model
        for part in parts[:-1]:
            module = getattr(module, part)
        last_part = parts[-1]
        setattr(module, last_part, nn.Identity())

        for param in self.model.parameters():
            param.requires_grad = False

        self.mlp = nn.Linear(in_feature, num_class)

    def forward(self, text):
        # Динамически определяем устройство
        if text.device != next(self.parameters()).device:
            text = text.to(next(self.parameters()).device)
        output = self.model(text)
        output = self.mlp(output.last_hidden_state[:, 0])
        return output

    def configure_metrics(self):
        pass


@st.cache_resource
def create_model_and_tokenizer():
    pre_model = AutoModel.from_pretrained(model_name)
    model = MyBert.load_from_checkpoint(
        model_path,
        pretrain_model=pre_model,
        name_mlp="classifier",
        num_class=154,
        weights_only=False,
        map_location=torch.device("cpu"),
    )

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    return model, tokenizer


model, tokenizer = create_model_and_tokenizer()


max_seq_len = 512


if "show_categories" not in st.session_state:
    st.session_state.show_categories = False

if st.session_state.show_categories:
    st.button(
        "Back", on_click=lambda: setattr(st.session_state, "show_categories", False)
    )

    category_names = list(ARXIV_CATEGORIES.keys())
    tabs = st.tabs(category_names)

    temp = 0
    for tab_idx, (cat_name, cat_items) in enumerate(list(ARXIV_CATEGORIES.items())):
        with tabs[tab_idx]:
            elements = list_arxiv[temp : temp + len(cat_items)]
            sub_cols = st.columns(3)
            for elem_idx, element in enumerate(elements):
                with sub_cols[elem_idx % 3]:
                    st.write(element)
        temp += len(cat_items)

else:
    if st.button("Show all possible classes"):
        st.session_state.show_categories = True
        st.rerun()

    title = st.text_input(label="Title", placeholder="Enter title article")
    caprion = st.text_input(label="Abstract", placeholder="Enter abstract article")

    if st.button("Predict"):
        if title == "" and caprion == "":
            st.text("Please enter title or abstrain")
        else:
            text = tokenizer(("Title: " + title + " Abstract: " + caprion))[
                "input_ids"
            ][:max_seq_len]
            text = torch.tensor(text, device=device)[None, :]
            pred = nn.functional.softmax(model(text)[0])
            top_k = pred.argsort(descending=True)
            temp = 0
            result = []
            index = 0
            while temp < 0.95:
                result.append(top_k[index])
                temp += pred[result[-1]]
                index += 1
            labels = []
            probabilities = []
            for i in result:
                labels.append(list_arxiv[i])
                probabilities.append(float(pred[i]))

            fig, ax = plt.subplots()
            ax.bar(labels, probabilities, width=0.5)
            ax.set_xticklabels(labels, rotation=45, ha="right")
            ax.set_ylabel("Probability")
            st.pyplot(fig)