predictor.py

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"}