mainchainge cgpt 2

app.py

@@ -5,24 +5,20 @@ from langchain.embeddings import HuggingFaceBgeEmbeddings
 from langchain_together import Together
 from langchain import hub
 from operator import itemgetter
-from langchain.schema import RunnableParallel, format_document
-from typing import List, Tuple
-from langchain.chains import LLMChain, ConversationalRetrievalChain
-from langchain.schema.output_parser import StrOutputParser
-from langchain.memory import StreamlitChatMessageHistory, ConversationBufferMemory, ConversationSummaryMemory
-from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate
-from langchain.schema import RunnableLambda, RunnablePassthrough
+from langchain.schema import format_document
+from langchain.prompts import ChatPromptTemplate, PromptTemplate
+from langchain.memory import StreamlitChatMessageHistory, ConversationBufferMemory
+from langchain_core.runnables import RunnableLambda, RunnableParallel, RunnablePassthrough
 
 # Load the embedding function
 model_name = "BAAI/bge-base-en"
 encode_kwargs = {'normalize_embeddings': True}
-
 embedding_function = HuggingFaceBgeEmbeddings(
     model_name=model_name,
     encode_kwargs=encode_kwargs
 )
 
-# Load the LLM
+# Initialize the LLMs
 llm = Together(
     model="mistralai/Mixtral-8x22B-Instruct-v0.1",
     temperature=0.2,
@@ -31,7 +27,6 @@ llm = Together(
     together_api_key=os.environ['pilotikval']
 )
 
-# Load the summarizeLLM
 llmc = Together(
     model="mistralai/Mixtral-8x22B-Instruct-v0.1",
     temperature=0.2,
@@ -40,19 +35,40 @@ llmc = Together(
     together_api_key=os.environ['pilotikval']
 )
 
+# Memory setup
 msgs = StreamlitChatMessageHistory(key="langchain_messages")
 memory = ConversationBufferMemory(chat_memory=msgs)
 
-DEFAULT_DOCUMENT_PROMPT = PromptTemplate.from_template(template="{page_content}")
+# Define the prompt templates
+CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(
+    """Given the following conversation and a follow-up question, rephrase the follow-up question to be a standalone question.
+    Chat History:
+    {chat_history}
+    Follow Up Input: {question}
+    Standalone question:"""
+)
+
+ANSWER_PROMPT = ChatPromptTemplate.from_template(
+    """You are helping a doctor. Answer based on the provided context:
+    {context}
+    Question: {question}"""
+)
 
-def _combine_documents(docs, document_prompt=DEFAULT_DOCUMENT_PROMPT, document_separator="\n\n"):
+# Function to combine documents
+def _combine_documents(docs, document_prompt=PromptTemplate.from_template("{page_content}"), document_separator="\n\n"):
     doc_strings = [format_document(doc, document_prompt) for doc in docs]
     return document_separator.join(doc_strings)
 
-chistory = []
+# Define the chain using LCEL
+condense_question_chain = RunnableLambda(lambda x: {"chat_history": chistory, "question": x}) | CONDENSE_QUESTION_PROMPT | llmc
+retriever_chain = RunnableLambda(lambda x: {"standalone_question": x}) | retriever | _combine_documents
+answer_chain = ANSWER_PROMPT | llm
 
-def store_chat_history(role: str, content: str):
-    chistory.append({"role": role, "content": content})
+conversational_qa_chain = RunnableParallel(
+    condense_question=condense_question_chain,
+    retrieve=retriever_chain,
+    generate_answer=answer_chain
+)
 
 # Define the Streamlit app
 def app():
@@ -63,67 +79,35 @@ def app():
             ('General Medicine', 'RespiratoryFishman', 'RespiratoryMurray', 'MedMRCP2', 'OldMedicine')
         )
         
-        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})
+        # Define retrievers based on option
+        persist_directory = {
+            'General Medicine': "./oxfordmedbookdir/",
+            'RespiratoryFishman': "./respfishmandbcud/",
+            'RespiratoryMurray': "./respmurray/",
+            'MedMRCP2': "./medmrcp2store/",
+            'OldMedicine': "./mrcpchromadb/"
+        }.get(option, "./mrcpchromadb/")
         
-        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})
+        collection_name = {
+            'General Medicine': "oxfordmed",
+            'RespiratoryFishman': "fishmannotescud",
+            'RespiratoryMurray': "respmurraynotes",
+            'MedMRCP2': "medmrcp2notes",
+            'OldMedicine': "mrcppassmednotes"
+        }.get(option, "mrcppassmednotes")
         
-        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})
-
+        vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function, collection_name=collection_name)
+        retriever = vectordb.as_retriever(search_kwargs={"k": 5})
+    
     if "messages" not in st.session_state:
         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("Ask Away!")
     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: