Update server.py

server.py

@@ -4,36 +4,44 @@ from fastapi import Depends, HTTPException
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 import os
 
-class EmbeddingAPI(ls.LitAPI):
+
+class ToxicAPI(ls.LitAPI):
     def setup(self, device):
+        model_name = "martin-ha/toxic-comment-model"
+
         tokenizer = AutoTokenizer.from_pretrained(
-             "martin-ha/toxic-comment-model",
-            device=device,
+            model_name,
+            use_auth_token=os.getenv("HF_TOKEN"),
             trust_remote_code=True,
-            token=os.getenv("HF_TOKEN")
         )
-        model = AutoTokenizer.from_pretrained(
-             "martin-ha/toxic-comment-model",
-            device=device,
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_name,
+            use_auth_token=os.getenv("HF_TOKEN"),
             trust_remote_code=True,
-            token=os.getenv("HF_TOKEN")
         )
-        self.pipeline = TextClassificationPipeline(model=model, tokenizer=tokenizer)
+
+        self.pipeline = TextClassificationPipeline(
+            model=model,
+            tokenizer=tokenizer,
+            device=-1 if device == "cpu" else 0,  # cpu = -1, gpu = 0
+        )
 
     def decode_request(self, request):
         return request
 
     def predict(self, query):
-            return self.pipeline(query)
+        return self.pipeline(query)
 
     def encode_response(self, output):
-        return output.tolist()
-    
+        return output  # already a JSON-serializable list of dicts
+
     def authorize(self, auth: HTTPAuthorizationCredentials = Depends(HTTPBearer())):
         if auth.scheme != "Bearer" or auth.credentials != os.getenv("auth_token"):
             raise HTTPException(status_code=401, detail="Bad token")
 
+
 if __name__ == "__main__":
-    api = EmbeddingAPI()
+    api = ToxicAPI()
     server = ls.LitServer(api, devices="cpu", accelerator="cpu")
     server.run(port=7860)
+