從0開始建構個人專屬的向量資料庫吧! Milvus 簡介、操作與實務教學

19 min readJan 31, 2024

Milvus向量数据库安装、使用全中文文档教程 - Milvus向量库中文文档

LangChain中文站，助力大语言模型LLM应用开发、chatGPT应用开发。

www.milvus-io.com

Milvus是在2019年創建的，其唯一目標是儲存、索引和管理由深度神經網路和其他機器學習（ML）模型產生的大規模嵌入向量。

隨著網路的發展和演變，非結構化資料變得越來越常見，包括電子郵件、論文、物聯網感測器資料、Facebook照片、蛋白質結構等等。為了使電腦能夠理解和處理非結構化數據，使用嵌入技術將它們轉換為向量。 Milvus儲存和索引這些向量。 Milvus能夠透過計算它們的相似距離來分析兩個向量之間的相關性。 如果兩個嵌入向量非常相似，則表示原始資料來源也很相似。

如果你來到這邊，你應該知道向量就是一個N維的數字字串，而如果要儲存這一大堆的向量的話，就需要透過向量資料庫(vector DB)來處理這些資料，眾多vector DB之中，跟Python合作最好，適合各種ML場景的Milvus就是多數人選擇的Vector DB。

我們可以用以下的code來連接到自訂的vector db:

from pymilvus import (
    connections,
    utility,
    FieldSchema, CollectionSchema, DataType,
    Collection,
)
import time
import numpy as np

connections.connect(
    alias="default",
    user='root',
    password='xxx',
    host='demodb.tbd.tw',
    port='443',
    secure=True,
)

# 以下是一些格式設定，看看就好
fmt = "\n=== {:30} ===\n"
search_latency_fmt = "search latency = {:.4f}s"

# 檢查某個db有沒有在我們的帳號裡面
has = utility.has_collection("customdb")
print(f"Does collection hello_milvus exist in Milvus: {has}")

透過Transformer將我們資料集裡面的文字轉換成向量

找到將中文轉成向量的工具!

我們在這邊忽略詳細的轉換流程，我們使用預設的bert-base-chinese來處理這項工作，讓我們友善地詢問ChatGPT老師我們該怎麼做：

bert-base-chinese · Hugging Face

We're on a journey to advance and democratize artificial intelligence through open source and open science.

huggingface.co

Question: 如何將中文句子轉換成向量的工具，並且進行相似度搜尋?

Answer:

將中文句子轉換成向量並進行相似度搜尋的方法通常涉及使用自然語言處理（NLP）和嵌入式表示（embedding）技術。以下是一個基本的步驟：

選擇嵌入模型： 選擇一個適合中文的嵌入模型，例如BERT、Word2Vec、或FastText。這些模型在處理中文文本上表現良好。
下載或訓練模型： 如果你選擇的模型需要訓練，你需要下載預訓練的權重或訓練模型。許多NLP模型的預訓練模型已經可以在互聯網上找到。我們去Huggingface上面找到Bert-base-Chinese的使用方式，來幫助我們斷詞(tokenizer)還有化成向量(transformer)
預處理文本： 將你的中文句子進行預處理，包括分詞、去除停用詞、標點符號等。這是為了讓模型更好地理解文本。(已經達成。)
使用模型轉換成向量： 將預處理的文本通過嵌入模型，獲得每個詞的嵌入向量，然後取這些向量的平均或加權平均得到整個句子的向量表示。
計算相似度：使用向量表示來計算句子之間的相似度。常見的方法包括餘弦相似度或歐氏距離。計算相似度也就是我們在使用Milvus VectorDB的時候最需要的東西

以下是使用Python的一個簡單範例，以使用Hugging Face的Transformers庫和BERT模型為例：

from transformers import BertTokenizer, BertModel
import torch
from sklearn.metrics.pairwise import cosine_similarity

# 載入BERT模型和分詞器
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
model = BertModel.from_pretrained('bert-base-chinese')

# 預處理文本
text1 = "這是一個示例句子。"
text2 = "這是另一個示例句子。"
tokens1 = tokenizer.tokenize(tokenizer.decode(tokenizer.encode(text1)))
tokens2 = tokenizer.tokenize(tokenizer.decode(tokenizer.encode(text2)))

# 將文本轉換成BERT的輸入格式
input_ids1 = tokenizer.encode(text1, return_tensors='pt')
input_ids2 = tokenizer.encode(text2, return_tensors='pt')

# 獲取BERT模型的輸出
output1 = model(input_ids1)
output2 = model(input_ids2)

# 取CLS位置的隱藏狀態作為句子向量
sentence_vector1 = output1.last_hidden_state.mean(dim=1).detach().numpy()
sentence_vector2 = output2.last_hidden_state.mean(dim=1).detach().numpy()

# 計算餘弦相似度
similarity = cosine_similarity(sentence_vector1, sentence_vector2)[0][0]
print("相似度:", similarity)

