Update app.py

app.py

@@ -2,39 +2,74 @@ import streamlit as st
 from transformers import pipeline
 import os
 
-# Set the cache directory
-CACHE_DIR = os.getenv('HF_HOME', os.path.expanduser('~/.cache/huggingface/hub'))
+# Set Hugging Face cache directory
+os.environ['TRANSFORMERS_CACHE'] = os.getenv('HF_HOME', os.path.expanduser('~/.cache/huggingface/hub'))
 
-# Function to load a single model
+# Function to load all three models
 @st.cache_resource
-def load_model(model_name):
-    if model_name == 'BART':
-        return pipeline("summarization", model="facebook/bart-large-cnn", cache_dir=CACHE_DIR)
-    elif model_name == 'T5':
-        return pipeline("summarization", model="t5-large", cache_dir=CACHE_DIR)
-    elif model_name == 'Pegasus':
-        return pipeline("summarization", model="google/pegasus-cnn_dailymail", cache_dir=CACHE_DIR)
+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)")
+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:
-    # 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}..."):
-            summarizer = load_model(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']
-        
-        summary_word_count = len(summary.split())
-        st.subheader(f"Summary using {model_choice}")
-        st.write(summary.replace('<n>', ''))
-        st.write(f"**Summary Word Count:** {summary_word_count}")
+    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.error("Please enter text to summarize.")
+    st.warning("Please enter text to summarize.")