快速搭建sentence_transformer方法

最新推荐文章于 2025-09-19 04:14:29 发布

原创最新推荐文章于 2025-09-19 04:14:29 发布 · 1.9k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#transformer #深度学习 #人工智能

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]])

可以看出来相似性还挺高的。