import streamlit as st
import torch 
from transformers import pipeline
from utils.process_data import generate_chunks, pdf_to_text

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

st.set_page_config(page_title="Summarizer", page_icon="⚖️")
st.title("Summarize Text")
st.subheader("🚗 🔗 Transformers Summarization Pipeline")

max = st.slider('Select max', 50, 500, step=10, value=150)
min = st.slider('Select min', 10, 450, step=10, value=50)
do_sample = st.checkbox("Do sample", value=False)

sentence = st.text_area('Please paste your article:', height=50)
button = st.button("Summarize")

@st.cache_data
def load_summarizer():
    model = pipeline("summarization", model=st.secrets["SUM_MODEL"], device=device)
    return model

with st.spinner("Generating Summary.."):
    if button and sentence:
        chunks = generate_chunks(sentence)
        summarizer = load_summarizer()
        res = summarizer(chunks,
                         max_length=max, 
                         min_length=min, 
                         do_sample=do_sample)
        text = ' '.join([summ['summary_text'] for summ in res])
        st.write(text)

st.divider()

st.subheader('🚙🔗 Summarize PDF')
pdf_path = st.file_uploader('Upload your PDF Document', type='pdf')
button2 = st.button("Summarize PDF")

if pdf_path is not None and button2:
    text = pdf_to_text(pdf_path)
    with st.spinner("Generating PDF Summary.."):
        chunks = generate_chunks(text)
        summarizer = load_summarizer()
        res = summarizer(chunks,
                        max_length=max, 
                        min_length=min, 
                        do_sample=do_sample)
        text_sum = ' '.join([summ['summary_text'] for summ in res])
        st.write(text_sum)