Update app.py

app.py

@@ -1,5 +1,5 @@
 """
-Enhanced ArXiv RAG System - Hugging Face Spaces Compatible Version (Fixed)
+Enhanced ArXiv RAG System - Hugging Face Spaces Compatible Version
 """
 
 import os
@@ -13,7 +13,6 @@ from datetime import datetime, timedelta
 import logging
 import tempfile
 import shutil
-import spaces
 
 # Core ML libraries
 import torch
@@ -31,14 +30,14 @@ from nltk.stem import PorterStemmer
 
 # Download required NLTK data
 try:
-    nltk.data.find('tokenizers/punkt')
+    nltk.data.find("tokenizers/punkt")
 except LookupError:
-    nltk.download('punkt')
+    nltk.download("punkt")
 
 try:
-    nltk.data.find('corpora/stopwords')
+    nltk.data.find("corpora/stopwords")
 except LookupError:
-    nltk.download('stopwords')
+    nltk.download("stopwords")
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -75,7 +74,7 @@ class BM25Retriever:
         self.avg_doc_length = 0
         self.stemmer = PorterStemmer()
         try:
-            self.stop_words = set(stopwords.words('english'))
+            self.stop_words = set(stopwords.words("english"))
         except:
             self.stop_words = set()
     
@@ -151,7 +150,7 @@ class SimpleVectorStore:
     def query(self, query_embedding: List[float], n_results: int = 10) -> Dict:
         """Query the vector store"""
         if not self.embeddings:
-            return {'ids': [[]], 'documents': [[]], 'metadatas': [[]]}
+            return {"ids": [[]], "documents": [[]], "metadatas": [[]]}
         
         # Calculate cosine similarities
         query_embedding = np.array(query_embedding)
@@ -168,25 +167,25 @@ class SimpleVectorStore:
         top_indices = np.argsort(similarities)[::-1][:n_results]
         
         return {
-            'ids': [[self.ids[i] for i in top_indices]],
-            'documents': [[self.documents[i] for i in top_indices]],
-            'metadatas': [[self.metadatas[i] for i in top_indices]]
+            "ids": [[self.ids[i] for i in top_indices]],
+            "documents": [[self.documents[i] for i in top_indices]],
+            "metadatas": [[self.metadatas[i] for i in top_indices]]
         }
     
     def get(self, ids: Optional[List[str]] = None) -> Dict:
         """Get documents by IDs or all documents"""
         if ids is None:
             return {
-                'ids': self.ids,
-                'documents': self.documents,
-                'metadatas': self.metadatas
+                "ids": self.ids,
+                "documents": self.documents,
+                "metadatas": self.metadatas
             }
         else:
             indices = [self.ids.index(id_) for id_ in ids if id_ in self.ids]
             return {
-                'ids': [self.ids[i] for i in indices],
-                'documents': [self.documents[i] for i in indices],
-                'metadatas': [self.metadatas[i] for i in indices]
+                "ids": [self.ids[i] for i in indices],
+                "documents": [self.documents[i] for i in indices],
+                "metadatas": [self.metadatas[i] for i in indices]
             }
     
     def clear(self):
@@ -202,16 +201,29 @@ class EnhancedArxivRAG:
     def __init__(self):
         logger.info("Initializing Enhanced ArXiv RAG System 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')
-        
-        # Initialize summarizer without GPU specification
-        self.summarizer = pipeline(
-            "summarization", 
-            model="facebook/bart-large-cnn",
-            device=-1  # Force CPU usage
-        )
+        # Determine device (GPU if available, else CPU)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Using device: {self.device}")
+
+        # Load models with appropriate device settings
+        try:
+            logger.info("Loading embedding model...")
+            self.embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device=self.device)
+            logger.info("Embedding model loaded.")
+
+            logger.info("Loading reranker model...")
+            self.reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-2-v2", device=self.device)
+            logger.info("Reranker model loaded.")
+
+            logger.info("Loading summarizer model...")
+            # For pipeline, device_map="auto" is often better for ZeroGPU
+            # If issues persist, try device=0 for the first GPU, or device=self.device
+            self.summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-6-6", device_map="auto")
+            logger.info("Summarizer model loaded.")
+
+        except Exception as e:
+            logger.error(f"Error loading models: {e}")
+            raise
         
         # Use simple vector store instead of ChromaDB for HF Spaces
         self.vector_store = SimpleVectorStore()
