Spaces:
Build error
Build error
File size: 4,514 Bytes
e479683 8d8846d e479683 8d8846d e479683 1b2eba5 8d8846d e479683 8d8846d |
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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
import streamlit as st
import torch
import joblib
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from huggingface_hub import hf_hub_download
@st.cache_resource
def load_model():
repo_id = "YakovPodlesnov/best-article-classifier"
model = AutoModelForSequenceClassification.from_pretrained(repo_id)
tokenizer = AutoTokenizer.from_pretrained(repo_id)
label_encoder_path = hf_hub_download(
repo_id=repo_id,
filename="label_encoder.joblib"
)
label_encoder = joblib.load(label_encoder_path)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
return model, tokenizer, label_encoder, device
def predict(text, model, tokenizer, label_encoder, device):
inputs = tokenizer(
text,
padding="max_length",
truncation=True,
max_length=512,
return_tensors="pt"
).to(device)
with torch.no_grad():
outputs = model(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=-1).cpu().numpy()[0]
sorted_indices = np.argsort(probs)[::-1]
cumulative = 0.0
results = []
for idx in sorted_indices:
if cumulative >= 0.95:
break
label = label_encoder.inverse_transform([idx])[0]
prob = float(probs[idx])
results.append((label, prob))
cumulative += prob
return results
st.set_page_config(page_title="arXiv Classifier", layout="wide")
st.title("Классификатор научных статей")
st.markdown("Определяем тематику статьи по заголовку и аннотации (arXiv)")
with st.spinner("Инициализация модели..."):
try:
model, tokenizer, label_encoder, device = load_model()
except Exception as e:
st.error(f"Ошибка загрузки модели: {str(e)}")
st.stop()
col1, col2 = st.columns([1, 1])
with col1:
title = st.text_input("Название статьи:",
placeholder="Attention Is All You Need",
help="Введите полное название статьи")
with col2:
abstract = st.text_area("Аннотация:",
height=150,
placeholder="We propose a new simple network architecture...",
help="Введите текст аннотации (необязательно)")
if st.button("Определить категории", type="primary"):
if not title:
st.error("Пожалуйста, введите название статьи")
else:
text = title + (" " + abstract if abstract else "")
with st.spinner("Анализируем текст..."):
try:
predictions = predict(text, model, tokenizer, label_encoder, device)
except Exception as e:
st.error(f"Ошибка предсказания: {str(e)}")
st.stop()
st.subheader(" Результаты классификации")
st.markdown("""
Расшифровку названий, например для cs, можно посмотреть здесь: https://arxiv.org/archive/cs
""")
if not predictions:
st.warning("Не удалось определить категории")
else:
cols = st.columns([2, 3])
with cols[0]:
st.markdown("**Топ-категории:**")
for label, prob in predictions:
st.markdown(f"▸ {label} ({prob:.1%})")
with cols[1]:
labels = [p[0] for p in predictions]
probs = [p[1] for p in predictions]
st.bar_chart(
dict(zip(labels, probs)),
use_container_width=True,
color="#FF4B4B"
)
st.markdown("---")
st.markdown("""
### Как это работает?
1. Введите **название** научной статьи
2. При необходимости добавьте **аннотацию**
3. Нажмите кнопку для классификации статьи
4. Получите предсказанные категории arXiv
**Модель: scibert_scivocab_uncased**
""")
if __name__ == "__main__":
pass |