File size: 1,010 Bytes
b82631d | 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 28 29 30 31 |
# 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
|