File size: 712 Bytes
97b056e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
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)}")