features updated

app.py

@@ -1,3 +1,87 @@
+# import os
+# import base64
+# import zipfile
+# from pathlib import Path
+# import streamlit as st
+# from byaldi import RAGMultiModalModel
+# from openai import OpenAI
+
+# import os
+# st.write("Current Working Directory:", os.getcwd())
+
+# # Function to unzip a folder if it does not exist
+# # def unzip_folder_if_not_exist(zip_path, extract_to):
+# #     if not os.path.exists(extract_to):
+# #         with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+# #             zip_ref.extractall(extract_to)
+
+# # # Example usage
+# # zip_path = 'medical_index.zip'
+# # extract_to = 'medical_index'
+# # unzip_folder_if_not_exist(zip_path, extract_to)
+
+# # Preload the RAGMultiModalModel
+# @st.cache_resource
+# def load_model():
+#     return RAGMultiModalModel.from_index("./medical_index")
+
+# RAG = load_model()
+
+# # OpenAI API key from environment
+# api_key = os.getenv("OPENAI_API_KEY")
+# client = OpenAI(api_key=api_key)
+
+# # Streamlit UI
+# st.title("Medical Diagnostic Assistant")
+# st.write("Enter a medical query and get diagnostic recommendations along with visual references.")
+
+# # User input
+# query = st.text_input("Query", "What should be the appropriate diagnostic test for peptic ulcer?")
+
+# if st.button("Submit"):
+#     if query:
+#         # Search using RAG model
+#         with st.spinner('Retrieving information...'):
+#             try:
+#                 returned_page = RAG.search(query, k=1)[0].base64
+
+#                 # Decode and display the retrieved image
+#                 image_bytes = base64.b64decode(returned_page)
+#                 filename = 'retrieved_image.jpg'
+#                 with open(filename, 'wb') as f:
+#                     f.write(image_bytes)
+
+#                 # Display image in Streamlit
+#                 st.image(filename, caption="Reference Image", use_column_width=True)
+
+#                 # Get model response
+#                 response = client.chat.completions.create(
+#                     model="gpt-4o-mini-2024-07-18",
+#                     messages=[
+#                         {"role": "system", "content": "You are a helpful assistant. You only answer the question based on the provided image"},
+#                         {
+#                             "role": "user",
+#                             "content": [
+#                                 {"type": "text", "text": query},
+#                                 {
+#                                     "type": "image_url",
+#                                     "image_url": {"url": f"data:image/jpeg;base64,{returned_page}"},
+#                                 },
+#                             ],
+#                         },
+#                     ],
+#                     max_tokens=300,
+#                 )
+                
+#                 # Display the response
+#                 st.success("Model Response:")
+#                 st.write(response.choices[0].message.content)
+#             except Exception as e:
+#                 st.error(f"An error occurred: {e}")
+#     else:
+#         st.warning("Please enter a query.")
+
+
 import os
 import base64
 import zipfile
@@ -6,24 +90,10 @@ import streamlit as st
 from byaldi import RAGMultiModalModel
 from openai import OpenAI
 
-import os
-st.write("Current Working Directory:", os.getcwd())
-
-# Function to unzip a folder if it does not exist
-# def unzip_folder_if_not_exist(zip_path, extract_to):
-#     if not os.path.exists(extract_to):
-#         with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-#             zip_ref.extractall(extract_to)
-
-# # Example usage
-# zip_path = 'medical_index.zip'
-# extract_to = 'medical_index'
-# unzip_folder_if_not_exist(zip_path, extract_to)
-
 # Preload the RAGMultiModalModel
 @st.cache_resource
 def load_model():
-    return RAGMultiModalModel.from_index("./medical_index")
+    return RAGMultiModalModel.from_index("/home/user/app/medical_index")
 
 RAG = load_model()
 
@@ -35,7 +105,11 @@ client = OpenAI(api_key=api_key)
 st.title("Medical Diagnostic Assistant")
 st.write("Enter a medical query and get diagnostic recommendations along with visual references.")
 
-# User input
+# User input for selecting the model
+model_options = ["gpt-4o", "gpt-4o-mini", "o1-preview", "o1-mini"]
+selected_model = st.selectbox("Choose a GPT model", model_options)
+
+# User input for query
 query = st.text_input("Query", "What should be the appropriate diagnostic test for peptic ulcer?")
 
 if st.button("Submit"):
@@ -43,40 +117,45 @@ if st.button("Submit"):
         # Search using RAG model
         with st.spinner('Retrieving information...'):
             try:
-                returned_page = RAG.search(query, k=1)[0].base64
-
-                # Decode and display the retrieved image
-                image_bytes = base64.b64decode(returned_page)
-                filename = 'retrieved_image.jpg'
-                with open(filename, 'wb') as f:
-                    f.write(image_bytes)
+                # Get top 10 images
+                returned_pages = RAG.search(query, k=10)
+                image_urls = []
+                for i, page in enumerate(returned_pages):
+                    image_bytes = base64.b64decode(page.base64)
+                    filename = f'retrieved_image_{i}.jpg'
+                    with open(filename, 'wb') as f:
+                        f.write(image_bytes)
+                    image_urls.append(f"data:image/jpeg;base64,{page.base64}")
 
-                # Display image in Streamlit
+                # Display the first image initially
                 st.image(filename, caption="Reference Image", use_column_width=True)
 
                 # Get model response
                 response = client.chat.completions.create(
-                    model="gpt-4o-mini-2024-07-18",
+                    model=selected_model,
                     messages=[
-                        {"role": "system", "content": "You are a helpful assistant. You only answer the question based on the provided image"},
+                        {"role": "system", "content": "You are a helpful assistant. You only answer the question based on the provided image and select the right option. You will need to provide the exaplanation from the context as well. DO NOT answer from your previous knowledge ; only answer from the images provided."},
                         {
                             "role": "user",
                             "content": [
                                 {"type": "text", "text": query},
-                                {
-                                    "type": "image_url",
-                                    "image_url": {"url": f"data:image/jpeg;base64,{returned_page}"},
-                                },
+                                *[{"type": "image_url", "image_url": {"url": url}} for url in image_urls],
                             ],
                         },
                     ],
                     max_tokens=300,
                 )
-                
+
                 # Display the response
                 st.success("Model Response:")
                 st.write(response.choices[0].message.content)
+
+                # Option to see all references
+                if st.button("Show References"):
+                    st.subheader("References")
+                    for i, page in enumerate(returned_pages):
+                        st.image(f'retrieved_image_{i}.jpg', caption=f"Reference Image {i+1}", use_column_width=True)
             except Exception as e:
                 st.error(f"An error occurred: {e}")
     else:
-        st.warning("Please enter a query.")
+        st.warning("Please enter a query.")