import streamlit as st

from category_classification.models import models as class_models
from common import Input

@st.cache_data
def load_class_model(name):
    model = class_models.get_model(name)
    return model


model_name = st.selectbox(
    "Select model",
    options=class_models.get_all_model_names()
)
title = st.text_area("Paper title")
abstract = st.text_area("Paper abstract")

if title:
    input = Input(title=title, abstract=abstract)
    model = load_class_model(model_name)
    out = model(input)
    st.json(out)