RAG fix

.streamlit/secrets.toml

@@ -1 +1,5 @@
 OPENAI_API_KEY = "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"
+QDRANT_API_KEY = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhY2Nlc3MiOiJtIn0.HXzezXdWMFeeR16F7zvqgjzsqrcm8hqa-StXdToFP9Q"
+COHERE_API_KEY = "PBEEQJ8n9sV2Xhpc7OMb8NRsBKtADvcEi9V0iPm5"
+GROQ_API_KEY = "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"
+GOOGLE_API_KEY = "AIzaSyD5pnSzkIuu86ByTPQewVKlh2zxOJI-f8M"

app.py

@@ -8,14 +8,23 @@ import torch.nn.functional as F
 from evo_vit import EvoViTModel
 import io
 import os
+import cohere
 from fpdf import FPDF
 from torchvision.models import resnet50
 import nest_asyncio
 from huggingface_hub import hf_hub_download
 from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 from SkinCancerDiagnosis import initialize_classifier
-from rag_pipeline import invoke_rag_chain
+from rag_pipeline import (
+    available_models,
+    initialize_llm,
+    load_rag_chain,
+    get_reranked_response,
+initialize_rag_components
+)
 from langchain_core.messages import HumanMessage, AIMessage
+from groq import Groq
+import google.generativeai as genai
 
 nest_asyncio.apply()
 device='cuda' if torch.cuda.is_available() else 'cpu'
@@ -25,33 +34,15 @@ st.set_page_config(page_title="DermBOT", page_icon="🧬", layout="centered")
 @st.cache_resource(show_spinner=False)
 def load_models():
     """Cache all models to load only once"""
-    with st.spinner("Loading AI models (one-time operation)..."):
-        classifier = initialize_classifier()
-        return classifier
-
-def initialize_llm(_model_name, _api_key):
-    """Initialize the LLM based on selection"""
-    print(f"Model name : {_model_name}")
-    if "OpenAI" in _model_name:
-        return ChatOpenAI(model="gpt-4o", temperature=0.2, api_key=_api_key)
-    elif "LLaMA" in _model_name:
-        st.warning("LLaMA integration is not implemented yet.")
-        st.stop()
-    elif "Gemini" in _model_name:
-        st.warning("Gemini integration is not implemented yet.")
-        st.stop()
-    else:
-        st.error("Unsupported model selected.")
-        st.stop()
-
-@st.cache_resource(show_spinner=False)
-def load_rag_chain(_model_name, _api_key):
-    """Initialize RAG chain only once"""
-    llm = initialize_llm(_model_name, _api_key)
-    return invoke_rag_chain(llm)
+    with st.spinner("Loading all AI models (one-time operation)..."):
+        models = {
+            'classifier': initialize_classifier(),
+            'rag_components': initialize_rag_components(),
+            'llm': initialize_llm(st.session_state["selected_model"])
+        }
+        models['rag_chain'] = load_rag_chain(models['llm'])
+        return models
 
-# === Model Selection ===
-available_models = ["OpenAI GPT-4o", "LLaMA 3", "Gemini Pro"]
 if "selected_model" not in st.session_state:
     st.session_state["selected_model"] = available_models[0]
 
@@ -63,7 +54,10 @@ st.session_state["selected_model"] = st.sidebar.selectbox(
     index=available_models.index(st.session_state["selected_model"])
 )
 
-OPENAI_API_KEY = st.secrets["OPENAI_API_KEY"]
+if 'app_models' not in st.session_state:
+    st.session_state.app_models = load_models()
+classifier = st.session_state.app_models['classifier']
+llm = st.session_state.app_models['llm']
 
 if "model_change_confirmed" not in st.session_state:
     st.session_state.model_change_confirmed = False
@@ -88,15 +82,12 @@ if st.session_state["selected_model"] != previous_model:
         st.session_state.model_change_confirmed = True
 
 if "model_change_confirmed" not in st.session_state or st.session_state.model_change_confirmed:
