背景

最近公司需要用到一個(gè)Bert模型，使用這個(gè)模型對一個(gè)短文本做實(shí)時(shí)的encode（也就是實(shí)現(xiàn)文本轉(zhuǎn)換成向量）。

因?yàn)槟Ｐ褪腔趐ython的transformers和sentence_transfromers。也就是只能使用python來做。

整體的數(shù)據(jù)流都是通過java來調(diào)用，而python這端只需要提供文本轉(zhuǎn)向量的接口即可。

因?yàn)橹熬捅容^喜歡使用fastapi，而且fastapi也比flask快得多。因此將fastapi結(jié)合sentence_transfromers是再正常不過的了。

過程

簡單版本

需要注意的是，這個(gè)代碼是cpu密集型的，非常吃cpu的計(jì)算。

要想實(shí)現(xiàn)這樣的一個(gè)功能，其實(shí)非常簡單，創(chuàng)建一個(gè)python文件叫nlp_api.py，填寫代碼如下：

# 導(dǎo)入包
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/模型部分/預(yù)訓(xùn)練模型/paraphrase-multilingual-MiniLM-L12-v2")

# 啟動(dòng)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"}
    
# 實(shí)現(xiàn)功能
@app.get('/get_vector_simple')
def sentenc2vector_simple(sentence):
    """
    這里是提供文本轉(zhuǎn)向量的接口，
    :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)

運(yùn)行這個(gè)代碼也是非常簡單，直接python運(yùn)行即可:python nlp_api.py。

上面代碼其實(shí)已經(jīng)實(shí)現(xiàn)了這個(gè)功能。但是要對這個(gè)接口 做壓測。這里使用的工具是wrk。我們設(shè)置了100個(gè)thread、1000個(gè)connection、時(shí)間是100秒。最終得到的結(jié)果是：

在1.67分鐘內(nèi)，接受了2529個(gè)請求；平均下來是每秒接受25.27個(gè)請求。這個(gè)時(shí)候我其實(shí)比較滿意了。但是技術(shù)不滿意，說這個(gè)太低了。

我使用htop查看了服務(wù)器的運(yùn)行情況，發(fā)現(xiàn)cpu基本上都是吃滿狀態(tài)。負(fù)荷非常高。

改進(jìn)

既然要考慮到每秒請求這么多的情況下，我用異步試一試。然后把上面的接口 做了異步處理。只要加上這個(gè)代碼就行。

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

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

這個(gè)時(shí)候，就創(chuàng)建了一個(gè)異步接口。然后我又使用wrk。設(shè)置了100個(gè)thread、1000個(gè)connection、時(shí)間是100秒。測試這個(gè)接口，最終得到的結(jié)果是：

在1.67分鐘內(nèi)，接受了7691個(gè)請求，平均下來說每秒接受76.84個(gè)請求。

我把這個(gè)給技術(shù)，技術(shù)那邊也基本是滿意了。這樣算下來，我平均一個(gè)句子轉(zhuǎn)向量的時(shí)間大概需要13ms。這個(gè)其實(shí)已經(jīng)非常高了。

對比

下圖就是一個(gè)接口對比：

• 最上面的框是同步接口效率展示
• 最下面的框是異步接口效率展示

在這次cpu密集型中，異步接口的效率是同步接口效率的3倍。我后來又測試了幾次，基本上都是在3倍以上。

在兩種不同的接口下，我使用htop查看了cpu的運(yùn)行情況：

• 同步接口被請求時(shí)，cpu負(fù)載情況：
• 異步接口被請求時(shí)，cpu負(fù)載情況：

可以看出來，相同的任務(wù)下，cpu的負(fù)載沒有那么高，但是效率反而還提高了。

總結(jié)

這個(gè)模型400MB，底層基于python，使用了pytorch、transformers、Fastapi等包，實(shí)現(xiàn)了文本轉(zhuǎn)向量功能，并且這個(gè)接口的效率可以達(dá)到每秒處理76條。折合每條的文本轉(zhuǎn)向量的時(shí)間只需要13ms左右，我還是很開心的。起碼不用去搞c++、TensorRT之類的東西。python yyds?。?/p>

但是我還沒搞懂為什么在cpu密集型的這種任務(wù)下，異步接口效率比同步接口效率高這么多，而且還降低了cpu的使用率。

這條路走通，起碼代表Fastapi一點(diǎn)也不差?。?！我后面如果要開別的接口，可能都會(huì)用這種方式試一試。

numpy這種，以及Sbert模型其實(shí)都不能異步操作的，但是我使用了asgiref.sync，這個(gè)可以將非異步的轉(zhuǎn)換成異步。非常方便。

作為Fastapi擁鱉，我還是很開心將他用在生產(chǎn)環(huán)境中。希望可以接受住考驗(yàn)！

以上為個(gè)人經(jīng)驗(yàn)，希望能給大家一個(gè)參考，也希望大家多多支持腳本之家。

參考資料：

• https://stackoverflow.com/questions/49822552/python-asyncio-typeerror-object-dict-cant-be-used-in-await-expression
• https://github.com/UKPLab/sentence-transformers/issues/628

最新評(píng)論