src/index_func.py

import os
import re
from typing import List, Optional, Any

from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from loguru import logger
from tqdm import tqdm

from src.config import local_embedding, retrieve_proxy, chunk_overlap, chunk_size, hf_emb_model_name
from src import shared
from src.utils import excel_to_string, get_files_hash, load_pkl, save_pkl

pwd_path = os.path.abspath(os.path.dirname(__file__))


class ChineseRecursiveTextSplitter(RecursiveCharacterTextSplitter):
    """Recursive text splitter for Chinese text.
    copy from: https://github.com/chatchat-space/Langchain-Chatchat/tree/master
    """

    def __init__(
            self,
            separators: Optional[List[str]] = None,
            keep_separator: bool = True,
            is_separator_regex: bool = True,
            **kwargs: Any,
    ) -> None:
        """Create a new TextSplitter."""
        super().__init__(keep_separator=keep_separator, **kwargs)
        self._separators = separators or [
            "\n\n",
            "\n",
            "。|！|？",
            "\.\s|\!\s|\?\s",
            "；|;\s",
            "，|,\s"
        ]
        self._is_separator_regex = is_separator_regex

    @staticmethod
    def _split_text_with_regex_from_end(
            text: str, separator: str, keep_separator: bool
    ) -> List[str]:
        # Now that we have the separator, split the text
        if separator:
            if keep_separator:
                # The parentheses in the pattern keep the delimiters in the result.
                _splits = re.split(f"({separator})", text)
                splits = ["".join(i) for i in zip(_splits[0::2], _splits[1::2])]
                if len(_splits) % 2 == 1:
                    splits += _splits[-1:]
            else:
                splits = re.split(separator, text)
        else:
            splits = list(text)
        return [s for s in splits if s != ""]

    def _split_text(self, text: str, separators: List[str]) -> List[str]:
        """Split incoming text and return chunks."""
        final_chunks = []
        # Get appropriate separator to use
        separator = separators[-1]
        new_separators = []
        for i, _s in enumerate(separators):
            _separator = _s if self._is_separator_regex else re.escape(_s)
            if _s == "":
                separator = _s
                break
            if re.search(_separator, text):
                separator = _s
                new_separators = separators[i + 1:]
                break

        _separator = separator if self._is_separator_regex else re.escape(separator)
        splits = self._split_text_with_regex_from_end(text, _separator, self._keep_separator)

        # Now go merging things, recursively splitting longer texts.
        _good_splits = []
        _separator = "" if self._keep_separator else separator
        for s in splits:
            if self._length_function(s) < self._chunk_size:
                _good_splits.append(s)
            else:
                if _good_splits:
                    merged_text = self._merge_splits(_good_splits, _separator)
                    final_chunks.extend(merged_text)
                    _good_splits = []
                if not new_separators:
                    final_chunks.append(s)
                else:
                    other_info = self._split_text(s, new_separators)
                    final_chunks.extend(other_info)
        if _good_splits:
            merged_text = self._merge_splits(_good_splits, _separator)
            final_chunks.extend(merged_text)
        return [re.sub(r"\n{2,}", "\n", chunk.strip()) for chunk in final_chunks if chunk.strip() != ""]


