added Topic Banner

main.py

@@ -10,7 +10,6 @@ class Guardrail:
     def __init__(self):
         tokenizer = AutoTokenizer.from_pretrained("ProtectAI/deberta-v3-base-prompt-injection")
         model = AutoModelForSequenceClassification.from_pretrained("ProtectAI/deberta-v3-base-prompt-injection")
-
         self.classifier = pipeline(
             "text-classification",
             model=model,
@@ -45,10 +44,6 @@ class ClassificationResult(BaseModel):
     level: int
     severity_label: str
 
-app = FastAPI()
-guardrail = Guardrail()
-toxicity_classifier = Detoxify('original')
-
 class ToxicityResult(BaseModel):
     toxicity: float
     severe_toxicity: float
@@ -57,6 +52,35 @@ class ToxicityResult(BaseModel):
     insult: float
     identity_attack: float
 
+class TopicBannerClassifier:
+    def __init__(self):
+        self.classifier = pipeline(
+            "zero-shot-classification",
+            model="MoritzLaurer/deberta-v3-large-zeroshot-v2.0",
+            device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        )
+        self.hypothesis_template = "This text is about {}"
+        self.classes_verbalized = ["politics", "economy", "entertainment", "environment"]
+
+    async def classify(self, text):
+        return await run_in_threadpool(
+            self.classifier,
+            text,
+            self.classes_verbalized,
+            hypothesis_template=self.hypothesis_template,
+            multi_label=False
+        )
+
+class TopicBannerResult(BaseModel):
+    sequence: str
+    labels: list
+    scores: list
+
+app = FastAPI()
+guardrail = Guardrail()
+toxicity_classifier = Detoxify('original')
+topic_banner_classifier = TopicBannerClassifier()
+
 @app.post("/api/models/toxicity/classify", response_model=ToxicityResult)
 async def classify_toxicity(text_prompt: TextPrompt):
     try:
@@ -83,6 +107,18 @@ async def classify_text(text_prompt: TextPrompt):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.post("/api/models/TopicBanner/classify", response_model=TopicBannerResult)
+async def classify_topic_banner(text_prompt: TextPrompt):
+    try:
+        result = await topic_banner_classifier.classify(text_prompt.prompt)
+        return {
+            "sequence": result["sequence"],
+            "labels": result["labels"],
+            "scores": result["scores"]
+        }
+    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)
+    uvicorn.run(app, host="0.0.0.0", port=8000)