Spaces:
Runtime error
Runtime error
import streamlit as st | |
import time | |
from transformers import pipeline | |
import torch | |
st.markdown('## Text-generation OPT from Meta ') | |
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' | |
with st.spinner('Loading Model... (This may take a while)'): | |
generator = get_model() | |
st.success('Model loaded correctly!') | |
gen = st.info('Generating text...') | |
answer = generator(prompt, | |
max_length=100, | |
no_repeat_ngram_size=3, early_stopping=True, num_beams=10, 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) |