File size: 6,765 Bytes
4b7893b 398db91 4b7893b 398db91 4b7893b 398db91 4b7893b cf7407a 4b7893b 4524c74 4b7893b 398db91 4b7893b d98df44 4b7893b 398db91 4b7893b 398db91 4b7893b 398db91 4b7893b 398db91 d98df44 4b7893b 398db91 4b7893b 398db91 4b7893b 398db91 4b7893b 398db91 4b7893b 398db91 |
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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
import streamlit as st
import os
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.llms import Together
from langchain import hub
from operator import itemgetter
from langchain.schema.runnable import RunnableParallel
from langchain.schema import format_document
from typing import List, Tuple
from langchain.chains import LLMChain
from langchain.chains import RetrievalQA
from langchain.schema.output_parser import StrOutputParser
from langchain_community.chat_message_histories import StreamlitChatMessageHistory
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationSummaryMemory
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate
from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
import time
# Load the embedding function
model_name = "BAAI/bge-base-en"
encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
embedding_function = HuggingFaceBgeEmbeddings(
model_name=model_name,
encode_kwargs=encode_kwargs
)
# Load the LLM
llm = Together(
model="mistralai/Mixtral-8x22B-Instruct-v0.1",
temperature=0.2,
max_tokens=19096,
top_k=10,
together_api_key=os.environ['pilotikval']
)
# Load the summarizeLLM
llmc = Together(
model="mistralai/Mixtral-8x7B-Instruct-v0.1",
temperature=0.2,
max_tokens=1024,
top_k=1,
together_api_key=os.environ['pilotikval']
)
msgs = StreamlitChatMessageHistory(key="langchain_messages")
memory = ConversationBufferMemory(chat_memory=msgs)
DEFAULT_DOCUMENT_PROMPT = PromptTemplate.from_template(template="{page_content}")
def _combine_documents(
docs, document_prompt=DEFAULT_DOCUMENT_PROMPT, document_separator="\n\n"
):
doc_strings = [format_document(doc, document_prompt) for doc in docs]
return document_separator.join(doc_strings)
chistory = []
def store_chat_history(role: str, content: str):
# Append the new message to the chat history
chistory.append({"role": role, "content": content})
# Define the Streamlit app
def app():
with st.sidebar:
st.title("dochatter")
# Create a dropdown selection box
option = st.selectbox(
'Which retriever would you like to use?',
('General Medicine', 'RespiratoryFishman', 'RespiratoryMurray', 'MedMRCP2', 'OldMedicine')
)
# Depending on the selected option, choose the appropriate retriever
if option == 'RespiratoryFishman':
persist_directory = "./respfishmandbcud/"
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function, collection_name="fishmannotescud")
retriever = vectordb.as_retriever(search_kwargs={"k": 5})
elif option == 'RespiratoryMurray':
persist_directory = "./respmurray/"
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function, collection_name="respmurraynotes")
retriever = vectordb.as_retriever(search_kwargs={"k": 5})
elif option == 'MedMRCP2':
persist_directory = "./medmrcp2store/"
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function, collection_name="medmrcp2notes")
retriever = vectordb.as_retriever(search_kwargs={"k": 5})
elif option == 'General Medicine':
persist_directory = "./oxfordmedbookdir/"
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function, collection_name="oxfordmed")
retriever = vectordb.as_retriever(search_kwargs={"k": 7})
else:
persist_directory = "./mrcpchromadb/"
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function, collection_name="mrcppassmednotes")
retriever = vectordb.as_retriever(search_kwargs={"k": 5})
# Session State
if "messages" not in st.session_state.keys():
st.session_state.messages = [{"role": "assistant", "content": "How may I help you?"}]
_template = """Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question which contains the themes of the conversation. Do not write the question. Do not write the answer.
Chat History:
{chat_history}
Follow Up Input: {question}
Standalone question:"""
CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)
template = """You are helping a doctor. Answer with what you know from the context provided. Please be as detailed and thorough. Answer the question based on the following context:
{context}
Question: {question}
"""
ANSWER_PROMPT = ChatPromptTemplate.from_template(template)
_inputs = RunnableParallel(
standalone_question=RunnablePassthrough.assign(
chat_history=lambda x: chistory
) | CONDENSE_QUESTION_PROMPT | llmc | StrOutputParser(),
)
_context = {
"context": itemgetter("standalone_question") | retriever | _combine_documents,
"question": lambda x: x["standalone_question"],
}
conversational_qa_chain = _inputs | _context | ANSWER_PROMPT | llm
st.header("Hello Doctor!")
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.write(message["content"])
store_chat_history(message["role"], message["content"])
prompts2 = st.chat_input("Say something")
if prompts2:
st.session_state.messages.append({"role": "user", "content": prompts2})
with st.chat_message("user"):
st.write(prompts2)
if st.session_state.messages[-1]["role"] != "assistant":
with st.chat_message("assistant"):
with st.spinner("Thinking..."):
for _ in range(3): # Retry up to 3 times
try:
response = conversational_qa_chain.invoke(
{
"question": prompts2,
"chat_history": chistory,
}
)
st.write(response)
message = {"role": "assistant", "content": response}
st.session_state.messages.append(message)
break
except Exception as e:
st.error(f"An error occurred: {e}")
time.sleep(2) # Wait 2 seconds before retrying
if __name__ == '__main__':
app()
|