Update app.py

app.py

@@ -30,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)
@@ -74,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()
     
@@ -150,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)
@@ -167,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):
@@ -201,29 +201,11 @@ class EnhancedArxivRAG:
     def __init__(self):
         logger.info("Initializing Enhanced ArXiv RAG System for HF Spaces...")
         
-        # 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 smaller, faster 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)
         
         # Use simple vector store instead of ChromaDB for HF Spaces
         self.vector_store = SimpleVectorStore()
@@ -256,9 +238,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),
@@ -383,12 +365,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
@@ -396,13 +378,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
         
@@ -413,12 +395,12 @@ 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
@@ -429,17 +411,17 @@ class EnhancedArxivRAG:
             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:
@@ -448,7 +430,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
@@ -461,12 +443,12 @@ 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]])
     
     def search_and_answer(self, query: str, max_papers: int = 15,
                          top_k_retrieval: int = 10, top_k_rerank: int = 5,
@@ -476,10 +458,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:
@@ -488,10 +470,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
@@ -509,36 +491,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
@@ -557,7 +539,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.", "", pd.DataFrame()
+        return "❌ Please enter a research topic or question.", "", ""
     
     try:
         # Initialize RAG system
@@ -566,7 +548,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(
@@ -579,33 +561,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
@@ -718,6 +700,5 @@ def create_interface():
 # Launch interface
 if __name__ == "__main__":
     interface = create_interface()
-    # Remove share=True for Hugging Face Spaces compatibility
     interface.launch()