背景

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

因為模型是基于python的transformers和sentence_transfromers。也就是只能使用python來做。

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

因為之前就比較喜歡使用fastapi，而且fastapi也比flask快得多。因此將fastapi結(jié)合sentence_transfromers是再正常不過的了。

過程

簡單版本

需要注意的是，這個代碼是cpu密集型的，非常吃cpu的計算。

要想實現(xiàn)這樣的一個功能，其實非常簡單，創(chuàng)建一個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")

# 啟動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"}
    
# 實現(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)

運行這個代碼也是非常簡單，直接python運行即可:python nlp_api.py。

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

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

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

改進(jìn)

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

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}

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

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

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

對比

下圖就是一個接口對比：

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

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

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

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

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

總結(jié)

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

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

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

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

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

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

參考資料：

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

最新評論