Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from PIL import Image
 import torch
+import cohere
 import torch.nn as nn
 from torchvision import transforms
 from torchvision.models import vit_b_16, vit_l_16, ViT_B_16_Weights, ViT_L_16_Weights
@@ -32,6 +33,8 @@ torch.cuda.empty_cache()
 
 import nest_asyncio
 nest_asyncio.apply()
+co = cohere.Client(st.secrets["COHERE_API_KEY"])
+
 
 st.set_page_config(page_title="DermBOT", page_icon="🧬", layout="centered")
 
@@ -97,6 +100,8 @@ vector_store = Qdrant(
 retriever = vector_store.as_retriever()
 
 
+
+
 # Dynamically initialize LLM based on selection
 OPENAI_API_KEY = st.secrets["OPENAI_API_KEY"]
 selected_model = st.session_state["selected_model"]
@@ -145,12 +150,12 @@ Your Response:
 
 prompt_template = PromptTemplate(template=AI_PROMPT_TEMPLATE, input_variables=["question", "context"])
 
-rag_chain = RetrievalQA.from_chain_type(
-    llm=llm,
-    retriever=retriever,
-    chain_type="stuff",
-    chain_type_kwargs={"prompt": prompt_template, "document_variable_name": "context"}
-)
+#rag_chain = RetrievalQA.from_chain_type(
+ #   llm=llm,
+  #  retriever=retriever,
+  #  chain_type="stuff",
+  #  chain_type_kwargs={"prompt": prompt_template, "document_variable_name": "context"}
+#)
 
 # === Class Names ===
 multilabel_class_names = [
@@ -268,6 +273,25 @@ def export_chat_to_pdf(messages):
     buf.seek(0)
     return buf
 
+
+#Reranker utility
+def rerank_with_cohere(query: str, documents: list, top_n: int = 5) -> list:
+    if not documents:
+        return []
+
+    raw_texts = [doc.page_content if hasattr(doc, "page_content") else str(doc) for doc in documents]
+    results = co.rerank(query=query, documents=raw_texts, top_n=min(top_n, len(raw_texts)), model="rerank-v3.5")
+    reranked_docs = [documents[result.index] for result in results]
+    return reranked_docs
+
+# Final answer generation using reranked context
+def get_reranked_response(query: str):
+    docs = retriever.get_relevant_documents(query)
+    reranked_docs = rerank_with_cohere(query, docs)
+    context = "\n\n".join([doc.page_content for doc in reranked_docs])
+    prompt = AI_PROMPT_TEMPLATE.format(question=query, context=context)
+    return llm.invoke([{"role": "system", "content": prompt}])
+
 # === App UI ===
 
 st.title("🧬 DermBOT — Skin AI Assistant")
@@ -278,22 +302,21 @@ if uploaded_file:
     st.image(uploaded_file, caption="Uploaded image", use_container_width=True)
     image = Image.open(uploaded_file).convert("RGB")
 
-
     predicted_multi, predicted_single = run_inference(image)
 
     # Show predictions clearly to the user
-    st.markdown(f" Skin Issues : {', '.join(predicted_multi)}")
-    st.markdown(f" Most Likely Diagnosis : {predicted_single}")
+    st.markdown(f"🧾 **Skin Issues**: {', '.join(predicted_multi)}")
+    st.markdown(f"📌 **Most Likely Diagnosis**: {predicted_single}")
 
     query = f"What are my treatment options for {predicted_multi} and {predicted_single}?"
     st.session_state.messages.append({"role": "user", "content": query})
 
-    with st.spinner("Analyzing the image and retrieving response..."):
-        response = rag_chain.invoke(query)
-        st.session_state.messages.append({"role": "assistant", "content": response['result']})
+    with st.spinner("🔎 Analyzing and retrieving context..."):
+        response = get_reranked_response(query)
+        st.session_state.messages.append({"role": "assistant", "content": response.content})
 
     with st.chat_message("assistant"):
-        st.markdown(response['result'])
+        st.markdown(response.content)
 
 # === Chat Interface ===
 if prompt := st.chat_input("Ask a follow-up..."):
@@ -301,7 +324,7 @@ if prompt := st.chat_input("Ask a follow-up..."):
     with st.chat_message("user"):
         st.markdown(prompt)
 
-    response = llm.invoke([{"role": m["role"], "content": m["content"]} for m in st.session_state.messages])
+    response = get_reranked_response(prompt)
     st.session_state.messages.append({"role": "assistant", "content": response.content})
     with st.chat_message("assistant"):
         st.markdown(response.content)