from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load model and tokenizer
MODEL_PATH = "./hate_speech_distilbert"  # Update with actual path
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

# Label Mapping
LABELS = {
    0: "Hate Speech",
    1: "Offensive Language",
    2: "NOT Hate Speech"
}

app = FastAPI()

class TextRequest(BaseModel):
    text: str

@app.get("/")
def greet_json():
    return {"Hello": "World!"}

@app.post("/predict")
async def predict(request: TextRequest):
    inputs = tokenizer(request.text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
    prediction = torch.argmax(outputs.logits, dim=1).item()
    return {"prediction": LABELS[prediction]}

# Example Usage
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)