Update handler.py

handler.py

@@ -16,35 +16,20 @@ from langchain.document_loaders import WebBaseLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from llm_for_langchain import LLM
 from langchain.chains.qa_with_sources import load_qa_with_sources_chain
-from langchain.chains.combine_documents import create_stuff_documents_chain
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.messages import HumanMessage
-from langchain_core.output_parsers import StrOutputParser
-from langchain_core.runnables import RunnableLambda, RunnableBranch, RunnablePassthrough
-from operator import itemgetter
-from langchain.schema import format_document
-from langchain.memory import ConversationBufferMemory
-from langchain_core.messages import AIMessage, HumanMessage, get_buffer_string
-
     
 class EndpointHandler():
     def __init__(self, path=""):
 
-        # Config LangChain
-        # os.environ["LANGCHAIN_TRACING_V2"] = "true"
-        # os.environ["LANGCHAIN_API_KEY"] = 
-
-        # Create LLM
-        chat = LLM(model_name_or_path=path, bit4=False)
-
-        # Create Text-Embedding Model
-        embedding_function = HuggingFaceBgeEmbeddings(
-            model_name="DMetaSoul/Dmeta-embedding",
+        self.llm = LLM(model_name_or_path=path, bit4=False)
+        
+        # Load Vector db
+        
+        self.embedding_function = HuggingFaceBgeEmbeddings(
+            model_name="BAAI/bge-large-zh",
             model_kwargs={'device': 'cuda'},
             encode_kwargs={'normalize_embeddings': True}
         )
-        
-        # Load Vector db
+
         urls = [
             "https://hk.on.cc/hk/bkn/cnt/news/20221019/bkn-20221019040039334-1019_00822_001.html",
             "https://www.hk01.com/%E7%A4%BE%E6%9C%83%E6%96%B0%E8%81%9E/822848/%E5%89%B5%E7%A7%91%E7%B2%BE%E8%8B%B1-%E5%87%BA%E6%88%B02022%E4%B8%96%E7%95%8C%E6%8A%80%E8%83%BD%E5%A4%A7%E8%B3%BD%E7%89%B9%E5%88%A5%E8%B3%BD",
@@ -58,87 +43,40 @@ class EndpointHandler():
         text_splitter = RecursiveCharacterTextSplitter(chunk_size = 1000, chunk_overlap = 16)
         all_splits = text_splitter.split_documents(data)
         
-        vectorstore = Chroma.from_documents(documents=all_splits, embedding=embedding_function)
-        retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
+        vectorstore = Chroma.from_documents(documents=all_splits, embedding=self.embedding_function)
         
-        compressor = LLMChainExtractor.from_llm(chat)
-        retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=retriever)
+        # vectorstore = Chroma(persist_directory="db", embedding_function=embedding_function)
 
-        _template = """Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question, in its original language.
+        compressor = LLMChainExtractor.from_llm(self.llm)
+        self.retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=vectorstore.as_retriever(search_kwargs={"k": 4}))
 
-        Chat History:
-        {chat_history}
-        Follow Up Input: {question}
-        Standalone question:"""
+        prompt_template = """<s>[INST] <<SYS>> You are a helpful assistant.
+        Use the following context to Answer the question below briefly: <<SYS>>
         
-        CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)
+        {history}
         
-        template = """Answer the question based only on the following context:
         {context}
         
-        Question: {question}
+        {question} [/INST] </s>
         """
-        ANSWER_PROMPT = ChatPromptTemplate.from_template(template)
-
-        self.memory = ConversationBufferMemory(
-            return_messages=True, output_key="answer", input_key="question"
-        )
-        
-        # First we add a step to load memory
-        # This adds a "memory" key to the input object
-        loaded_memory = RunnablePassthrough.assign(
-            chat_history=RunnableLambda(self.memory.load_memory_variables) | itemgetter("history"),
-        )
-        # Now we calculate the standalone question
-        standalone_question = {
-            "standalone_question": {
-                "question": lambda x: x["question"],
-                "chat_history": lambda x: get_buffer_string(x["chat_history"]),
-            }
-            | CONDENSE_QUESTION_PROMPT
-            | chat(temperature=0)
-            | StrOutputParser(),
-        }
 
-        DEFAULT_DOCUMENT_PROMPT = PromptTemplate.from_template(template="{page_content}")
+        prompt = PromptTemplate(input_variables=["history", "context", "question"], template=prompt_template)
+        memory = ConversationBufferMemory(input_key='question', memory_key='history', return_messages=True)
 
-        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)
-  
-        # Now we retrieve the documents
-        retrieved_documents = {
-            "docs": itemgetter("standalone_question") | retriever,
-            "question": lambda x: x["standalone_question"],
-        }
-        # Now we construct the inputs for the final prompt
-        final_inputs = {
-            "context": lambda x: _combine_documents(x["docs"]),
-            "question": itemgetter("question"),
-        }
-        # And finally, we do the part that returns the answers
-        answer = {
-            "answer": final_inputs | ANSWER_PROMPT | chat,
-            "docs": itemgetter("docs"),
-        }
-        # And now we put it all together!
-        self.final_chain = loaded_memory | standalone_question | retrieved_documents | answer
+        self.qa_chain = RetrievalQA.from_chain_type(
+            self.llm,
+            chain_type="stuff",
+            retriever=self.retriever,
+            chain_type_kwargs={"prompt": prompt, "memory": memory}
+        )
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         # pseudo
         # self.model(input)
         inputs = data.pop("inputs", data)
-        result = self.final_chain.invoke(inputs)
-        print(result['answer'])
-
-        # Note that the memory does not save automatically
-        # This will be improved in the future
-        # For now you need to save it yourself
-        self.memory.save_context(inputs, {"answer": result["answer"].content})
-        self.memory.load_memory_variables({})
+        output = self.qa_chain(inputs)
+        print(output)
         
-        return result
+        return output