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)
def get_model():
    return pipeline('text-generation', model=model, skip_special_tokens=True)

col1, col2 = st.beta_columns([2,1])

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'

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

with col1:
    prompt= st.text_area('Your prompt here',
        '''AI will help humanity?''') 
        
    answer = generator(prompt, max_length=60, no_repeat_ngram_size=2, early_stopping=True, num_beams=5)
    lst = answer[0]['generated_text']

    # answer = ['ciao come stai stutto bene']
    # lst = ''.join(answer)
    
t = st.empty()
for i in range(len(lst)):
    t.markdown(" %s..." % lst[0:i])
    time.sleep(0.04)