Upload handler.py

handler.py

@@ -0,0 +1,39 @@
+import torch
+from transformers import AutoModelForVision2Seq, AutoTokenizer
+
+class ModelHandler:
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+
+    def initialize(self, context):
+        """ Load the model and tokenizer """
+        model_dir = context.system_properties.get("model_dir")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        self.model = AutoModelForVision2Seq.from_pretrained(model_dir)
+        self.model.eval()
+
+    def preprocess(self, data):
+        """ Preprocess the input data before passing it to the model """
+        inputs = self.tokenizer(data, return_tensors="pt")
+        return inputs
+
+    def inference(self, inputs):
+        """ Run the forward pass of the model """
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        return outputs
+
+    def postprocess(self, outputs):
+        """ Post-process the output data from the model """
+        return outputs
+
+# This is required for the Hugging Face inference endpoints
+_handler = ModelHandler()
+
+def handle(data, context):
+    if not _handler.model:
+        _handler.initialize(context)
+    inputs = _handler.preprocess(data)
+    outputs = _handler.inference(inputs)
+    return _handler.postprocess(outputs)