feat: add BART

app.py

@@ -37,11 +37,8 @@ def fetch_arxiv_papers(query, max_results=5):
     """Fetch papers from arXiv"""
     client = arxiv.Client()
     
-    # Ensure query includes autism-related terms
-    if 'autism' not in query.lower():
-        search_query = f"(ti:{query} OR abs:{query}) AND (ti:autism OR abs:autism) AND cat:q-bio"
-    else:
-        search_query = f"(ti:{query} OR abs:{query}) AND cat:q-bio"
+    # Always include autism in the search query
+    search_query = f"(ti:autism OR abs:autism) AND (ti:\"{query}\" OR abs:\"{query}\") AND cat:q-bio"
     
     # Search arXiv
     search = arxiv.Search(
@@ -52,12 +49,18 @@ def fetch_arxiv_papers(query, max_results=5):
     
     papers = []
     for result in client.results(search):
-        papers.append({
-            'title': result.title,
-            'abstract': result.summary,
-            'url': result.pdf_url,
-            'published': result.published.strftime("%Y-%m-%d")
-        })
+        # Only include papers that mention autism in title or abstract
+        if ('autism' in result.title.lower() or 
+            'asd' in result.title.lower() or 
+            'autism' in result.summary.lower() or 
+            'asd' in result.summary.lower()):
+            papers.append({
+                'title': result.title,
+                'abstract': result.summary,
+                'url': result.pdf_url,
+                'published': result.published.strftime("%Y-%m-%d"),
+                'relevance_score': 1 if 'autism' in result.title.lower() else 0.5
+            })
     
     return papers
 
@@ -65,11 +68,8 @@ def fetch_pubmed_papers(query, max_results=5):
     """Fetch papers from PubMed"""
     base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
     
-    # Ensure query includes autism-related terms
-    if 'autism' not in query.lower():
-        search_term = f"({query}) AND (autism[Title/Abstract] OR ASD[Title/Abstract])"
-    else:
-        search_term = query
+    # Always include autism in the search term
+    search_term = f"(autism[Title/Abstract] OR ASD[Title/Abstract]) AND ({query}[Title/Abstract])"
     
     # Search for papers
     search_url = f"{base_url}/esearch.fcgi"
@@ -109,12 +109,19 @@ def fetch_pubmed_papers(query, max_results=5):
             pmid = article.find('.//PMID')
             
             if title is not None and abstract is not None:
-                papers.append({
-                    'title': title.text,
-                    'abstract': abstract.text,
-                    'url': f"https://pubmed.ncbi.nlm.nih.gov/{pmid.text}/",
-                    'published': year.text if year is not None else 'Unknown'
-                })
+                title_text = title.text.lower()
+                abstract_text = abstract.text.lower()
+                
+                # Only include papers that mention autism
+                if ('autism' in title_text or 'asd' in title_text or 
+                    'autism' in abstract_text or 'asd' in abstract_text):
+                    papers.append({
+                        'title': title.text,
+                        'abstract': abstract.text,
+                        'url': f"https://pubmed.ncbi.nlm.nih.gov/{pmid.text}/",
+                        'published': year.text if year is not None else 'Unknown',
+                        'relevance_score': 1 if ('autism' in title_text or 'asd' in title_text) else 0.5
+                    })
             
     except Exception as e:
         st.error(f"Error fetching PubMed papers: {str(e)}")
@@ -130,14 +137,28 @@ def search_research_papers(query):
     all_papers = []
     for paper in arxiv_papers + pubmed_papers:
         if paper['abstract'] and len(paper['abstract'].strip()) > 0:
-            all_papers.append({
-                'title': paper['title'],
-                'text': f"Title: {paper['title']}\n\nAbstract: {paper['abstract']}",
-                'url': paper['url'],
-                'published': paper['published']
-            })
+            # Check if the paper is actually about autism
+            if ('autism' in paper['title'].lower() or 
+                'asd' in paper['title'].lower() or 
+                'autism' in paper['abstract'].lower() or 
+                'asd' in paper['abstract'].lower()):
+                all_papers.append({
+                    'title': paper['title'],
+                    'text': f"Title: {paper['title']}\n\nAbstract: {paper['abstract']}",
+                    'url': paper['url'],
+                    'published': paper['published'],
+                    'relevance_score': paper.get('relevance_score', 0.5)
+                })
     
-    return pd.DataFrame(all_papers)
+    # Sort papers by relevance score and convert to DataFrame
+    all_papers.sort(key=lambda x: x['relevance_score'], reverse=True)
+    df = pd.DataFrame(all_papers)
+    
+    if df.empty:
+        st.warning("No autism-related papers found. Please try a different search term.")
+        return pd.DataFrame(columns=['title', 'text', 'url', 'published', 'relevance_score'])
+    
+    return df
 
 def generate_answer(question, context, max_length=512):
     """Generate a comprehensive answer using the local model"""
@@ -146,21 +167,24 @@ def generate_answer(question, context, max_length=512):
     if model is None or tokenizer is None:
         return "Error: Could not load the model. Please try again later."
     
+    if not context or len(context.strip()) == 0:
+        return "No relevant autism research papers found to answer your question. Please try rephrasing your question."
+    
     # Format the context as a structured query
-    prompt = f"""Summarize the following research about autism and answer the question.
+    prompt = f"""Based on the following autism research papers, provide a detailed answer:
 
 Research Context:
 {context}
 
 Question: {question}
 
-Instructions: Based on the research context above, provide a comprehensive answer that covers:
-1. Main findings from the research
-2. Research methods used
-3. Clinical implications
-4. Limitations of the studies
+Instructions: Provide a comprehensive answer that covers:
+1. Main findings about autism from the research
+2. Research methods used in autism studies
+3. Clinical implications for autism treatment
+4. Limitations of the autism studies
 
-If the research doesn't address the question directly, explain what information is missing."""
+If the research papers don't directly address autism and the question, explain what information is missing."""
     
     try:
         # Generate response
@@ -182,13 +206,13 @@ If the research doesn't address the question directly, explain what information
         
         # If response is too short or empty, provide a fallback message
         if len(response.strip()) < 50:
-            return """I apologize, but I couldn't generate a specific answer from the research papers provided. 
+            return """I apologize, but I couldn't find specific information about autism in the research papers provided. 
             This might be because:
-            1. The research papers don't directly address your question
-            2. The context needs more specific information
-            3. The question might need to be more specific
+            1. The research papers don't directly address autism
+            2. The context needs more specific information about autism
+            3. The question might need to be more focused on autism-specific aspects
             
-            Please try rephrasing your question or ask about a more specific aspect of autism."""
+            Please try asking a more specific question about autism."""
         
         # Format the response for better readability
         formatted_response = response.replace(". ", ".\n").replace("• ", "\n• ")