Update app.py

app.py

@@ -1,346 +1,147 @@
-import os
-import re
-import json
-import math
-import requests
-import threading
-from concurrent.futures import ThreadPoolExecutor, as_completed
-
 import streamlit as st
-import pandas as pd
-
-# Set page config FIRST, before any other Streamlit calls:
-st.set_page_config(page_title="Enhanced RAG + PubMed", layout="wide")
-
-# NLP
-import nltk
-nltk.download('punkt')
-from nltk.tokenize import sent_tokenize
-
-# Transformers for summarization
-from transformers import pipeline
-
-# Optional: OpenAI and Google Generative AI
-import openai
-import google.generativeai as genai
-
-###############################################################################
-#                              CONFIG & ENV                                   #
-###############################################################################
-
-
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
-GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "")
-MY_PUBMED_EMAIL = os.getenv("MY_PUBMED_EMAIL", "[email protected]")
-
-# Configure OpenAI if key is provided
-if OPENAI_API_KEY:
-    openai.api_key = OPENAI_API_KEY
+import os
 
-# Configure Google PaLM / Gemini if key is provided
-if GEMINI_API_KEY:
-    genai.configure(api_key=GEMINI_API_KEY)
+from config import (
+    OPENAI_API_KEY, 
+    GEMINI_API_KEY, 
+    DEFAULT_CHUNK_SIZE
+)
+from models import configure_llms, openai_chat, gemini_chat
+from pubmed_utils import (
+    search_pubmed, 
+    fetch_pubmed_abstracts, 
+    chunk_and_summarize
+)
+from image_pipeline import load_image_model, analyze_image
 
 ###############################################################################
-#                           SUMMARIZATION PIPELINE                            #
+#                           PAGE CONFIG FIRST                                 #
 ###############################################################################
-@st.cache_resource
-def load_summarizer():
-    """
-    Load a summarization model (e.g., BART, PEGASUS, T5).
-    For a more concise summarization, consider 'google/pegasus-xsum'.
-    For a balanced approach, 'facebook/bart-large-cnn' is popular.
-    """
-    return pipeline(
-        "summarization",
-        model="facebook/bart-large-cnn",
-        tokenizer="facebook/bart-large-cnn"
-    )
-
-summarizer = load_summarizer()
+st.set_page_config(page_title="RAG + Image: Production Scenario", layout="wide")
 
 ###############################################################################
-#                      PUBMED RETRIEVAL (NCBI E-utilities)                    #
+#                           INITIALIZE & LOAD MODELS                          #
 ###############################################################################
-def search_pubmed(query, max_results=3):
-    """
-    Searches PubMed for PMIDs matching the query.
-    Includes recommended 'tool' and 'email' in the request.
-    """
-    base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
-    params = {
-        "db": "pubmed",
-        "term": query,
-        "retmax": max_results,
-        "retmode": "json",
-        "tool": "ElysiumRAG",
-        "email": MY_PUBMED_EMAIL
-    }
-    resp = requests.get(base_url, params=params)
-    resp.raise_for_status()
-    data = resp.json()
-    id_list = data.get("esearchresult", {}).get("idlist", [])
-    return id_list
 
-def fetch_one_abstract(pmid):
+def initialize_app():
     """
-    Fetches a single abstract for a given PMID using EFetch.
-    Returns (pmid, text).
+    Configures LLMs, loads image model, etc.
+    Cache these calls for performance in HF Spaces.
     """
-    base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
-    params = {
-        "db": "pubmed",
-        "retmode": "text",
-        "rettype": "abstract",
-        "id": pmid,
-        "tool": "ElysiumRAG",
-        "email": MY_PUBMED_EMAIL
-    }
-    resp = requests.get(base_url, params=params)
-    resp.raise_for_status()
-    raw_text = resp.text.strip()
+    configure_llms()  # sets openai.api_key and genai.configure if keys are present
+    image_model = load_image_model()
+    return image_model
 
