import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

tokenizer = AutoTokenizer.from_pretrained("impyadav/GPT2-FineTuned-Hinglish-Song-Generation")

model = AutoModelForCausalLM.from_pretrained("impyadav/GPT2-FineTuned-Hinglish-Song-Generation")

def get_song(line):
    lyricist = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer
    )
    return lyricist(line, max_length=150, num_return_sequences=3)
    

if __name__ == '__main__':
    st.title('AI Lyricist')
    st.write('Transformer Architecture : {}'.format('gpt-2'))
    st.subheader("Input")
    #st.write('Paste your query act here:')
    user_input = st.text_area('', height=25)  # height in pixel
    # st.markdown('<style>red{color:red}</style>')
    result = get_song(user_input)
    if st.button('Run'):
        st.write(result)