File size: 2,822 Bytes
b947dc9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
from langchain_community.vectorstores import FAISS
from config import embeddings, llm, llm2
from pprint import pprint
from langchain.retrievers.document_compressors import FlashrankRerank
from langchain.retrievers import ContextualCompressionRetriever
from langchain_core.prompts import ChatPromptTemplate
compressor = FlashrankRerank()
def get_vectoreDB(file_path:str):
path = f"vectors/{file_path}".replace(".pdf","").replace("data/","")
vectoreStore = FAISS.load_local(
path,
embeddings,
allow_dangerous_deserialization=True
)
vectoreStore_as_retriever = vectoreStore.as_retriever(
search_type="similarity",
search_kwargs={
"k": 10,
"include_metadata": True
}
)
#ReRanker
compression_retriever = ContextualCompressionRetriever(
base_compressor=compressor,
base_retriever=vectoreStore_as_retriever
)
return compression_retriever
system_prompt = """
You are a Medical Chatbot named as Medi. Your task is to assist as a doctor.
Your tasks are given below:
- Answer the given question from the given context. If the question is out of context from medical feild, just tell the user "I don't know about it".
- If the question is related to medical field and not relevent to context, then generate by your own self
- If the question is about the context given, explain the user according to the most relevent context.
- Explain the user a proper short Answer, don't copy the whole context, explain it by your self.
- If someone ask your name, so tell them
Remember to answer in a precise way and to the point as a doctor.
query: '''{query}'''
"""
PROMPT = ChatPromptTemplate.from_messages(
[
("system", system_prompt),
("human", "context: '''{context}'''")
]
)
llm_chain = PROMPT | llm
mini_llm_chain = PROMPT | llm2
vectoreDB = get_vectoreDB("data/The_GALE_ENCYCLOPEDIA_of_MEDICINE_SECOND.pdf")
def ask_ai(query:str, model_name:str):
docs = vectoreDB.invoke(query)
context_text = "\n\n---\n\n".join([doc.page_content for doc in docs])
if model_name == "4o":
llm = llm_chain
elif model_name == "4o-mini":
llm = mini_llm_chain
else:
return "No model found"
response = llm.invoke({
"context": context_text,
"query" : query
}
)
return response.content |