lenyamal/hftuning

1import streamlit as st2import torch3import numpy as np4from transformers import AutoTokenizer, AutoModelForSequenceClassification5import os6 7dir = os.path.dirname(os.path.abspath(__file__))8 9@st.cache_resource10def load_model():11    model = AutoModelForSequenceClassification.from_pretrained(os.path.join(dir, "arxiv_classifier_final"))12    tokenizer = AutoTokenizer.from_pretrained(os.path.join(dir, "arxiv_classifier_final"))13    id2label = model.config.id2label14    return model, tokenizer, id2label15 16model, tokenizer, id2label = load_model()17 18st.title("📄 arXiv Paper Classifier")19 20title = st.text_input("Название статьи")21abstract = st.text_area("Абстракт (опционально)")22 23if st.button("Классифицировать"):24    if not title:25        st.error("Введите название статьи")26    else:27        text = f"Title: {title}\nAbstract: {abstract}" if abstract else title28        29        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)30        31        with torch.no_grad():32            outputs = model(**inputs)33            probs = torch.nn.functional.softmax(outputs.logits[0], dim=-1).numpy()34        35        indices = np.argsort(probs)[::-1]36        cumulative = 0.037        result = []38        for idx in indices:39            prob = probs[idx]40            cumulative += prob41            result.append((id2label[idx], prob))42            if cumulative >= 0.95:43                break44        45        st.subheader("Предсказанные категории (top-95%):")46        for label, prob in result:47            st.write(f"- **{label}**: {prob*100:.1f}%")