Upload handler.py

handler.py

@@ -0,0 +1,41 @@
+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)