【langchain 创建向量数据库非常美满的代码】

 - 支持faiss chroma两种数据库(faiss-cpu 支持旧数据库的合并)
- 支持制止重复文件embedding（hash）
- 支持浩繁文件格式
- 支持huggingface的embedding模型
- 优化了切分chunk的策略
- 支持多线程处理惩罚文件

   修改自langchain-chatchat， 增加了一些功能， 优化了splitter的策略

1. import os
2. import glob
3. import hashlib
4. from typing import List
5. from functools import partial
6. from tqdm import tqdm
8. from multiprocessing import Pool
9. from langchain.document_loaders import (
10.     CSVLoader,
11.     EverNoteLoader,
12.     PDFMinerLoader,
13.     TextLoader,
14.     UnstructuredEmailLoader,
15.     UnstructuredEPubLoader,
16.     UnstructuredHTMLLoader,
17.     UnstructuredMarkdownLoader,
18.     UnstructuredODTLoader,
19.     UnstructuredPowerPointLoader,
20.     UnstructuredWordDocumentLoader,
21.     UnstructuredExcelLoader,
22. )
23. from langchain.text_splitter import RecursiveCharacterTextSplitter, TextSplitter
24. from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
25. from langchain.docstore.document import Document
30. CONFIG = {
31.     "doc_source": "./docs",  # 需要向量化的文档
32.     "embedding_model": "hugging_models/text2vec-base-chinese_config",  # embeding模型
33.     "db_source": "./db",  # 向量化数据库
34.     "db_type": "faiss",  #
35.     "chunk_size": 200,  # 块词量
36.     "chunk_overlap": 20,  # 交集范围
37.     "k": 3,  # 查询文档量
38.     "merge_rows": 5,  # 合并表格的行数,
39.     "hash_file_path": "hash_file.txt",  #
40. }
43. # 设置目录和embedding基础变量
44. source_directory = CONFIG["doc_source"]
45. embeddings_model_name = CONFIG["embedding_model"]
46. chunk_size = CONFIG["chunk_size"]
47. chunk_overlap = CONFIG["chunk_overlap"]
48. output_dir = CONFIG["db_source"]
49. k = CONFIG["k"]
50. merge_rows = CONFIG["merge_rows"]
51. hash_file_path = CONFIG["hash_file_path"]
52. db_type = CONFIG["db_type"]
55. # Custom document loaders 自定义文档加载
56. class MyElmLoader(UnstructuredEmailLoader):
57.     def load(self) -> List[Document]:
58.         """Wrapper adding fallback for elm without html"""
59.         try:
60.             try:
61.                 doc = UnstructuredEmailLoader.load(self)
62.             except ValueError as e:
63.                 if "text/html content not found in email" in str(e):
64.                     # Try plain text
65.                     self.unstructured_kwargs["content_source"] = "text/plain"
66.                     doc = UnstructuredEmailLoader.load(self)
67.                 else:
68.                     raise
69.         except Exception as e:
70.             # Add file_path to exception message
71.             raise type(e)(f"{self.file_path}: {e}") from e
73.         return doc
76. # Map file extensions to document loaders and their arguments
77. # GBK2312 GB18030
78. LOADER_MAPPING = {
79.     ".csv": (CSVLoader, {}),
80.     ".doc": (UnstructuredWordDocumentLoader, {}),
81.     ".docx": (UnstructuredWordDocumentLoader, {}),
82.     ".enex": (EverNoteLoader, {}),
83.     ".eml": (MyElmLoader, {}),
84.     ".epub": (UnstructuredEPubLoader, {}),
85.     ".html": (UnstructuredHTMLLoader, {}),
86.     ".md": (UnstructuredMarkdownLoader, {}),
87.     ".odt": (UnstructuredODTLoader, {}),
88.     ".pdf": (PDFMinerLoader, {}),
89.     ".ppt": (UnstructuredPowerPointLoader, {}),
90.     ".pptx": (UnstructuredPowerPointLoader, {}),
91.     ".txt": (TextLoader, {"encoding": "utf8"}),
92.     ".xls": (UnstructuredExcelLoader, {}),
93.     ".xlsx": (UnstructuredExcelLoader, {}),
94. }
97. def read_hash_file(path="hash_file.txt"):
98.     hash_file_list = []
99.     if os.path.exists(path):
100.         with open(path, "r") as f:
101.             hash_file_list = [i.strip() for i in f.readlines()]
102.     return hash_file_list
105. def save_hash_file(hash_list, path="hash_file.txt"):
106.     with open(path, "w") as f:
107.         f.write("\n".join(hash_list))
110. def get_hash_from_file(path):
111.     with open(path, "rb") as f:
112.         readable_hash = hashlib.md5(f.read()).hexdigest()
113.     return readable_hash
116. def load_single_document(
117.     file_path: str, splitter: TextSplitter, merge_rows: int
118. ) -> List[Document]:
119.     ext = "." + file_path.rsplit(".", 1)[-1]
120.     if ext in LOADER_MAPPING:
121.         loader_class, loader_args = LOADER_MAPPING[ext]
122.         loader = loader_class(file_path, **loader_args)
123.         docs = loader.load()
125.         # 针对不同的文件类型分别进行处理
126.         if not file_path.endswith((".xlsx", "xls", "csv")):
127.             # 合并一个文件中的所有page_content
128.             tmp = [i.page_content for i in docs]
129.             docs = Document(
130.                 "".join(tmp).strip(),
131.                 metadata={"source": docs[0].metadata["source"], "pages": len(tmp)},
132.             )
133.             # 进行split
134.             docs = splitter.split_documents([docs])
135.         else:
136.             # 表格数据，合并多个行
137.             merge_n = len(docs) // merge_rows + bool(len(docs) % merge_rows)
138.             _docs = []
139.             for i in range(merge_n):
140.                 tmp = "\n\n".join(
141.                     [
142.                         d.page_content
143.                         for d in docs[i * merge_rows : (i + 1) * merge_rows]
144.                     ]
145.                 )
146.                 _docs.append(Document(tmp, metadata=dict(source=docs[0]["source"])))
147.             docs = _docs
149.         return docs
151.     raise ValueError(f"Unsupported file extension '{ext}'")
154. def load_documents(source_dir: str, ignored_files: List[str] = []) -> List[Document]:
155.     """
156.     Loads all documents from the source documents directory, ignoring specified files
157.     """
158.     all_files = []
159.     for ext in LOADER_MAPPING:
160.         all_files.extend(
161.             glob.glob(os.path.join(source_dir, f"**/*{ext}"), recursive=True)
162.         )
163.     filtered_files = [
164.         file_path for file_path in all_files if file_path not in ignored_files
165.     ]
167.     # hash filter
168.     if os.path.exists(hash_file_path):
169.         hash_file_list = read_hash_file(hash_file_path)
170.         if hash_file_list:
171.             tmp = []
172.             for file in filtered_files:
173.                 hash = get_hash_from_file(file)
174.                 if hash not in hash_file_list:
175.                     tmp.append(file)
176.                     hash_file_list(hash)
177.             filtered_files = tmp
179.             save_hash_file(hash_file_list, hash_file_path)
181.     # splitter
182.     splitter = RecursiveCharacterTextSplitter(
183.         chunk_size=chunk_size, chunk_overlap=chunk_overlap
184.     )
185.     load_document = partial(
186.         load_single_document, splitter=splitter, merge_rows=merge_rows
187.     )
189.     # load
190.     with Pool(processes=os.cpu_count()) as pool:
191.         results = []
192.         with tqdm(
193.             total=len(filtered_files), desc="Loading new documents", ncols=80
194.         ) as pbar:
195.             for i, docs in enumerate(
196.                 pool.imap_unordered(load_document, filtered_files)
197.             ):
198.                 results.extend(docs)
199.                 pbar.update()
201.     return results
204. def process_documents(ignored_files: List[str] = []) -> List[Document]:
205.     """
206.     Load documents and split in chunks
207.     """
208.     print(f"Loading documents from {source_directory}")
209.     documents = load_documents(source_directory, ignored_files)
210.     if not documents:
211.         print("No new documents to load")
212.         exit(0)
213.     print(
214.         f"Loaded {len(documents)} new documents from {source_directory}."
215.         f"\nSplit into {len(documents)} chunks of text (max. {chunk_size} tokens each)"
216.     )
217.     return documents
220. def main():
221.     # Create embeddings
222.     # print(torch.cuda.is_available())
223.     # Create and store locally vectorstore
224.     print("Creating new vectorstore")
225.     documents = process_documents()
226.     print(f"Creating embeddings. May take some minutes...")
227.     embedding_function = SentenceTransformerEmbeddings(model_name=embeddings_model_name)
229.     if db_type == "chroma":
230.         from langchain.vectorstores import Chroma
232.         db = Chroma.from_documents(
233.             documents, embedding_function, persist_directory=output_dir
234.         )
235.         db.persist()
236.         db = None
237.     elif db_type == "faiss":
238.         from langchain.vectorstores import FAISS
240.         print("创建新数据库")
241.         db = FAISS.from_documents(documents, embedding_function)
243.         # 读取之前的db data(GPU版本的不支持)
244.         if os.path.exists(output_dir):
245.             try:
246.                 print("读取旧数据库")
247.                 old_db = FAISS.load_local(output_dir, embedding_function)
249.                 print("融合新旧数据库")
250.                 db.merge_from(old_db)
251.             except Exception as e:
252.                 print(e)
254.         print("保存")
255.         db.save_local(output_dir)
256.         db = None
257.     else:
258.         raise NotImplementedError(f'未定义数据库 {db_type} 的实现.')
260. if __name__ == "__main__":
261.     main()

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。