API для обслуживания Roberta ClassificationModel с FastAPI

Я обучил преобразователь, используя модель простых преобразователей на colab, загрузил модель с поиском, и у меня есть небольшие проблемы с ее использованием для вывода. Загрузка модели в модель на 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}

2 ответа

Не могли бы вы указать ошибку, которую вы получаете, иначе ее довольно сложно увидеть отладку

Также кажется, что model.predict функция в коде jupyter получает массив в качестве входных данных, а в коде fastapi вы передаете строку непосредственно в эту функцию.

Так может попробовать

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

Без ошибки сложно сказать.

Если это поможет, вы можете взглянуть на эту статью, в которой показано, как построить классификацию с помощью преобразователей Hugging Face (модель Bart Large MNLI) и FastAPI: https://nlpcloud.io/nlp-machine-learning-classification-api-production-fastapi-transformers-nlpcloud.html

Другие вопросы по тегам