|
|
|
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 |
|
} |
