from t5 import t5
import torch
import streamlit as st
from streamlit import session_state as ses
import timeit


if "page" not in ses:
    ses["page"] = "run"
    ses.out = ""
    ses.time = ""


def __run_t5():
    try:
        ses.out = t5(ses.input_text)
    except:
        ses.out = "error"


def __test_t5():
    try:
        ses.time = timeit.timeit(__run_t5(), globals=globals(), number=20)
    except:
        ses.out = "error"


st.text_input("enter text", key="input_text", on_change=__run_t5)
st.code(ses.out, language="json")
if st.button("average time", on_click=__test_t5):
    st.metric("time", ses.time)