import streamlit as st
import time
from transformers import pipeline
import torch
trust_remote_code=True
st.markdown('## Text-generation gpt Muse 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)

    top_k = st.slider('top_k', 1, 10,  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)

    st.markdown('## how to convert it into midi. go to this site https://mrcheeze.github.io/musenet-midi/ and then paste the numbers/musenet encoders you get from the ai into the big box and then click export midi')

    
with col1:
    prompt= st.text_area('Your prompt here',
        '''2623 2619 3970 3976 2607 3973 2735 3973 2598 3985 2726 3973 2607 4009 2735 3973 2598 3973 2726 3973 2607 3973 2735 4009''') 
        
with col2:
    select_model = st.radio(
        "Select the model to use:",
        ('MusePy', 'MuseMini', 'MusePy-1-1', 'MuseCan', 'MuseCan-1-2'), index = 4)

    if select_model == 'MusePy':
        model = 'breadlicker45/MusePy'
    elif select_model == 'MuseNeo':
        model = 'BreadAi/MuseMini'
    elif select_model == 'MusePy-1-1':
        model = 'BreadAi/MusePy-1-1'
    elif select_model == 'MuseCan':
        model = 'BreadAi/MuseCan'
    elif select_model == 'MuseCan-1-2':
        model = 'BreadAi/MuseCan-1-2'

    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, top_k=top_k)                      
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)