關于SPTAG
SPTAG是微軟開發(fā)的一款近似最近鄰搜索( approximate nearest neighbor search)的庫,可以用它來做dense vector的索引。
最常見的就是圖像搜索這樣的應用,當然文本檢索做語義匹配也可以用到。
Docker 環(huán)境安裝
因為SPTAG目前不支持mac版本,所以安裝在docker里面就好了。我試了一下官方的dockerfile寫的有點兒問題,我沒運行起來。于是打算直接建一個docker鏡像安裝。
docker run -it ubuntu:18.04 "/bin/bash"
進入docker后運行:
apt update
apt -y install wget build-essential openjdk-8-jdk python3-pip swig
apt -y install software-properties-common git
安裝boost1.67
add-apt-repository ppa:mhier/libboost-latest
apt update
apt install libboost1.67 libboost1.67-dev
安裝cmake3.15.5
wget https://cmake.org/files/v3.15/cmake-3.15.5.tar.gz
tar zvxf cmake-3.15.5.tar.gz
cd cmake-3.15.5
./bootstrap
make -j2 && make install
編譯SPTAG
git clone https://github.com/microsoft/SPTAG.git
cd SPTAG && mkdir build && cd build
cmake ..
make -j2
到這里環(huán)境就算搞好了。把編譯好的結果放到/app下
mkdir /app
mv Release /app
裝一下測試用到的python庫
pip3 install numpy rpyc
Docker鏡像的保存
這里算是裝好了基本的環(huán)境,這里再把環(huán)境保存一下。
先查看一下自己的container id。
docker ps

我這里是:
46b0c72411dc
docker commit -m "create SPTAG env" 46b0c72411dc nladuo/sptag-rpc-server:latest
再看下當前的鏡像。

編寫rpc服務
因為SPTAG不支持Mac,所以為了能在Mac上訪問,這里可以編寫一個簡單的Rpc Demo服務,將接口稍微封裝一下即可。
這里代碼放到了:https://github.com/nladuo/SPTAG-rpc-service
SPTAG_rpc_demo_server.py需要放到docker中,SPTAG_rpc_demo_client.py則直接import到自己的包里即可。
這里我們先把之前的鏡像停掉,重新開一個帶端口映射的容器(我這里用的8888端口)。
docker run -p 8888:8888 -t -i nladuo/sptag-rpc-server:latest "/bin/bash"
這里先把SPTAG_rpc_demo_server.py拷貝到新的docker容器中(注意容器的id的變化)。
docker cp SPTAG_rpc_demo_server.py 25042d741f07:/app/Release/
然后通過python運行起來:

到這里SPTAG的rpc服務算是搞好了,我們可以摁下Ctrl+p 然后再摁下Ctrl+q把服務放到后臺運行。
測試Demo API
添加索引測試
import numpy as np
from SPTAG_rpc_demo_client import SPTAG_RpcDemoClient, DataBean
client = SPTAG_RpcDemoClient("127.0.0.1", "8888")
beans = []
for i in range(5):
vec = i * np.ones((10,), dtype=np.float32)
beans.append(DataBean(_id=f"s{i}", vec=vec))
index_name = "test"
print("Adding Data:", client.add_data(index_name, beans))
這里添加了5個向量,分別是10個0,10個1,..., 10個4。

搜索測試
print("Test Search")
q = DataBean(_id=f"s{0}", vec=0 * np.ones((10,), dtype=np.float32))
print(client.search(index_name, [q], 3))
然后測試下搜索,我們搜索10個0的向量,可以看到返回的10個0(本身)的距離是0,10個1的距離為=10,10個2的距離為
=40。沒有問題

刪除數(shù)據(jù)測試
print("*"*100)
print("Test Delete:", client.delete_data(index_name, [q]))
print("*"*100)
print("Test Search After Deletion")
print(client.search(index_name, [q], 3))
刪除之后,本身不在了。第三近的變成了10個3,=90

刪除索引測試
print("*"*100)
print("Test Delete Index:", client.delete_index(index_name))
最后是刪除索引,可以看到返回結果為True,刪除成功。
真實場景下的接口
在真實場景中,
- 1 . 索引文件可能非常大,我們不會一批一批的添加數(shù)據(jù),也不會使用網(wǎng)絡開銷很大的數(shù)據(jù)傳輸。
- 2 . 每次搜索時候不會每次都重新導入index,然后再調(diào)用搜索。
基于以上兩個缺陷,這里我想到了以下解決方法:
- 1 . 針對索引的很大的問題,直接使用SPTAG提供的indexbuilder工具建立索引。
- 2 . 對于搜索問題,直接編寫一個專門的搜索服務接口,在剛啟動的時候就導入index。(如有需要可以對索引做定期更新,而非每次都重新加載)
通過indexbuilder建立索引
這里先導出一批測試數(shù)據(jù)到test_index_input.txt中
import numpy as np
with open("test_index_input.txt", "w") as f:
for i in range(5):
_id = f"s{i}"
vec = i * np.ones((10,), dtype=np.float32)
vec_str = "|".join([str(i) for i in vec.tolist()])
f.write(f"{_id}\t{vec_str}\n")
然后把test_index_input.txt放到docker里面

然后進入到/app/Release目錄建立索引
cd /app/Release
./indexbuilder -d 10 -v Float -i ./test_index_input.txt -o data/test_index -a BKT -t 2 Index.DistCalcMethod=L2
測試搜索服務
然后和上面的demo API類似,這里我編寫了個search API。還是在這個項目里:https://github.com/nladuo/SPTAG-rpc-service
SPTAG_rpc_search_server.py需要放到docker中,SPTAG_rpc_search_client.py則直接import到自己的包里即可。
這里把SPTAG_rpc_search_server.py放到/app/Release目錄后,啟動起來。
python3 SPTAG_rpc_search_server.py -i test_index

然后再測一下搜索客戶端。

和之前的一樣,沒問題。
到這里SPTAG的安裝和測試就結束了。