nn200433/text2vec-server repository overview

FROM nn200433/text2vec-base-env:3.9.8

# 拷贝文件
COPY text2vec /home/text2vec

# 指定工作空间
WORKDIR /home/text2vec

# 设置脚本为可执行
RUN chmod +x run.sh

# 端口变量
ENV PORT=8001
ENV MODEL_PATH=/home/text2vec/model

# 暴露端口
EXPOSE 8001

ENTRYPOINT ["/home/text2vec/run.sh"]

# -*- coding: utf-8 -*-
"""
@author:XuMing([email protected])
@description: pip install fastapi uvicorn
"""
import argparse
import uvicorn
import sys
import os
from fastapi import FastAPI, Query
from starlette.middleware.cors import CORSMiddleware
import torch
from loguru import logger

sys.path.append('..')
from text2vec import SentenceModel, Similarity, EncoderType

# 判断是否支持cuda（显卡）
use_cuda = torch.cuda.is_available()
logger.info(f'use_cuda:{use_cuda}')

# 默认参数
default_model_path = "text2vec_base_chinese"
default_port = 8001

# 获取模型路径
env_port = int(os.environ.get('PORT', default_port))
env_model_path = os.environ.get('MODEL_PATH')

parser = argparse.ArgumentParser()
parser.add_argument("--model_path", type=str, default=default_model_path, help="模型保存位置或模型名称")
parser.add_argument("--port", type=int, default=default_port, help="端口号")
args = parser.parse_args()
args_port = args.port
args_model_path = args.model_path

# 最终参数获取
result_model_path = env_model_path if env_model_path else args_model_path
result_port = env_port if env_port is not None else args_port

# 初始化模型
s_model = SentenceModel(result_model_path, encoder_type="FIRST_LAST_AVG")
sim_model = Similarity(result_model_path, encoder_type=EncoderType.FIRST_LAST_AVG)

# 定义应用程序
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"])

@app.get('/')
async def index():
    return {"message": f"启动成功，使用模型 {result_model_path}"}


@app.get('/vec')
async def vec(q: str = Query(..., min_length=1, max_length=512, title='待转换为向量的文本')):
    try:
        embeddings = s_model.encode(q)
        return {'status': True, "data" : embeddings.tolist(), 'msg': "成功"}
    except Exception as e:
        logger.error(e)
        return {'status': False, 'msg': e}, 400

@app.get('/sim')
async def sim(q1: str = Query(..., min_length=1, max_length=512, title='文本1'),
              q2: str = Query(..., min_length=1, max_length=512, title='文本2')):
    try:
        cosine_scores = sim_model.get_score(q1, q2)
        return {'status': True, "data" : cosine_scores, 'msg': "成功"}
    except Exception as e:
        logger.error(e)
        return {'status': False, 'msg': e}, 400

if __name__ == '__main__':
    uvicorn.run(app=app, host='0.0.0.0', port=result_port)

#!/bin/bash

echo "即将启动text2vec...."
python WebApi.py --model_path=./model

version: '3.3'
services:
  text2vec-server:
    environment:
      - PORT=9999
      #- MODEL_PATH=/home/text2vec/build/text2vec/model
    volumes:
      # 模型路径演示，可选（内置了 text2vec-base-chinese）
      - ./build/text2vec/model:/home/text2vec/model
    ports:
      - 9999:9999
    container_name: text2vecServer
    image: nn200433/text2vec-server:cn_base_py398_v129

# 没东西,就看下是否启动了
curl http://127.0.0.1:9999/

# 计算相似度
# http://127.0.0.1:9999/sim?q1=俗人&q2=凡人
curl http://127.0.0.1:9999/sim?q1=%E4%BF%97%E4%BA%BA\&q2=%E5%87%A1%E4%BA%BA

# 文本转向量
# http:127.0.0.1:9999/vec?q=test
curl http:127.0.0.1:9999/vec?q=%E6%B5%8B%E8%AF%95