Initial Commit

.gitignore

@@ -0,0 +1,2 @@
+venv
+.env

app.py

@@ -0,0 +1,72 @@
+import streamlit as st
+from transformers import pipeline
+from PIL import Image
+from rag_model import *
+
+
+@st.cache_resource
+def load_image_model():
+    return pipeline("image-classification", model="Heem2/wound-image-classification")
+
+pipeline = load_image_model()
+
+st.title("FirstAid-AI")
+
+# Initialize chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Upload an image of a wound
+file = st.file_uploader("Upload an image of your wound")
+
+# Reset chat history if no file is uploaded
+if file is None:
+    st.session_state.messages = []
+
+if file is not None:
+    # Display the image and predictions
+    col1, col2 = st.columns(2)
+    image = Image.open(file)
+    col1.image(image, use_container_width=True)
+
+    # Classify the image
+    predictions = pipeline(image)
+    detected_wound = predictions[0]['label']
+    col2.header("Detected Wound")
+    for p in predictions:
+        col2.subheader(f"{p['label']}: {round(p['score'] * 100, 1)}%")
+
+    # Initial advice for wound advice
+    if not st.session_state.messages:
+        initial_query = f"Provide treatment advice for a {detected_wound} wound"
+        initial_response = rag_chain.invoke(initial_query)
+        st.session_state.messages.append({"role": "assistant", "content": initial_response})
+
+# Display chat messages from history
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Accept user input if an image is uploaded
+if file is not None and (prompt := st.chat_input("Ask a follow-up question or continue the conversation:")):
+    # Display user message in chat
+    with st.chat_message("user"):
+        st.markdown(prompt)
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+
+    # Prepare the conversation history for rag_chain
+    conversation_history = "\n".join(
+        f"{message['role']}: {message['content']}" for message in st.session_state.messages
+    )
+
+    # Generate response from rag_chain
+    query = f"Context:\n{conversation_history}\n\nAssistant, respond to the user's latest query: {prompt}"
+    response = rag_chain.invoke(query)
+
+    # Display assistant response in chat message container
+    with st.chat_message("assistant"):
+        st.markdown(response)
+        
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": response})

rag_model.py

@@ -0,0 +1,132 @@
+# Import libraries.
+
+# File loading and environment variables.
+import os
+from dotenv import load_dotenv
+
+# HuggingFace LLM.
+from huggingface_hub import InferenceClient
+
+# Langchain.
+from langchain.document_loaders import TextLoader
+from langchain.prompts import PromptTemplate
+from langchain.schema.runnable import RunnablePassthrough, RunnableLambda
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import MongoDBAtlasVectorSearch
+from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
+
+# MongoDB.
+from pymongo import MongoClient
+
+# Function type hints.
+from typing import Dict, Any
+
+# Streamlit
+import streamlit as st
+
+# Load the environment variables from the .env file
+load_dotenv()
+
+# Load the environment variables
+MONGO_URI = os.getenv("MONGO_URI")
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+# Setup Vector Store and MongoDB Atlas connection.
+
+# Connect to MongoDB Atlas cluster using the connection string.
+DB_NAME = "ericguan04"
+COLLECTION_NAME = "first_aid_intents"
+vector_search_index = "vector_index"
+
+@st.cache_resource
+def get_mongodb_collection():
+    # Connect to MongoDB Atlas cluster using the connection string.
+    cluster = MongoClient(MONGO_URI)
+    # Connect to the specific collection in the database.
+    return cluster[DB_NAME][COLLECTION_NAME]
+
+MONGODB_COLLECTION = get_mongodb_collection()
+
+@st.cache_resource
+def load_embedding_model():
+    return HuggingFaceInferenceAPIEmbeddings(
+        api_key=HF_TOKEN, model_name="sentence-transformers/all-mpnet-base-v2"
+    )
+
+embedding_model = load_embedding_model()
+
+vector_search = MongoDBAtlasVectorSearch.from_connection_string(
+    connection_string=MONGO_URI,
+    namespace=f"{DB_NAME}.{COLLECTION_NAME}",
+    embedding=embedding_model,
+    index_name=vector_search_index,
+)
+
+# k to search for only the X most relevant documents.
+k = 10
+
+# score_threshold to use only documents with a relevance score above 0.80.
+score_threshold = 0.80
+
+# Build your retriever
+retriever_1 = vector_search.as_retriever(
+   search_type = "similarity", # similarity, mmr, similarity_score_threshold. https://api.python.langchain.com/en/latest/vectorstores/langchain_core.vectorstores.VectorStore.html#langchain_core.vectorstores.VectorStore.as_retriever
+   search_kwargs = {"k": k, "score_threshold": score_threshold}
+)
+
+
+# Initialize Hugging Face client
+hf_client = InferenceClient(api_key=HF_TOKEN)
+
+# Define the prompt template
+prompt = PromptTemplate.from_template(
+    """Use the following pieces of context to answer the question at the end.
+
+    START OF CONTEXT:
+    {context}
+    END OF CONTEXT:
+
+    START OF QUESTION:
+    {question}
+    END OF QUESTION:
+
+    If you do not know the answer, just say that you do not know.
+    NEVER assume things.
+    """
+)
+
+
+# Formatting the retrieved documents before inserting them in the system prompt template.
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
+
+@st.cache_resource
+def generate_response(input_dict: Dict[str, Any]) -> str:
+    formatted_prompt = prompt.format(**input_dict)
+    # print(formatted_prompt)
+
+    ## THIS IS YOUR LLM
+    response = hf_client.chat.completions.create(
+        model="Qwen/Qwen2.5-1.5B-Instruct",
+        messages=[{
+            "role": "system",
+            "content": formatted_prompt
+        },{
+            "role": "user",
+            "content": input_dict["question"]
+        }],
+        max_tokens=1000,
+        temperature=0.2,
+    )
+
+    return response.choices[0].message.content
+
+
+# Build the chain with retriever_1.
+rag_chain = (
+    {
+        "context": retriever_1 | RunnableLambda(format_docs),
+        "question": RunnablePassthrough()
+    }
+    | RunnableLambda(generate_response)
+)

