Docker下SPTAG的安裝與測試

關于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

docker ps

我這里是:46b0c72411dc

docker commit -m "create SPTAG env" 46b0c72411dc nladuo/sptag-rpc-server:latest

再看下當前的鏡像。


docker images

編寫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*(1-0)^2=10,10個2的距離為10*(2-0)^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,10*(3-0)^2=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的安裝和測試就結束了。

最后編輯于
?著作權歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成,瀏覽時請結合常識與多方信息審慎甄別。
平臺聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點,簡書系信息發(fā)布平臺,僅提供信息存儲服務。

友情鏈接更多精彩內(nèi)容