nlp - FastAPI を使用して roberta ClassificationModel を提供する API

Question

colab で simpletransformers モデルを使用してトランスフォーマーをトレーニングし、シリアル化されたモデルをダウンロードしましたが、それを使用して推論を行う際に問題はほとんどありません。モデルをjupyterのモデルにロードすると機能しますが、fastapiで使用するとエラーが発生しますこれは私がjupyterで使用している方法です:

from scipy.special import softmax
label_cols = ['art', 'politics', 'health', 'tourism']
model = torch.load("model.bin")
pred = model.predict(['i love politics'])[1]
preds = softmax(pred,axis=1)
preds

次の結果が得られます: array([[0.00230123, 0.97465035, 0.00475409, 0.01829433]])

次のように fastapi を使用しようとしましたが、エラーが発生し続けます。

from pydantic import BaseModel
class Message(BaseModel):
    text : str
model = torch.load("model.bin")
@app.post("/predict")
def predict_health(data: Message):
    prediction = model.predict(data.text)[1]
    preds = softmax(prediction, axis=1)
    return {"results": preds}