| # handler.py | |
| import json | |
| import os | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline | |
| class EndpointHandler: | |
| def __init__(self, path=""): | |
| # Load the fine-tuned model from the specified path | |
| self.tokenizer = AutoTokenizer.from_pretrained(path) | |
| self.model = AutoModelForSequenceClassification.from_pretrained(path) | |
| # Create a pipeline for text classification | |
| self.pipeline = pipeline( | |
| "text-classification", | |
| model=self.model, | |
| tokenizer=self.tokenizer | |
| ) | |
| def __call__(self, data): | |
| # The data parameter is a dictionary with a 'inputs' key containing the text(s) | |
| inputs = data.get("inputs", data) | |
| if isinstance(inputs, str): | |
| inputs = [inputs] # Wrap single string in a list | |
| # Perform inference | |
| predictions = self.pipeline(inputs) | |
| # The output should be a list of dictionaries, one for each input | |
| return predictions | |