ColossalAI/applications/ColossalQA/colossalqa/retrieval_conversation_universal.py

"""
Multilingual retrieval based conversation system
"""
from typing import List

from colossalqa.data_loader.document_loader import DocumentLoader
from colossalqa.mylogging import get_logger
from colossalqa.retrieval_conversation_en import EnglishRetrievalConversation
from colossalqa.retrieval_conversation_zh import ChineseRetrievalConversation
from colossalqa.retriever import CustomRetriever
from colossalqa.text_splitter import ChineseTextSplitter
from colossalqa.utils import detect_lang_naive
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter, TextSplitter

logger = get_logger()


class UniversalRetrievalConversation:
    """
    Wrapper class for bilingual retrieval conversation system
    """

    def __init__(
        self,
        embedding_model_path: str = "moka-ai/m3e-base",
        embedding_model_device: str = "cpu",
        zh_model_path: str = None,
        zh_model_name: str = None,
        en_model_path: str = None,
        en_model_name: str = None,
        sql_file_path: str = None,
        files_zh: List[List[str]] = None,
        files_en: List[List[str]] = None,
        text_splitter_chunk_size=100,
        text_splitter_chunk_overlap=10,
    ) -> None:
        """
        Wrapper for multilingual retrieval qa class (Chinese + English)
        Args:
            embedding_model_path: local or huggingface embedding model
            embedding_model_device:
            files_zh: [[file_path, name_of_file, separator],...] defines the files used as supporting documents for Chinese retrieval QA
            files_en: [[file_path, name_of_file, separator],...] defines the files used as supporting documents for English retrieval QA
        """
        self.embedding = HuggingFaceEmbeddings(
            model_name=embedding_model_path,
            model_kwargs={"device": embedding_model_device},
            encode_kwargs={"normalize_embeddings": False},
        )
        print("Select files for constructing Chinese retriever")
        docs_zh = self.load_supporting_docs(
            files=files_zh,
            text_splitter=ChineseTextSplitter(
                chunk_size=text_splitter_chunk_size, chunk_overlap=text_splitter_chunk_overlap
            ),
        )
        # Create retriever
        self.information_retriever_zh = CustomRetriever(
            k=3, sql_file_path=sql_file_path.replace(".db", "_zh.db"), verbose=True
        )
        self.information_retriever_zh.add_documents(
            docs=docs_zh, cleanup="incremental", mode="by_source", embedding=self.embedding
        )

        print("Select files for constructing English retriever")
        docs_en = self.load_supporting_docs(
            files=files_en,
            text_splitter=RecursiveCharacterTextSplitter(
                chunk_size=text_splitter_chunk_size, chunk_overlap=text_splitter_chunk_overlap
            ),
        )
        # Create retriever
        self.information_retriever_en = CustomRetriever(
            k=3, sql_file_path=sql_file_path.replace(".db", "_en.db"), verbose=True
        )
        self.information_retriever_en.add_documents(
            docs=docs_en, cleanup="incremental", mode="by_source", embedding=self.embedding
        )

        self.chinese_retrieval_conversation = ChineseRetrievalConversation.from_retriever(
            self.information_retriever_zh, model_path=zh_model_path, model_name=zh_model_name
        )
        self.english_retrieval_conversation = EnglishRetrievalConversation.from_retriever(
            self.information_retriever_en, model_path=en_model_path, model_name=en_model_name
        )
        self.memory = None

    def load_supporting_docs(self, files: List[List[str]] = None, text_splitter: TextSplitter = None):
        """
        Load supporting documents, currently, all documents will be stored in one vector store
        """
        documents = []
        if files:
            for file in files:
                retriever_data = DocumentLoader([[file["data_path"], file["name"]]]).all_data
                splits = text_splitter.split_documents(retriever_data)
                documents.extend(splits)
        else:
            while True:
                file = input("Select a file to load or press Enter to exit:")
                if file == "":
                    break
                data_name = input("Enter a short description of the data:")
                separator = input(
                    "Enter a separator to force separating text into chunks, if no separator is given, the default separator is '\\n\\n', press ENTER directly to skip:"
                )
                separator = separator if separator != "" else "\n\n"
                retriever_data = DocumentLoader([[file, data_name.replace(" ", "_")]]).all_data

                # Split
                splits = text_splitter.split_documents(retriever_data)
                documents.extend(splits)
        return documents

    def start_test_session(self):
        """
        Simple multilingual session for testing purpose, with naive language selection mechanism
        """
        while True:
            user_input = input("User: ")
            lang = detect_lang_naive(user_input)
            if "END" == user_input:
                print("Agent: Happy to chat with you ：)")
                break
            agent_response = self.run(user_input, which_language=lang)
            print(f"Agent: {agent_response}")

    def run(self, user_input: str, which_language=str):
        """
        Generate the response given the user input and a str indicates the language requirement of the output string
        """
        assert which_language in ["zh", "en"]
        if which_language == "zh":
            agent_response, self.memory = self.chinese_retrieval_conversation.run(user_input, self.memory)
        else:
            agent_response, self.memory = self.english_retrieval_conversation.run(user_input, self.memory)
        return agent_response.split("\n")[0]