File size: 2,767 Bytes
c81207b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | """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"]
|