Update main.py

main.py

@@ -2,6 +2,8 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 import torch
+from detoxify import Detoxify
+
 
 
 class Guardrail:
@@ -48,6 +50,30 @@ class ClassificationResult(BaseModel):
 
 app = FastAPI()
 guardrail = Guardrail()
+toxicity_classifier = Detoxify('original')
+
+class ToxicityResult(BaseModel):
+    toxicity: float
+    severe_toxicity: float
+    obscene: float
+    threat: float
+    insult: float
+    identity_attack: float
+
+@app.post("/api/models/toxicity/classify", response_model=ToxicityResult)
+def classify_toxicity(text_prompt: TextPrompt):
+    try:
+        result = toxicity_classifier.predict(text_prompt.prompt)
+        return {
+            "toxicity": result['toxicity'],
+            "severe_toxicity": result['severe_toxicity'],
+            "obscene": result['obscene'],
+            "threat": result['threat'],
+            "insult": result['insult'],
+            "identity_attack": result['identity_attack']
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 
 @app.post("/api/models/PromptInjection/classify", response_model=ClassificationResult)