-    llm = initialize_llm(st.session_state["selected_model"], OPENAI_API_KEY)
-    rag_chain = load_rag_chain(st.session_state["selected_model"], OPENAI_API_KEY)
-    st.session_state.llm = llm
-    st.session_state.rag_chain = rag_chain
+    st.session_state.app_models['llm'] = initialize_llm(st.session_state["selected_model"])
+    st.session_state.app_models['rag_chain'] = load_rag_chain(st.session_state.app_models['llm'])
+    llm = st.session_state.app_models['llm']
 else:
-    llm = st.session_state.get("llm", initialize_llm(previous_model, OPENAI_API_KEY))
-    rag_chain = st.session_state.get("rag_chain", load_rag_chain(previous_model, OPENAI_API_KEY))
+    pass
 
-classifier = load_models()
 
 # === Session Init ===
 if "messages" not in st.session_state:
@@ -149,10 +140,10 @@ if uploaded_file is not None and uploaded_file != st.session_state.current_image
 
     initial_query = f"What are my treatment options for {predicted_label}?"
     st.session_state.messages.append({"role": "user", "content": initial_query})
-
     with st.spinner("Retrieving medical information..."):
-        response = rag_chain.invoke(initial_query)
-        st.session_state.messages.append({"role": "assistant", "content": response['result']})
+        response = get_reranked_response(initial_query, st.session_state.app_models['llm'], st.session_state.app_models['rag_components'])
+        st.session_state.messages.append({"role": "assistant", "content": response})
+
 
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
@@ -165,32 +156,6 @@ if prompt := st.chat_input("Ask a follow-up question..."):
         st.markdown(prompt)
 
     with st.chat_message("assistant"):
-        # with st.spinner("Thinking..."):
-            # Convert messages to LangChain format
-            # chat_history = []
-            # for msg in st.session_state.messages[:-1]:  # Exclude the current prompt
-            #     if msg["role"] == "user":
-            #         chat_history.append(HumanMessage(content=msg["content"]))
-            #     else:
-            #         chat_history.append(AIMessage(content=msg["content"]))
-            #
-            # # Get response
-            # response = llm.invoke([HumanMessage(content=prompt)] + chat_history)
-            # assistant_response = response.content
-            #
-            # st.markdown(assistant_response)
-            # st.session_state.messages.append({"role": "assistant", "content": assistant_response})
-        # with st.spinner("Thinking..."):
-        #     if len(st.session_state.messages) > 1:
-        #         response = llm.invoke([{"role": m["role"], "content": m["content"]} for m in st.session_state.messages])
-        #         response = response.content
-        #     else:
-        #         response = rag_chain.invoke(prompt)
-        #         response = response['result']
-        #
-        #     st.markdown(response)
-        #     st.session_state.messages.append({"role": "assistant", "content": response})
-
         with st.spinner("Thinking..."):
             if len(st.session_state.messages) > 1:
                 conversation_context = "\n".join(
@@ -201,14 +166,12 @@ if prompt := st.chat_input("Ask a follow-up question..."):
                     f"Conversation history:\n{conversation_context}\n\n"
                     f"Current question: {prompt}"
                 )
-                response = rag_chain.invoke({"query": augmented_prompt})
-                assistant_response = response['result']
+                response = get_reranked_response(augmented_prompt, st.session_state.app_models['llm'], st.session_state.app_models['rag_components'])
             else:
-                response = rag_chain.invoke({"query": prompt})
-                assistant_response = response['result']
+                response = get_reranked_response(prompt, st.session_state.app_models['llm'], st.session_state.app_models['rag_components'])
 
-            st.markdown(assistant_response)
-            st.session_state.messages.append({"role": "assistant", "content": assistant_response})
+            st.markdown(response)
+            st.session_state.messages.append({"role": "assistant", "content": response})
 
 if st.session_state.messages and st.button("📄 Download Chat as PDF"):
     pdf_file = export_chat_to_pdf(st.session_state.messages)

