File size: 800 Bytes
2601994 ed8a7fb 2601994 6dbd488 c5d740a 6dbd488 c5d740a 59fca60 c5d740a 187df69 c5d740a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
# Load model directly
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
class EndpointHandler():
def __init__(self, path=""):
tokenizer = AutoTokenizer.from_pretrained(path)
model = AutoModelForSequenceClassification.from_pretrained(path)
self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
def __call__(self, data):
inputs = data.pop("inputs", data)
def iterator():
for i in inputs:
yield i
labels = []
for out in self.pipeline(iterator(), padding=True, truncation=True, max_length=253):
labels.append(int(out["label"][-1]))
return {
"pairs": inputs,
"evaluations": labels
} |