import streamlit as st
import time
from multipage import MultiPage
from transformers import pipeline
import torch


def app():
    st.markdown('## Text Generation task')
    st.write('Write something and AI will continue the sentence ')
    st.markdown('## ')

    @st.cache(allow_output_mutation=True, suppress_st_warning =True, show_spinner=False)
    def get_model():
        return pipeline('text-generation', model=model, do_sample=True, skip_special_tokens=True)
        
    col1, col2 = st.columns([2,1])

        
    with col1:
        prompt= st.text_area('Your prompt here',
            '''Who is Elon Musk?''') 
            
    with col2:
        select_model = st.radio(
            "Select the model to use:",
            ('OPT-125m', 'OPT-350m'), index = 1)
        if select_model == 'OPT-350m':
            model = 'facebook/opt-350m'
        elif select_model == 'OPT-125m':
            model = 'facebook/opt-125m'

        with st.spinner('Loading Model... (This may take a while)'):
            generator = get_model()    
            st.success('Model loaded correctly!')

    with col1:    
        gen = st.info('Generating text...')
        answer = generator(prompt,
                                max_length=80, no_repeat_ngram_size=2,
                                 early_stopping=True, num_beams=8,
                                   skip_special_tokens=True)                      
        gen.empty()                      
                        
    lst = answer[0]['generated_text']
    
    t = st.empty()
    for i in range(len(lst)):
        t.markdown("#### %s..." % lst[0:i])
        time.sleep(0.04)