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

st.markdown('## Text-generation OPT from Meta ')

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

with col1:
    prompt= st.text_area('Your prompt here',
        '''AI will help humanity? ''') 
        
with col2:
    select_model = st.radio(
        "Select the model to use:",
        ('OPT-125m', 'OPT-350m', 'OPT-1.3b'), index = 1)

    if select_model == 'OPT-1.3b':
        model = 'facebook/opt-1.3b'
    elif select_model == 'OPT-350m':
        model = 'facebook/opt-350m'
    elif select_model == 'OPT-125m':
        model = 'facebook/opt-125m'

    if select_model:
        with st.spinner('Loading Model... (This may take a while)'):
            generator = get_model()    
            st.success('Model loaded correctly!')
     
answer = generator(prompt,
                       max_length=100,
                       no_repeat_ngram_size=3, early_stopping=True, num_beams=10, skip_special_tokens=True)
                       
lst = answer[0]['generated_text']
   
t = st.empty()
for i in range(len(lst)):
    t.markdown(" %s..." % lst[0:i])
    time.sleep(0.04)