File size: 720 Bytes
cf509fc 9a4ae09 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("text-classification", model="Maxim01/intent-classification")
@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
|