rag_pipeline.py

@@ -1,3 +1,4 @@
+import streamlit as st
 from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
 from sentence_transformers import SentenceTransformer
@@ -5,10 +6,89 @@ from qdrant_client import QdrantClient
 from langchain_qdrant import Qdrant
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.embeddings import SentenceTransformerEmbeddings
+from transformers import pipeline
 import os
 import torch
+from groq import Groq
+import google.generativeai as genai
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+import cohere
 
-def invoke_rag_chain(llm):
+available_models = ["OpenAI GPT-4o", "LLaMA 3", "Gemini Pro", "Ensemble"]
+AI_PROMPT_TEMPLATE = """You are an AI-assisted Dermatology Chatbot, specializing in diagnosing and educating users about skin diseases.
+    You provide accurate, compassionate, and detailed explanations while using correct medical terminology.
+
+    Guidelines:
+    1. Symptoms - Explain in simple terms with proper medical definitions.
+    2. Causes - Include genetic, environmental, and lifestyle-related risk factors.
+    3. Medications & Treatments - Provide common prescription and over-the-counter treatments.
+    4. Warnings & Emergencies - Always recommend consulting a licensed dermatologist.
+    5. Emergency Note - If symptoms worsen or include difficulty breathing, **advise calling 911 immediately.
+
+    Query: {question}
+    Relevant Information: {context}
+    Answer:
+    """
+
+
+@st.cache_resource(show_spinner=False)
+def initialize_rag_components():
+    components = {
+        'cohere_client': cohere.Client(st.secrets["COHERE_API_KEY"]),
+        'pair_ranker': pipeline("text-classification",
+                            model="llm-blender/PairRM",
+                            tokenizer="llm-blender/PairRM",
+                            return_all_scores=True
+                        ),
+        'gen_fuser': pipeline("text-generation",
+                        model="llm-blender/gen_fuser_3b",
+                        tokenizer="llm-blender/gen_fuser_3b",
+                        max_length=2048,
+                        do_sample=False
+                    ),
+        'retriever': get_retriever()
+    }
+    return components
+
+class AllModelsWrapper:
+    def invoke(self, messages):
+        prompt = messages[0]["content"]
+        rag_components = st.session_state.app_models['rag_components']  # Get components
+        responses = get_all_responses(prompt)
+        fused = rank_and_fuse(prompt, responses, rag_components)
+        return type('obj', (object,), {'content': fused})()
+
+def get_all_responses(prompt):
+    # Get responses from all models
+    openai_resp = ChatOpenAI(model="gpt-4o", temperature=0.2,
+                             api_key=st.secrets["OPENAI_API_KEY"]).invoke(
+        [{"role": "user", "content": prompt}]).content
+
+    gemini = genai.GenerativeModel("gemini-2.5-pro-exp-03-25")
+    gemini_resp = gemini.generate_content(prompt).text
+
+    llama = Groq(api_key=st.secrets["GROQ_API_KEY"])
+    llama_resp = llama.chat.completions.create(
+        model="meta-llama/llama-4-maverick-17b-128e-instruct",
+        messages=[{"role": "user", "content": prompt}],
+        temperature=1, max_completion_tokens=1024, top_p=1, stream=False
+    ).choices[0].message.content
+
+    return [openai_resp, gemini_resp, llama_resp]
+
+
+def rank_and_fuse(prompt, responses, rag_components):
+    ranked = [(resp, rag_components['pair_ranker'](f"{prompt}\n\n{resp}")[0][1]['score'])
+              for resp in responses]
+    ranked.sort(key=lambda x: x[1], reverse=True)
+
+    # Fuse top responses
+    fusion_input = "\n\n".join([f"[Answer {i + 1}]: {ans}" for i, (ans, _) in enumerate(ranked[:2])])
+    return rag_components['gen_fuser'](f"Fuse these responses:\n{fusion_input}",
+                     return_full_text=False)[0]['generated_text']
+
+
+def get_retriever():
     # === Qdrant DB Setup ===
     qdrant_client = QdrantClient(
         url="https://2715ddd8-647f-40ee-bca4-9027d193e8aa.us-east-1-0.aws.cloud.qdrant.io",
@@ -30,30 +110,74 @@ def invoke_rag_chain(llm):
         collection_name=collection_name,
         embeddings=local_embedding
     )
