binuser007/Toxic_comment_classification_using_Bert

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