Upload 2 files

src/label_encoder.json

@@ -0,0 +1,87 @@
+{
+  "0": "astro-ph.CO",
+  "1": "astro-ph.IM",
+  "2": "astro-ph.SR",
+  "3": "cmp-lg",
+  "4": "cond-mat.dis-nn",
+  "5": "cond-mat.mtrl-sci",
+  "6": "cond-mat.stat-mech",
+  "7": "cs.AI",
+  "8": "cs.AR",
+  "9": "cs.CC",
+  "10": "cs.CE",
+  "11": "cs.CG",
+  "12": "cs.CL",
+  "13": "cs.CR",
+  "14": "cs.CV",
+  "15": "cs.CY",
+  "16": "cs.DB",
+  "17": "cs.DC",
+  "18": "cs.DL",
+  "19": "cs.DM",
+  "20": "cs.DS",
+  "21": "cs.ET",
+  "22": "cs.FL",
+  "23": "cs.GR",
+  "24": "cs.GT",
+  "25": "cs.HC",
+  "26": "cs.IR",
+  "27": "cs.IT",
+  "28": "cs.LG",
+  "29": "cs.LO",
+  "30": "cs.MA",
+  "31": "cs.MM",
+  "32": "cs.MS",
+  "33": "cs.NA",
+  "34": "cs.NE",
+  "35": "cs.NI",
+  "36": "cs.OH",
+  "37": "cs.PF",
+  "38": "cs.PL",
+  "39": "cs.RO",
+  "40": "cs.SD",
+  "41": "cs.SE",
+  "42": "cs.SI",
+  "43": "cs.SY",
+  "44": "econ.EM",
+  "45": "eess.AS",
+  "46": "eess.IV",
+  "47": "eess.SP",
+  "48": "eess.SY",
+  "49": "gr-qc",
+  "50": "hep-ph",
+  "51": "math.CO",
+  "52": "math.DS",
+  "53": "math.FA",
+  "54": "math.LO",
+  "55": "math.NA",
+  "56": "math.OC",
+  "57": "math.PR",
+  "58": "math.ST",
+  "59": "nlin.AO",
+  "60": "physics.ao-ph",
+  "61": "physics.bio-ph",
+  "62": "physics.chem-ph",
+  "63": "physics.comp-ph",
+  "64": "physics.data-an",
+  "65": "physics.flu-dyn",
+  "66": "physics.geo-ph",
+  "67": "physics.med-ph",
+  "68": "physics.optics",
+  "69": "physics.soc-ph",
+  "70": "q-bio.BM",
+  "71": "q-bio.GN",
+  "72": "q-bio.MN",
+  "73": "q-bio.NC",
+  "74": "q-bio.PE",
+  "75": "q-bio.QM",
+  "76": "q-fin.CP",
+  "77": "q-fin.PM",
+  "78": "q-fin.ST",
+  "79": "q-fin.TR",
+  "80": "quant-ph",
+  "81": "stat.AP",
+  "82": "stat.CO",
+  "83": "stat.ME",
+  "84": "stat.ML"
+}

src/streamlit_app.py

@@ -1,40 +1,128 @@
-import altair as alt
+import json
+import re
+import html
+import unicodedata
 import numpy as np
-import pandas as pd
 import streamlit as st