def get_documents(file_paths):
    text_splitter = ChineseRecursiveTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)

    documents = []
    logger.debug("Loading documents...")
    logger.debug(f"file_paths: {file_paths}")
    for file in file_paths:
        filepath = file.name
        filename = os.path.basename(filepath)
        file_type = os.path.splitext(filename)[1]
        logger.info(f"loading file: {filename}")
        texts = None
        try:
            if file_type == ".pdf":
                import PyPDF2
                logger.debug("Loading PDF...")
                try:
                    from src.pdf_func import parse_pdf
                    from src.config import advance_docs

                    two_column = advance_docs["pdf"].get("two_column", False)
                    pdftext = parse_pdf(filepath, two_column).text
                except:
                    pdftext = ""
                    with open(filepath, "rb") as pdfFileObj:
                        pdfReader = PyPDF2.PdfReader(pdfFileObj)
                        for page in tqdm(pdfReader.pages):
                            pdftext += page.extract_text()
                texts = [Document(page_content=pdftext,
                                  metadata={"source": filepath})]
            elif file_type == ".docx":
                logger.debug("Loading Word...")
                from langchain.document_loaders import UnstructuredWordDocumentLoader
                loader = UnstructuredWordDocumentLoader(filepath)
                texts = loader.load()
            elif file_type == ".pptx":
                logger.debug("Loading PowerPoint...")
                from langchain.document_loaders import UnstructuredPowerPointLoader
                loader = UnstructuredPowerPointLoader(filepath)
                texts = loader.load()
            elif file_type == ".epub":
                logger.debug("Loading EPUB...")
                from langchain.document_loaders import UnstructuredEPubLoader
                loader = UnstructuredEPubLoader(filepath)
                texts = loader.load()
            elif file_type == ".xlsx":
                logger.debug("Loading Excel...")
                text_list = excel_to_string(filepath)
                texts = []
                for elem in text_list:
                    texts.append(Document(page_content=elem,
                                          metadata={"source": filepath}))
            else:
                logger.debug("Loading text file...")
                from langchain_community.document_loaders import TextLoader
                loader = TextLoader(filepath, "utf8")
                texts = loader.load()
            logger.debug(f"text size: {len(texts)}, text top3: {texts[:3]}")
        except Exception as e:
            logger.error(f"Error loading file: {filename}, {e}")

        if texts is not None:
            texts = text_splitter.split_documents(texts)
            documents.extend(texts)
    logger.debug(f"Documents loaded. documents size: {len(documents)}, top3: {documents[:3]}")
    return documents


def construct_index(api_key, files, load_from_cache_if_possible=True):
    from langchain_community.vectorstores import FAISS
    from langchain.embeddings.huggingface import HuggingFaceEmbeddings

    if api_key:
        os.environ["OPENAI_API_KEY"] = api_key
    else:
        os.environ["OPENAI_API_KEY"] = "sk-xxxxxxx"

    index_name = get_files_hash(files)
    index_dir = os.path.join(pwd_path, 'index')
    index_path = os.path.join(index_dir, index_name)
    doc_file = os.path.join(index_path, 'docs.pkl')

    if local_embedding:
        embeddings = HuggingFaceEmbeddings(model_name=hf_emb_model_name)
    else:
        from langchain_community.embeddings import OpenAIEmbeddings
        if os.environ.get("OPENAI_API_TYPE", "openai") == "openai":
            embeddings = OpenAIEmbeddings(
                openai_api_base=shared.state.openai_api_base,
                openai_api_key=os.environ.get("OPENAI_EMBEDDING_API_KEY", api_key)
            )
        else:
            embeddings = OpenAIEmbeddings(
                deployment=os.environ["AZURE_EMBEDDING_DEPLOYMENT_NAME"],
                openai_api_key=os.environ["AZURE_OPENAI_API_KEY"],
                model=os.environ["AZURE_EMBEDDING_MODEL_NAME"],
                openai_api_base=os.environ["AZURE_OPENAI_API_BASE_URL"],
                openai_api_type="azure"
            )

    # 确保索引路径存在
    os.makedirs(index_dir, exist_ok=True)

    if os.path.exists(index_path) and load_from_cache_if_possible:
        try:
            logger.info("找到了缓存的索引文件，加载中……")
            index = FAISS.load_local(index_path, embeddings)
            documents = load_pkl(doc_file)
            return index, documents
        except (FileNotFoundError, RuntimeError) as e:
            logger.error(f"加载缓存的索引文件失败，重新构建索引…… 错误: {e}")

    try:
        documents = get_documents(files)
        logger.info("构建索引中……")
        with retrieve_proxy():
            index = FAISS.from_documents(documents, embeddings)
        logger.debug("索引构建完成！")
        os.makedirs(index_path, exist_ok=True)
        index.save_local(index_path)
        logger.debug("索引已保存至本地!")
        save_pkl(documents, doc_file)
        logger.debug("索引文档已保存至本地!")
        return index, documents
    except Exception as e:
        logger.error(f"索引构建失败！错误: {e}")
        return None