Delete num max classes

app.py

@@ -1,5 +1,3 @@
-import os
-import json
 import numpy as np
 import pandas as pd
 import streamlit as st
@@ -7,9 +5,9 @@ import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import plotly.express as px
 
-MODEL_REPO = os.getenv("MODEL_REPO", "ChocoLord/paper-classifier-model")
-MAX_LENGTH = int(os.getenv("MAX_LENGTH", "512"))
-TOP_P = float(os.getenv("TOP_P", "0.95"))
+MODEL_REPO = "ChocoLord/paper-classifier-model"
+MAX_LENGTH = 512
+TOP_P = 0.95
 
 st.set_page_config(page_title="Paper classifier", layout="wide")
 st.title("Paper classifier")
@@ -19,20 +17,12 @@ def load_artifacts():
     tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
     model = AutoModelForSequenceClassification.from_pretrained(MODEL_REPO)
     model.eval()
+    return tokenizer, model
 
-    id2label = model.config.id2label
-    if id2label is None or len(id2label) == 0:
-        raise ValueError("Model config must contain id2label.")
-
-    id2label = {int(k): v for k, v in id2label.items()} if not isinstance(list(id2label.keys())[0], int) else id2label
-    return tokenizer, model, id2label
-
-tokenizer, model, id2label = load_artifacts()
+tokenizer, model = load_artifacts()
 
 def predict(title: str, summary: str):
-    title = title or ""
-    summary = summary or ""
-    text = f"{title}\n{summary}".strip()
+    text = f"{title or ''}\n{summary or ''}".strip()
 
     inputs = tokenizer(
         text,
@@ -46,7 +36,8 @@ def predict(title: str, summary: str):
         logits = model(**inputs).logits
         probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
 
-    labels = [id2label[i] for i in range(len(probs))]
+    labels = [model.config.id2label[i] for i in range(len(probs))]
+
     df = pd.DataFrame({
         "class_name": labels,
         "predicted_proba": probs,
@@ -61,8 +52,6 @@ def predict(title: str, summary: str):
 title = st.text_input("Title")
 summary = st.text_area("Summary", height=250)
 
-n_value = st.number_input("Max classes to display in text output", min_value=1, max_value=100, value=20, step=1)
-
 if st.button("Classify", type="primary"):
     if not title.strip() and not summary.strip():
         st.warning("Enter title and/or summary.")
@@ -70,15 +59,9 @@ if st.button("Classify", type="primary"):
         df, selected_df = predict(title, summary)
 
         st.subheader("Selected classes")
-        st.write(
-            f"Top classes whose cumulative predicted probability reaches at least {TOP_P:.2f}. "
-            f"Selected {len(selected_df)} classes with total probability {selected_df['predicted_proba'].sum():.4f}."
-        )
-
-        text_df = selected_df.head(int(n_value)).copy()
         lines = [
             f"{i+1}. {row.class_name} — {row.predicted_proba:.4f}"
-            for i, row in text_df.iterrows()
+            for i, row in selected_df.iterrows()
         ]
         st.text("\n".join(lines))
 
@@ -87,7 +70,6 @@ if st.button("Classify", type="primary"):
             df,
             x="class_name",
             y="predicted_proba",
-            hover_data=["cumsum"],
         )
         fig.update_layout(
             xaxis_title="Class",