-    retriever = vector_store.as_retriever()
+    return vector_store.as_retriever()
 
+def initialize_llm(_model_name):
+    """Initialize the LLM based on selection"""
+    print(f"Model name : {_model_name}")
+    if "OpenAI" in _model_name:
+        return ChatOpenAI(model="gpt-4o", temperature=0.2, api_key=st.secrets["OPENAI_API_KEY"])
+    elif "LLaMA" in _model_name:
+        client = Groq(api_key=st.secrets["GROQ_API_KEY"])
+        def get_llama_response(prompt):
+            completion = client.chat.completions.create(
+                model="meta-llama/llama-4-maverick-17b-128e-instruct",
+                messages=[{"role": "user", "content": prompt}],
+                temperature=1,
+                max_completion_tokens=1024,
+                top_p=1,
+                stream=False
+            )
+            return completion.choices[0].message.content
+        return type('obj', (object,), {'invoke': lambda self, x: get_llama_response(x[0]["content"])})()
+
+    elif "Gemini" in _model_name:
+        genai.configure(api_key=st.secrets["GEMINI_API_KEY"])
+        gemini_model = genai.GenerativeModel("gemini-2.5-pro-exp-03-25")
+        def get_gemini_response(prompt):
+            response = gemini_model.generate_content(prompt)
+            return response.text
+        return type('obj', (object,), {'invoke': lambda self, x: get_gemini_response(x[0]["content"])})()
+
+    elif "Ensemble" in _model_name:
+        return AllModelsWrapper()
+    else:
+        raise ValueError("Unsupported model selected")
+
+
+def load_rag_chain(llm):
 
-    AI_PROMPT_TEMPLATE = """You are an AI-assisted Dermatology Chatbot, specializing in diagnosing and educating users about skin diseases.
-    You provide accurate, compassionate, and detailed explanations while using correct medical terminology.
-    
-    Guidelines:
-    1. Symptoms - Explain in simple terms with proper medical definitions.
-    2. Causes - Include genetic, environmental, and lifestyle-related risk factors.
-    3. Medications & Treatments - Provide common prescription and over-the-counter treatments.
-    4. Warnings & Emergencies - Always recommend consulting a licensed dermatologist.
-    5. Emergency Note - If symptoms worsen or include difficulty breathing, **advise calling 911 immediately.
-    
-    Query: {question}
-    Relevant Information: {context}
-    Answer:
-    """
     prompt_template = PromptTemplate(template=AI_PROMPT_TEMPLATE, input_variables=["question", "context"])
 
     rag_chain = RetrievalQA.from_chain_type(
         llm=llm,
-        retriever=retriever,
+        retriever=get_retriever(),
         chain_type="stuff",
         chain_type_kwargs={"prompt": prompt_template, "document_variable_name": "context"}
     )
 
