Skip to content

Instantly share code, notes, and snippets.

@caleb-kaiser
Last active Oct 9, 2020
Embed
What would you like to do?
import torch
from transformers import BertTokenizerFast, BertModel
class PythonPredictor:
def __init__(self, config):
self.device = "cpu"
self.tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
self.model = BertModel.from_pretrained("bert-base-uncased")
def predict(self, payload):
inputs = self.tokenizer(payload["text"], return_tensors="pt")
predictions = self.model(**inputs)
return {"request": "done"}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment