Classification / server /classification_server.py
Ravi Pandey
Docker deploy
d737375
Raw
History Blame Contribute Delete
1.27 kB
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from transformers import pipeline
from fastapi import FastAPI, Request
import uvicorn
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["https://www.google.com"], # or specify other origins
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
classifier = pipeline("zero-shot-classification", model="MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli")
@app.post("/classify_text")
async def classify_text(request: Request):
try:
body = await request.json()
texts = body.get('texts', [])
labels = body.get('labels', ["Violent", "Neutral", "Sexually explicit"])
results = []
# Classify the selected caption and send response to server
for text in texts:
result = classifier(text, labels , max_new_tokens=50)
results.append({
"text": text,
"category": result["labels"][0],
"scores": result["scores"][0]
})
return results
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8001)