import streamlit as st

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import nltk
nltk.download('punkt')

#tokenizer = AutoTokenizer.from_pretrained("fabiochiu/t5-small-medium-title-generation")
#model = AutoModelForSeq2SeqLM.from_pretrained("fabiochiu/t5-small-medium-title-generation")

tokenizer = AutoTokenizer.from_pretrained("Soooma/titles_gen")
model = AutoModelForSeq2SeqLM.from_pretrained("Soooma/titles_gen")

text = st.text_area('Enter an abstract to summerize, i.e. generate a title!', height=440)

max_input_length = 512

if text:
    inputs = ["summarize: " + text]
    inputs = tokenizer(inputs, max_length=max_input_length, truncation=True, return_tensors="pt")
    output = model.generate(**inputs, num_beams=8, do_sample=True, min_length=10, max_length=64)
    decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
    predicted_title = nltk.sent_tokenize(decoded_output.strip())[0]
    
    html_string = f"<h4>The predicted title is:</h4> \'{predicted_title}\'"

    st.markdown(html_string, unsafe_allow_html=True)