Update app.py

app.py

@@ -1,18 +1,54 @@
+import os
 import streamlit as st
 from transformers import pipeline
 import traceback
 
+# Use Hugging Face Spaces' recommended persistent storage
+CACHE_DIR = os.path.join(os.getcwd(), "model_cache")
+
+def ensure_cache_dir():
+    """
+    Ensure the cache directory exists.
+    
+    Returns:
+        str: Path to the cache directory
+    """
+    os.makedirs(CACHE_DIR, exist_ok=True)
+    return CACHE_DIR
+
 def load_pipelines():
     """
-    Load summarization pipelines with error handling.
+    Load summarization pipelines with persistent caching.
     
     Returns:
         dict: Dictionary of model pipelines
     """
     try:
-        bart_pipeline = pipeline("summarization", model="facebook/bart-large-cnn")
-        t5_pipeline = pipeline("summarization", model="t5-large")
-        pegasus_pipeline = pipeline("pegasus-cnn_dailymail")
+        # Ensure cache directory exists
+        cache_dir = ensure_cache_dir()
+        
+        # Define model paths within the cache directory
+        bart_cache = os.path.join(cache_dir, "bart-large-cnn")
+        t5_cache = os.path.join(cache_dir, "t5-large")
+        pegasus_cache = os.path.join(cache_dir, "pegasus-cnn_dailymail")
+        
+        # Load pipelines with explicit cache directories
+        bart_pipeline = pipeline(
+            "summarization", 
+            model="facebook/bart-large-cnn", 
+            cache_dir=bart_cache
+        )
+        t5_pipeline = pipeline(
+            "summarization", 
+            model="t5-large", 
+            cache_dir=t5_cache
+        )
+        pegasus_pipeline = pipeline(
+            "summarization", 
+            model="google/pegasus-cnn_dailymail", 
+            cache_dir=pegasus_cache
+        )
+        
         return {
             'BART': bart_pipeline, 
             'T5': t5_pipeline, 
@@ -52,6 +88,9 @@ def generate_summary(pipeline, text, model_name):
 def main():
     st.title("Text Summarization with Pre-trained Models")
     
+    # Display cache directory info (optional)
+    st.info(f"Models will be cached in: {CACHE_DIR}")
+    
     # Text input
     text_input = st.text_area("Enter text to summarize:")