Update app.py

app.py

@@ -261,18 +261,22 @@ class OptimizedRagSystem:
                     logger.warning("Query too short, using default search")
                     search_query = "machine learning"
                 
-                # Enhance transformer-related queries for better results
-                transformer_keywords = ["transformer", "attention", "bert", "gpt", "llm", "language model"]
-                if any(keyword in search_query.lower() for keyword in transformer_keywords):
-                    # Add related terms to improve relevance
-                    enhanced_terms = []
-                    if "attention" in search_query.lower():
-                        enhanced_terms.extend(["self-attention", "multi-head attention", "scaled dot-product"])
-                    if "transformer" in search_query.lower():
-                        enhanced_terms.extend(["encoder", "decoder", "positional encoding"])
-                    
-                    if enhanced_terms:
-                        search_query = f"({search_query}) OR ({' OR '.join(enhanced_terms)})"
+                # Simplify transformer queries for better ArXiv results
+                # ArXiv search works better with simple, specific terms
+                if "attention" in search_query.lower() and "transformer" in search_query.lower():
+                    search_query = "attention mechanism transformer"
+                elif "transformer" in search_query.lower():
+                    search_query = "transformer neural network"
+                elif "attention" in search_query.lower():
+                    search_query = "attention mechanism"
+                
+                logger.info(f"Simplified query: '{search_query}'")
+                
+                # Handle categories - for transformers, default to relevant categories
+                if "attention" in search_query.lower() or "transformer" in search_query.lower():
+                    if not categories or len(categories) == 0:
+                        categories = ["cs.CL", "cs.LG", "cs.AI"]  # Default to relevant categories
+                        logger.info(f"Added default categories for transformer search: {categories}")
                 
                 if categories and len(categories) > 0:
                     category_filter = " OR ".join([f"cat:{cat.strip()}" for cat in categories if cat.strip()])
@@ -352,35 +356,46 @@ class OptimizedRagSystem:
         # If all attempts failed, try a simple fallback search
         logger.warning("All search attempts failed, trying fallback search...")
         try:
-            fallback_search = arxiv.Search(
-                query="artificial intelligence",  # Simple fallback
-                max_results=5,
-                sort_by=arxiv.SortCriterion.SubmittedDate,
-                sort_order=arxiv.SortOrder.Descending
-            )
+            # Try a specific search that should return transformer papers
+            fallback_queries = [
+                "attention is all you need",
+                "transformer attention mechanism",
+                "BERT language representation",
+                "GPT generative pretrained",
+                "artificial intelligence"
+            ]
             
-            papers = []
-            for i, result in enumerate(fallback_search.results()):
-                if i >= 5:  # Limit fallback results
-                    break
-                try:
-                    paper = Paper(
-                        id=result.entry_id.split('/')[-1],
-                        title=result.title,
-                        abstract=result.summary,
-                        authors=[author.name for author in result.authors],
-                        categories=result.categories,
-                        published=result.published,
-                        url=result.entry_id
-                    )
-                    papers.append(paper)
-                except Exception as e:
-                    logger.warning(f"Error in fallback paper processing: {e}")
-                    continue
-                    
-            if papers:
-                logger.info(f"🔄 Fallback search returned {len(papers)} papers")
-                return papers
+            for fallback_query in fallback_queries:
+                logger.info(f"Trying fallback: '{fallback_query}'")
+                fallback_search = arxiv.Search(
+                    query=fallback_query,
+                    max_results=5,
+                    sort_by=arxiv.SortCriterion.Relevance,
+                    sort_order=arxiv.SortOrder.Descending
+                )
+            
+                papers = []
+                for i, result in enumerate(fallback_search.results()):
+                    if i >= 5:  # Limit fallback results
+                        break
+                    try:
+                        paper = Paper(
+                            id=result.entry_id.split('/')[-1],
+                            title=result.title,
+                            abstract=result.summary,
+                            authors=[author.name for author in result.authors],
+                            categories=result.categories,
+                            published=result.published,
+                            url=result.entry_id
+                        )
+                        papers.append(paper)
+                    except Exception as e:
+                        logger.warning(f"Error in fallback paper processing: {e}")
+                        continue
+                        
+                if papers:
+                    logger.info(f"🔄 Fallback search '{fallback_query}' returned {len(papers)} papers")
+                    return papers
                 
         except Exception as e:
             logger.error(f"Even fallback search failed: {e}")
@@ -546,10 +561,14 @@ class OptimizedRagSystem:
             with torch.cuda.amp.autocast() if DEVICE == "cuda" else torch.no_grad():
                 rerank_scores = self.reranker.predict(pairs, show_progress_bar=False)
             
-            # Combine with original scores
+            # Combine with original scores (ensure positive scores)
             reranked_results = []
             for i, (chunk, original_score) in enumerate(results):
-                combined_score = 0.6 * float(rerank_scores[i]) + 0.4 * original_score
+                # Normalize rerank scores to [0, 1] and ensure positive
+                rerank_score = float(rerank_scores[i])
+                rerank_score = max(0.0, min(1.0, (rerank_score + 1) / 2))  # Convert from [-1,1] to [0,1]
+                
+                combined_score = 0.6 * rerank_score + 0.4 * max(0.0, original_score)
                 reranked_results.append((chunk, combined_score))
             
             # Sort by new scores