Update app.py

app.py

@@ -1,11 +1,73 @@
-import streamlit as st
-import PyPDF2
-from huggingface_hub import InferenceClient
-
-# Initialize the Inference Client
-client = InferenceClient("meta-llama/Llama-3.2-3B-Instruct")
-
-
+import gradio as gr
+from langchain_community.vectorstores import FAISS
+from langchain_community.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.chains import ConversationalRetrievalChain
+from langchain_community.llms import HuggingFaceEndpoint
+from langchain.chains import ConversationChain
+from langchain.memory import ConversationBufferMemory
+import os
+
+api_token = os.getenv("HF_TOKEN")
+
+# List of LLMs
+list_llm = ["meta-llama/Meta-Llama-3-8B-Instruct", "mistralai/Mistral-7B-Instruct-v0.2"]
+list_llm_simple = [os.path.basename(llm) for llm in list_llm]
+
+# Load and split PDF documents
+def load_doc(list_file_path):
+    loaders = [PyPDFLoader(x) for x in list_file_path]
+    pages = []
+    for loader in loaders:
+        pages.extend(loader.load())
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=64)
+    doc_splits = text_splitter.split_documents(pages)
+    return doc_splits
+
+# Create vector database
+def create_db(splits):
+    embeddings = HuggingFaceEmbeddings()
+    vectordb = FAISS.from_documents(splits, embeddings)
+    return vectordb
+
+# Initialize LLM chain
+def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db):
+    if llm_model == "meta-llama/Meta-Llama-3-8B-Instruct":
+        llm = HuggingFaceEndpoint(
+            repo_id=llm_model,
+            huggingfacehub_api_token=api_token,
+            temperature=temperature,
+            max_new_tokens=max_tokens,
+            top_k=top_k,
+        )
+    else:
+        llm = HuggingFaceEndpoint(
+            huggingfacehub_api_token=api_token,
+            repo_id=llm_model,
+            temperature=temperature,
+            max_new_tokens=max_tokens,
+            top_k=top_k,
+        )
+    
+    memory = ConversationBufferMemory(
+        memory_key="chat_history",
+        output_key='answer',
+        return_messages=True
+    )
+
+    retriever = vector_db.as_retriever()
+    qa_chain = ConversationalRetrievalChain.from_llm(
+        llm,
+        retriever=retriever,
+        chain_type="stuff",
+        memory=memory,
+        return_source_documents=True,
+        verbose=False,
+    )
+    return qa_chain
+
+# Function to handle chatbot responses
 def respond(
     message,
     history: list[tuple[str, str]],
@@ -13,149 +75,79 @@ def respond(
     max_tokens,
     temperature,
     top_p,
-    uploaded_pdf=None
+    vector_db,
+    llm_model,
 ):
-    messages = [{"role": "system", "content": system_message}]
-
-    # Add previous conversation history to the messages
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    # If a new message is entered, add it to the conversation history
-    messages.append({"role": "user", "content": message})
-
-    # If a PDF is uploaded, process its content
-    if uploaded_pdf is not None:
-        file_content = extract_pdf_text(uploaded_pdf)
-        if file_content:
-            messages.append({"role": "user", "content": f"Document Content: {file_content}"})
-
-    # Get response from the model
-    response = ""
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-        response += token
-        yield response
-
-
-def extract_pdf_text(file):
-    """Extract text from a PDF file."""
-    try:
-        reader = PyPDF2.PdfReader(file)
-        text = ""
-        for page in reader.pages:
-            text += page.extract_text()
-        return text
-    except Exception as e:
-        return f"Error extracting text from PDF: {str(e)}"
-
-
-# Streamlit UI
-st.set_page_config(page_title="Health Assistant", layout="wide")
-
-# Custom CSS for Streamlit app
-st.markdown(
-    """
-    <style>
-    body {
-        background-color: #1e2a38; /* Dark blue background */
-        color: #ffffff; /* White text for readability */
-        font-family: 'Arial', sans-serif; /* Clean and modern font */
-    }
-    .stButton button {
-        background-color: #42B3CE !important; /* Light blue button */
-        color: #2e3b4e !important; /* Dark text for contrast */
-        border: none !important;
-        padding: 10px 20px !important;
-        border-radius: 8px !important;
-        font-size: 16px;
-        font-weight: bold;
-        transition: background-color 0.3s ease, transform 0.2s ease;
-    }
-    .stButton button:hover {
-        background-color: #3189A2 !important; /* Darker blue on hover */
-        transform: scale(1.05);
-    }
-    .stTextInput input {
-        background-color: #2f3b4d;
-        color: white;
-        border: 2px solid #42B3CE;
-        padding: 12px;
-        border-radius: 8px;
-        font-size: 16px;
-        transition: border 0.3s ease;
-    }
-    .stTextInput input:focus {
-        border-color: #3189A2;
-    }
-    </style>
-    """,
-    unsafe_allow_html=True,
+    # Initialize LLM chain if not already initialized
+    if not hasattr(respond, 'qa_chain'):
+        respond.qa_chain = initialize_llmchain(llm_model, temperature, max_tokens, top_p, vector_db)
+
+    # Format chat history
+    formatted_chat_history = []
+    for user_message, bot_message in history:
+        formatted_chat_history.append(f"User: {user_message}")
+        formatted_chat_history.append(f"Assistant: {bot_message}")
+    formatted_chat_history.append(f"User: {message}")
+
+    # Generate response using QA chain
+    response = respond.qa_chain.invoke({"question": message, "chat_history": formatted_chat_history})
+    response_answer = response["answer"]
+    if response_answer.find("Helpful Answer:") != -1:
+        response_answer = response_answer.split("Helpful Answer:")[-1]
+
+    return response_answer
+
+# CSS for styling the interface
+css = """
+body {
+    background-color: #06688E; /* Dark background */
+    color: white; /* Text color for better visibility */
+}
+.gr-button {
+    background-color: #42B3CE !important; /* White button color */
+    color: black !important; /* Black text for contrast */
+    border: none !important;
+    padding: 8px 16px !important;
+    border-radius: 5px !important;
+}
+.gr-button:hover {
+    background-color: #e0e0e0 !important; /* Slightly lighter button on hover */
+}
+.gr-slider-container {
+    color: white !important; /* Slider labels in white */
+}
+"""
+
+# Initialize database and LLM chain
+def initialize_database_and_llm(list_file_obj, llm_option, max_tokens, temperature, top_p):
+    list_file_path = [x.name for x in list_file_obj if x is not None]
+    doc_splits = load_doc(list_file_path)
+    vector_db = create_db(doc_splits)
+    llm_name = list_llm[llm_option]
+    return vector_db, llm_name
+
+# Gradio interface
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Files(file_count="multiple", file_types=["pdf"], label="Upload PDF documents", visible=False),
+        gr.Radio(list_llm_simple, label="Available LLMs", value=list_llm_simple, visible=False),
+        gr.Slider(minimum=128, maximum=9192, value=4096, step=128, label="Max new tokens", visible=False),
+        gr.Slider(minimum=0.01, maximum=1.0, value=0.5, step=0.1, label="Temperature", visible=False),
+        gr.Slider(minimum=1, maximum=10, value=3, step=1, label="Top-k", visible=False),
+    ],
+    css=css,
+    title="RAG PDF Chatbot",
+    description="Query your PDF documents using a Retrieval Augmented Generation (RAG) chatbot.",
 )
 
