Rksafaryan's picture
Update app.py
4972136
import streamlit as st
st.markdown("### Article Classifier")
st.markdown("<img width=200px src='https://media.istockphoto.com/photos/funny-cat-is-studying-chemistry-picture-id526831620'>", unsafe_allow_html=True)
st.markdown("This is a tool for classifying article category by it's title and summary. \n Follow the instructions below")
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import torch
import numpy as np
tokenizer = AutoTokenizer.from_pretrained("Wi/arxiv-topics-distilbert-base-cased")
model = AutoModelForSequenceClassification.from_pretrained("Wi/arxiv-topics-distilbert-base-cased")
title = st.text_area("Put the title here")
abstract = st.text_area("Put the abstract here")
if st.button('Press when ready'):
text = 'Title:' + title + '\n' + 'Abstract:' + abstract
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
logits = model(**inputs).logits
probs = logits.softmax(dim=-1).detach().cpu().flatten().numpy().tolist()
order = np.argsort(probs)[::-1]
i = 0
sum = 0
predicted_class_id = []
while sum < 0.95:
predicted_class_id.append(order[i])
sum += probs[order[i]]
i+=1
for id in predicted_class_id:
st.markdown(model.config.id2label[id])