-    if not raw_text:
-        return (pmid, "No abstract text found.")
-    return (pmid, raw_text)
-
-def fetch_pubmed_abstracts(pmids):
-    """
-    Parallel fetching of multiple PMIDs to reduce overall latency.
-    Returns {pmid: abstract_text}.
-    """
-    abstracts_map = {}
-    if not pmids:
-        return abstracts_map
-
-    with ThreadPoolExecutor(max_workers=min(len(pmids), 5)) as executor:
-        future_to_pmid = {executor.submit(fetch_one_abstract, pmid): pmid for pmid in pmids}
-        for future in as_completed(future_to_pmid):
-            pmid = future_to_pmid[future]
-            try:
-                pmid_result, text = future.result()
-                abstracts_map[pmid_result] = text
-            except Exception as e:
-                abstracts_map[pmid] = f"Error fetching abstract: {str(e)}"
-    return abstracts_map
+image_model = initialize_app()
 
 ###############################################################################
-#                  ABSTRACT CHUNKING + SUMMARIZATION LOGIC                    #
+#                    HELPER: BUILD SYSTEM PROMPT WITH REFERENCES             #
 ###############################################################################
-def chunk_and_summarize(abstract_text, chunk_size=512):
+def build_system_prompt_with_refs(pmids, summaries):
     """
-    Splits a large abstract into manageable chunks (by sentences),
-    then summarizes each chunk with the Hugging Face pipeline.
-    Returns a combined summary for the entire abstract.
+    Creates a system prompt that includes references [Ref1], [Ref2], etc.
     """
-    sentences = sent_tokenize(abstract_text)
-    chunks = []
-
-    current_chunk = []
-    current_length = 0
-    for sent in sentences:
-        tokens_in_sent = len(sent.split())
-        # If adding this sentence exceeds the chunk_size limit, finalize the chunk
-        if current_length + tokens_in_sent > chunk_size:
-            chunks.append(" ".join(current_chunk))
-            current_chunk = []
-            current_length = 0
-
-        current_chunk.append(sent)
-        current_length += tokens_in_sent
-
-    # Final chunk if it exists
-    if current_chunk:
-        chunks.append(" ".join(current_chunk))
-
-    summarized_pieces = []
-    for c in chunks:
-        summary_out = summarizer(
-            c,
-            max_length=100,   # Tweak for desired summary length
-            min_length=30,
-            do_sample=False
-        )
-        summarized_pieces.append(summary_out[0]['summary_text'])
-
-    final_summary = " ".join(summarized_pieces)
-    return final_summary.strip()
-
-###############################################################################
-#                      LLM CALLS (OpenAI / Gemini)                            #
-###############################################################################
-def openai_chat(system_prompt, user_message, model="gpt-3.5-turbo", temperature=0.3):
-    """
-    Basic ChatCompletion with a system + user role for OpenAI.
-    """
-    if not OPENAI_API_KEY:
-        return "Error: OpenAI API key not provided."
-    try:
-        response = openai.ChatCompletion.create(
-            model=model,
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": user_message}
-            ],
-            temperature=temperature
-        )
-        return response.choices[0].message["content"].strip()
-    except Exception as e:
-        return f"Error calling OpenAI: {str(e)}"
-
-def gemini_chat(system_prompt, user_message, model_name="models/chat-bison-001", temperature=0.3):
-    """
-    Basic PaLM2/Gemini chat call using google.generativeai.
-    """
-    if not GEMINI_API_KEY:
-        return "Error: Gemini API key not provided."
-    try:
-        model = genai.GenerativeModel(model_name=model_name)
-        chat_session = model.start_chat(history=[("system", system_prompt)])
-        reply = chat_session.send_message(user_message, temperature=temperature)
-        return reply.text
-    except Exception as e:
-        return f"Error calling Gemini: {str(e)}"
-
-###############################################################################
-#                         BUILD REFERENCES FOR ANSWER                         #
-###############################################################################
-def build_system_prompt_with_refs(pmids, summarized_map):
-    """
-    Creates a system prompt that includes the summarized abstracts alongside
-    labeled references (e.g., [Ref1]) so the LLM can cite them in the final answer.
-    """
-    system_context = (
-        "You have access to the following summarized PubMed articles. "
-        "When relevant, cite them using their reference label.\n\n"
-    )
+    system_context = "You have access to the following summarized PubMed articles:\n\n"
     for idx, pmid in enumerate(pmids, start=1):
         ref_label = f"[Ref{idx}]"
