Spaces:
Running
Running
File size: 1,067 Bytes
c0921a2 |
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 |
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)
|