Spaces:
Sleeping
Sleeping
File size: 2,568 Bytes
e8e0226 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | from pathlib import Path
import streamlit as st
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
MODEL_DIR = Path(__file__).resolve().parent / "model"
DEMO_LABELS = [
"cs.CL",
"math.NT",
"astro-ph.GA",
"q-bio.NC",
"physics.optics",
"econ.EM",
]
def make_text(title: str, abstract: str) -> str:
title = title.strip()
abstract = abstract.strip()
if abstract:
return f"Title: {title}\nAbstract: {abstract}"
return f"Title: {title}"
def top_95(labels, probabilities):
order = torch.argsort(probabilities, descending=True)
result = []
total = 0.0
for idx in order:
probability = float(probabilities[idx])
result.append((labels[int(idx)], probability))
total += probability
if total >= 0.95:
break
return result
@st.cache_resource(show_spinner=False)
def load_model():
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, use_fast=False)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
model.eval()
labels = [model.config.id2label[i] for i in range(model.config.num_labels)]
return tokenizer, model, labels
def predict(text: str):
tokenizer, model, labels = load_model()
encoded = tokenizer(
text,
truncation=True,
max_length=384,
padding=True,
return_tensors="pt",
)
with torch.inference_mode():
logits = model(**encoded).logits[0]
probabilities = torch.softmax(logits, dim=-1)
return top_95(labels, probabilities)
st.set_page_config(page_title="ArXiv Topic Classifier", page_icon=":books:", layout="centered")
st.title("ArXiv Topic Classifier")
st.caption("Классификатор тематики научной статьи по названию и abstract.")
title = st.text_input("Название статьи", placeholder="Attention Is All You Need")
abstract = st.text_area(
"Abstract",
height=180,
placeholder="Вставьте abstract. Если его нет, модель использует только название.",
)
if st.button("Классифицировать", type="primary", use_container_width=True):
if not title.strip() and not abstract.strip():
st.warning("Введите название, abstract или оба поля.")
else:
st.subheader("Топ-95% тематик")
for label, probability in predict(make_text(title, abstract)):
st.progress(probability, text=f"{label}: {probability:.1%}")
|