-        system_context += f"{ref_label} (PMID {pmid}): {summarized_map[pmid]}\n\n"
-
-    system_context += "Use this contextual info to provide a concise, evidence-based answer."
+        system_context += f"{ref_label} (PMID {pmid}): {summaries[pmid]}\n\n"
+    system_context += (
+        "Use this info to answer the user's question. Cite references as needed."
+    )
     return system_context
 
 ###############################################################################
-#                                STREAMLIT APP                                #
+#                                 MAIN APP                                    #
 ###############################################################################
 def main():
-    # From here on, we do NOT call st.set_page_config() again (to avoid the error).
-    st.title("Enhanced RAG + PubMed: Production-Ready Medical Insights")
+    st.title("RAG + Image: Production-Ready Medical AI")
 
     st.markdown("""
-    **Welcome** to an advanced demonstration of **Retrieval-Augmented Generation (RAG)**
-    using PubMed E-utilities, Hugging Face Summarization, and optional LLM calls 
-    (OpenAI or Gemini).
-    
-    This version includes:
-    - **Parallel** fetching for multiple PMIDs
-    - Advanced **chunking & summarization** of large abstracts
-    - **Reference labeling** in the final answer
-    - Clear disclaimers & best-practice structures
-    
-    ---
-    **Disclaimer**: This is a demonstration prototype for educational or research purposes.
-    It is *not* a substitute for professional medical advice. Always consult a qualified
-    healthcare provider for personal health decisions.
-    """)
+    **Features**:
+    1. *PubMed RAG Pipeline*: Search, fetch, summarize, then generate a final answer with LLM.
+    2. *Optional Image Analysis*: Upload an image for a simple caption or interpretive text.
+    3. *Separation of Concerns*: Each major function is in its own module for maintainability.
 
