關(guān)于SPTAG
SPTAG是微軟開發(fā)的一款近似最近鄰搜索( approximate nearest neighbor search)的庫(kù)生巡,可以用它來做dense vector的索引。
最常見的就是圖像搜索這樣的應(yīng)用,當(dāng)然文本檢索做語(yǔ)義匹配也可以用到聊闯。
Docker 環(huán)境安裝
因?yàn)镾PTAG目前不支持mac版本史简,所以安裝在docker里面就好了。我試了一下官方的dockerfile寫的有點(diǎn)兒?jiǎn)栴}顽频,我沒運(yùn)行起來藤肢。于是打算直接建一個(gè)docker鏡像安裝。
docker run -it ubuntu:18.04 "/bin/bash"
進(jìn)入docker后運(yùn)行:
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)境就算搞好了糯景。把編譯好的結(jié)果放到/app下
mkdir /app
mv Release /app
裝一下測(cè)試用到的python庫(kù)
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
再看下當(dāng)前的鏡像最住。
編寫rpc服務(wù)
因?yàn)镾PTAG不支持Mac,所以為了能在Mac上訪問怠惶,這里可以編寫一個(gè)簡(jiǎn)單的Rpc Demo服務(wù)涨缚,將接口稍微封裝一下即可。
這里代碼放到了:https://github.com/nladuo/SPTAG-rpc-service
SPTAG_rpc_demo_server.py需要放到docker中策治,SPTAG_rpc_demo_client.py則直接import到自己的包里即可脓魏。
這里我們先把之前的鏡像停掉,重新開一個(gè)帶端口映射的容器(我這里用的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運(yùn)行起來:
到這里SPTAG的rpc服務(wù)算是搞好了,我們可以摁下Ctrl+p 然后再摁下Ctrl+q把服務(wù)放到后臺(tái)運(yùn)行履腋。
測(cè)試Demo API
添加索引測(cè)試
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個(gè)向量珊燎,分別是10個(gè)0,10個(gè)1,..., 10個(gè)4俐末。
搜索測(cè)試
print("Test Search")
q = DataBean(_id=f"s{0}", vec=0 * np.ones((10,), dtype=np.float32))
print(client.search(index_name, [q], 3))
然后測(cè)試下搜索料按,我們搜索10個(gè)0的向量,可以看到返回的10個(gè)0(本身)的距離是0卓箫,10個(gè)1的距離為=10载矿,10個(gè)2的距離為=40。沒有問題
刪除數(shù)據(jù)測(cè)試
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個(gè)3,=90
刪除索引測(cè)試
print("*"*100)
print("Test Delete Index:", client.delete_index(index_name))
最后是刪除索引旅急,可以看到返回結(jié)果為True逢勾,刪除成功。
真實(shí)場(chǎng)景下的接口
在真實(shí)場(chǎng)景中藐吮,
- 1 . 索引文件可能非常大溺拱,我們不會(huì)一批一批的添加數(shù)據(jù),也不會(huì)使用網(wǎng)絡(luò)開銷很大的數(shù)據(jù)傳輸谣辞。
- 2 . 每次搜索時(shí)候不會(huì)每次都重新導(dǎo)入index迫摔,然后再調(diào)用搜索。
基于以上兩個(gè)缺陷泥从,這里我想到了以下解決方法:
- 1 . 針對(duì)索引的很大的問題句占,直接使用SPTAG提供的indexbuilder工具建立索引。
- 2 . 對(duì)于搜索問題躯嫉,直接編寫一個(gè)專門的搜索服務(wù)接口纱烘,在剛啟動(dòng)的時(shí)候就導(dǎo)入index。(如有需要可以對(duì)索引做定期更新祈餐,而非每次都重新加載)
通過indexbuilder建立索引
這里先導(dǎo)出一批測(cè)試數(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里面
然后進(jìn)入到/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
測(cè)試搜索服務(wù)
然后和上面的demo API類似擂啥,這里我編寫了個(gè)search API。還是在這個(gè)項(xiàng)目里: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
目錄后,啟動(dòng)起來舱痘。
python3 SPTAG_rpc_search_server.py -i test_index
然后再測(cè)一下搜索客戶端变骡。
和之前的一樣,沒問題芭逝。
到這里SPTAG的安裝和測(cè)試就結(jié)束了塌碌。