| 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, data): | |
| inputs = data.pop("inputs", data) | |
| 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 |