sentence transformer模型可以将句子进行Embedding。这里使用docker快速构建一个sentence transformer环境,小试牛刀。
准备镜像
这里Docker的安装就不介绍了,我们使用Dockerfile
来构建我们的镜像,Dockerfile
内容如下:
FROM python:3.10-slim
WORKDIR /app
RUN 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 SentenceTransformer
model = 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, util
app = Flask(__name__)
model = SentenceTransformer('/app/distiluse-base-multilingual-cased-v1')
def encode_text(texts):
# Encode the texts and convert to list
embeddings = 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.json
sentence_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]])
可以看出来相似性还挺高的。