import streamlit as st
from transformers import pipeline
import os

# Set Hugging Face cache directory
os.environ['TRANSFORMERS_CACHE'] = os.getenv('HF_HOME', os.path.expanduser('~/.cache/huggingface/hub'))

# Function to load all three models
@st.cache_resource
def load_models():
    bart_summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
    t5_summarizer = pipeline("summarization", model="t5-large")
    pegasus_summarizer = pipeline("summarization", model="google/pegasus-cnn_dailymail")
    return bart_summarizer, t5_summarizer, pegasus_summarizer

# Streamlit app layout
st.title("Text Summarization with Pre-trained Models: BART, T5, Pegasus")

# Load models
with st.spinner("Loading models..."):
    bart_model, t5_model, pegasus_model = load_models()

# Input text
text_input = st.text_area("Enter text to summarize:")

# User input for min and max words
st.sidebar.header("Summary Length Settings")
min_words = st.sidebar.slider("Minimum words in summary:", 10, 100, 50, step=5)
max_words = st.sidebar.slider("Maximum words in summary:", min_words + 10, 300, 150, step=10)

if text_input:
    word_count = len(text_input.split())
    st.write(f"**Input Word Count:** {word_count}")

    if st.button("Generate Summaries"):
        with st.spinner("Generating summaries..."):
            # Generate summaries with dynamic length constraints
            bart_summary = bart_model(
                text_input,
                max_length=max_words,
                min_length=min_words,
                num_beams=4,
                early_stopping=True
            )[0]['summary_text']

            t5_summary = t5_model(
                text_input,
                max_length=max_words,
                min_length=min_words,
                num_beams=4,
                early_stopping=True
            )[0]['summary_text']

            pegasus_summary = pegasus_model(
                text_input,
                max_length=max_words,
                min_length=min_words,
                num_beams=4,
                early_stopping=True
            )[0]['summary_text']

        # Display summaries
        st.subheader("BART Summary")
        st.write(bart_summary)
        st.write(f"**Word Count:** {len(bart_summary.split())}")

        st.subheader("T5 Summary")
        st.write(t5_summary)
        st.write(f"**Word Count:** {len(t5_summary.split())}")

        st.subheader("Pegasus Summary")
        st.write(pegasus_summary)
        st.write(f"**Word Count:** {len(pegasus_summary.split())}")
else:
    st.warning("Please enter text to summarize.")