Update src/streamlit_app.py

src/streamlit_app.py

@@ -33,10 +33,15 @@ nltk.download('punkt_tab', download_dir=nltk_data_dir, quiet=True)
 
 
 # Load models
-summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=torch.device("cpu"))
+# summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=torch.device("cpu"))
+summarizer = pipeline(
+    "summarization",
+    model="sshleifer/distilbart-cnn-12-6",
+    from_flax=True,
+    device=-1  # CPU mode, use device=0 for GPU
+)
 qa_pipeline = pipeline("question-answering", model="distilbert-base-cased-distilled-squad", device=torch.device("cpu"))
 embedding_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-
 vector_dbs = {}  # Dictionary to store multiple vector databases, keyed by document title
 extracted_texts = {} # Dictionary to store extracted text, keyed by document title
 current_doc_title = None