langchain Chroma 构建本地向量数据库

langchain Chroma 构建本地向量数据库

python 复制代码
# import
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings.sentence_transformer import (
    SentenceTransformerEmbeddings,
)
from langchain_community.embeddings import HuggingFaceEmbeddings 
from langchain_community.vectorstores import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders.word_document import Docx2txtLoader

import glob
import os

# 数据库路径
db_dir = "./db"
# 文档路径
source_directory = "./docs"
# 文件后缀
file_ext = '*.docx'

# create the open-source embedding function
# embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
# 使用中文嵌入层编码器
ebd_function = HuggingFaceEmbeddings(model_name="shibing624/text2vec-base-chinese")

def add_files_to_db(filepath:str="",file_ext:str=""):
    docx_files = glob.glob(os.path.join(source_directory, file_ext))
    text_list=[]
    for file_name in docx_files:
        print(file_name)
        loader = Docx2txtLoader(file_name)
        documents = loader.load()
        text_list.extend(documents)

    # split it into chunks
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    docs = text_splitter.split_documents(text_list)

    # load it into Chroma
    db = Chroma.from_documents(docs, ebd_function, persist_directory=db_dir)
    # save db to disk
    db.persist()


def query_db(db:Chroma,query:str=""):

    # query it
    docs = db.similarity_search(query)

    # print results
    print(docs[0].page_content)
    print("-----------------------------------------")


    
 
if __name__=="__main__":

    # 只需执行一次
    # add_files_to_db(source_directory,file_ext)

    db = Chroma(persist_directory=db_dir,embedding_function=ebd_function)
    query = "怎么治疗骨质疏松症?"
    query_db(db,query)
    query = "怎么治疗鼻炎?"
    query_db(db,query)
    db = None
    pass

文档在当前代码目录下

./docs/第十六章-感染性疾病.docx

./docs/第八章-骨骼关节和肌肉疾病.docx

./docs/第十九章-耳鼻咽喉疾病.docx

相关推荐
寻星探路12 分钟前
数据库造神计划第九天---增删改查(CRUD)(5)
数据库
Alan5215944 分钟前
🚀 阿里云 ECS + MySQL 环境搭建全流程(用于个人博客系统开发)
数据库·程序员
Huhbbjs1 小时前
SQL 核心概念与实践总结
开发语言·数据库·sql
wuyunhang1234561 小时前
Redis---集群模式
数据库·redis·缓存
sensenlin911 小时前
Mybatis中SQL全大写或全小写影响执行性能吗
数据库·sql·mybatis
IAtlantiscsdn2 小时前
Redis Stack扩展功能
java·数据库·redis
没有bug.的程序员2 小时前
Redis 大 Key 与热 Key:生产环境的风险与解决方案
java·数据库·redis·缓存·热key·大key
王维志2 小时前
LiteDB详解
数据库·后端·mongodb·sqlite·c#·json·database
2301_815357703 小时前
parameterType和@Param注解的区别
java·开发语言·数据库
零雲3 小时前
除了缓存,我们还可以用redis做什么?
数据库·redis·缓存