async

main.py

@@ -3,8 +3,8 @@ from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 import torch
 from detoxify import Detoxify
-
-
+import asyncio
+from fastapi.concurrency import run_in_threadpool
 
 class Guardrail:
     def __init__(self):
@@ -20,8 +20,8 @@ class Guardrail:
             device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
         )
 
-    def guard(self, prompt):
-        return self.classifier(prompt)
+    async def guard(self, prompt):
+        return await run_in_threadpool(self.classifier, prompt)
 
     def determine_level(self, label, score):
         if label == "SAFE":
@@ -36,18 +36,15 @@ class Guardrail:
             else:
                 return 1, "very low"
 
-
 class TextPrompt(BaseModel):
     prompt: str
 
-
 class ClassificationResult(BaseModel):
     label: str
     score: float
     level: int
     severity_label: str
 
-
 app = FastAPI()
 guardrail = Guardrail()
 toxicity_classifier = Detoxify('original')
@@ -61,9 +58,9 @@ class ToxicityResult(BaseModel):
     identity_attack: float
 
 @app.post("/api/models/toxicity/classify", response_model=ToxicityResult)
-def classify_toxicity(text_prompt: TextPrompt):
+async def classify_toxicity(text_prompt: TextPrompt):
     try:
-        result = toxicity_classifier.predict(text_prompt.prompt)
+        result = await run_in_threadpool(toxicity_classifier.predict, text_prompt.prompt)
         return {
             "toxicity": result['toxicity'],
             "severe_toxicity": result['severe_toxicity'],
@@ -75,11 +72,10 @@ def classify_toxicity(text_prompt: TextPrompt):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-
 @app.post("/api/models/PromptInjection/classify", response_model=ClassificationResult)
-def classify_text(text_prompt: TextPrompt):
+async def classify_text(text_prompt: TextPrompt):
     try:
-        result = guardrail.guard(text_prompt.prompt)
+        result = await guardrail.guard(text_prompt.prompt)
         label = result[0]['label']
         score = result[0]['score']
         level, severity_label = guardrail.determine_level(label, score)
@@ -87,8 +83,6 @@ def classify_text(text_prompt: TextPrompt):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-
 if __name__ == "__main__":
     import uvicorn
-
     uvicorn.run(app, host="0.0.0.0", port=8000)