MayZhou commited on
Commit
2424f5b
·
verified ·
1 Parent(s): 718e2d5

Upload handler.py

Browse files
Files changed (1) hide show
  1. handler.py +23 -0
handler.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import transformers
3
+ from transformers import AutoModelForSequenceClassification, AutoTokenizer
4
+
5
+
6
+
7
+ class EndpointHandler:
8
+ def __init__(self, path=""):
9
+ tokenizer = AutoTokenizer.from_pretrained(path)
10
+ model = AutoModelForSequenceClassification.from_pretrained(path)
11
+ model.eval()
12
+ self.pipeline = transformers.pipeline(
13
+ "text-classification", model=model, tokenizer=tokenizer
14
+ )
15
+
16
+ def __call__(self, inputs):
17
+ result = self.pipeline(inputs, truncation=True, padding=False, max_length=512)
18
+ for item in result:
19
+ if item['label'] == 'LABEL_0':
20
+ item['label'] = 'human-written'
21
+ else:
22
+ item['label'] = 'AI-generated'
23
+ return result