ColossalAI/applications/ColossalQA/examples/retrieval_conversation_chatgpt.py

"""
Multilingual retrieval based conversation system backed by ChatGPT
"""

import argparse
import os

from colossalqa.data_loader.document_loader import DocumentLoader
from colossalqa.memory import ConversationBufferWithSummary
from colossalqa.retriever import CustomRetriever
from langchain import LLMChain
from langchain.chains import RetrievalQA
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.llms import OpenAI
from langchain.prompts.prompt import PromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Multilingual retrieval based conversation system backed by ChatGPT")
    parser.add_argument("--open_ai_key_path", type=str, default=None, help="path to the model")
    parser.add_argument(
        "--sql_file_path", type=str, default=None, help="path to the a empty folder for storing sql files for indexing"
    )

    args = parser.parse_args()

    if not os.path.exists(args.sql_file_path):
        os.makedirs(args.sql_file_path)

    # Setup openai key
    # Set env var OPENAI_API_KEY or load from a file
    openai_key = open(args.open_ai_key_path).read()
    os.environ["OPENAI_API_KEY"] = openai_key

    llm = OpenAI(temperature=0.6)

    information_retriever = CustomRetriever(k=3, sql_file_path=args.sql_file_path, verbose=True)
    # VectorDB
    embedding = HuggingFaceEmbeddings(
        model_name="moka-ai/m3e-base", model_kwargs={"device": "cpu"}, encode_kwargs={"normalize_embeddings": False}
    )

    # Define memory with summarization ability
    memory = ConversationBufferWithSummary(llm=llm)

    # Load data to vector store
    print("Select files for constructing retriever")
    documents = []
    while True:
        file = input("Enter a file path or press Enter directory without input to exit:").strip()
        if file == "":
            break
        data_name = input("Enter a short description of the data:")
        retriever_data = DocumentLoader([[file, data_name.replace(" ", "_")]]).all_data

        # Split
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=0)
        splits = text_splitter.split_documents(retriever_data)
        documents.extend(splits)
    # Create retriever
    information_retriever.add_documents(docs=documents, cleanup="incremental", mode="by_source", embedding=embedding)

    prompt_template = """Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
    If the answer cannot be inferred based on the given context, please don't share false information.
    Use the context and chat history to respond to the human's input at the end or carry on the conversation. You should generate one response only. No following up is needed.

    context:
    {context}

    chat history
    {chat_history}

    Human: {question}
    Assistant:"""

    prompt_template_disambiguate = """You are a helpful, respectful and honest assistant. You always follow the instruction.
    Please replace any ambiguous references in the given sentence with the specific names or entities mentioned in the chat history or just output the original sentence if no chat history is provided or if the sentence doesn't contain ambiguous references. Your output should be the disambiguated sentence itself (in the same line as "disambiguated sentence:") and contain nothing else.

    Here is an example:
    Chat history:
    Human: I have a friend, Mike. Do you know him?
    Assistant: Yes, I know a person named Mike

    sentence: What's his favorite food?
    disambiguated sentence: What's Mike's favorite food?
    END OF EXAMPLE

    Chat history:
    {chat_history}

    sentence: {input}
    disambiguated sentence:"""

    PROMPT = PromptTemplate(template=prompt_template, input_variables=["question", "chat_history", "context"])

    memory.initiate_document_retrieval_chain(
        llm,
        PROMPT,
        information_retriever,
        chain_type_kwargs={
            "chat_history": "",
        },
    )

    PROMPT_DISAMBIGUATE = PromptTemplate(
        template=prompt_template_disambiguate, input_variables=["chat_history", "input"]
    )

    llm_chain = RetrievalQA.from_chain_type(
        llm=llm,
        verbose=False,
        chain_type="stuff",
        retriever=information_retriever,
        chain_type_kwargs={"prompt": PROMPT, "memory": memory},
    )
    llm_chain_disambiguate = LLMChain(llm=llm, prompt=PROMPT_DISAMBIGUATE)

    def disambiguity(input):
        out = llm_chain_disambiguate.run({"input": input, "chat_history": memory.buffer})
        return out.split("\n")[0]

    information_retriever.set_rephrase_handler(disambiguity)

    while True:
        user_input = input("User: ")
        if " end " in user_input:
            print("Agent: Happy to chat with you ：)")
            break
        agent_response = llm_chain.run(user_input)
        agent_response = agent_response.split("\n")[0]
        print(f"Agent: {agent_response}")