import streamlit as st
import joblib,torch
import time
device = 'cuda' if torch.cuda.is_available() else 'cpu'
loaded_tokenizer = joblib.load("finalized_tokenizer.sav")
loaded_model = joblib.load("finalized_model.sav")

st.title('Text Summarization using Pegasus')

txt = st.text_area('Enter Text to summarize here', '')

if st.button('Summarize'):
    with st.spinner('Summarizing..'):
        batch = loaded_tokenizer(txt, truncation=True, padding='longest', return_tensors="pt").to(device)
        translated = loaded_model.generate(**batch)
        tgt_text = loaded_tokenizer.batch_decode(translated, skip_special_tokens=True)
    st.success('Summarized Text')
    st.subheader(tgt_text[0])