aaaaasss/0406

1import os2import logging3 4from llama_index import download_loader5from llama_index import (6    Document,7    LLMPredictor,8    PromptHelper,9    QuestionAnswerPrompt,10    RefinePrompt,11)12import colorama13import PyPDF214from tqdm import tqdm15 16from modules.presets import *17from modules.utils import *18 19def get_index_name(file_src):20    file_paths = [x.name for x in file_src]21    file_paths.sort(key=lambda x: os.path.basename(x))22 23    md5_hash = hashlib.md5()24    for file_path in file_paths:25        with open(file_path, "rb") as f:26            while chunk := f.read(8192):27                md5_hash.update(chunk)28 29    return md5_hash.hexdigest()30 31def block_split(text):32    blocks = []33    while len(text) > 0:34        blocks.append(Document(text[:1000]))35        text = text[1000:]36    return blocks37 38def get_documents(file_src):39    documents = []40    logging.debug("Loading documents...")41    logging.debug(f"file_src: {file_src}")42    for file in file_src:43        filepath = file.name44        filename = os.path.basename(filepath)45        file_type = os.path.splitext(filepath)[1]46        logging.info(f"loading file: {filename}")47        if file_type == ".pdf":48            logging.debug("Loading PDF...")49            try:50                from modules.pdf_func import parse_pdf51                from modules.config import advance_docs52                two_column = advance_docs["pdf"].get("two_column", False)53                pdftext = parse_pdf(filepath, two_column).text54            except:55                pdftext = ""56                with open(filepath, 'rb') as pdfFileObj:57                    pdfReader = PyPDF2.PdfReader(pdfFileObj)58                    for page in tqdm(pdfReader.pages):59                        pdftext += page.extract_text()60            text_raw = pdftext61        elif file_type == ".docx":62            logging.debug("Loading Word...")63            DocxReader = download_loader("DocxReader")64            loader = DocxReader()65            text_raw = loader.load_data(file=filepath)[0].text66        elif file_type == ".epub":67            logging.debug("Loading EPUB...")68            EpubReader = download_loader("EpubReader")69            loader = EpubReader()70            text_raw = loader.load_data(file=filepath)[0].text71        elif file_type == ".xlsx":72            logging.debug("Loading Excel...")73            text_raw = excel_to_string(filepath)74        else:75            logging.debug("Loading text file...")76            with open(filepath, "r", encoding="utf-8") as f:77                text_raw = f.read()78        text = add_space(text_raw)79        # text = block_split(text)80        # documents += text81        documents += [Document(text)]82    logging.debug("Documents loaded.")83    return documents84 85 86def construct_index(87        api_key,88        file_src,89        max_input_size=4096,90        num_outputs=5,91        max_chunk_overlap=20,92        chunk_size_limit=600,93        embedding_limit=None,94        separator=" "95):96    from langchain.chat_models import ChatOpenAI97    from llama_index import GPTSimpleVectorIndex, ServiceContext98 99    os.environ["OPENAI_API_KEY"] = api_key100    chunk_size_limit = None if chunk_size_limit == 0 else chunk_size_limit101    embedding_limit = None if embedding_limit == 0 else embedding_limit102    separator = " " if separator == "" else separator103 104    llm_predictor = LLMPredictor(105        llm=ChatOpenAI(model_name="gpt-3.5-turbo-0301", openai_api_key=api_key)106    )107    prompt_helper = PromptHelper(max_input_size = max_input_size, num_output = num_outputs, max_chunk_overlap = max_chunk_overlap, embedding_limit=embedding_limit, chunk_size_limit=600, separator=separator)108    index_name = get_index_name(file_src)109    if os.path.exists(f"./index/{index_name}.json"):110        logging.info("找到了缓存的索引文件，加载中……")111        return GPTSimpleVectorIndex.load_from_disk(f"./index/{index_name}.json")112    else:113        try:114            documents = get_documents(file_src)115            logging.info("构建索引中……")116            with retrieve_proxy():117                service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper, chunk_size_limit=chunk_size_limit)118                index = GPTSimpleVectorIndex.from_documents(119                    documents,  service_context=service_context120                )121            logging.debug("索引构建完成！")122            os.makedirs("./index", exist_ok=True)123            index.save_to_disk(f"./index/{index_name}.json")124            logging.debug("索引已保存至本地!")125            return index126 127        except Exception as e:128            logging.error("索引构建失败！", e)129            print(e)130            return None131 132 133def add_space(text):134    punctuations = {"，": "， ", "。": "。 ", "？": "？ ", "！": "！ ", "：": "： ", "；": "； "}135    for cn_punc, en_punc in punctuations.items():136        text = text.replace(cn_punc, en_punc)137    return text138