Create new file

pipeline.py

@@ -0,0 +1,21 @@
+from typing import Dict, List, Any
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+
+
+class PreTrainedPipeline():
+    def __init__(self, path=""):
+      self.processor = TrOCRProcessor.from_pretrained(path)
+      self.model = VisionEncoderDecoderModel.from_pretrained(path)
+    
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        image = data.pop("inputs", data)
+
+        # process image
+        pixel_values = self.processor(images=image, return_tensors="pt").pixel_values
+
+        # run prediction
+        generated_ids = self.model.generate(pixel_values)
+        
+        # decode output
+        prediction = generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)
+        return {"text":prediction[0]}