handler.py

from typing import Dict, List, Any
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

class EndpointHandler():
 def __init__(self, path=""):
 quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)
 # load the optimized model
 tokenizer = AutoTokenizer.from_pretrained(path)
 model = AutoModelForCausalLM.from_pretrained(
 path,
 quantization_config=quantization_config,
 device_map="auto",
 torch_dtype='auto'
 ).eval()
 # create inference pipeline
 self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)


 def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
 """
 Args:
 data (:obj:):
 includes the input data and the parameters for the inference.
 Return:
 A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
 - "label": A string representing what the label/class is. There can be multiple labels.
 - "score": A score between 0 and 1 describing how confident the model is for this label/class.
 """
 inputs = data.pop("inputs", data)
 parameters = data.pop("parameters", None)

 # pass inputs with all kwargs in data
 if parameters is not None:
 prediction = self.pipeline(inputs, **parameters)
 else:
 prediction = self.pipeline(inputs)
 # postprocess the prediction
 return prediction