GitHub - UKPLab/sentence-transformers: Multilingual Sentence & Image Embeddings with BERT

Multilingual Sentence & Image Embeddings with BERT - GitHub - UKPLab/sentence-transformers: Multilingual Sentence &…

github.com

這個看起來也不錯，等等再試試看，雖然我印象中我有看過transformer效度排行榜之類的東西，但是我沒找到，反正如果在huggingface的框架之下都可以替換使用的model。

創立Schema

Schema就是資料庫的欄位，假設創立了三個資料庫欄位，之後每筆資料(entity)都會有相對應的屬性，直接來看Code：

#################################################################################
# 2. create collection
# We're going to create a collection with 3 fields.
# +-+------------+------------+------------------+------------------------------+
# | | field name | field type | other attributes |       field description      |
# +-+------------+------------+------------------+------------------------------+
# |1|    "pk"    |   Int64    |  is_primary=True |      "primary field"         |
# | |            |            |   auto_id=True   |                              |
# +-+------------+------------+------------------+------------------------------+
# |2|  "info"    |Varchar(20000)|                |      "a double field"        |
# +-+------------+------------+------------------+------------------------------+
# |3|"embeddings"| FloatVector|     dim=768      |  "float vector with dim 768"   |
# +-+------------+------------+------------------+------------------------------+
fields = [
    FieldSchema(name="pk", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="info", dtype=DataType.VARCHAR,max_length=20000),
    FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=768)
]
schema = CollectionSchema(fields, "VectorDB_StockAI is the smartest stock robot!")
print(fmt.format("Create collection `VectorDB_StockAI`"))

VectorDB_StockAI = Collection("VectorDB_StockAI", schema)

這樣的設計是正確的，embeddings 字段顯然是用來存儲句子的向量資訊的。然而，需要注意一點的是，Milvus 的 FLOAT_VECTOR 支援固定維度的向量，確保每個向量都有相同的維度。在你的代碼中，embeddings 字段的維度為 768，這是一個合理的選擇，因為 BERT 模型常用的向量維度就是 768。

正式把所有的資料丟進去

我把所有的資料都丟進去，並且設定overlap的框架，因為bert最多只能接受512個token，接者我們把所有的資料餵給我們的資料庫。

df = pd.read_csv('VectorDB.csv', encoding='utf-8')

vectors_list = []
sentences_list = []

max_len = 512  # BERT的最大輸入長度
stride = 256  # 步長

# 循環處理每個句子
for index, row in df.iterrows():
    # 得到句子
    sentence = row['content']  # 假設句子在content欄位中，請根據實際情況調整
    if (index%3 == 0):
        print(index, sentence)
    # 分詞
    tokens = tokenizer.tokenize(sentence)

    # 分段
    for i in range(0, len(tokens), stride):
        segment = tokens[i:i + max_len]
        # 把分段後的句子重新組合起來，丟進sentences_list當中
        sentences_list.append("".join(segment))
        # 將文本轉換成BERT的輸入格式
        input_ids = tokenizer.convert_tokens_to_ids(segment)
        input_ids = torch.tensor([input_ids])

        # 獲取BERT模型的輸出
        output = model(input_ids)

        # 取CLS位置的隱藏狀態作為句子向量
        sentence_vector = output.last_hidden_state.mean(dim=1).detach().numpy()

        # 將向量添加到列表中
        vectors_list.append(sentence_vector[0])

# 打印第一個句子的向量和對應的文字
print("第一個句子的向量:", vectors_list[0])
print("對應的文字:", sentences_list[0])索引類型

entities = [
    # provide the pk field because `auto_id` is set to False
    sentences_list, vectors_list]
insert_result = VectorDB_StockAI.insert(entities)
VectorDB_StockAI.flush()
print(f"Number of entities in Milvus: {VectorDB_StockAI.num_entities}")  # check the num_entites

搜尋方式概要

大多數Milvus支援的向量索引類型使用近似最近鄰搜尋（ANNS），包括：

FLAT：FLAT最適合在小規模，百萬級資料集上尋求完全準確和精確的搜尋結果的場景。我們的資料集用FLAT即可。

IVF_FLAT：IVF_FLAT是一種量化索引，最適合在精度和查詢速度之間尋求理想平衡的場景。IVE_FLAT也可以

IVF_SQ8：IVF_SQ8是一種量化索引，最適合在磁碟、CPU和GPU記憶體消耗非常有限的場景中顯著減少資源消耗。

IVF_PQ：IVF_PQ是一種量化索引，最適合在高查詢速度的情況下以犧牲精確度為代價的場景。

HNSW：HNSW是一種基於圖形的索引，最適合對搜尋效率有高需求的場景。

ANNOY：ANNOY是基於樹狀結構的索引，最適合尋求高召回率的場景。

我們這次採用IVE_FLAT作為搜尋的基準。

