datafreak/hatespeech-detect

1from fastapi import FastAPI
2from pydantic import BaseModel
3from transformers import AutoTokenizer, AutoModelForSequenceClassification
4import torch
5
6# Load model and tokenizer
7MODEL_PATH = "./hate_speech_distilbert"  # Update with actual path
8model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
9tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
10
11# Label Mapping
12LABELS = {
13    0: "Hate Speech",
14    1: "Offensive Language",
15    2: "NOT Hate Speech"
16}
17
18app = FastAPI()
19
20class TextRequest(BaseModel):
21    text: str
22
23@app.get("/")
24def greet_json():
25    return {"Hello": "World!"}
26
27@app.post("/predict")
28async def predict(request: TextRequest):
29    inputs = tokenizer(request.text, return_tensors="pt", truncation=True, padding=True)
30    with torch.no_grad():
31        outputs = model(**inputs)
32    prediction = torch.argmax(outputs.logits, dim=1).item()
33    return {"prediction": LABELS[prediction]}
34
35# Example Usage
36if __name__ == "__main__":
37    import uvicorn
38    uvicorn.run(app, host="0.0.0.0", port=8000)
39