| from transformers import pipeline | |
| class EndpointHandler: | |
| def __init__(self, path=""): | |
| self.pipeline = pipeline( | |
| "text-classification", | |
| model=path, | |
| tokenizer=path, | |
| top_k=None, | |
| truncation=True, | |
| max_length=512, | |
| ) | |
| def __call__(self, data): | |
| inputs = data.get("inputs", []) | |
| if isinstance(inputs, str): | |
| inputs = [inputs] | |
| results = self.pipeline(inputs) | |
| return results | |