Chroma向量資料庫使用案例

轉載自：https://blog.csdn.net/xzq_qzx_/article/details/136535125

基於ChromaDB與BGEToward-VectorModel的本地私有化向量檢索之路

安裝基礎環境

包括Python安裝、pip安裝及映象設定等網上提供很多資料，這裡就不介紹了。

安裝chromaDB環境安裝

pip install chromadb

安裝pdf解析庫

pip install pdfminer.six

安裝模型庫

pip install sentence_transformers

下載bge-large-zh-v1.5向量模型

如果能訪問huggingface執行程式時自動下載，如果不能訪問huggingface，請點選以下網盤連結進行下載：

連結:百度網盤請輸入提取碼提取碼: fpej

手動下載模型需解壓到專案工程目錄，即與MyVectorDb.py在同一目錄。

實踐

在最後的環節，我們即將把理論化為實踐。我們將用Python編寫出一套基於chromadb向量資料庫和bge-large-zh-v1.5向量模型實現本地向量檢索的程式碼。

chromadb向量資料部分程式碼示例

引用

import chromadb
from chromadb.config import Settings

建立資料庫物件

chroma_client = chromadb.Client(Settings(allow_reset=True))

建立一個 collection，即建立一個名為“demo”的資料庫

collection = chroma_client.get_or_create_collection(name="demo")

向 collection 中新增文件與向量

collection.add(
embeddings= embedding_fn(documents), # 每個文件的向量，這裡呼叫了向量模型，將文件轉化成向量
documents=documents, # 文件的原文
ids=[f"id{i}" for i in range(len(documents))] # 每個文件的 id
)

檢索向量資料庫

collection.query(
query_embeddings=embedding_fn([query]), # 檢索內容的向量，這裡呼叫了向量模型，將文件轉化成向量
n_results=top_n
)

beg模型程式碼示例

引入

from sentence_transformers import SentenceTransformer

建立模型

model = SentenceTransformer('BAAI/bge-large-zh-v1.5')

內容向量化embedding_fn

doc_vecs = [
model.encode(doc, normalize_embeddings=True).tolist()
for doc in documents
]

完整原始碼：

import chromadb
from chromadb.config import Settings
from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer
from sentence_transformers import SentenceTransformer
class MyVectorDB:
'''
私有化向量資料庫
\n作者：George
\n時間：2024年3月7日
'''
def __init__(self, collection_name):
chroma_client = chromadb.Client(Settings(allow_reset=True))
model = SentenceTransformer('BAAI/bge-large-zh-v1.5')
# 為了演示，實際不需要每次 reset()
chroma_client.reset()
# 建立一個 collection
self.collection = chroma_client.get_or_create_collection(name=collection_name)
self.bge_model = model
def add_documents(self, filename, page_numbers=None, min_line_length=1, metadata={}):
paragraphs = self.extract_text_from_pdf(filename, page_numbers, min_line_length)
'''向 collection 中新增文件與向量'''
self.collection.add(
embeddings=self.embedding_fn(paragraphs), # 每個文件的向量
documents=paragraphs, # 文件的原文
ids=[f"id{i}" for i in range(len(paragraphs))] # 每個文件的 id
)
def search(self, query, top_n):
'''檢索向量資料庫'''
results = self.collection.query(
query_embeddings=self.embedding_fn([query]),
n_results=top_n
)
return results
def embedding_fn(self, paragraphs):
'''文字向量化'''
doc_vecs = [
self.bge_model.encode(doc, normalize_embeddings=True).tolist()
for doc in paragraphs
]
return doc_vecs
def extract_text_from_pdf(self, filename, page_numbers=None, min_line_length=1):
'''從 PDF 檔案中（按指定頁碼）提取文字'''
paragraphs = []
buffer = ''
full_text = ''
# 提取全部文字
for i, page_layout in enumerate(extract_pages(filename)):
# 如果指定了頁碼範圍，跳過範圍外的頁
if page_numbers is not None and i not in page_numbers:
continue
for element in page_layout:
if isinstance(element, LTTextContainer):
full_text += element.get_text() + '\n'
# 按空行分隔，將文字重新組織成段落
lines = full_text.split('\n')
for text in lines:
if len(text) >= min_line_length:
buffer += (' '+text) if not text.endswith('-') else text.strip('-')
elif buffer:
paragraphs.append(buffer)
buffer = ''
if buffer:
paragraphs.append(buffer)
return paragraphs
if "__main__" == __name__:
# 建立一個向量資料庫物件
vector_db = MyVectorDB("demo")
# 向向量資料庫中新增文件
vector_db.add_documents("llama2.pdf", page_numbers=[
2, 3], min_line_length=10)
user_query = "Llama 2有多少引數"
results = vector_db.search(user_query, 2)
for para in results['documents'][0]:
print(para+"\n")

總結：

這只是一個簡單的演示樣例，方便大家進一步理解和操作Chroma資料庫，也希望大家一起進步，有問題也可以評論相互學習！

Chroma向量資料庫使用案例

轉載自：https://blog.csdn.net/xzq_qzx_/article/details/136535125

​​​​​基於ChromaDB與BGEToward-VectorModel的本地私有化向量檢索之路

安裝基礎環境

實踐

總結：

相關文章

基於ChromaDB與BGEToward-VectorModel的本地私有化向量檢索之路