Update app.py

app.py

@@ -1,5 +1,5 @@
 """
-Enhanced ArXiv RAG System - Hugging Face Spaces Compatible Version
+Enhanced ArXiv RAG System - Hugging Face Spaces Compatible Version (Fixed)
 """
 
 import os
@@ -13,6 +13,7 @@ from datetime import datetime, timedelta
 import logging
 import tempfile
 import shutil
+import spaces
 
 # Core ML libraries
 import torch
@@ -201,11 +202,16 @@ class EnhancedArxivRAG:
     def __init__(self):
         logger.info("Initializing Enhanced ArXiv RAG System for HF Spaces...")
         
-        # Use smaller, faster models for HF Spaces
+        # Use CPU-friendly models for HF Spaces
         self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-        self.reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-2-v2')  # Smaller reranker
-        self.summarizer = pipeline("summarization", model="facebook/bart-large-cnn", 
-                                  device=0 if torch.cuda.is_available() else -1)
+        self.reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-2-v2')
+        
+        # Initialize summarizer without GPU specification
+        self.summarizer = pipeline(
+            "summarization", 
+            model="facebook/bart-large-cnn",
+            device=-1  # Force CPU usage
+        )
         
         # Use simple vector store instead of ChromaDB for HF Spaces
         self.vector_store = SimpleVectorStore()
@@ -312,6 +318,7 @@ class EnhancedArxivRAG:
         
         return chunks
     
+    @spaces.GPU(duration=60)  # GPU decorator for processing
     def process_and_store(self, papers: List[Paper]):
         """Process papers and store in vector store"""
         logger.info("Processing and storing papers...")
@@ -405,6 +412,7 @@ class EnhancedArxivRAG:
         
         return final_results
     
+    @spaces.GPU(duration=30)  # GPU decorator for reranking
     def rerank_results(self, query: str, results: List[Dict], top_k: int = 5) -> List[Dict]:
         """Rerank results using cross-encoder"""
         if not results:
@@ -450,6 +458,7 @@ class EnhancedArxivRAG:
             return f"Based on the retrieved papers about '{query}', here are the key findings:\n\n" + \
                    "\n\n".join([chunk['document'][:150] + "..." for chunk in context_chunks[:2]])
     
+    @spaces.GPU(duration=120)  # Main GPU decorator for the full pipeline
     def search_and_answer(self, query: str, max_papers: int = 15,
                          top_k_retrieval: int = 10, top_k_rerank: int = 5,
                          categories: Optional[List[str]] = None,
@@ -700,4 +709,4 @@ def create_interface():
 # Launch interface
 if __name__ == "__main__":
     interface = create_interface()
-    interface.launch(share=True)
+    interface.launch()