added labels as in input

main.py

@@ -5,6 +5,7 @@ import torch
 from detoxify import Detoxify
 import asyncio
 from fastapi.concurrency import run_in_threadpool
+from typing import List
 
 class Guardrail:
     def __init__(self):
@@ -60,17 +61,20 @@ class TopicBannerClassifier:
             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):
+    async def classify(self, text, labels):
         return await run_in_threadpool(
             self.classifier,
             text,
-            self.classes_verbalized,
+            labels,
             hypothesis_template=self.hypothesis_template,
             multi_label=False
         )
 
+class TopicBannerRequest(BaseModel):
+    prompt: str
+    labels: List[str]
+
 class TopicBannerResult(BaseModel):
     sequence: str
     labels: list
@@ -108,9 +112,9 @@ async def classify_text(text_prompt: TextPrompt):
         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):
+async def classify_topic_banner(request: TopicBannerRequest):
     try:
-        result = await topic_banner_classifier.classify(text_prompt.prompt)
+        result = await topic_banner_classifier.classify(request.prompt, request.labels)
         return {
             "sequence": result["sequence"],
             "labels": result["labels"],