Mirageinv/Article_classifier

1import json2import streamlit as st3import torch4import torch.nn.functional as F5from transformers import AutoTokenizer, DistilBertForSequenceClassification6 7CHECKPOINT_PATH = "checkpoints/epoch_8.pt"8LABELS_PATH = "checkpoints/labels_info.json"9 10with open(LABELS_PATH, 'r') as f:11    LABELS = json.load(f)12 13print(len(LABELS))14BASE_MODEL = "distilbert-base-cased"15 16@st.cache_resource17def load_model():18    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)19    # The same model20    model = DistilBertForSequenceClassification.from_pretrained(BASE_MODEL, num_labels=len(LABELS))21    state_dict = torch.load(CHECKPOINT_PATH, map_location=torch.device('cpu'))22    model.load_state_dict(state_dict)23    model.eval()24 25    return tokenizer, model26 27tokenizer, model = load_model()28 29st.title("Классификатор научных статей по заголовку и описанию")30st.write("Введите название и аннотацию статьи для предсказания её тематики по таксономии arxiv.org")31 32title = st.text_input("Название статьи:")33abstract = st.text_area("Аннотация (abstract):")34 35if st.button("Классифицировать"):36    if not title and not abstract:37        st.warning("Введите хотя бы название статьи.")38    else:39        text = title if not abstract else f"{title} {abstract}"40 41        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=256)42 43        with torch.no_grad():44            outputs = model(**inputs)45            probs = F.softmax(outputs.logits, dim=1).squeeze()46 47        label_probs = [(label, prob.item()) for label, prob in zip(list(LABELS.values()), probs)]48 49        # Sorting for getting 95% afterwards50        label_probs.sort(key=lambda x: x[1], reverse=True)51 52        cumulative = 0.053        top_labels = []54        for label, prob in label_probs:55            cumulative += prob56            top_labels.append((label, prob))57            if cumulative >= 0.95:58                break59 60        # Вывод61        st.subheader("Наиболее вероятные тематики (суммарно ≥95%):")62        for label, prob in top_labels:63            st.write(f"**{label}** — {prob * 100:.2f}%")64