Update handler.py

handler.py

@@ -16,20 +16,58 @@ 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=""):
 
-        self.llm = LLM(model_name_or_path=path, bit4=False)
+        # Config LangChain
+        # os.environ["LANGCHAIN_TRACING_V2"] = "true"
+        # os.environ["LANGCHAIN_API_KEY"] = 
+
+        # Create LLM
+
+        # load the tokenizer and the quantized mistral model
+        model = AutoModelForCausalLM.from_pretrained(
+                     path, 
+                     device_map="auto")
         
-        # Load Vector db
+        tokenizer = AutoTokenizer.from_pretrained(path)
         
-        self.embedding_function = HuggingFaceBgeEmbeddings(
-            model_name="BAAI/bge-large-zh",
+        # using HuggingFace's pipeline
+        pipeline = pipeline(
+               "text-generation",
+               model=model,
+               tokenizer=tokenizer,
+               use_cache=True,
+               device_map="auto",
+               max_new_tokens=5000,
+               do_sample=True,
+               top_k=1,
+               temperature = 0.01,
+               num_return_sequences=1,
+               eos_token_id=tokenizer.eos_token_id,
+               pad_token_id=tokenizer.eos_token_id,
+        )
+        chat = HuggingFacePipeline(pipeline=pipeline)
+
+        # Create Text-Embedding Model
+        embedding_function = HuggingFaceBgeEmbeddings(
+            model_name="DMetaSoul/Dmeta-embedding",
             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",
@@ -43,40 +81,84 @@ 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=self.embedding_function)
+        vectorstore = Chroma.from_documents(documents=all_splits, embedding=embedding_function)
+        retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
         
-        # vectorstore = Chroma(persist_directory="db", embedding_function=embedding_function)
+        compressor = LLMChainExtractor.from_llm(chat)
+        retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=retriever)
 
-        compressor = LLMChainExtractor.from_llm(self.llm)
-        self.retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=vectorstore.as_retriever(search_kwargs={"k": 4}))
-
-        prompt_template = """<s>[INST] <<SYS>> You are a helpful assistant.
-        Use the following context to Answer the question below briefly: <<SYS>>
+        _template = """Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question, in its original language.
+        Chat History:
+        {chat_history}
+        Follow Up Input: {question}
+        Standalone question:"""
         
-        {history}
+        CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)
         
+        template = """Answer the question based only on the following context:
         {context}
         
-        {question} [/INST] </s>
+        Question: {question}
         """
+        ANSWER_PROMPT = ChatPromptTemplate.from_template(template)
 
-        prompt = PromptTemplate(input_variables=["history", "context", "question"], template=prompt_template)
-        memory = ConversationBufferMemory(input_key='question', memory_key='history', return_messages=True)
-
-        self.qa_chain = RetrievalQA.from_chain_type(
-            self.llm,
-            chain_type="stuff",
-            retriever=self.retriever,
-            chain_type_kwargs={"prompt": prompt, "memory": memory}
+        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}")
+
+        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
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         # pseudo
         # self.model(input)
         inputs = data.pop("inputs", data)
-        output = self.qa_chain(inputs)
-        print(output)
-        
-        return output
-
+        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({})
+        
+        return result