import streamlit as st

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

title = {en: "Papers classification", ru: "Классификация статей"}
model_label = {en: "Select model", ru: "Выберете модель"}
abstract_label = {en: "Abstract", ru: "Аннотация"}


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


lang = st.pills(label=langs_str, options=langs, default=en)
st.title(title[lang])
model_name = st.selectbox(
    model_label[lang], options=class_models.get_model_names_by_lang(lang)
)
abstract = st.text_area(abstract_label[lang])

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