File size: 876 Bytes
dacf14a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from fastapi import APIRouter, HTTPException
from .schemas import QARequest
from transformers import AutoModelForSequenceClassification
from transformers import AutoModelForQuestionAnswering, AutoTokenizer
from app.models import predict_sentiment

router = APIRouter()

model_name = "fine-tuned-model"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_safetensors=True)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

@router.post("/")
def get_answer(request: QARequest):
    try:
        sentiment, probs = predict_sentiment(request.command)
        probs = [float(prob) for prob in probs]
        print({"review": request.command, "sentiment": sentiment, "probs": probs})
        return {"review": request.command, "sentiment": sentiment, "probs": probs}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))