from io import BytesIO
from typing import Dict, Any
from transformers import NougatProcessor, VisionEncoderDecoderModel
from transformers.image_utils import base64
from PIL import Image
import torch

class EndpointHandler():
    def __init__(self, path="facebook/nougat-base") -> None:
        self.processor = NougatProcessor.from_pretrained(path)
        self.model = VisionEncoderDecoderModel.from_pretrained(path)

        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model.to(self.device)

    def __call__(self, data: Dict[str, Any]) -> str:
        image = data.pop("inputs", data)
        image_data = Image.open(BytesIO(base64.b64decode(image)))
        pixel_values = self.processor(image_data, return_tensors="pt").pixel_values

        outputs = self.model.generate(
            pixel_values.to(self.device),
            min_length=1,
            max_new_tokens=30,
            bad_words_ids=[[self.processor.tokenizer.unk_token_id]]
        )

        text = self.processor.batch_decode(outputs, skip_special_tokens=True)[0]
        text = self.processor.post_process_generation(text, fix_markdown=False)

        return text