Update backend.py

backend.py

@@ -1,3 +1,10 @@
+"""
+backend.py
+----------
+This module handles the core processing of clinical queries. It retrieves relevant PubMed abstracts
+via the retrieval pipeline and then uses BioGPT to generate a clinical answer. Designed for speed and clarity.
+"""
+
 from transformers import pipeline
 from retrieval import get_relevant_pubmed_docs
 
@@ -5,43 +12,43 @@ from retrieval import get_relevant_pubmed_docs
 MODEL_NAME = "microsoft/BioGPT-Large-PubMedQA"
 qa_pipeline = pipeline("text-generation", model=MODEL_NAME)
 
-# In-memory cache for documents (used later in the knowledge graph).
+# In-memory cache for retrieved documents (used for knowledge graph visualization).
 docs_cache = {}
 
 def process_medical_query(query: str):
     """
-    Processes a clinical query in two steps:
+    Processes a clinical query by:
     
-    1. Retrieve relevant PubMed abstracts using a retrieval pipeline.
-    2. Generate an answer using BioGPT, leveraging the retrieved abstracts as context.
+    1. Retrieving relevant PubMed abstracts.
+    2. Generating an answer using BioGPT.
     
     Returns:
-        - The generated answer.
-        - An empty list for sub-questions (omitted for speed).
-        - The initial answer (same as final in this simplified pipeline).
-        - An empty critique (omitted).
-    
-    Designed for clinical use with clear, concise responses.
+      - final_answer: The generated answer.
+      - sub_questions: (Empty list; previously used for self-critique but omitted for speed.)
+      - initial_answer: The same as final_answer in this streamlined version.
+      - critique: (Empty string; omitted for performance.)
     """
-    # Retrieve relevant documents via PubMed and Chroma.
+    # Retrieve relevant documents using the retrieval pipeline.
     relevant_docs = get_relevant_pubmed_docs(query)
     docs_cache[query] = relevant_docs
 
     if not relevant_docs:
         return ("No documents found for this query.", [], "", "")
 
-    # Combine the retrieved abstracts as context.
+    # Combine abstracts into a context.
     context_text = "\n\n".join(relevant_docs)
     prompt = f"Question: {query}\nContext: {context_text}\nAnswer:"
 
-    # Generate an answer using BioGPT.
-    generation = qa_pipeline(prompt, max_new_tokens=100, truncation=True)
+    try:
+        generation = qa_pipeline(prompt, max_new_tokens=100, truncation=True)
+    except Exception as e:
+        generation = None
+        print(f"[ERROR] BioGPT generation failed: {e}")
+
     if generation and isinstance(generation, list):
         answer = generation[0]["generated_text"]
     else:
         answer = "No answer found."
 
-    # In this streamlined pipeline, we return the same answer for all outputs.
-    sub_questions = []  # No sub-questions generated.
-    critique = ""       # No self-critique performed.
-    return answer, sub_questions, answer, critique
+    # Return the answer along with empty placeholders for sub-questions and critique.
+    return answer, [], answer, ""