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)