PaperClassifier / app.py
oracat's picture
Initial commit
97b056e
raw
history blame
712 Bytes
from transformers import AutoTokenizer, AutoModelForMaskedLM, pipeline
import streamlit as st
@st.cache_data
def prepare():
tokenizer = AutoTokenizer.from_pretrained("oracat/bert-paper-classifier")
model = AutoModelForSequenceClassification.from_pretrained("oracat/bert-paper-classifier")
def process(text):
pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)
result = pipe(text)[0]
return results['label']
prepare()
st.markdown("### Hello, paper classifier!")
title = st.text_input("Enter the title...")
abstract = st.text_area("... and maybe the abstract of the paper you want to classify")
text = "\n".join([title, abstract])
st.markdown(f"{process(text)}")