API to serve roberta ClassificationModel with FastAPI

Question

I have trained transformer using simpletransformers model on colab ,downloaded the searialized model and i have little issues on using it to make inferences. Loading the model on model on jupyter works but while using it with fastapi gives an error This is how I,m using it on 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

It gives the following result:array([[0.00230123, 0.97465035, 0.00475409, 0.01829433]])

I have tried using fastapi as follows but keeps on getting an error:

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}

Answer 1

Could you please specify the error you get, otherwise its quite hard to see debug

Also it seems that the model.predict function in the jupyter code gets an array as input, while in your fastapi code you are passing a string directly to that function.

So maybe try

...
prediction = model.predict([data.text])[1]
...

Answer 2

It's hard to say without the error.

In case it helps you could have a look at this article that shows how to build classification with Hugging Face transformers (Bart Large MNLI model) and FastAPI: https://nlpcloud.io/nlp-machine-learning-classification-api-production-fastapi-transformers-nlpcloud.html