Spaces:
Sleeping
Sleeping
File size: 3,310 Bytes
e47a867 a7e2485 78d506d a7e2485 431d797 78d506d a7e2485 431d797 a7e2485 0ae5cd0 a7e2485 9fef900 0ae5cd0 ce37e77 a7e2485 |
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 |
import streamlit as st
import torch
import torch.serialization
from transformers import AutoModelForSequenceClassification, AutoTokenizer
@st.cache_resource
def load_model():
#trained_model = 'TinyBERT_cls_model.pt'
#base_model = 'huawei-noah/TinyBERT_General_4L_312D'
trained_model = 'distilbert-base_cls_model.pt'
base_model = 'distilbert-base-uncased'
checkpoint = torch.load(trained_model,
map_location='cpu',
weights_only=False)
model = AutoModelForSequenceClassification.from_pretrained(
base_model,
num_labels=len(checkpoint['idx_to_category'])
)
model.load_state_dict(checkpoint['model_state_dict'])
tokenizer = checkpoint['tokenizer']
idx_to_category = checkpoint['idx_to_category']
return model, tokenizer, idx_to_category
def predict(title, abstract, model, tokenizer, idx_to_category, threshold=0.95):
text = f"{title} /n {abstract}" if abstract else title
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
outputs = model(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
sorted_probs, sorted_indices = torch.sort(probs, descending=True)
results = []
cumulative_prob = 0.0
for i in range(len(sorted_probs)):
if cumulative_prob >= threshold:
break
prob = sorted_probs[i].item()
results.append({
"category": idx_to_category[sorted_indices[i].item()],
"probability": prob
})
cumulative_prob += prob
return results, cumulative_prob
def main():
model, tokenizer, idx_to_category = load_model()
st.title("Классификатор статей")
st.markdown("Определение тематики научных статей по названию и аннотации")
with st.form("input_form"):
title = st.text_input("Название статьи*", placeholder="Введите название...")
abstract = st.text_area("Аннотация", placeholder="Введите текст аннотации (необязательно)...", height=150)
submitted = st.form_submit_button("Классифицировать")
if submitted and not title:
st.error("Пожалуйста, введите название статьи")
if submitted and title:
with st.spinner("Анализируем статью..."):
results, total_prob = predict(
title=title,
abstract=abstract,
model=model,
tokenizer=tokenizer,
idx_to_category=idx_to_category
)
st.success("Результаты классификации:")
st.metric("Общая вероятность", f"{total_prob*100:.1f}%")
for i, res in enumerate(results, 1):
col1, col2 = st.columns([1, 4])
with col1:
st.metric(f"Топ {i}", f"{res['probability']*100:.1f}%")
with col2:
st.progress(res['probability'], text=res['category'])
if __name__ == "__main__":
main()
|