Update server.py

server.py

@@ -1,37 +1,22 @@
-from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
+from transformers import pipeline
 import litserve as ls
 from fastapi import Depends, HTTPException
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 import os
 
 
-class BiasAPI(ls.LitAPI):
+class TopicBanningAPI(ls.LitAPI):
     def setup(self, device):
-        model_name = "d4data/bias-detection-model"
-
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_name,
-            token=os.getenv("HF_TOKEN"),
-            trust_remote_code=True,
-        )
-        model = TFAutoModelForSequenceClassification.from_pretrained(
-            model_name,
-            token=os.getenv("HF_TOKEN"),
-            trust_remote_code=True,
-        )
-
-        self.pipeline = TextClassificationPipeline(
-            'text-classification',
-            model=model,
-            tokenizer=tokenizer,
-            device=-1 if device == "cpu" else 0,  # cpu = -1, gpu = 0
-        )
+        model_name = "facebook/bart-large-mnli"
+        self.pipeline = pipeline("zero-shot-classification", model=model, token=os.getenv("HF_TOKEN"))
 
     def decode_request(self, request):
         return request
 
     def predict(self, query):
-        return self.pipeline(query)
+        text = query["text"]
+        topics = query["topics"]
+        return self.pipeline(text, topics)
 
     def encode_response(self, output):
         return output 
@@ -42,7 +27,6 @@ class BiasAPI(ls.LitAPI):
 
 
 if __name__ == "__main__":
-    api = BiasAPI()
+    api = TopicBanningAPI()
     server = ls.LitServer(api, devices="cpu", accelerator="cpu")
     server.run(port=7860)
-