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

相关推荐
Wyc724095 分钟前
Maven
java·数据库·maven
程序猿小D8 分钟前
[附源码+数据库+毕业论文]基于Spring+MyBatis+MySQL+Maven+jsp实现的电影小说网站管理系统,推荐!
java·数据库·mysql·spring·毕业设计·ssm框架·电影小说网站
羊小猪~~14 分钟前
数据库学习笔记(十七)--触发器的使用
数据库·人工智能·后端·sql·深度学习·mysql·考研
背太阳的牧羊人1 小时前
Neo4j 的向量搜索(Neo4jVector)和常见的向量数据库(比如 Milvus、Qdrant)之间的区别与联系
数据库·neo4j·milvus
liulun1 小时前
在浏览器中使用SQLite(官方sqlite3.wasm)
数据库·sqlite·wasm
IT项目管理2 小时前
达梦数据库DMHS介绍及安装部署
linux·数据库
你都会上树?2 小时前
MySQL MVCC 详解
数据库·mysql
大春儿的试验田3 小时前
高并发收藏功能设计:Redis异步同步与定时补偿机制详解
java·数据库·redis·学习·缓存
Ein hübscher Kerl.3 小时前
虚拟机上安装 MariaDB 及依赖包
数据库·mariadb
醇醛酸醚酮酯4 小时前
Qt项目锻炼——TODO清单(二)
开发语言·数据库·qt