File size: 682 Bytes
cf509fc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
from transformers import pipeline
import os
from fastapi import APIRouter
from models.prediction import Prediction
from services.answers_service import AnswersService
router = APIRouter(
prefix='/prediction',
tags=['prediction'],
)
nlp_task = os.getenv('NLP_TASK')
model_name = os.getenv('MODEL')
model = pipeline(nlp_task, model=model_name)
@router.post("/get_prediction", name='Получение предсказания модели')
def get_prediction(message_from_user: str):
output = model(message_from_user)
prediction = Prediction.from_output(dictionary=output[0])
answer = AnswersService.get_answer_by_id(prediction=prediction)
return answer
|