【山东大学项目实训】4

这两周主要进行了知识库的搭建和修改

我主要负责知识数据库的创建和管理。

导入模块：

import os
import sys
import time
import tempfile
from dotenv import load_dotenv, find_dotenv
from embedding.call_embedding import get_embedding
from langchain.document_loaders import UnstructuredFileLoader
from langchain.document_loaders import UnstructuredMarkdownLoader
from langchain.document_loaders import UnstructuredWordDocumentLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import PyMuPDFLoader
from langchain.vectorstores import Chroma

• os 和 sys：用于与操作系统进行交互和系统级操作。
• time：用于时间相关的操作。
• tempfile：用于创建临时文件和目录。
• dotenv：用于从环境文件中加载环境变量。
• get_embedding：自定义函数，用于获取嵌入。
• UnstructuredFileLoader、UnstructuredMarkdownLoader、UnstructuredWordDocumentLoader、PyMuPDFLoader：用于加载不同类型的文档。
• RecursiveCharacterTextSplitter：用于将文档拆分成小段文本。
• Chroma：用于创建和管理文档的嵌入向量数据库。

选择 Chroma 是因为它轻量级且数据存储在内存中，这使得它非常容易启动和开始使用。

设置默认变量：定义数据库文件和向量数据存储的默认路径 

获取文件列表：

def get_files(dir_path):
    file_list = []
    for filepath, dirnames, filenames in os.walk(dir_path):
        for filename in filenames:
            file_list.append(os.path.join(filepath, filename))
    return file_list

递归获取目录下所有文件的路径

文件加载：

def file_loader(file, loaders):
    if isinstance(file, tempfile._TemporaryFileWrapper):
        file = file.name
    if not os.path.isfile(file):
        [file_loader(os.path.join(file, f), loaders) for f in  os.listdir(file)]
        return
    file_type = file.split('.')[-1]
    if file_type == 'pdf':
        loaders.append(PyMuPDFLoader(file))
    elif file_type == 'md':
        loaders.append(UnstructuredMarkdownLoader(file))
    elif file_type == 'txt':
        loaders.append(UnstructuredFileLoader(file))
    elif file_type == 'docx':
        loaders.append(UnstructuredWordDocumentLoader(file))
    return

根据文件类型选择合适的加载器，加载文件并将其添加到加载器列表中。

创建数据库信息：

def create_db_info(files=DEFAULT_DB_PATH, embeddings="zhipuai", persist_directory=DEFAULT_PERSIST_PATH):
    vectordb = create_db(files, persist_directory, embeddings)
    return ""

创建数据库：

def create_db(files=DEFAULT_DB_PATH, persist_directory=DEFAULT_PERSIST_PATH, embeddings="zhipuai"):
    
    if files == None:
        files=DEFAULT_DB_PATH
    if type(files) != list:
        files = [files]
    loaders = []
    [file_loader(file, loaders) for file in files]
    docs = []
    for loader in loaders:
        if loader is not None:
            docs.extend(loader.load())
    # 切分文档
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500, chunk_overlap=150)
    split_docs = text_splitter.split_documents(docs[:10])

    # 定义持久化路径
    persist_directory = DEFAULT_PERSIST_PATH
    if type(embeddings) == str:
        embeddings = get_embedding(embedding=embeddings)

    # 加载数据库
    vectordb = Chroma.from_documents(
        documents=split_docs,
        embedding=embeddings,
        persist_directory=persist_directory  # 允许我们将persist_directory目录保存到磁盘上
    )
    vectordb.persist()
    return vectordb

加载文件，切分文档，生成嵌入向量，并创建向量数据库。

持久化数据库：

def presit_knowledge_db(vectordb):
    ...

加载数据库：

def load_knowledge_db(path, embeddings):
    ...
 

通过搭建并使用向量数据库，将文本数据转化为嵌入向量，并将这些向量存储在一个数据库中，以便后续的文本分析、检索或其他自然语言处理任务中使用。