@@ -244,9 +256,9 @@ class EnhancedArxivRAG:
             papers = []
             for result in search.results():
                 paper = Paper(
-                    id=result.entry_id.split('/')[-1],
-                    title=result.title.strip().replace('\n', ' '),
-                    abstract=result.summary.strip().replace('\n', ' '),
+                    id=result.entry_id.split("/")[-1],
+                    title=result.title.strip().replace("\n", " "),
+                    abstract=result.summary.strip().replace("\n", " "),
                     authors=[author.name for author in result.authors],
                     categories=result.categories,
                     published=result.published.replace(tzinfo=None),
@@ -318,7 +330,6 @@ 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...")
@@ -372,12 +383,12 @@ class EnhancedArxivRAG:
             bm25_scores = self.bm25_retriever.score(query, top_k * 2)
             
             for idx, score in bm25_scores:
-                if idx < len(all_docs['ids']):
+                if idx < len(all_docs["ids"]):
                     bm25_results.append({
-                        'id': all_docs['ids'][idx],
-                        'document': all_docs['documents'][idx],
-                        'metadata': all_docs['metadatas'][idx],
-                        'score': score
+                        "id": all_docs["ids"][idx],
+                        "document": all_docs["documents"][idx],
+                        "metadata": all_docs["metadatas"][idx],
+                        "score": score
                     })
         
         # Combine results using RRF
@@ -385,13 +396,13 @@ class EnhancedArxivRAG:
         bm25_weight = 1.0 - semantic_weight
         
         # Add semantic scores
-        for i, doc_id in enumerate(semantic_results['ids'][0]):
+        for i, doc_id in enumerate(semantic_results["ids"][0]):
             rank = i + 1
             combined_scores[doc_id] = combined_scores.get(doc_id, 0) + semantic_weight / rank
         
         # Add BM25 scores
         for i, result in enumerate(bm25_results):
-            doc_id = result['id']
+            doc_id = result["id"]
             rank = i + 1
             combined_scores[doc_id] = combined_scores.get(doc_id, 0) + bm25_weight / rank
         
@@ -402,34 +413,33 @@ class EnhancedArxivRAG:
         final_results = []
         for doc_id, score in sorted_results[:top_k]:
             doc_result = self.vector_store.get(ids=[doc_id])
-            if doc_result['ids']:
+            if doc_result["ids"]:
                 final_results.append({
-                    'id': doc_id,
-                    'document': doc_result['documents'][0],
-                    'metadata': doc_result['metadatas'][0],
-                    'combined_score': score
+                    "id": doc_id,
+                    "document": doc_result["documents"][0],
+                    "metadata": doc_result["metadatas"][0],
+                    "combined_score": score
                 })
         
         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:
             return results
         
         # Prepare query-document pairs
-        query_doc_pairs = [(query, result['document']) for result in results]
+        query_doc_pairs = [(query, result["document"]) for result in results]
         
         # Get reranking scores
         rerank_scores = self.reranker.predict(query_doc_pairs)
         
         # Add rerank scores to results
         for i, result in enumerate(results):
-            result['rerank_score'] = float(rerank_scores[i])
+            result["rerank_score"] = float(rerank_scores[i])
         
         # Sort by rerank score
-        reranked_results = sorted(results, key=lambda x: x['rerank_score'], reverse=True)
+        reranked_results = sorted(results, key=lambda x: x["rerank_score"], reverse=True)
         return reranked_results[:top_k]
     
     def generate_answer(self, query: str, context_chunks: List[Dict]) -> str:
@@ -438,7 +448,7 @@ class EnhancedArxivRAG:
             return "No relevant information found to answer your query."
         
         # Combine context from top chunks
-        context_texts = [chunk['document'] for chunk in context_chunks[:3]]
+        context_texts = [chunk["document"] for chunk in context_chunks[:3]]
         combined_context = "\n\n".join(context_texts)
         
         # Limit context length
@@ -451,14 +461,13 @@ class EnhancedArxivRAG:
             summary = self.summarizer(summary_input, 
                                     max_length=120, 
                                     min_length=30, 
-                                    do_sample=False)[0]['summary_text']
+                                    do_sample=False)[0]["summary_text"]
             return summary
         except Exception as e:
             logger.error(f"Error generating summary: {e}")