-    return rag_chain
+    return rag_chain
+
+def rerank_with_cohere(query, documents, co, top_n=5):
+    if not documents:
+        return []
+    raw_texts = [doc.page_content for doc in documents]
+    results = co.rerank(query=query, documents=raw_texts, top_n=min(top_n, len(raw_texts)), model="rerank-v3.5")
+    return [documents[result.index] for result in results]
+
+
+def get_reranked_response(query, llm, rag_components):
+    """Get response with reranking"""
+    docs = rag_components['retriever'].get_relevant_documents(query)
+    reranked_docs = rerank_with_cohere(query, docs, rag_components['cohere_client'])
+    context = "\n\n".join([doc.page_content for doc in reranked_docs])
+
+    if isinstance(llm, (ChatOpenAI, AllModelsWrapper)):
+        return load_rag_chain(llm).invoke({"query": query, "context": context})['result']
+    else:
+        prompt = AI_PROMPT_TEMPLATE.format(question=query, context=context)
+        return llm.invoke([{"role": "user", "content": prompt}]).content
+
+
+if __name__ == "__main__":
+    print("This is a module - import it instead of running directly")

requirements.txt

@@ -17,4 +17,7 @@ nest_asyncio
 sentence_transformers
 langchain-qdrant
 huggingface_hub
-langchain_core
+langchain_core
+groq
+google
+cohere

test.py

