Spaces:
Runtime error
Runtime error
| 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('## ') | |
| 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) | |