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)}")