Update app.py

app.py

@@ -2,145 +2,141 @@ import streamlit as st
 import os
 
 from config import (
-    OPENAI_API_KEY, 
-    GEMINI_API_KEY, 
-    DEFAULT_CHUNK_SIZE
+    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, 
+    search_pubmed,
+    fetch_pubmed_abstracts,
     chunk_and_summarize
 )
 from image_pipeline import load_image_model, analyze_image
 
 ###############################################################################
-#                           PAGE CONFIG FIRST                                 #
+#                             STREAMLIT PAGE CONFIG                           #
 ###############################################################################
-st.set_page_config(page_title="RAG + Image: Production Scenario", layout="wide")
+st.set_page_config(page_title="RAG + Image Captioning Demo", layout="wide")
 
 ###############################################################################
 #                           INITIALIZE & LOAD MODELS                          #
 ###############################################################################
-
+@st.cache_resource
 def initialize_app():
     """
-    Configures LLMs, loads image model, etc.
-    Cache these calls for performance in HF Spaces.
+    Configure LLMs (OpenAI/Gemini) and load the image captioning model once.
     """
-    configure_llms()  # sets openai.api_key and genai.configure if keys are present
-    image_model = load_image_model()
-    return image_model
+    configure_llms()
+    model = load_image_model()
+    return model
 
 image_model = initialize_app()
 
 ###############################################################################
-#                    HELPER: BUILD SYSTEM PROMPT WITH REFERENCES             #
+#                    HELPER: BUILD SYSTEM PROMPT WITH REFS                    #
 ###############################################################################
 def build_system_prompt_with_refs(pmids, summaries):
     """
-    Creates a system prompt that includes references [Ref1], [Ref2], etc.
+    Creates a system prompt for the LLM that includes references [Ref1], [Ref2], etc.
     """
     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}): {summaries[pmid]}\n\n"
-    system_context += (
-        "Use this info to answer the user's question. Cite references as needed."
-    )
+    system_context += "Use this info to answer the user's question, citing references if needed."
     return system_context
 
 ###############################################################################
-#                                 MAIN APP                                    #
+#                                MAIN APP                                     #
 ###############################################################################
 def main():
-    st.title("RAG + Image: Production-Ready Medical AI")
+    st.title("RAG + Image Captioning: Production Demo")
 
     st.markdown("""
-    **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.
-
-    **Disclaimer**: Not a substitute for professional medical advice. 
+    This demonstration shows:
+    1. **PubMed RAG**: Retrieve abstracts, summarize, and feed them into an LLM.
+    2. **Image Captioning**: Upload an image for analysis using a known stable model.
+    3. **Separation of Concerns**: Each pipeline is in its own module.
+    
+    **Note**: If you previously encountered KeyError: 'idefics3', it's because the 
+    `SmolVLM-500M-Instruct` model was incompatible with your Transformers version. 
+    Here, we use a supported model such as `nlpconnect/vit-gpt2-image-captioning`.
+    
+    **Disclaimer**: This is a prototype, not medical advice.
     """)
 
-    # Section A: Image pipeline
-    st.subheader("Image Analysis")
-    uploaded_image = st.file_uploader("Upload an image (optional)", type=["png", "jpg", "jpeg"])
-    if uploaded_image:
+    # Section A: Image Upload / Caption
+    st.subheader("Image Captioning")
+    uploaded_img = st.file_uploader("Upload an image (optional)", type=["png", "jpg", "jpeg"])
+    if uploaded_img:
         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)
+            caption = analyze_image(uploaded_img, image_model)
+        st.image(uploaded_img, use_column_width=True)
+        st.write("**Caption**:", 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:
-        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", ["OpenAI GPT-3.5", "Gemini PaLM2"])
-    with col3:
-        temperature = st.slider("LLM Temperature", 0.0, 1.0, 0.3, 0.1)
+    st.subheader("PubMed RAG Pipeline")
+    user_query = st.text_input("Enter a medical question:", "What are the latest treatments for type 2 diabetes?")
+    
+    c1, c2, c3 = st.columns([2,1,1])
+    with c1:
+        st.markdown("**Parameters**:")
+        max_papers = st.slider("Number of Articles", 1, 10, 3)
+        chunk_size = st.slider("Chunk Size", 128, 1024, DEFAULT_CHUNK_SIZE)
+    with c2:
+        llm_choice = st.selectbox("Choose LLM", ["OpenAI: GPT-3.5", "Gemini: PaLM2"])
+    with c3:
+        temperature = st.slider("LLM Temperature", 0.0, 1.0, 0.3, step=0.1)
 
     if st.button("Run RAG Pipeline"):
         if not user_query.strip():
-            st.warning("Please enter a question.")
+            st.warning("Please enter a query.")
             return
-
-        # 1) PubMed retrieval
+        
         with st.spinner("Searching PubMed..."):
-            pmids = search_pubmed(user_query, max_results=max_papers)
-
+            pmids = search_pubmed(user_query, max_papers)
+        
         if not pmids:
-            st.error("No relevant results found. Try a different query.")
+            st.error("No PubMed results. Try a different query.")
             return
-
-        # 2) Fetch & Summarize
-        with st.spinner("Fetching & Summarizing abstracts..."):
-            abs_map = fetch_pubmed_abstracts(pmids)
+        
+        with st.spinner("Fetching & Summarizing..."):
+            abstracts_map = fetch_pubmed_abstracts(pmids)
             summarized_map = {}
-            for pmid, text in abs_map.items():
+            for pmid, text in abstracts_map.items():
                 if text.startswith("Error:"):
                     summarized_map[pmid] = text
                 else:
                     summarized_map[pmid] = chunk_and_summarize(text, chunk_size=chunk_size)
 
-        # 3) Display Summaries
         st.subheader("Retrieved & Summarized PubMed Articles")
         for idx, pmid in enumerate(pmids, start=1):
             st.markdown(f"**[Ref{idx}] PMID {pmid}**")
             st.write(summarized_map[pmid])
             st.write("---")
 
-        # 4) Final LLM Answer
         st.subheader("RAG-Enhanced Final Answer")
         system_prompt = build_system_prompt_with_refs(pmids, summarized_map)
-        with st.spinner("Generating answer..."):
-            if selected_llm == "OpenAI GPT-3.5":
+        with st.spinner("Generating LLM response..."):
+            if llm_choice == "OpenAI: GPT-3.5":
                 answer = openai_chat(system_prompt, user_query, temperature=temperature)
             else:
                 answer = gemini_chat(system_prompt, user_query, temperature=temperature)
 
         st.write(answer)
-        st.success("RAG Pipeline Complete.")
+        st.success("Pipeline Complete.")
 
-    # Production tips
     st.markdown("---")
     st.markdown("""
-    ### 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
+    **Production Tips**:
+    - Vector DB for advanced retrieval
+    - Precise citation parsing
+    - Rate limiting on PubMed
+    - Multi-lingual expansions
+    - Logging & monitoring
+    - Security & privacy compliance
     """)
 
 if __name__ == "__main__":