-# Title and description
-st.title("Health Assistant Chat")
-st.subheader("Chat with your health assistant and upload a document for analysis")
-
-# System message for health-related responses
-system_message = (
-    "You are a virtual health assistant designed to provide accurate and reliable information "
-    "related to health, wellness, and medical topics. Your primary goal is to assist users with "
-    "their health-related queries, offer general guidance, and suggest when to consult a licensed "
-    "medical professional. If a user asks a question that is unrelated to health, wellness, or medical "
-    "topics, respond politely but firmly with: 'I'm sorry, I can't help with that because I am a virtual "
-    "health assistant designed to assist with health-related needs. Please let me know if you have any health-related questions.'"
+# Preprocessing events
+demo.preprocess(
+    initialize_database_and_llm,
+    inputs=["document", "llm_btn", "slider_maxtokens", "slider_temperature", "slider_topk"],
+    outputs=["vector_db", "llm_model"],
+    api_name="initialize",
 )
 
-# Upload a PDF file
-uploaded_pdf = st.file_uploader("Upload a PDF file (Optional)", type="pdf")
-
-# User input message
-message = st.text_input("Type your health-related question:")
-
-# History for conversation tracking
-if 'history' not in st.session_state:
-    st.session_state['history'] = []
-
-# Collect and display previous conversation history
-history = st.session_state['history']
-for user_message, assistant_message in history:
-    st.markdown(f"**You:** {user_message}")
-    st.markdown(f"**Assistant:** {assistant_message}")
-
-# Max tokens, temperature, and top-p sliders
-max_tokens = st.slider("Max new tokens", min_value=1, max_value=2048, value=512)
-temperature = st.slider("Temperature", min_value=0.1, max_value=4.0, value=0.7, step=0.1)
-top_p = st.slider("Top-p (nucleus sampling)", min_value=0.1, max_value=1.0, value=0.95, step=0.05)
-
-# Button to generate response
-if st.button("Generate Response"):
-    if message:
-        # Append the user's question to the conversation history
-        st.session_state.history.append((message, ""))
-        # Generate the response based on the user's input and any uploaded document
-        response = respond(
-            message,
-            st.session_state.history,
-            system_message,
-            max_tokens,
-            temperature,
-            top_p,
-            uploaded_pdf
-        )
-        # Display the response
-        for resp in response:
-            st.markdown(f"**Assistant:** {resp}")
-        # Update the conversation history with the assistant's response
-        st.session_state.history[-1] = (message, resp)
-    else:
-        st.error("Please enter a question to proceed.")
+if __name__ == "__main__":
+    demo.launch(share=True)