Aimff/Diesel_Fault_Diagnosis_System

1from utils import *2 3import os4from glob import glob#遍历文件5from langchain.vectorstores.chroma import Chroma6from langchain.document_loaders import CSVLoader, PyMuPDFLoader, TextLoader7from langchain.text_splitter import RecursiveCharacterTextSplitter8 9def doc2vec():10    # 定义文本分割器11    text_splitter = RecursiveCharacterTextSplitter(12        chunk_size = 300,13        chunk_overlap = 5014    )15 16    # 读取并分割文件17    dir_path = os.path.join(os.path.dirname(__file__), './data/inputs/')18 19    documents = []20    for file_path in glob(dir_path + '*.*'):21        loader = None22        if '.csv' in file_path:23            loader = CSVLoader(file_path)24        if '.pdf' in file_path:25            loader = PyMuPDFLoader(file_path)26        if '.txt' in file_path:27            loader = TextLoader(file_path)28        if loader:29            documents += loader.load_and_split(text_splitter)30            31    # 向量化并存储32    if documents:33        vdb = Chroma.from_documents(34            documents = documents, 35            embedding = get_embeddings_model(),36            persist_directory = os.path.join(os.path.dirname(__file__), './data/db/')37        )38        vdb.persist()39 40 41if __name__ == '__main__':42    doc2vec()