cguynup/flash_gen_bert_para

1# Load model directly2from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline3 4 5class EndpointHandler():6    def __init__(self, path=""):7        tokenizer = AutoTokenizer.from_pretrained(path)8        model = AutoModelForSequenceClassification.from_pretrained(path)9        10        self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)11 12 13    def __call__(self, data):14        inputs = data.pop("inputs", data)15 16        def iterator():17            for i in inputs:18                yield i19 20        labels = []21        for out in self.pipeline(iterator(), padding=True, truncation=True, max_length=253):22            labels.append(int(out["label"][-1]))23 24        return {25            "pairs": inputs,26            "evaluations": labels27        }