smvaideesh/MedicalAIChatbot

1import gradio as gr2from huggingface_hub import InferenceClient3from langchain.embeddings import HuggingFaceEmbeddings4from langchain.vectorstores import FAISS5from langchain.text_splitter import CharacterTextSplitter6from langchain.document_loaders import PyPDFLoader7import os8 9 10# Load the model client11client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")12 13# Initialize vector store14vector_store = None15 16# Preload and process the PDF document17#PDF_PATH = "general symptoms.pdf"  # Path to the pre-defined PDF document18 19#PDF_PATH = "general symptoms.pdf" 20PDF_PATH = "general symptoms.pdf" 21 22def preload_pdf():23    global vector_store24 25    # Load PDF and extract text26    loader = PyPDFLoader(PDF_PATH)27    documents = loader.load()28 29    # Split the text into smaller chunks for retrieval30    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)31    docs = text_splitter.split_documents(documents)32 33    # Compute embeddings for the chunks34    embeddings = HuggingFaceEmbeddings()35    vector_store = FAISS.from_documents(docs, embeddings)36 37    print(f"PDF '{PDF_PATH}' loaded and indexed successfully.")38 39# Response generation40def respond(41    message,42    history: list[tuple[str, str]],43    system_message,44    max_tokens,45    temperature,46    top_p,47):48    global vector_store49 50    if vector_store is None:51        return "The PDF document is not loaded. Please check the code setup."52 53    # Retrieve relevant chunks from the PDF54    relevant_docs = vector_store.similarity_search(message, k=3)55    context = "\n".join([doc.page_content for doc in relevant_docs])56 57    # Combine system message, context, and user message58    full_system_message = (59        f"{system_message}\n\nContext from the document:\n{context}\n\n"60    )61 62    messages = [{"role": "system", "content": full_system_message}]63 64    for val in history:65        if val[0]:66            messages.append({"role": "user", "content": val[0]})67        if val[1]:68            messages.append({"role": "assistant", "content": val[1]})69 70    messages.append({"role": "user", "content": message})71 72    response = ""73 74    for message in client.chat_completion(75        messages,76        max_tokens=max_tokens,77        stream=True,78        temperature=temperature,79        top_p=top_p,80    ):81        token = message.choices[0].delta.content82        response += token83        yield response84 85# Gradio interface86#demo = gr.Blocks()87 88demo = gr.Blocks(css="""89 90.gr-chat-container {91    display: flex;92    background-color: skyblue;93    justify-content: center;94    align-items: center;95    height: 90vh;96    padding: 20px;97}98 99.gr-chat {100    height: 80vh;101    justify-content: center;102    align-items: center;103    border: 1px solid #ccc;104    padding: 10px;105    box-shadow: 2px 2px 10px rgba(0, 0, 0, 0.1);106}107""")108 109 110with demo:111    with gr.Row(elem_classes=["gr-chat-container"]):112    #with gr.Row():113        with gr.Column(elem_classes=["gr-chat"]):114        #with gr.Column():115            chatbot = gr.ChatInterface(116                respond,117                additional_inputs=[118                    gr.Textbox(119                        value=(120                            "You are going to act like a medical practitioner. Hear the symptoms, "121                            "diagnose the disease, mention the disease in seperate line, suggest tips to overcome the issue and suggest some good habits "122                            "to overcome the issue. Base your answers on the provided document. limit the response to 5 to 6 sentence point by point" 123                        ),visible=False,124                        label="system_message",125                    ),126                    gr.Slider(minimum=1, maximum=2048, value=512, step=1,visible=False, label="Max new tokens"),127                    gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, visible=False,label="Temperature"),128                    gr.Slider(minimum=0.1,maximum=1.0,value=0.95,step=0.05,visible=False,label="Top-p (nucleus sampling)", ),129                ],130                examples=[131                    ["I am not well and feeling feverish, tired?"],132                    ["Can you guide me through quick health tips?"],133                    ["How do I stop worrying about things I can't control?"],134                ],135                title="Diagnify 🕊️",136            )137    138    139if __name__ == "__main__":140    preload_pdf()141    demo.launch()142