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