sentence transformer模型可以將句子進行Embedding。這里使用docker快速構建一個sentence transformer環境,小試牛刀。
準備鏡像
這里Docker的安裝就不介紹了,我們使用Dockerfile
來構建我們的鏡像,Dockerfile
內容如下:
FROM python:3.10-slimWORKDIR /appRUN pip install transformers \&& pip3 install torch --index-url https://download.pytorch.org/whl/cpu \&& pip install sentence-transformers \&& pip install Flask
Dockerfile
里面的內容很簡單,主要是安裝pytorch
以及sentence-transformers
庫
構建鏡像:docker build -t ddy_sentence_transformer .
下載模型
在官網https://www.sbert.net/
的介紹中,通過下面的方式使用模型,如果模型不存在,則會自動下載:
from sentence_transformers import SentenceTransformermodel = SentenceTransformer("hkunlp/instructor-large")
然而在使用的過程中,發現根本下載不下來,因此使用離線的方式下載模型,在https://public.ukp.informatik.tu-darmstadt.de/reimers/sentence-transformers/v0.2/
中,找到我們需要的模型然后下載。
這里我們使用較小的模型,distiluse-base-multilingual-cased-v1.zip
,模型支持中文。
將模型進行解壓,得到:0_Transformer
1_Pooling
2_Dense
config.json
modules.json
構建一個簡單的demo
新建一個main.py
用于測試sentence_transformer:
from flask import Flask, request, jsonify
from sentence_transformers import SentenceTransformer, utilapp = Flask(__name__)
model = SentenceTransformer('/app/distiluse-base-multilingual-cased-v1')def encode_text(texts): # Encode the texts and convert to listembeddings = model.encode(texts, convert_to_tensor=True, show_progress_bar=False)embeddings_list = embeddings.tolist()return embeddings_list@app.route('/cal_distance', methods=['POST'])
def cal_distance():params = request.jsonsentence_1 = params['sentence_1']sentence_2 = params['sentence_2']embedding_1 = encode_text([sentence_1])[0]embedding_2 = encode_text([sentence_2])[0]distance = util.pytorch_cos_sim(embedding_1, embedding_2)print('sentence_1:', sentence_1, '|| sentence_2:' , sentence_2 , '===distance:', distance)return jsonify({'status': 'ok'})
if __name__ == '__main__':app.run(host='0.0.0.0', port=5000)
程序的主要作用就是判斷兩句話的相似度
運行程序
docker run \-p 5000:5000 \-v /home/my_user/lib/ddy/sentence_transformer/distiluse-base-multilingual-cased-v1/:/app/distiluse-base-multilingual-cased-v1/ \-v /home/my_user/lib/ddy/sentence_transformer/code/main.py:/app/ \--name ddy \ddy_sentence_transformer \python -u main.py
測試
data = {"sentence_1": '教師', "sentence_2": '老師'}
requests.post('http://10.26.120.58:5000/cal_distance', json=data)
可以看出程序的輸出結果:
sentence_1: 教師 || sentence_2: 老師 ===distance: tensor([[0.9625]])
可以看出來相似性還挺高的。