processing_molparser_mobile.py

"""Processor that combines MolParser Mobile image preprocessing and tokenizer."""

from __future__ import annotations

from pathlib import Path
from typing import Sequence

from .image_processing_molparser_mobile import MolParserImageProcessor
from .tokenization_molparser_mobile import MolParserTokenizer


class MolParserProcessor:
    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "MolParserImageProcessor"
    tokenizer_class = "MolParserTokenizer"

    def __init__(
        self,
        image_processor: MolParserImageProcessor | None = None,
        tokenizer: MolParserTokenizer | None = None,
    ):
        self.image_processor = image_processor or MolParserImageProcessor()
        self.tokenizer = tokenizer

    @classmethod
    def register_for_auto_class(cls, auto_class: str = "AutoProcessor"):
        cls._auto_class = auto_class

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs) -> "MolParserProcessor":
        path = str(pretrained_model_name_or_path)
        image_processor = MolParserImageProcessor.from_pretrained(path, **kwargs)
        tokenizer = MolParserTokenizer.from_pretrained(path, **kwargs)
        return cls(image_processor=image_processor, tokenizer=tokenizer)

    def save_pretrained(self, save_directory: str, **kwargs):
        Path(save_directory).mkdir(parents=True, exist_ok=True)
        image_files = self.image_processor.save_pretrained(save_directory, **kwargs)
        tokenizer_files = ()
        if self.tokenizer is not None:
            tokenizer_files = self.tokenizer.save_pretrained(save_directory, **kwargs)
        return tuple(image_files) + tuple(tokenizer_files)

    def __call__(
        self,
        images=None,
        text: str | Sequence[str] | None = None,
        return_tensors: str | None = None,
        **kwargs,
    ):
        encoded = {}
        if images is not None:
            encoded.update(self.image_processor(images=images, return_tensors=return_tensors, **kwargs))
        if text is not None:
            if self.tokenizer is None:
                raise ValueError("MolParserProcessor was created without a tokenizer.")
            encoded.update(self.tokenizer(text, return_tensors=return_tensors, **kwargs))
        return encoded

    def decode(self, *args, **kwargs):
        if self.tokenizer is None:
            raise ValueError("MolParserProcessor was created without a tokenizer.")
        return self.tokenizer.decode(*args, **kwargs)

    def batch_decode(self, *args, **kwargs):
        if self.tokenizer is None:
            raise ValueError("MolParserProcessor was created without a tokenizer.")
        return self.tokenizer.batch_decode(*args, **kwargs)


__all__ = ["MolParserProcessor"]