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