pages/Text_Generation.py

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

import streamlit as st
import AI.text_generation as ai

def run():

    st.set_page_config(page_title="Text Generation", page_icon="📈")

    st.markdown("# Text Generation")
    st.write('Generate some text.')
    input = st.text_input('Enter your prompt 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.generate(model, input)
            progress_bar.progress(90)
            status_text.text("Clearing Cache")
            ai.clearCache("cpu", model)
            progress_bar.progress(100)
            status_text.text("Done")
            st.code(output[0].get('generated_text'))
            # st.success('Done')

if __name__ == '__main__':
    run()