Upload handler.py

handler.py

@@ -0,0 +1,23 @@
+import torch
+import transformers
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+ 
+ 
+class EndpointHandler:
+    def __init__(self, path=""):
+        tokenizer = AutoTokenizer.from_pretrained(path)
+        model = AutoModelForSequenceClassification.from_pretrained(path)
+        model.eval()
+        self.pipeline = transformers.pipeline(
+           "text-classification", model=model, tokenizer=tokenizer
+        )
+ 
+    def __call__(self, inputs):
+        result = self.pipeline(inputs, truncation=True, padding=False, max_length=512)
+        for item in result:
+            if item['label'] == 'LABEL_0':
+                item['label'] = 'human-written'
+            else:
+                item['label'] = 'AI-generated'
+        return result