bash
from fastapi import FastAPI
from transformers import AutoModel, AutoTokenizer
import numpy as np
from starlette.responses import JSONResponse
app = FastAPI()
加载模型和分词器
bash
model = AutoModel.from_pretrained("distilbert-base-uncased")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
异步函数用于将输入文本转换为模型需要的格式
bash
async def prepare_input_for_model(text: str):
inputs = tokenizer.encode(text, return_tensors='pt')
return inputs
异步函数用于模型预测
bash
async def get_prediction(inputs):
outputs = model(inputs)
return outputs.logits
异步接口用于处理HTTP请求并返回预测结果
bash
@app.post("/predict")
async def predict(text: str):
inputs = await prepare_input_for_model(text)
outputs = await get_prediction(inputs)
predictions = np.argmax(outputs.numpy(), axis=-1)
return JSONResponse(content={"prediction": predictions[0]})
这段代码展示了如何使用FastAPI框架的异步功能来提高性能。通过异步函数prepare_input_for_model和get_prediction,我们能够处理并行任务,有效利用服务器资源。这样的设计模式对于需要处理大量并发请求的应用程序非常有用。