Update main.py

main.py

@@ -53,6 +53,10 @@ class ToxicityResult(BaseModel):
     insult: float
     identity_attack: float
 
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**{k: float(v) for k, v in data.items()})
+
 class TopicBannerClassifier:
     def __init__(self):
         self.classifier = pipeline(
@@ -99,14 +103,7 @@ topic_banner_classifier = TopicBannerClassifier()
 async def classify_toxicity(text_prompt: TextPrompt):
     try:
         result = await run_in_threadpool(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']
-        }
+        return ToxicityResult.from_dict(result)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))