requirements.txt

@@ -0,0 +1,87 @@
+aiohappyeyeballs==2.4.3
+aiohttp==3.11.5
+aiosignal==1.3.1
+altair==5.4.1
+annotated-types==0.7.0
+anyio==4.6.2.post1
+attrs==24.2.0
+blinker==1.9.0
+cachetools==5.5.0
+certifi==2024.8.30
+charset-normalizer==3.4.0
+click==8.1.7
+dataclasses-json==0.6.7
+dnspython==2.7.0
+filelock==3.16.1
+frozenlist==1.5.0
+fsspec==2024.10.0
+gitdb==4.0.11
+GitPython==3.1.43
+h11==0.14.0
+httpcore==1.0.7
+httpx==0.27.2
+httpx-sse==0.4.0
+huggingface-hub==0.26.2
+idna==3.10
+Jinja2==3.1.4
+jsonpatch==1.33
+jsonpointer==3.0.0
+jsonschema==4.23.0
+jsonschema-specifications==2024.10.1
+langchain==0.3.7
+langchain-community==0.3.7
+langchain-core==0.3.19
+langchain-text-splitters==0.3.2
+langsmith==0.1.143
+markdown-it-py==3.0.0
+MarkupSafe==3.0.2
+marshmallow==3.23.1
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.1.0
+mypy-extensions==1.0.0
+narwhals==1.14.1
+networkx==3.4.2
+numpy==1.26.4
+orjson==3.10.11
+packaging==24.2
+pandas==2.2.3
+pillow==11.0.0
+propcache==0.2.0
+protobuf==5.28.3
+pyarrow==18.0.0
+pydantic==2.9.2
+pydantic-settings==2.6.1
+pydantic_core==2.23.4
+pydeck==0.9.1
+Pygments==2.18.0
+pymongo==4.10.1
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+pytz==2024.2
+PyYAML==6.0.2
+referencing==0.35.1
+regex==2024.11.6
+requests==2.32.3
+requests-toolbelt==1.0.0
+rich==13.9.4
+rpds-py==0.21.0
+safetensors==0.4.5
+six==1.16.0
+smmap==5.0.1
+sniffio==1.3.1
+SQLAlchemy==2.0.35
+streamlit==1.40.1
+sympy==1.13.1
+tenacity==9.0.0
+tokenizers==0.20.3
+toml==0.10.2
+torch==2.5.1
+tornado==6.4.1
+tqdm==4.67.0
+transformers==4.46.3
+typing-inspect==0.9.0
+typing_extensions==4.12.2
+tzdata==2024.2
+urllib3==2.2.3
+yarl==1.17.2