解读FastAPI异步化为transformers模型打造高性能接口

# 导入包
import numpy as np
import pandas as pd
import torch as t
from tqdm import tqdm
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from sentence_transformers import SentenceTransformer as SBert
import uvicorn
from asgiref.sync import sync_to_async
import asyncio

# 加载模型
model = SBert("/home/dataai/文档/huzheng/模型部分/预训练模型/paraphrase-multilingual-MiniLM-L12-v2")

# 启动app
app = FastAPI()

# 让app可以跨域
origins = ["*"]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/")
async def main():
    return {"message": "Hello World"}
    
# 实现功能
@app.get('/get_vector_simple')
def sentenc2vector_simple(sentence):
    """
    这里是提供文本转向量的接口，
    :param sentence: 文本字符串
    :return: 文本向量
    """
    encode1 = model.encode(sentence)
    encode1 = encode1.flatten().tolist()
    return {'vector': encode1}

if __name__ == '__main__':
    uvicorn.run(app='nlp_api:app', host="0.0.0.0",
                port=8000, reload=True, debug=True)

async def encode2list(encode):
    return encode.flatten().tolist()

@app.get('/get_vector_async')
async def sentenc2vector_async(sentence):
    """
    异步版本
    这里是提供文本转向量的接口，
    :param sentence: 文本字符串
    :return: 文本向量
    """
    encode1 = await sync_to_async(model.encode)(sentence)
    encode1 = await encode2list(encode1)
    return {'vector': encode1}