Update app.py

app.py

@@ -1,39 +1,40 @@
 import streamlit as st
 from transformers import pipeline
+import time
 
-# Load models and tokenizers using Hugging Face's pipeline
-def load_pipelines():
-    bart_pipeline = pipeline("summarization", model="facebook/bart-large-cnn")
-    t5_pipeline = pipeline("summarization", model="t5-large")
-    pegasus_pipeline = pipeline("summarization", model="google/pegasus-cnn_dailymail")
-    
-    return {
-        'BART': bart_pipeline,
-        'T5': t5_pipeline,
-        'Pegasus': pegasus_pipeline,
-    }
-
-prompt = """
-Summarize the below paragraph
-"""
+# Load selected model's pipeline
+@st.cache_resource
+def load_pipeline(model_name):
+    if model_name == 'BART':
+        return pipeline("summarization", model="facebook/bart-large-cnn")
+    elif model_name == 'T5':
+        return pipeline("summarization", model="t5-large")
+    elif model_name == 'Pegasus':
+        return pipeline("summarization", model="google/pegasus-cnn_dailymail")
 
 # Streamlit app layout
 st.title("Text Summarization with Pre-trained Models (BART, T5, Pegasus)")
 
 text_input = st.text_area("Enter text to summarize:")
 
-if st.button("Generate Summary"):
-    if text_input:
-        pipelines = load_pipelines()
-        summaries = {}
-        for model_name, pipeline in pipelines.items():
-            summary = pipeline(f"{prompt}\n{text_input}", max_length=150, min_length=50, length_penalty=2.0, num_beams=4, early_stopping=True)[0]['summary_text']
-            summaries[model_name] = summary
+if text_input:
+    # Display word count of input text
+    word_count = len(text_input.split())
+    st.write(f"**Word Count:** {word_count}")
+
+    model_choice = st.selectbox("Choose a model:", ['BART', 'T5', 'Pegasus'])
+
+    if st.button("Generate Summary"):
+        with st.spinner(f"Generating summary using {model_choice}..."):
+            start_time = time.time()
+            summarizer = load_pipeline(model_choice)
+            summary = summarizer(text_input, max_length=150, min_length=50, length_penalty=2.0, num_beams=4, early_stopping=True)[0]['summary_text']
+            end_time = time.time()
+            summary_word_count = len(summary.split())
 
-        st.subheader("Summaries")
-        for model_name, summary in summaries.items():
-            st.write(f"**{model_name}**")
-            st.write(summary.replace('<n>', ''))
-            st.write("---")
-    else:
-        st.error("Please enter text to summarize.")
+        st.subheader(f"Summary using {model_choice}")
+        st.write(summary.replace('<n>', ''))
+        st.write(f"**Summary Word Count:** {summary_word_count}")
+        st.write(f"**Time Taken:** {end_time - start_time:.2f} seconds")
+else:
+    st.error("Please enter text to summarize.")