import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
#from summarizer import Summarizer

tokenizer = AutoTokenizer.from_pretrained("DemocracyStudio/generate_nft_content")
model = AutoModelForCausalLM.from_pretrained("DemocracyStudio/generate_nft_content")

#summarize=Summarizer()

st.title("Text generation for the marketing content of NFTs")
st.subheader("Course project 'NLP with transformers' at opencampus.sh, Spring 2022")
    
st.sidebar.image("bayc crown.png", use_column_width=True)
topics=["NFT", "Blockchain", "Metaverse"]
choice = st.sidebar.selectbox("Select one topic", topics)
    
if choice == 'NFT':
    keywords=st.text_area("Input 4 keywords here: (optional)")
    length=st.text_area("How long should be your text? (default: 512 words)")

    if st.button("Generate"):
        prompt = "<|startoftext|>"
        generated = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)        
        output = model.generate(generated, do_sample=True, top_k=50, max_length = 512, top_p=0.95, num_return_sequences=1)
        generated_text = tokenizer.decode(output, skip_special_tokens=True)
        #summary = summarize(generated_text, num_sentences=1)
      
        #st.text("Keywords: {}\n".format(keywords))
        #st.text("Length in number of words: {}\n".format(length))
        st.text("This is your tailored blog article:", generated_text)
        #st.text("This is a tweet-sized summary of your article: ", summary)
else:
    st.write("Topic not available yet")