|
|
from fastapi import FastAPI |
|
|
from pydantic import BaseModel |
|
|
from transformers import pipeline |
|
|
import os |
|
|
|
|
|
|
|
|
os.environ["HF_HOME"] = "/tmp" |
|
|
|
|
|
spam = pipeline("text-classification", model="valurank/distilroberta-spam-comments-detection") |
|
|
|
|
|
toxic = pipeline("text-classification", model="s-nlp/roberta_toxicity_classifier") |
|
|
|
|
|
sentiment = pipeline("text-classification", model = "nlptown/bert-base-multilingual-uncased-sentiment") |
|
|
|
|
|
nsfw = pipeline("text-classification", model = "michellejieli/NSFW_text_classifier") |
|
|
|
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
@app.get("/") |
|
|
def root(): |
|
|
return {"status": "ok"} |
|
|
|
|
|
class Query(BaseModel): |
|
|
text: str |
|
|
|
|
|
@app.post("/spam") |
|
|
def predict_spam(query: Query): |
|
|
result = spam(query.text)[0] |
|
|
return {"label": result["label"], "score": result["score"]} |
|
|
|
|
|
@app.post("/toxic") |
|
|
def predict_toxic(query: Query): |
|
|
result = toxic(query.text)[0] |
|
|
return {"label": result["label"], "score": result["score"]} |
|
|
|
|
|
@app.post("/sentiment") |
|
|
def predict_sentiment(query: Query): |
|
|
result = sentiment(query.text)[0] |
|
|
return {"label": result["label"], "score": result["score"]} |
|
|
|
|
|
@app.post("/nsfw") |
|
|
def predict_nsfw(query: Query): |
|
|
result = nsfw(query.text)[0] |
|
|
return {"label": result["label"], "score": result["score"]} |
