predictor.py

from transformers import BertTokenizerFast
import numpy


class ONNXPredictor:
    def __init__(self, onnx_client, config):
        self.device = "cpu"
        self.client = onnx_client
        self.tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")

    def predict(self, payload):
        tokens = self.tokenizer.encode_plus(payload["text"])
        tokens = {name: numpy.atleast_2d(value) for name, value in tokens.items()}
        prediction = self.client.predict(tokens)
        return {"request": "done"}