Initial commit

app.py

@@ -0,0 +1,63 @@
+from indexer import (
+    load_raw_dataset,
+    create_vector_database,
+    get_llm,
+    get_prompt_template,
+)
+import gradio as gr
+
+
+def format_contexts(contexts):
+    return "\n".join(
+        [
+            f"Reference {i+1}:\n{doc.metadata['question']}\n{doc.metadata['answer']}"
+            for i, doc in enumerate(contexts)
+        ]
+    )
+
+
+class CustomRAG:
+    def __init__(self, vector_db, llm, prompt_template):
+        self.vector_db = vector_db
+        self.llm = llm
+        self.prompt_template = prompt_template
+
+    def run(self, query):
+        retriever = self.vector_db.as_retriever(search_kwargs={"k": 3})
+        contexts = retriever.get_relevant_documents(query)
+        formatted_context = format_contexts(contexts)
+        prompt = self.prompt_template.format(context=formatted_context, question=query)
+        return self.llm.invoke(prompt), contexts
+
+
+def answer_question(query):
+    docs = load_raw_dataset()
+    rag = CustomRAG(
+        create_vector_database(docs, "all-MiniLM-L6-v2"),
+        get_llm("FreedomIntelligence/HuatuoGPT-o1-7B"),
+        get_prompt_template(),
+    )
+    response, _ = rag.run(query)
+
+    return response
+
+
+demo = gr.Interface(
+    fn=answer_question,
+    inputs=[
+        gr.Textbox(
+            label="Describe your medical concern",
+            placeholder="e.g. I've been feeling tired and dizzy lately.",
+            lines=3,
+        ),
+    ],
+    outputs="text",
+    title="Medical Assistant – Powered by AI & RAG",
+    description=(
+        "Get helpful insights based on your described symptoms. "
+        "This assistant uses medical reference data to provide informative responses. "
+        "Note: This is not a substitute for professional medical advice."
+    ),
+)
+
+demo.launch()

indexer.py

@@ -0,0 +1,74 @@
+from datasets import load_dataset
+import pandas as pd
+from langchain.schema import Document
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from langchain.llms import HuggingFacePipeline
+from langchain.prompts import PromptTemplate
+
+
+def load_raw_dataset():
+    dataset = load_dataset("lavita/ChatDoctor-HealthCareMagic-100k")
+
+    df = pd.DataFrame(dataset["train"])
+
+    df["combined"] = df["input"] + " " + df["output"]
+
+    docs = [
+        Document(
+            page_content=row["combined"],
+            metadata={"question": row["input"], "answer": row["output"]},
+        )
+        for _, row in df.iterrows()
+    ]
+
+    return docs
+
+
+def create_vector_database(docs, model_name):
+    embedding_model = HuggingFaceEmbeddings(model_name=model_name)
+    vectorstore = FAISS.from_documents(docs, embedding_model)
+    return vectorstore
+
+
+def get_llm(model_name):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name, torch_dtype="auto", device_map="auto"
+    )
+
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=512,
+        temperature=0.7,
+        do_sample=True,
+    )
+
+    llm = HuggingFacePipeline(pipeline=pipe)
+    return llm
+
+
+def get_prompt_template():
+    prompt_template = PromptTemplate(
+        input_variables=["context", "question"],
+        template="""Based on the following references and your medical knowledge, provide a detailed response:
+
+        References:
+        {context}
+
+        Question: {question}
+
+        By considering:
+        1. The key medical concepts in the question.
+        2. How the reference cases relate to this question.
+        3. What medical principles should be applied.
+        4. Any potential complications or considerations.
+
+        Give the final response:
+        """,
+    )
+
+    return prompt_template

requirements.txt

@@ -0,0 +1,10 @@
+gradio
+transformers
+sentence-transformers
+torch
+langchain
+faiss-cpu
+huggingface-hub
+praw
+langchain-community
+accelerate