# Author: Ricardo Lisboa Santos
# Creation date: 2024-01-10

import streamlit as st
import AI.zero_shot_classification as ai

def run():

    st.set_page_config(page_title="Zero Shot Classification", page_icon="📈")

    st.markdown("# Zero Shot Classification")
    st.write('Add some labels and then add a text to classify.')
    labels_box = st.text_area('Enter your labels here.')
    st.markdown('<p class="breadcrumb">{}</p>'.format(labels_box))
    text = st.text_input('Enter your text here.')
    if st.button('Click me to run'):
        progress_bar = st.sidebar.progress(0)
        status_text = st.sidebar.empty()
        with st.spinner(text='Loading Model'):
            status_text.text("Getting Device")
            device = ai.getDevice("cpu")
            progress_bar.progress(30)
            status_text.text("Loading Model")
            model = ai.loadGenerator(device)
            progress_bar.progress(60)
            status_text.text("Generating Answer")
            output = ai.classify(model, question, context)
            progress_bar.progress(90)
            status_text.text("Clearing Cache")
            ai.clearCache("cpu", model)
            progress_bar.progress(100)
            status_text.text("Done")
            st.code(output.get('answer'))
            # st.success('Done')

if __name__ == '__main__':
    run()