File size: 1,852 Bytes
dec266f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Dict
import torch
from src.preprocessing.text_processor import TextPreprocessor
from src.models.toxic_classifier import ToxicClassifier
app = FastAPI()
class CommentRequest(BaseModel):
text: str
class ToxicityResponse(BaseModel):
toxic: float
severe_toxic: float
obscene: float
threat: float
insult: float
identity_hate: float
confidence: float
@app.post("/predict", response_model=ToxicityResponse)
async def predict_toxicity(comment: CommentRequest):
try:
# Preprocess text
preprocessor = TextPreprocessor()
processed_text = preprocessor.process(comment.text)
# Tokenize for BERT
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
encoded = tokenizer(
processed_text,
padding=True,
truncation=True,
max_length=128,
return_tensors='pt'
)
# Get model prediction
model.eval()
with torch.no_grad():
outputs = model(
encoded['input_ids'].to(device),
encoded['attention_mask'].to(device)
)
predictions = outputs[0].cpu().numpy()
confidence = float(outputs.max())
return ToxicityResponse(
toxic=float(predictions[0]),
severe_toxic=float(predictions[1]),
obscene=float(predictions[2]),
threat=float(predictions[3]),
insult=float(predictions[4]),
identity_hate=float(predictions[5]),
confidence=confidence
)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e)) |