import streamlit as st
from keybert import KeyBERT
from transformers.pipelines import pipeline

st.title('Title Summarization')


path='checkpoint-210000'
hf_model = pipeline("feature-extraction",model=path,tokenizer='bert-base-uncased' )
etsy_model = KeyBERT(model=hf_model)


text = st.text_area("Enter title here for the summarization:")

if st.button('Generate'):
    out = etsy_model.extract_keywords(text, keyphrase_ngram_range=(1,3), stop_words='english', top_n=20)
    if len(out[0][0]) < 20:
        out = etsy_model.extract_keywords(text, keyphrase_ngram_range=(1,4), stop_words='english', top_n=20)

    if len(out)==20:
        show=[out[0],out[4],out[9],out[14],out[19]]
    else:
        show=out
    st.json(show)