# 4. create index
# We are going to create an IVF_FLAT index for hello_milvus collection.
# create_index() can only be applied to `FloatVector` and `BinaryVector` fields.
print(fmt.format("Start Creating index IVF_FLAT"))
index = {
    "index_type": "IVF_FLAT",
    "metric_type": "L2",
    "params": {"nlist": 128},
}

VectorDB_StockAI.create_index("embeddings", index)

我們終於可以開始查詢了！

我們可以藉由以下三種方式來把玩我們的向量資料庫：

透過相似度來搜尋相關的資訊 (search based on vector similarity)

每次開始前，我們都需要把我們的資料放進MEMORY之中，我們才可以開始進行搜尋作業。

# - search based on vector similarity
# Before conducting a search or a query, you need to load the data in `VectorDB_StockAI` into memory.
print(fmt.format("Start loading"))
VectorDB_StockAI.load()

接者，我們就可以進行查詢了。我假設我想要查詢的句子是我剛剛輸入進資料庫的最後一個資料，並且限制系統給我前三相似的資料。可想而知，相似度最高的結果的距離為0 (因為我輸入的字詞跟資料庫裏面有的資料一模一樣)。以下為範例code:

print(fmt.format("Start searching based on vector similarity"))
#entity 最後一個項目
vectors_to_search = entities[-1][-1:]

#這裡是系統預設的 我就不改了
search_params = {
    "metric_type": "L2",
    "params": {"nprobe": 10},
}

#可以注意到後面兩個 limit是限制輸出三個最相近的項目，output_fields是顯示輸出時我可以看到的內容
start_time = time.time()
result = VectorDB_StockAI.search(vectors_to_search, "embeddings", search_params, limit=3, output_fields=["info"])
end_time = time.time()

# 處理搜索結果
for hits in result:
    for hit in hits:
        info_data = hit.entity.get('info')
        print(f"hit: {hit}, info: {info_data}")

print(search_latency_fmt.format(end_time - start_time))

=== Start searching based on vector similarity ===

hit: id: 446486765575146598, distance: 0.0, entity: {‘info’: ‘幫助雲端環境建置與維運的系統整合業者、設備安裝業者，以及顧問諮詢業者。’}
hit: id: 446486765575146596, distance: 93.17062377929688, entity: {‘info’: ‘雲端產業可以區分為三個部分，分別是硬體設施、軟體工具以及應用暨支援服務。其中硬體設施供應商是支持雲端應用與服務的基礎設施提供者……’}
hit: id: 446486765575146597, distance: 95.15864562988281, entity: {‘info’: ‘軟體工具：雲端產業中的軟體工具，是提供支持硬體設施運作、相關資源管理以及雲端服務基礎環境，其業者涵蓋設備管理、營運管理、虛擬化、雲端作業系統及資安防護軟體等軟體供應商……’}
search latency = 0.1412s

可以看到我們輸出的結果都跟雲端環境建置有關。接者我們再來問問看一個問題：

Question: 我應該如何研究買賣超動態，來了解籌碼變化?

hit: id: 446486765575146545, distance: 125.75736999511719, entity: {‘info’: ‘三大法人近幾日的買賣超佔個股股本多少我們可以透過「法人買賣超佔股本」，搭配大盤行情，判斷個股的進出場點三大法人買賣超佔比＞##０，是法人買超的意思；三大法人買賣超佔比＜##０，則是法人賣超大盤多頭漲勢時，我們應該關注排名前面的個股，準備進場；大盤空頭下跌時，應該注意排名前面的個股，賣出持股’}
hit: id: 446486765575146562, distance: 126.03663635253906, entity: {‘info’: ‘短期均線＞中期均線＞長期均線，我們可以透過「多頭排列」，找出不論是短期或中長期投資人，都看好未來行情的個股從收盤價介於哪兩個均線之間，我們也能判斷目前市場多空或整盤趨勢多設定一個[UNK]遞增條件，篩選出短期均線仍是向上走的個股’}
hit: id: 446486765575146535, distance: 128.12078857421875, entity: {‘info’: ‘利用個股在近[UNK]日獲利最多的分點，對個股的買賣超情形創造最高投資報酬率的前10名券商分點，這些分點對於此檔股票的判斷有較高的準確度我們可以透過「股懂券商買賣超」的買賣超動態，來獲知籌碼變化，確認主力的進出動作’}
search latency = 0.1024s

可以看到大致符合我們想問的問題

(query based on scalar filtering(boolean, int, etc.))

(hybrid search)

之後補充

刪除資料表

沒什麼特別的，直接改依

###############################################################################
# 6. delete entities by PK
# You can delete entities by their PK values using boolean expressions.
ids = insert_result.primary_keys
expr = f'pk in ["{ids[0]}" , "{ids[1]}"]'
result = VectorDB_StockAI.query(expr=expr, output_fields=["info", "embeddings"])
VectorDB_StockAI.delete(expr)

###############################################################################
# 7. drop collection
# Finally, drop the VectorDB_StockAI collection
print(fmt.format("Drop collection `VectorDB_StockAI`"))
utility.drop_collection("VectorDB_StockAI")