import streamlit as st import time from transformers import pipeline import torch trust_remote_code=True st.markdown('## story-generation from Breadlicker45') use_auth_token=True @st.cache(allow_output_mutation=True, suppress_st_warning =True, show_spinner=False) def get_model(): return pipeline('text-generation', model=model, do_sample=True) col1, col2 = st.columns([2,1]) with st.sidebar: st.markdown('## Model Parameters') max_length = st.slider('Max text length', 80, 2000, 80) min_length = st.slider('Min text length', 80, 500, 80) num_beams = st.slider('N° tree beams search', 1, 15, 1) temperature = st.slider('temperature', 0.0, 1.0, 0.5, 0.1) early_stopping = st.selectbox( 'Early stopping text generation', ('True', 'False'), key={'True' : True, 'False': False}, index=0) no_ngram_repeat = st.slider('Max repetition limit', 1, 3, 1) with col1: prompt= st.text_area('Your prompt here', '''in a world''') with col2: select_model = st.radio( "Select the model to use:", ('StoryPy', 'null'), index = 0) if select_model == 'StoryPy': model = 'BreadAi/StoryPy' elif select_model == 'null': model = 'BreadAi/StoryPy' elif select_model == 'MuseNeo': model = 'breadlicker45/MuseNeo' elif select_model == 'MusePy-1-1': model = 'BreadAi/MusePy-1-1' elif select_model == 'MuseCan': model = 'BreadAi/MuseCan' 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=max_length, no_repeat_ngram_size=no_ngram_repeat, early_stopping=early_stopping, num_beams=num_beams, min_length=min_length, temperature=temperature) 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)