Update app.py

app.py

@@ -1,28 +1,19 @@
 import streamlit as st
 from transformers import pipeline
-from concurrent.futures import ThreadPoolExecutor
-import time
+import os
+
+# Set the cache directory
+CACHE_DIR = os.getenv('HF_HOME', os.path.expanduser('~/.cache/huggingface/hub'))
 
 # Function to load a single model
+@st.cache_resource
 def load_model(model_name):
     if model_name == 'BART':
-        return pipeline("summarization", model="facebook/bart-large-cnn")
+        return pipeline("summarization", model="facebook/bart-large-cnn", cache_dir=CACHE_DIR)
     elif model_name == 'T5':
-        return pipeline("summarization", model="t5-large")
+        return pipeline("summarization", model="t5-large", cache_dir=CACHE_DIR)
     elif model_name == 'Pegasus':
-        return pipeline("summarization", model="google/pegasus-cnn_dailymail")
-
-# Function to load all models concurrently and cache them
-@st.cache_resource
-def load_all_models():
-    model_names = ['BART', 'T5', 'Pegasus']
-    models = {}
-    with ThreadPoolExecutor() as executor:
-        futures = {executor.submit(load_model, name): name for name in model_names}
-        for future in futures:
-            model_name = futures[future]
-            models[model_name] = future.result()
-    return models
+        return pipeline("summarization", model="google/pegasus-cnn_dailymail", cache_dir=CACHE_DIR)
 
 # Streamlit app layout
 st.title("Text Summarization with Pre-trained Models (BART, T5, Pegasus)")
@@ -34,23 +25,16 @@ if text_input:
     word_count = len(text_input.split())
     st.write(f"**Word Count:** {word_count}")
 
-    if st.button("Generate Summaries"):
-        with st.spinner("Loading models and generating summaries..."):
-            start_time = time.time()
-            models = load_all_models()
-            summaries = {}
-            for model_name, model in models.items():
-                summary = model(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
-            end_time = time.time()
+    model_choice = st.selectbox("Choose a model:", ['BART', 'T5', 'Pegasus'])
 
-        st.subheader("Summaries")
-        for model_name, summary in summaries.items():
-            summary_word_count = len(summary.split())
-            st.write(f"**{model_name}**")
-            st.write(summary.replace('<n>', ''))
-            st.write(f"**Summary Word Count:** {summary_word_count}")
-            st.write("---")
-        st.write(f"**Total Time Taken:** {end_time - start_time:.2f} seconds")
+    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}")
 else:
     st.error("Please enter text to summarize.")