+import torch
+from transformers import AutoTokenizer, T5ForSequenceClassification
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+MODEL_DIR = "best_model"
+LABEL_ENCODER = "label_encoder.json"
+MAX_LENGTH = 256
+TOP_PROB = 0.95
+DEVICE = torch.device("cpu")
+MIN_CHARS = 20
+MAX_CHARS = 5000
+
+def clean_text(text: str) -> str:
+    text = html.unescape(text)
+    text = unicodedata.normalize("NFKC", text)
+    text = re.sub(r"\$.*?\$", "", text)
+    text = re.sub(r"\\[a-zA-Z]+\{.*?\}", "", text)
+    text = re.sub(r"\s+", " ", text)
+    return text.strip()
+
+
+def validate(title: str, abstract: str):
+    if len(title.strip()) < MIN_CHARS:
+        return f"Title too short (at least {MIN_CHARS} characters)"
+    if len(abstract.strip()) < MIN_CHARS:
+        return f"Abstract too short (at least {MIN_CHARS} characters)"
+    if len(title) > MAX_CHARS:
+        return f"Title too long (max {MAX_CHARS} characters)"
+    if len(abstract) > MAX_CHARS:
+        return f"Abstract too long (max {MAX_CHARS} characters)"
+    return None
+
+
+@st.cache_resource(show_spinner="Loading model…")
+def load_model():
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
+        model = T5ForSequenceClassification.from_pretrained(MODEL_DIR)
+        model.to(DEVICE)
+        model.eval()
+        with open(LABEL_ENCODER) as f:
+            id2label = {int(k): v for k, v in json.load(f).items()}
+        return tokenizer, model, id2label
+    except FileNotFoundError as e:
+        st.error(f"Model files not found: {e}")
+        st.stop()
+    except Exception as e:
+        st.error(f"Failed to load model: {e}")
+        st.stop()
+
+
+@st.cache_data(show_spinner=False)
+def predict(title: str, abstract: str):
+    tokenizer, model, id2label = load_model()
+    if title and abstract:
+        text = f"classify: {clean_text(title)} [SEP] {clean_text(abstract)[:1000]}"
+    elif title:
+        text = f"classify: {clean_text(title)}"
+    else:
+        text = f"classify: {clean_text(abstract)[:1000]}"
+    eos_id = tokenizer.eos_token_id
+
+    enc = tokenizer(text, max_length=MAX_LENGTH, truncation=True, return_tensors="pt")
+    ids = enc["input_ids"][0].tolist()
+
+    if eos_id is not None and ids.count(eos_id) > 1:
+        ids = [t for t in ids[:-1] if t != eos_id] + [eos_id]
+
+    input_ids = torch.tensor([ids], dtype=torch.long).to(DEVICE)
+    attention_mask = torch.ones_like(input_ids)
+
+    try:
+        with torch.no_grad():
+            logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
+    except Exception as e:
+        st.error(f"Model inference failed: {e}")
+        st.stop()
+
+    probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
+    sorted_idx = np.argsort(probs)[::-1]
+    cumsum = np.cumsum(probs[sorted_idx])
+
+    return [
+        {"label": id2label.get(int(i), f"class_{i}"), "probability": float(probs[i])}
+        for i in sorted_idx[:int(np.searchsorted(cumsum, TOP_PROB)) + 1]
+    ]
+
+
+st.set_page_config(page_title="Paper Classifier", page_icon="📄", layout="centered")
+st.title("📄 Research Paper Classifier")
+st.caption("Predicts arXiv categories from a paper's title and abstract.")
+
+load_model()
+
+title = st.text_input("Title *", placeholder="e.g. Attention Is All You Need", max_chars=MAX_CHARS)
+abstract = st.text_area("Abstract *", placeholder="Paste the abstract here…", height=200, max_chars=MAX_CHARS)
+
+col1, col2 = st.columns([1, 4])
+classify_btn = col1.button("Classify", type="primary", use_container_width=True)
+col2.button("Clear", on_click=lambda: None, use_container_width=True)
+
+if classify_btn:
+    error = validate(title, abstract)
+    if error:
+        st.warning(error)
+    if not clean_text(title) and not clean_text(abstract):
+        st.warning("Please enter at least a title or an abstract.")
+    else:
+        with st.spinner("Classifying…"):
+            results = predict(title, abstract)
+
+        st.divider()
+        st.subheader("Predicted categories")
+        st.caption(
+            f"Top-95% probability set * {len(results)} label(s) * "
+            f"combined probability {sum(r['probability'] for r in results):.1%}"
+        )
+
+        for r in results:
+            col1, col2 = st.columns([3, 1])
+            col1.markdown(f"**{r['label']}**")
+            col2.markdown(f"`{r['probability']:.1%}`")
+            st.progress(float(r["probability"]))