基於GPT搭建私有知識庫聊天機器人（三）向量資料訓練

在前面的文章中，我們介紹了實現原理和基本環境安裝。本文將重點介紹資料訓練的流程，以及如何載入、切割、訓練資料，並使用向量資料庫Milvus進行資料儲存。

1. 資料訓練依賴於向量資料庫

在本文中，我們使用了Milvus作為向量資料庫。讀者可以參考之前的文章《基於GPT搭建私有知識庫聊天機器人（二）環境安裝》來準備其他基礎環境。

2. 資料訓練流程

資料訓練的流程包括準備PDF檔案、上傳至系統檔案目錄、開始訓練、載入檔案內容、內容切割和儲存至向量資料庫。下面是整個流程的流程圖：

3. 程式碼展示

3.1 上傳檔案至系統檔案目錄

@app.route('/upload', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        # 獲取文字內容
        text = request.form.get('name')
        # 獲取檔案內容
        file = request.files.get('file')
        if file:
            # 儲存檔案到伺服器
            filename = file.filename
            file.save(os.path.join(KNOWLEDGE_FOLDER, text, filename))
            file_path = os.path.join(KNOWLEDGE_FOLDER, text, filename)
        else:
            file_path = None

        return jsonify({'message': '上傳成功', 'fileServicePath': file_path})

    return render_template('index.html')

3.2 載入檔案內容

# 對映檔案載入
LOADER_MAPPING = {
    ".csv": (CSVLoader, {}),
    ".docx": (Docx2txtLoader, {}),
    ".doc": (UnstructuredWordDocumentLoader, {}),
    ".docx": (UnstructuredWordDocumentLoader, {}),
    ".enex": (EverNoteLoader, {}),
    ".eml": (MyElmLoader, {}),
    ".epub": (UnstructuredEPubLoader, {}),
    ".html": (UnstructuredHTMLLoader, {}),
    ".md": (UnstructuredMarkdownLoader, {}),
    ".odt": (UnstructuredODTLoader, {}),
    ".pdf": (PDFMinerLoader, {}),
    ".ppt": (UnstructuredPowerPointLoader, {}),
    ".pptx": (UnstructuredPowerPointLoader, {}),
    ".txt": (TextLoader, {"encoding": "utf8"}),
}

def load_single_document(file_path: str) -> List[Document]:
    ext = "." + file_path.rsplit(".", 1)[-1]
    if ext in LOADER_MAPPING:
        loader_class, loader_args = LOADER_MAPPING[ext]
        loader = loader_class(file_path, **loader_args)
        return loader.load()

    raise ValueError(f"檔案不存在 '{ext}'")

# 載入檔案
def load_documents_knowledge(source_dir: str, secondary_directories: str) -> List[Document]:
    """
    Loads all documents from the source documents directory, ignoring specified files
    """
    all_files = []
    for ext in LOADER_MAPPING:
        all_files.extend(

            glob.glob(os.path.join(source_dir, secondary_directories, f"**/*{ext}"), recursive=True)
        )
    filtered_files = [file_path for file_path in all_files if file_path]

    with Pool(processes=os.cpu_count()) as pool:
        results = []
        with tqdm(total=len(filtered_files), desc='Loading new documents', ncols=80) as pbar:
            for i, docs in enumerate(pool.imap_unordered(load_single_document, filtered_files)):
                results.extend(docs)
                pbar.update()

    return results

3.3 內容切割

text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
texts = text_splitter.split_documents(documents)

3.4 儲存至向量資料庫

Milvus.from_documents(
        texts,
        collection_name=collection_name,
        embedding=embeddings,
        connection_args={"host": MILVUS_HOST, "port": MILVUS_PORT}
    )

3.5 全部程式碼

#!/usr/bin/env python3
import glob
import os
import shutil
from multiprocessing import Pool
from typing import List

from dotenv import load_dotenv
from langchain.docstore.document import Document
from langchain.document_loaders import (
    CSVLoader,
    EverNoteLoader,
    PDFMinerLoader,
    TextLoader,
    UnstructuredEmailLoader,
    UnstructuredEPubLoader,
    UnstructuredHTMLLoader,
    UnstructuredMarkdownLoader,
    UnstructuredODTLoader,
    UnstructuredPowerPointLoader,
    UnstructuredWordDocumentLoader, )
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Milvus
from tqdm import tqdm

load_dotenv(".env")

MILVUS_HOST = os.environ.get('MILVUS_HOST')
MILVUS_PORT = os.environ.get('MILVUS_PORT')
source_directory = os.environ.get('SOURCE_DIRECTORY', 'source_documents')
KNOWLEDGE_FOLDER = os.environ.get('KNOWLEDGE_FOLDER')
KNOWLEDGE_FOLDER_BK = os.environ.get('KNOWLEDGE_FOLDER_BK')
chunk_size = 500
chunk_overlap = 50


# Custom document loaders
class MyElmLoader(UnstructuredEmailLoader):
    """在預設值不起作用時回退到文字純"""

    def load(self) -> List[Document]:
        """EMl沒有 html 使用text/plain"""
        try:
            try:
                doc = UnstructuredEmailLoader.load(self)
            except ValueError as e:
                if 'text/html content not found in email' in str(e):
                    # Try plain text
                    self.unstructured_kwargs["content_source"] = "text/plain"
                    doc = UnstructuredEmailLoader.load(self)
                else:
                    raise
        except Exception as e:
            # Add file_path to exception message
            raise type(e)(f"{self.file_path}: {e}") from e

        return doc


# 對映檔案載入
LOADER_MAPPING = {
    ".csv": (CSVLoader, {}),
    # ".docx": (Docx2txtLoader, {}),
    ".doc": (UnstructuredWordDocumentLoader, {}),
    ".docx": (UnstructuredWordDocumentLoader, {}),
    ".enex": (EverNoteLoader, {}),
    ".eml": (MyElmLoader, {}),
    ".epub": (UnstructuredEPubLoader, {}),
    ".html": (UnstructuredHTMLLoader, {}),
    ".md": (UnstructuredMarkdownLoader, {}),
    ".odt": (UnstructuredODTLoader, {}),
    ".pdf": (PDFMinerLoader, {}),
    ".ppt": (UnstructuredPowerPointLoader, {}),
    ".pptx": (UnstructuredPowerPointLoader, {}),
    ".txt": (TextLoader, {"encoding": "utf8"}),
}


def load_single_document(file_path: str) -> List[Document]:
    ext = "." + file_path.rsplit(".", 1)[-1]
    if ext in LOADER_MAPPING:
        loader_class, loader_args = LOADER_MAPPING[ext]
        loader = loader_class(file_path, **loader_args)
        return loader.load()

    raise ValueError(f"檔案不存在 '{ext}'")


def load_documents_knowledge(source_dir: str, secondary_directories: str) -> List[Document]:
    """
    Loads all documents from the source documents directory, ignoring specified files
    """
    all_files = []
    for ext in LOADER_MAPPING:
        all_files.extend(

            glob.glob(os.path.join(source_dir, secondary_directories, f"**/*{ext}"), recursive=True)
        )
    filtered_files = [file_path for file_path in all_files if file_path]

    with Pool(processes=os.cpu_count()) as pool:
        results = []
        with tqdm(total=len(filtered_files), desc='Loading new documents', ncols=80) as pbar:
            for i, docs in enumerate(pool.imap_unordered(load_single_document, filtered_files)):
                results.extend(docs)
                pbar.update()

    return results


def process_documents_knowledge(secondary_directories: str) -> List[Document]:
    """
    載入檔案並拆分為塊
    """
    print(f"載入檔案目錄: {KNOWLEDGE_FOLDER}")
    documents = load_documents_knowledge(KNOWLEDGE_FOLDER, secondary_directories)
    if not documents:
        print("沒有檔案需要載入")
        exit(0)
    print(f"載入 {len(documents)} 檔案從 {KNOWLEDGE_FOLDER}")
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    texts = text_splitter.split_documents(documents)
    print(f"切割 {len(texts)} 文字塊 (最大. {chunk_size} tokens 令牌)")
    return texts


def main_knowledge(collection_name: str):
    # Create embeddings
    embeddings = OpenAIEmbeddings()

    texts = process_documents_knowledge(collection_name)

    Milvus.from_documents(
        texts,
        collection_name=collection_name,
        embedding=embeddings,
        connection_args={"host": MILVUS_HOST, "port": MILVUS_PORT}
    )

4. 總結

在本文中，我們詳細介紹了基於GPT搭建私有知識庫聊天機器人的資料訓練過程，包括資料訓練的依賴、流程和程式碼展示。資料訓練是搭建聊天機器人的重要步驟，希望本文能對讀者有所幫助。在下一篇文章中，我們將介紹如何使用訓練好的模型進行聊天機器人的測試和使用。