langchain Chroma 构建本地向量数据库

langchain Chroma 构建本地向量数据库

1. # import
2. from langchain_community.document_loaders import TextLoader
3. from langchain_community.embeddings.sentence_transformer import (
4.     SentenceTransformerEmbeddings,
5. )
6. from langchain_community.embeddings import HuggingFaceEmbeddings
7. from langchain_community.vectorstores import Chroma
8. from langchain_text_splitters import RecursiveCharacterTextSplitter
9. from langchain_community.document_loaders.word_document import Docx2txtLoader
11. import glob
12. import os
14. # 数据库路径
15. db_dir = "./db"
16. # 文档路径
17. source_directory = "./docs"
18. # 文件后缀
19. file_ext = '*.docx'
21. # create the open-source embedding function
22. # embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
23. # 使用中文嵌入层编码器
24. ebd_function = HuggingFaceEmbeddings(model_name="shibing624/text2vec-base-chinese")
26. def add_files_to_db(filepath:str="",file_ext:str=""):
27.     docx_files = glob.glob(os.path.join(source_directory, file_ext))
28.     text_list=[]
29.     for file_name in docx_files:
30.         print(file_name)
31.         loader = Docx2txtLoader(file_name)
32.         documents = loader.load()
33.         text_list.extend(documents)
35.     # split it into chunks
36.     text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
37.     docs = text_splitter.split_documents(text_list)
39.     # load it into Chroma
40.     db = Chroma.from_documents(docs, ebd_function, persist_directory=db_dir)
41.     # save db to disk
42.     db.persist()
45. def query_db(db:Chroma,query:str=""):
47.     # query it
48.     docs = db.similarity_search(query)
50.     # print results
51.     print(docs[0].page_content)
52.     print("-----------------------------------------")
55.    
57. if __name__=="__main__":
59.     # 只需执行一次
60.     # add_files_to_db(source_directory,file_ext)
62.     db = Chroma(persist_directory=db_dir,embedding_function=ebd_function)
63.     query = "怎么治疗骨质疏松症?"
64.     query_db(db,query)
65.     query = "怎么治疗鼻炎?"
66.     query_db(db,query)
67.     db = None
68.     pass

复制代码

文档在当前代码目录下
./docs/第十六章-感染性疾病.docx
./docs/第八章-骨骼关节和肌肉疾病.docx
./docs/第十九章-耳鼻咽喉疾病.docx

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