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向量資料部分程式碼示例

引用

1. import chromadb
2. 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 中新增文件與向量

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

檢索向量資料庫

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

• beg模型程式碼示例

引入

from sentence_transformers import SentenceTransformer

建立模型

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

內容向量化embedding_fn

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

完整原始碼：

1. import chromadb
2. from chromadb.config import Settings
3. from pdfminer.high_level import extract_pages
4. from pdfminer.layout import LTTextContainer
5. from sentence_transformers import SentenceTransformer
7. class MyVectorDB:
8. '''
9. 私有化向量資料庫
10. \n作者：George
11. \n時間：2024年3月7日
12. '''
15. def __init__(self, collection_name):
16. chroma_client = chromadb.Client(Settings(allow_reset=True))
17. model = SentenceTransformer('BAAI/bge-large-zh-v1.5')
19. # 為了演示，實際不需要每次 reset()
20. chroma_client.reset()
22. # 建立一個 collection
23. self.collection = chroma_client.get_or_create_collection(name=collection_name)
24. self.bge_model = model
27. def add_documents(self, filename, page_numbers=None, min_line_length=1, metadata={}):
28. paragraphs = self.extract_text_from_pdf(filename, page_numbers, min_line_length)
29. '''向 collection 中新增文件與向量'''
30. self.collection.add(
31. embeddings=self.embedding_fn(paragraphs), # 每個文件的向量
32. documents=paragraphs, # 文件的原文
33. ids=[f"id{i}" for i in range(len(paragraphs))] # 每個文件的 id
34. )
36. def search(self, query, top_n):
37. '''檢索向量資料庫'''
38. results = self.collection.query(
39. query_embeddings=self.embedding_fn([query]),
40. n_results=top_n
41. )
42. return results
44. def embedding_fn(self, paragraphs):
45. '''文字向量化'''
46. doc_vecs = [
47. self.bge_model.encode(doc, normalize_embeddings=True).tolist()
48. for doc in paragraphs
49. ]
50. return doc_vecs
52. def extract_text_from_pdf(self, filename, page_numbers=None, min_line_length=1):
53. '''從 PDF 檔案中（按指定頁碼）提取文字'''
54. paragraphs = []
55. buffer = ''
56. full_text = ''
57. # 提取全部文字
58. for i, page_layout in enumerate(extract_pages(filename)):
59. # 如果指定了頁碼範圍，跳過範圍外的頁
60. if page_numbers is not None and i not in page_numbers:
61. continue
62. for element in page_layout:
63. if isinstance(element, LTTextContainer):
64. full_text += element.get_text() + '\n'
65. # 按空行分隔，將文字重新組織成段落
66. lines = full_text.split('\n')
67. for text in lines:
68. if len(text) >= min_line_length:
69. buffer += (' '+text) if not text.endswith('-') else text.strip('-')
70. elif buffer:
71. paragraphs.append(buffer)
72. buffer = ''
73. if buffer:
74. paragraphs.append(buffer)
75. return paragraphs
78. if "__main__" == __name__:
79. # 建立一個向量資料庫物件
80. vector_db = MyVectorDB("demo")
81. # 向向量資料庫中新增文件
82. vector_db.add_documents("llama2.pdf", page_numbers=[
83. 2, 3], min_line_length=10)
85. user_query = "Llama 2有多少引數"
86. results = vector_db.search(user_query, 2)
88. for para in results['documents'][0]:
89. print(para+"\n")

總結：

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