-            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]])
+            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,
@@ -467,10 +476,10 @@ class EnhancedArxivRAG:
         
         if not query.strip():
             return {
-                'answer': "Please enter a valid research query.",
-                'papers': [],
-                'retrieved_chunks': [],
-                'search_stats': {'papers_found': 0, 'chunks_retrieved': 0}
+                "answer": "Please enter a valid research query.",
+                "papers": [],
+                "retrieved_chunks": [],
+                "search_stats": {"papers_found": 0, "chunks_retrieved": 0}
             }
         
         try:
@@ -479,10 +488,10 @@ class EnhancedArxivRAG:
             
             if not papers:
                 return {
-                    'answer': "No papers found for your query. Please try different keywords.",
-                    'papers': [],
-                    'retrieved_chunks': [],
-                    'search_stats': {'papers_found': 0, 'chunks_retrieved': 0}
+                    "answer": "No papers found for your query. Please try different keywords.",
+                    "papers": [],
+                    "retrieved_chunks": [],
+                    "search_stats": {"papers_found": 0, "chunks_retrieved": 0}
                 }
             
             # Process and store papers
@@ -500,36 +509,36 @@ class EnhancedArxivRAG:
             # Prepare unique papers
             unique_papers = {}
             for chunk in reranked_results:
-                paper_id = chunk['id'].split('_')[0]
+                paper_id = chunk["id"].split("_")[0]
                 if paper_id in self.papers_cache and paper_id not in unique_papers:
                     paper = self.papers_cache[paper_id]
                     unique_papers[paper_id] = {
-                        'title': paper.title,
-                        'authors': paper.authors,
-                        'abstract': paper.abstract,
-                        'url': paper.url,
-                        'categories': paper.categories,
-                        'published': paper.published.strftime('%Y-%m-%d')
+                        "title": paper.title,
+                        "authors": paper.authors,
+                        "abstract": paper.abstract,
+                        "url": paper.url,
+                        "categories": paper.categories,
+                        "published": paper.published.strftime("%Y-%m-%d")
                     }
             
             return {
-                'answer': answer,
-                'papers': list(unique_papers.values()),
-                'retrieved_chunks': reranked_results,
-                'search_stats': {
-                    'papers_found': len(papers),
-                    'chunks_retrieved': len(reranked_results),
-                    'unique_papers_in_results': len(unique_papers)
+                "answer": answer,
+                "papers": list(unique_papers.values()),
+                "retrieved_chunks": reranked_results,
+                "search_stats": {
+                    "papers_found": len(papers),
+                    "chunks_retrieved": len(reranked_results),
+                    "unique_papers_in_results": len(unique_papers)
                 }
             }
             
         except Exception as e:
             logger.error(f"Error in search_and_answer: {e}")
             return {
-                'answer': f"An error occurred while processing your query: {str(e)}",
-                'papers': [],
-                'retrieved_chunks': [],
-                'search_stats': {'papers_found': 0, 'chunks_retrieved': 0}
+                "answer": f"An error occurred while processing your query: {str(e)}",
+                "papers": [],
+                "retrieved_chunks": [],
+                "search_stats": {"papers_found": 0, "chunks_retrieved": 0}
             }
 
 # Global RAG instance
@@ -548,7 +557,7 @@ def search_papers(query: str, max_papers: int = 15, top_k_retrieval: int = 10,
     """Main search function for Gradio interface"""
     
     if not query.strip():
-        return "❌ Please enter a research topic or question.", "", ""
+        return "❌ Please enter a research topic or question.", "", pd.DataFrame()
     
     try:
         # Initialize RAG system
@@ -557,7 +566,7 @@ def search_papers(query: str, max_papers: int = 15, top_k_retrieval: int = 10,
         # Parse categories
         category_list = None
         if categories.strip():
-            category_list = [cat.strip() for cat in categories.split(',') if cat.strip()]
+            category_list = [cat.strip() for cat in categories.split(",") if cat.strip()]
         
         # Perform search
         result = rag.search_and_answer(
@@ -570,33 +579,33 @@ def search_papers(query: str, max_papers: int = 15, top_k_retrieval: int = 10,
         )
         
         # Format answer
-        answer = f"## 🤖 AI-Generated Answer\n\n{result['answer']}\n\n"
+        answer = f"## 🤖 AI-Generated Answer\n\n{result["answer"]}\n\n"
         answer += f"**Search Statistics:**\n"
-        answer += f"- Papers found: {result['search_stats']['papers_found']}\n"
-        answer += f"- Chunks retrieved: {result['search_stats']['chunks_retrieved']}\n"
-        answer += f"- Unique papers in results: {result['search_stats']['unique_papers_in_results']}\n\n"
+        answer += f"- Papers found: {result["search_stats"]["papers_found"]}\n"
+        answer += f"- Chunks retrieved: {result["search_stats"]["chunks_retrieved"]}\n"
+        answer += f"- Unique papers in results: {result["search_stats"]["unique_papers_in_results"]}\n\n"
         
         # Format papers
         papers_md = "## 📚 Relevant Papers\n\n"
-        for i, paper in enumerate(result['papers'], 1):
-            papers_md += f"### {i}. {paper['title']}\n\n"
-            papers_md += f"**Authors:** {', '.join(paper['authors'][:3])}{'...' if len(paper['authors']) > 3 else ''}\n\n"
-            papers_md += f"**Categories:** {', '.join(paper['categories'])}\n\n"
-            papers_md += f"**Published:** {paper['published']}\n\n"
-            papers_md += f"**Abstract:** {paper['abstract'][:250]}{'...' if len(paper['abstract']) > 250 else ''}\n\n"
-            papers_md += f"**URL:** [{paper['url']}]({paper['url']})\n\n"
+        for i, paper in enumerate(result["papers"], 1):
+            papers_md += f"### {i}. {paper["title"]}\n\n"
+            papers_md += f"**Authors:** {", ".join(paper["authors"][:3])}{"..." if len(paper["authors"]) > 3 else ""}\n\n"
+            papers_md += f"**Categories:** {", ".join(paper["categories"])}\n\n"
+            papers_md += f"**Published:** {paper["published"]}\n\n"
+            papers_md += f"**Abstract:** {paper["abstract"][:250]}{"..." if len(paper["abstract"]) > 250 else ""}\n\n"
+            papers_md += f"**URL:** [{paper["url"]}]({paper["url"]})\n\n"
             papers_md += "---\n\n"
         
         # Create papers dataframe
         papers_df = pd.DataFrame([
             {
-                'Title': paper['title'][:50] + '...' if len(paper['title']) > 50 else paper['title'],
-                'Authors': ', '.join(paper['authors'][:2]) + ('...' if len(paper['authors']) > 2 else ''),
-                'Categories': ', '.join(paper['categories'][:2]),
-                'Published': paper['published'],
-                'URL': paper['url']
+                "Title": paper["title"][:50] + "..." if len(paper["title"]) > 50 else paper["title"],
+                "Authors": ", ".join(paper["authors"][:2]) + ("..." if len(paper["authors"]) > 2 else ""),
+                "Categories": ", ".join(paper["categories"][:2]),
+                "Published": paper["published"],
+                "URL": paper["url"]
             }
-            for paper in result['papers']
+            for paper in result["papers"]
         ])
         
         return answer, papers_md, papers_df
@@ -709,4 +718,6 @@ def create_interface():
 # Launch interface
 if __name__ == "__main__":
     interface = create_interface()
-    interface.launch()
+    # Remove share=True for Hugging Face Spaces compatibility
+    interface.launch()
+