import streamlit as st

from back import load_setup, predict
from model import DistillBERTClass


st.markdown("""# Hi, let's find some arxiv tags of papers! <img width=50px src='https://tahayassine.me/media/icon_huecf151563ab8d381649e162e44411916_11613_512x512_fill_lanczos_center_3.png'>""",
            unsafe_allow_html=True)

title = st.text_area("$\LARGE \\text{Title of the paper:}$", placeholder='')
abstract = st.text_area("$\LARGE \\text{Abstract of the paper(if any):}$", placeholder='')

if st.button('Generate tags!'):
    if title == '':
        st.write('# Please, enter title!')
    else:
        text = title + '[SEP]' + abstract
        model, tokenizer = load_setup(path_to_model='distilbert1.bin', path_to_vocab='vocab_distilbert1.bin')
        prediction = predict(model, tokenizer, text)
        st.markdown(f'## Predicted tag(s): {prediction}')
else:
    pass