| import torch | |
| from transformers import AutoModel, AutoTokenizer | |
| class ModelHandler: | |
| def __init__(self): | |
| self.model = None | |
| self.tokenizer = None | |
| def initialize(self, model_path): | |
| """Load model and tokenizer.""" | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_path) | |
| self.model = AutoModel.from_pretrained(model_path) | |
| def preprocess(self, data): | |
| """Preprocess input data.""" | |
| text = data.get("text", "") | |
| inputs = self.tokenizer(text, return_tensors="pt") | |
| return inputs | |
| def inference(self, inputs): | |
| """Run inference on the model.""" | |
| outputs = self.model(**inputs) | |
| return outputs | |
| def postprocess(self, outputs): | |
| """Postprocess model output.""" | |
| return {"output": outputs.logits.tolist()} | |
| _handler = ModelHandler() | |
| def handle(data, context): | |
| if not _handler.model: | |
| model_path = context.system_properties.get("model_dir") | |
| _handler.initialize(model_path) | |
| if data is None: | |
| return {"error": "No input data"} | |
| inputs = _handler.preprocess(data[0]) | |
| outputs = _handler.inference(inputs) | |
| return _handler.postprocess(outputs) | |