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