embeding/faissdb.py

from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.document_loaders import TextLoader,UnstructuredCSVLoader, UnstructuredPDFLoader,UnstructuredWordDocumentLoader,UnstructuredExcelLoader,UnstructuredMarkdownLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
import shutil
import os
from .asr_utils import get_spk_txt


class FaissDB():
    def __init__(self, embedding="mofanke/acge_text_embedding:latest", persist_directory="./Faiss_db/"):

        self.embedding = OllamaEmbeddings(model=embedding)
        self.persist_directory = persist_directory
        self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50, add_start_index=True)

    def parse_data(self,file):
        if "txt" in file.lower() or "csv" in file.lower():
            try:
                loaders = UnstructuredCSVLoader(file)
                data = loaders.load()
            except:
                loaders = TextLoader(file,encoding="utf-8")
                data = loaders.load()
        if ".doc" in file.lower() or ".docx" in file.lower():
            loaders = UnstructuredWordDocumentLoader(file)
            data = loaders.load()
        if "pdf" in file.lower():
            loaders = UnstructuredPDFLoader(file)
            data = loaders.load()
        if ".xlsx" in file.lower():
            loaders = UnstructuredExcelLoader(file)
            data = loaders.load()
        if ".md" in file.lower():
            loaders = UnstructuredMarkdownLoader(file)
            data = loaders.load()
        if "mp3" in file.lower() or "mp4" in file.lower() or "wav" in file.lower():
            # 语音解析成文字
            fw = get_spk_txt(file)
            loaders = UnstructuredCSVLoader(fw)
            data = loaders.load()
            tmp = []
            for i in data:
                i.metadata["source"] = file
                tmp.append(i)
            data = tmp
        return data

    # 创建 新的collection 并且初始化
    def create_collection(self, files, c_name,chunk_size=200, chunk_overlap=50):
        self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        print("开始创建数据库 ....")
        tmps = []
        for file in files:
            data = self.parse_data(file)
            tmps.extend(data)

        splits = self.text_splitter.split_documents(tmps)

        vectorstore = FAISS.from_documents(documents=splits,
                                           embedding=self.embedding)
        vectorstore.save_local(self.persist_directory + c_name)
        print("数据块总量:", vectorstore.index.ntotal)

        return vectorstore

    # 添加 数据到已有数据库
    def add_chroma(self, files, c_name,chunk_size=200, chunk_overlap=50):
        self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        print("开始添加文件...")
        tmps = []
        for file in files:
            data = self.parse_data(file)
            tmps.extend(data)

        splits = self.text_splitter.split_documents(tmps)

        vectorstore = FAISS.load_local(folder_path=self.persist_directory + c_name, embeddings=self.embedding,
                                       allow_dangerous_deserialization=True)
        vectorstore.add_documents(documents=splits)
        vectorstore.save_local("Faiss_db/" + c_name)
        print("数据块总量:", vectorstore.index.ntotal)

        return vectorstore

    # 删除 某个collection中的 某个文件
    def del_files(self, del_files_name, c_name):

        vectorstore = FAISS.load_local(folder_path=self.persist_directory + c_name, embeddings=self.embedding,
                                       allow_dangerous_deserialization=True)
        del_ids = []
        vec_dict = vectorstore.docstore._dict
        for id, md in vec_dict.items():
            for dl in del_files_name:
                if dl in md.metadata["source"]:
                    del_ids.append(id)
        vectorstore.delete(ids=del_ids)
        vectorstore.save_local(self.persist_directory + c_name)
        print("数据块总量:", vectorstore.index.ntotal)

        return vectorstore

    # 删除某个 知识库 collection
    def delete_collection(self, c_name):
        shutil.rmtree(self.persist_directory + c_name)

    # 获取目前所有 collection
    def get_all_collections_name(self):
        cl_names = [i for i in os.listdir(self.persist_directory) if os.path.isdir(self.persist_directory+i)]

        return cl_names

    # 获取 collection中的所有文件
    def get_collcetion_content_files(self, c_name):
        vectorstore = FAISS.load_local(folder_path=self.persist_directory + c_name, embeddings=self.embedding,
                                       allow_dangerous_deserialization=True)
        c_files = []
        vec_dict = vectorstore.docstore._dict
        for _, md in vec_dict.items():
            c_files.append(md.metadata["source"])

        return list(set(c_files))


# if __name__ == "__main__":
#     chromadb = FaissDB()
#     c_name = "sss3"
#
#     print(chromadb.get_all_collections_name())
#     chromadb.create_collection(["data/jl.txt", "data/jl.pdf"], c_name=c_name)
#     print(chromadb.get_all_collections_name())
#     chromadb.add_chroma(["data/tmp.txt"], c_name=c_name)
#     print(c_name, "包含的文件:", chromadb.get_collcetion_content_files(c_name))
#     chromadb.del_files(["data/tmp.txt"], c_name=c_name)
#     print(c_name, "包含的文件:", chromadb.get_collcetion_content_files(c_name))
#     print(chromadb.get_all_collections_name())
#     chromadb.delete_collection(c_name=c_name)
#     print(chromadb.get_all_collections_name())