-    user_query = st.text_area(
-        "Enter your medical question or topic:",
-        placeholder="e.g., 'What are the latest treatments for type 2 diabetes complications?'",
-        height=120
-    )
+    **Disclaimer**: Not a substitute for professional medical advice. 
+    """)
 
-    col1, col2 = st.columns(2)
+    # Section A: Image pipeline
+    st.subheader("Image Analysis")
+    uploaded_image = st.file_uploader("Upload an image (optional)", type=["png", "jpg", "jpeg"])
+    if uploaded_image:
+        with st.spinner("Analyzing image..."):
+            caption = analyze_image(uploaded_image, image_model)
+        st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
+        st.write("**Model Output:**", caption)
+        st.write("---")
+
+    # Section B: PubMed-based RAG
+    st.subheader("PubMed Retrieval & Summarization")
+    user_query = st.text_input("Enter your medical question:", "What are the latest treatments for type 2 diabetes complications?")
+
+    col1, col2, col3 = st.columns([2, 1, 1])
     with col1:
-        max_papers = st.slider(
-            "Number of PubMed Articles to Retrieve",
-            min_value=1,
-            max_value=10,
-            value=3,
-            help="Number of articles to fetch & summarize."
-        )
+        st.markdown("**Set Pipeline Params**")
+        max_papers = st.slider("PubMed Articles to Retrieve", 1, 10, 3)
+        chunk_size = st.slider("Summarization Chunk Size", 256, 1024, DEFAULT_CHUNK_SIZE)
     with col2:
-        selected_llm = st.selectbox(
-            "Select LLM for Final Generation",
-            ["OpenAI: GPT-3.5", "Gemini: PaLM2"],
-            help="Choose which large language model to finalize the answer."
-        )
-
-    # Additional advanced parameter: chunk size
-    chunk_size = st.slider(
-        "Summarization Chunk Size (words)",
-        min_value=256,
-        max_value=1024,
-        value=512,
-        help=(
-            "Larger chunks produce fewer summarization calls, but risk token limits. "
-            "Smaller chunks produce more robust summaries."
-        )
-    )
+        selected_llm = st.selectbox("Select LLM", ["OpenAI GPT-3.5", "Gemini PaLM2"])
+    with col3:
+        temperature = st.slider("LLM Temperature", 0.0, 1.0, 0.3, 0.1)
 
-    if st.button("Run Enhanced RAG Pipeline"):
+    if st.button("Run RAG Pipeline"):
         if not user_query.strip():
-            st.warning("Please enter a query before running RAG.")
+            st.warning("Please enter a question.")
             return
 
-        # 1. PubMed Search
+        # 1) PubMed retrieval
         with st.spinner("Searching PubMed..."):
-            pmids = search_pubmed(query=user_query, max_results=max_papers)
+            pmids = search_pubmed(user_query, max_results=max_papers)
 
         if not pmids:
-            st.error("No matching PubMed results. Try a different query.")
+            st.error("No relevant results found. Try a different query.")
             return
 
-        # 2. Fetch & Summarize
-        with st.spinner("Fetching and summarizing abstracts..."):
-            abstracts_map = fetch_pubmed_abstracts(pmids)
+        # 2) Fetch & Summarize
+        with st.spinner("Fetching & Summarizing abstracts..."):
+            abs_map = fetch_pubmed_abstracts(pmids)
             summarized_map = {}
-            for pmid, abstract_text in abstracts_map.items():
-                if "Error fetching" in abstract_text:
-                    summarized_map[pmid] = abstract_text
+            for pmid, text in abs_map.items():
+                if text.startswith("Error:"):
+                    summarized_map[pmid] = text
                 else:
-                    summarized_map[pmid] = chunk_and_summarize(abstract_text, chunk_size=chunk_size)
+                    summarized_map[pmid] = chunk_and_summarize(text, chunk_size=chunk_size)
 
-        # 3. Display Summaries
+        # 3) Display Summaries
         st.subheader("Retrieved & Summarized PubMed Articles")
         for idx, pmid in enumerate(pmids, start=1):
-            ref_label = f"[Ref{idx}]"
-            st.markdown(f"**{ref_label} PMID {pmid}**")
+            st.markdown(f"**[Ref{idx}] PMID {pmid}**")
             st.write(summarized_map[pmid])
             st.write("---")
 
-        # 4. Build Prompt & Generate Final Answer
+        # 4) Final LLM Answer
         st.subheader("RAG-Enhanced Final Answer")
         system_prompt = build_system_prompt_with_refs(pmids, summarized_map)
-
-        with st.spinner("Generating final answer..."):
-            if selected_llm == "OpenAI: GPT-3.5":
-                answer = openai_chat(system_prompt=system_prompt, user_message=user_query)
+        with st.spinner("Generating answer..."):
+            if selected_llm == "OpenAI GPT-3.5":
+                answer = openai_chat(system_prompt, user_query, temperature=temperature)
             else:
-                answer = gemini_chat(system_prompt=system_prompt, user_message=user_query)
+                answer = gemini_chat(system_prompt, user_query, temperature=temperature)
 
         st.write(answer)
         st.success("RAG Pipeline Complete.")
 
-    # Production notes:
+    # Production tips
     st.markdown("---")
     st.markdown("""
-    ### Production-Ready Enhancements
-    1. **Vector Databases & Advanced Retrieval**  
-       - For large-scale usage, index PubMed articles in a vector DB to quickly retrieve relevant passages.  
-    2. **Citation Parsing**  
-       - Automatically detect which chunk or article contributed to each sentence for more precise referencing.  
-    3. **Multi-Lingual**  
-       - Integrate translation pipelines for non-English queries or abstracts to expand global reach.  
-    4. **Rate Limiting**  
-       - Respect NCBI's ~3 requests/sec guideline if scaling up usage.  
-    5. **Logging & Monitoring**  
-       - In production, set up robust logging/observability for success/failure rates.  
-    6. **Security & Privacy**  
-       - Currently only uses public info. If patient data is included, ensure HIPAA/GDPR compliance.
+    ### Production Enhancements
+    - **Vector Database** for advanced retrieval
+    - **Citation Parsing** for accurate referencing
+    - **Multi-Lingual** expansions
+    - **Rate Limiting** for PubMed (max ~3 requests/sec)
+    - **Robust Logging / Monitoring**
+    - **Security & Privacy** if patient data is integrated
     """)
 
-
 if __name__ == "__main__":
     main()