import os import streamlit as st from PIL import Image import torch from langchain_community.llms import LlamaCpp from langchain.chains import ConversationChain from langchain.memory import ConversationBufferMemory import os.path # Set page configuration st.set_page_config( page_title="MedGenius Assistant", page_icon="🩺", layout="wide" ) # Check if model exists locally, otherwise provide download instructions MODEL_DIR = "./models" MODEL_FILENAME = "MedGenius_LLaMA-3.2B.Q8_0.gguf" MODEL_PATH = os.path.join(MODEL_DIR, MODEL_FILENAME) # Initialize session state for conversation history if "messages" not in st.session_state: st.session_state.messages = [] # Sidebar for navigation st.sidebar.title("MedGenius Assistant") page = st.sidebar.radio("Navigation", ["Chat", "Image Analysis", "Symptom Analysis"]) # Model setup instructions in sidebar with st.sidebar.expander("Model Setup Instructions", expanded=not os.path.exists(MODEL_PATH)): st.markdown(""" ### Setup Instructions: 1. Run the download script to get the model: ``` python download_model.py ``` 2. This will download the model from tatendachirume/zems to your local machine 3. Refresh this page after downloading """) # Initialize LLM if model exists if os.path.exists(MODEL_PATH) and "llm" not in st.session_state: try: @st.cache_resource def load_llm(): return LlamaCpp( model_path=MODEL_PATH, temperature=0.7, max_tokens=2000, top_p=0.95, verbose=True, n_ctx=4096 ) st.session_state.llm = load_llm() st.session_state.memory = ConversationBufferMemory() st.session_state.conversation = ConversationChain( llm=st.session_state.llm, memory=st.session_state.memory, verbose=True ) except Exception as e: st.sidebar.error(f"Error loading the model: {e}") # Header st.title("🩺 MedGenius Assistant") # If model doesn't exist, show download message if not os.path.exists(MODEL_PATH): st.warning(f"Model file not found at {MODEL_PATH}. Please follow the setup instructions in the sidebar.") # Main application if page == "Chat": st.subheader("Medical Chat Assistant") # Display chat messages for message in st.session_state.messages: with st.chat_message(message["role"]): st.write(message["content"]) # Chat input if prompt := st.chat_input("Ask me about medical topics...", disabled=not os.path.exists(MODEL_PATH)): # Add user message to chat history st.session_state.messages.append({"role": "user", "content": prompt}) # Display user message with st.chat_message("user"): st.write(prompt) # Generate response with st.chat_message("assistant"): with st.spinner("Thinking..."): try: response = st.session_state.conversation.predict(input=prompt) st.write(response) # Add assistant response to chat history st.session_state.messages.append({"role": "assistant", "content": response}) except Exception as e: st.error(f"Error generating response: {e}") elif page == "Image Analysis": st.subheader("Medical Image Analysis") if not os.path.exists(MODEL_PATH): st.warning("Please download the model first to use this feature.") else: uploaded_file = st.file_uploader("Upload a medical image for analysis", type=["jpg", "jpeg", "png"]) if uploaded_file is not None: image = Image.open(uploaded_file) col1, col2 = st.columns(2) with col1: st.image(image, caption="Uploaded Image", use_column_width=True) with col2: st.write("Image analysis:") with st.spinner("Analyzing image..."): # For demonstration, we'll create a prompt about the image image_description = "a medical image" prompt = f"""This is {image_description}. I need an analysis of this medical image. Consider potential abnormalities, findings that might be relevant, and provide a detailed but focused assessment of what might be seen in this image. What are the key features that would be important for a medical professional to note?""" try: analysis = st.session_state.conversation.predict(input=prompt) st.write(analysis) except Exception as e: st.error(f"Error analyzing image: {e}") st.info("Note: Full image analysis requires additional integration with vision models.") elif page == "Symptom Analysis": st.subheader("Symptom Analysis") if not os.path.exists(MODEL_PATH): st.warning("Please download the model first to use this feature.") else: with st.form("symptom_form"): st.write("Please describe your symptoms:") symptoms = st.text_area("Symptoms", height=150) col1, col2 = st.columns(2) with col1: age = st.number_input("Age", min_value=0, max_value=120, value=30) with col2: gender = st.selectbox("Gender", ["Male", "Female", "Other"]) medical_history = st.text_area("Any relevant medical history?", height=100) submit_button = st.form_submit_button("Analyze Symptoms") if submit_button: prompt = f""" Patient Information: - Age: {age} - Gender: {gender} - Symptoms: {symptoms} - Medical History: {medical_history} Based on this information, what could be potential causes of these symptoms? What are recommended next steps or additional tests? Please note any warning signs that would require immediate medical attention. """ with st.spinner("Analyzing symptoms..."): try: analysis = st.session_state.conversation.predict(input=prompt) st.write(analysis) st.warning("Note: This is not a substitute for professional medical advice. Please consult with a healthcare provider for proper diagnosis and treatment.") except Exception as e: st.error(f"Error analyzing symptoms: {e}") # Add a footer st.markdown("---") st.markdown("**Disclaimer:** This application is for informational purposes only and is not a substitute for professional medical advice, diagnosis, or treatment.")