import streamlit as st
from transformers import pipeline

@st.cache_resource
def summarization_model():
    model_name = "google/pegasus-xsum"
    summarizer = pipeline(
        model=model_name, 
        tokenizer=model_name, 
        task="summarization"
    )
    return summarizer

def summarization_main():
    st.markdown("<h2 style='text-align: center; color:grey;'>Text Summarization</h2>", unsafe_allow_html=True)
    st.markdown("<h3 style='text-align: left; color:#F63366; font-size:18px;'><b>What is text summarization about?<b></h3>", unsafe_allow_html=True)
    st.write("Text summarization is producing a shorter version of a given text while preserving its important information.")
    st.markdown('___')
    source = st.radio("How would you like to start? Choose an option below", ["I want to input some text", "I want to upload a file"])
    if source == "I want to input some text":
        sample_text = ""
        text = st.text_area("Input a text in English (10,000 characters max) or use the example below", value=sample_text, max_chars=10000, height=330)
        
        button = st.button("Get summary")
        if button:
            with st.spinner(text="Loading summarization model..."):
                summarizer = summarization_model()
            with st.spinner(text="Summarizing text..."):
                summary = summarizer(text, max_length=130, min_length=30)
                st.text(summary[0]["summary_text"])

    elif source == "I want to upload a file":
        uploaded_file = st.file_uploader("Choose a .txt file to upload", type=["txt"])
        if uploaded_file is not None:
            raw_text = str(uploaded_file.read(),"utf-8")
            text = st.text_area("", value=raw_text, height=330)
            button = st.button("Get summary")
            if button:
                with st.spinner(text="Loading summarization model..."):
                    summarizer = summarization_model()
                with st.spinner(text="Summarizing text..."):
                    summary = summarizer(text, max_length=130, min_length=30)
                    st.text(summary[0]["summary_text"])