| from fastapi import FastAPI |
| from pydantic import BaseModel |
| from transformers import pipeline |
|
|
| app = FastAPI() |
|
|
| classifier = pipeline( |
| "text-classification", |
| model="LokeshDevCreates/tone-baseline-v3", |
| top_k=None |
| ) |
|
|
| class TextRequest(BaseModel): |
| text: str |
|
|
| @app.post("/predict") |
| def predict_tone(req: TextRequest): |
| results = classifier(req.text)[0] |
| results = sorted(results, key=lambda x: x["score"], reverse=True) |
|
|
| return { |
| "detected_tone": results[0]["label"], |
| "confidence": round(results[0]["score"], 4), |
| "all_probs": {r["label"]: round(r["score"], 4) for r in results} |
| } |
|
|
