Skip to content

Instantly share code, notes, and snippets.

@caleb-kaiser
Last active Oct 9, 2020
Embed
What would you like to do?
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"}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment