|
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 |
|
|