@@ -0,0 +1,416 @@
+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
+import pandas as pd
+from huggingface_hub import hf_hub_download
+from langchain_huggingface import HuggingFaceEmbeddings
+import io
+import os
+import base64
+from fpdf import FPDF
+from sqlalchemy import create_engine
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import Distance, VectorParams
+from sentence_transformers import SentenceTransformer
+# from langchain_community.vectorstores.pgvector import PGVector
+# from langchain_postgres import PGVector
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+from langchain_community.vectorstores import Qdrant
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.embeddings import SentenceTransformerEmbeddings
+from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM
+import nest_asyncio
+
+torch.cuda.empty_cache()
+nest_asyncio.apply()
+co = cohere.Client(st.secrets["COHERE_API_KEY"])
+
+st.set_page_config(page_title="DermBOT", page_icon="🧬", layout="centered")
+
+# === Model Selection ===
+available_models = ["GPT-4o", "LLaMA 4 Maverick", "Gemini 2.5 Pro", "All"]
+st.session_state["selected_model"] = st.sidebar.selectbox("Select LLM Model", available_models)
+
+# === Qdrant DB Setup ===
+qdrant_client = QdrantClient(
+    url="https://2715ddd8-647f-40ee-bca4-9027d193e8aa.us-east-1-0.aws.cloud.qdrant.io",
+    api_key="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhY2Nlc3MiOiJtIn0.HXzezXdWMFeeR16F7zvqgjzsqrcm8hqa-StXdToFP9Q"
+)
+collection_name = "ks_collection_1.5BE"
+# embedding_model = SentenceTransformer("D:\DR\RAG\gte-Qwen2-1.5B-instruct", trust_remote_code=True)
+# embedding_model.max_seq_length = 8192
+# local_embedding = SentenceTransformerEmbeddings(model=embedding_model)
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def get_safe_embedding_model():
+    model_name = "Alibaba-NLP/gte-Qwen2-1.5B-instruct"
+
+    try:
+        print("Trying to load embedding model on CUDA...")
+        embedding = HuggingFaceEmbeddings(
+            model_name=model_name,
+            model_kwargs={
+                "trust_remote_code": True,
+                "device": "cuda"
+            }
+        )
+        print("Loaded embedding model on GPU.")
+        return embedding
+    except RuntimeError as e:
+        if "CUDA out of memory" in str(e):
+            print("CUDA OOM. Falling back to CPU.")
+        else:
+            print(" Error loading model on CUDA:", str(e))
+        print("Loading embedding model on CPU...")
+        return HuggingFaceEmbeddings(
+            model_name=model_name,
+            model_kwargs={
+                "trust_remote_code": True,
+                "device": "cpu"
+            }
+        )
+
+
+# Replace your old local_embedding line with this
+local_embedding = get_safe_embedding_model()
+
+print(" Qwen2-1.5B local embedding model loaded.")
+
+vector_store = Qdrant(
+    client=qdrant_client,
+    collection_name=collection_name,
+    embeddings=local_embedding
+)
+retriever = vector_store.as_retriever()
+
+pair_ranker = pipeline(
+    "text-classification",
+    model="llm-blender/PairRM",
+    tokenizer="llm-blender/PairRM",
+    return_all_scores=True
+)
+
+gen_fuser = pipeline(
+    "text-generation",
+    model="llm-blender/gen_fuser_3b",
+    tokenizer="llm-blender/gen_fuser_3b",
+    max_length=2048,
+    do_sample=False
+)
+
+# selected_model = st.session_state["selected_model"]
+
+if "OpenAI" in selected_model:
+    from langchain_openai import ChatOpenAI
+
+    llm = ChatOpenAI(model="gpt-4o", temperature=0.2, api_key=st.secrets["OPENAI_API_KEY"])
+
+elif "LLaMA" in selected_model:
+    from groq import Groq
+
+    client = Groq(api_key=st.secrets["GROQ_API_KEY"])  # Store in `.streamlit/secrets.toml`
+
+
+    def get_llama_response(prompt):
+        completion = client.chat.completions.create(
+            model="meta-llama/llama-4-maverick-17b-128e-instruct",
+            messages=[{"role": "user", "content": prompt}],
+            temperature=1,
+            max_completion_tokens=1024,
+            top_p=1,
+            stream=False
+        )
+        return completion.choices[0].message.content
+
+
+    llm = get_llama_response  # use this in place of llm.invoke()
+
+elif "Gemini" in selected_model:
+    import google.generativeai as genai
+
+    genai.configure(api_key=st.secrets["GEMINI_API_KEY"])  # Store in `.streamlit/secrets.toml`
+
+    gemini_model = genai.GenerativeModel("gemini-2.5-pro-exp-03-25")
+
+
+    def get_gemini_response(prompt):
+        response = gemini_model.generate_content(prompt)
+        return response.text
+
+
+    llm = get_gemini_response
+
+elif "All" in selected_model:
+
+    from groq import Groq
+    import google.generativeai as genai
+
+    genai.configure(api_key=st.secrets["GEMINI_API_KEY"])
+
+
+    def get_all_model_responses(prompt):
+        openai_resp = ChatOpenAI(model="gpt-4o", temperature=0.2, api_key=st.secrets["OPENAI_API_KEY"]).invoke(
+            [{"role": "system", "content": prompt}]).content
+
+        gemini = genai.GenerativeModel("gemini-2.5-pro-exp-03-25")
+        gemini_resp = gemini.generate_content(prompt).text
+
+        llama = Groq(api_key=st.secrets["GROQ_API_KEY"])
+        llama_resp = llama.chat.completions.create(
+            model="meta-llama/llama-4-maverick-17b-128e-instruct",
+            messages=[{"role": "user", "content": prompt}],
+            temperature=1, max_completion_tokens=1024, top_p=1, stream=False
+        ).choices[0].message.content
+
+        return [openai_resp, gemini_resp, llama_resp]
+
+
+    def rank_and_fuse(prompt, responses):
+        ranked = [(resp, pair_ranker(f"{prompt}\n\n{resp}")[0][1]['score']) for resp in responses]
+        ranked.sort(key=lambda x: x[1], reverse=True)
+        fusion_input = "\n\n".join([f"[Answer {i + 1}]: {ans}" for i, (ans, _) in enumerate(ranked)])
+        return gen_fuser(f"Fuse these responses:\n{fusion_input}", return_full_text=False)[0]['generated_text']
+
+
+else:
+    st.error("Unsupported model selected.")
+    st.stop()
+
+# retriever = vector_store.as_retriever()
+
+AI_PROMPT_TEMPLATE = """
+You are DermBOT, a compassionate and knowledgeable AI Dermatology Assistant designed to educate users about skin-related health concerns with clarity, empathy, and precision.
+
+Your goal is to respond like a well-informed human expert—balancing professionalism with warmth and reassurance.
+
+When crafting responses:
+- Begin with a clear, engaging summary of the condition or concern.
+- Use short paragraphs for readability.
+- Include bullet points or numbered lists where appropriate.
+- Avoid overly technical terms unless explained simply.
+- End with a helpful next step, such as lifestyle advice or when to see a doctor.
+
+🩺 Response Structure:
+1. **Overview** — Briefly introduce the condition or concern.
+2. **Common Symptoms** — Describe noticeable signs in simple terms.
+3. **Causes & Risk Factors** — Include genetic, lifestyle, and environmental aspects.
+4. **Treatment Options** — Outline common OTC and prescription treatments.
+5. **When to Seek Help** — Warn about symptoms that require urgent care.
+
+Always encourage consulting a licensed dermatologist for personal diagnosis and treatment. For any breathing difficulties, serious infections, or rapid symptom worsening, advise calling emergency services immediately.
+
+---
+
+Query: {question}
+Relevant Context: {context}
+
+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"}
+# )
+
+# === Class Names ===
+multilabel_class_names = [
+    "Vesicle", "Papule", "Macule", "Plaque", "Abscess", "Pustule", "Bulla", "Patch",
+    "Nodule", "Ulcer", "Crust", "Erosion", "Excoriation", "Atrophy", "Exudate", "Purpura/Petechiae",
+    "Fissure", "Induration", "Xerosis", "Telangiectasia", "Scale", "Scar", "Friable", "Sclerosis",
+    "Pedunculated", "Exophytic/Fungating", "Warty/Papillomatous", "Dome-shaped", "Flat topped",
+    "Brown(Hyperpigmentation)", "Translucent", "White(Hypopigmentation)", "Purple", "Yellow",
+    "Black", "Erythema", "Comedo", "Lichenification", "Blue", "Umbilicated", "Poikiloderma",
+    "Salmon", "Wheal", "Acuminate", "Burrow", "Gray", "Pigmented", "Cyst"
+]
+
+multiclass_class_names = [
+    "systemic", "hair", "drug_reactions", "uriticaria", "acne", "light",
+    "autoimmune", "papulosquamous", "eczema", "skincancer",
+    "benign_tumors", "bacteria_parasetic_infections", "fungal_infections", "viral_skin_infections"
+]
+
+
+# === Load Models ===
+class SkinViT(nn.Module):
+    def __init__(self, num_classes):
+        super(SkinViT, self).__init__()
+        self.model = vit_b_16(weights=ViT_B_16_Weights.DEFAULT)
+        in_features = self.model.heads.head.in_features
+        self.model.heads.head = nn.Linear(in_features, num_classes)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class DermNetViT(nn.Module):
+    def __init__(self, num_classes):
+        super(DermNetViT, self).__init__()
+        self.model = vit_l_16(weights=ViT_L_16_Weights.DEFAULT)
+        in_features = self.model.heads[0].in_features
+        self.model.heads[0] = nn.Sequential(
+            nn.Dropout(0.3),
+            nn.Linear(in_features, num_classes)
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+
+# multilabel_model = torch.load("D:/DR/RAG/BestModels2703/skin_vit_fold10.pth", map_location='cpu')
+# multiclass_model = torch.load("D:/DR/RAG/BestModels2703/best_dermnet_vit.pth", map_location='cpu')
+
+# === Load Model State Dicts ===
+multilabel_model_path = hf_hub_download(repo_id="santhoshraghu/DermBOT", filename="skin_vit_fold10_sd.pth")
+multiclass_model_path = hf_hub_download(repo_id="santhoshraghu/DermBOT", filename="best_dermnet_vit_sd.pth")
+
+
+def load_model_with_fallback(model_class, weight_path, num_classes, model_name):
+    try:
+        print(f"🔍 Loading {model_name} on GPU...")
+        model = model_class(num_classes)
+        model.load_state_dict(torch.load(weight_path, map_location="cuda"))
+        model.to("cuda")
+        print(f"✅ {model_name} loaded on GPU.")
+        return model
+    except RuntimeError as e:
+        if "CUDA out of memory" in str(e):
+            print(f"⚠️ {model_name} OOM. Falling back to CPU.")
+        else:
+            print(f"❌ Error loading {model_name} on CUDA: {e}")
+        print(f"🔄 Loading {model_name} on CPU...")
+        model = model_class(num_classes)
+        model.load_state_dict(torch.load(weight_path, map_location="cpu"))
+        model.to("cpu")
+        return model
+
+
+# Load both models with fallback
+multilabel_model = load_model_with_fallback(SkinViT, multilabel_model_path, len(multilabel_class_names), "SkinViT")
+multiclass_model = load_model_with_fallback(DermNetViT, multiclass_model_path, len(multiclass_class_names),
+                                            "DermNetViT")
+
+multilabel_model.eval()
+multiclass_model.eval()
+
+# === Session Init ===
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+
+# === Image Processing Function ===
+def run_inference(image):
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5])
+    ])
+    input_tensor = transform(image).unsqueeze(0)
+
+    # Automatically match model device (GPU or CPU)
+    model_device = next(multilabel_model.parameters()).device
+    input_tensor = input_tensor.to(model_device)
+
+    with torch.no_grad():
+        probs_multi = torch.sigmoid(multilabel_model(input_tensor)).squeeze().cpu().numpy()
+        pred_idx = torch.argmax(multiclass_model(input_tensor), dim=1).item()
+        predicted_multi = [multilabel_class_names[i] for i, p in enumerate(probs_multi) if p > 0.5]
+        predicted_single = multiclass_class_names[pred_idx]
+
+    return predicted_multi, predicted_single
+
+
+# === PDF Export ===
+def export_chat_to_pdf(messages):
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_font("Arial", size=12)
+    for msg in messages:
+        role = "You" if msg["role"] == "user" else "AI"
+        pdf.multi_cell(0, 10, f"{role}: {msg['content']}\n")
+    buf = io.BytesIO()
+    pdf.output(buf)
+    buf.seek(0)
+    return buf
+
+
+# Reranker utility
+def rerank_with_cohere(query, documents, top_n=5):
+    if not documents:
+        return []
+    raw_texts = [doc.page_content for doc in documents]
+    results = co.rerank(query=query, documents=raw_texts, top_n=min(top_n, len(raw_texts)), model="rerank-v3.5")
+    return [documents[result.index] for result in results]
+
+
+# Final answer generation using reranked context
+def get_reranked_response(query):
+    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)
+
+    if selected_model == "All":
+        responses = get_all_model_responses(prompt)
+        fused = rank_and_fuse(prompt, responses)
+        return type("Obj", (), {"content": fused})
+
+    if callable(llm):
+        return type("Obj", (), {"content": llm(prompt)})
+    else:
+        return llm.invoke([{"role": "system", "content": prompt}])
+
+
+# === App UI ===
+
+st.title("🧬 DermBOT — Skin AI Assistant")
+st.caption(f"🧠 Using model: {selected_model}")
+uploaded_file = st.file_uploader("Upload a skin image", type=["jpg", "jpeg", "png"])
+
+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}")
+
+    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 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.content)
+
+# === Chat Interface ===
+if prompt := st.chat_input("Ask a follow-up..."):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    response = get_reranked_response(prompt)
+    st.session_state.messages.append({"role": "assistant", "content": response.content})
+    with st.chat_message("assistant"):
+        st.markdown(response.content)
+
+# === PDF Button ===
+if st.button("📄 Download Chat as PDF"):
+    pdf_file = export_chat_to_pdf(st.session_state.messages)
+    st.download_button("Download PDF", data=pdf_file, file_name="chat_history.pdf", mime="application/pdf")