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