optimum-rbln 0.1.13__py3-none-any.whl → 0.2.0__py3-none-any.whl

optimum/rbln/transformers/modeling_generic.py

@@ -0,0 +1,385 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+"""
+This file defines generic base classes for various RBLN models,
+such as Question Answering, Image Classification, Audio Classification,
+Sequence Classification, and Masked Language Modeling. These classes
+implement common functionalities and configurations to be used across
+different model architectures.
+"""
+
+import inspect
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+
+import transformers
+from transformers import (
+    AutoModelForAudioClassification,
+    AutoModelForImageClassification,
+    AutoModelForMaskedLM,
+    AutoModelForQuestionAnswering,
+    AutoModelForSequenceClassification,
+    PretrainedConfig,
+)
+
+from ..modeling import RBLNModel
+from ..modeling_config import RBLNCompileConfig, RBLNConfig
+from ..utils.logging import get_logger
+
+
+if TYPE_CHECKING:
+    from transformers import (
+        AutoFeatureExtractor,
+        AutoProcessor,
+        AutoTokenizer,
+    )
+
+logger = get_logger()
+
+
+class RBLNModelForQuestionAnswering(RBLNModel):
+    auto_model_class = AutoModelForQuestionAnswering
+    rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_kwargs: Dict[str, Any] = {},
+    ) -> RBLNConfig:
+        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+        rbln_model_input_names = rbln_kwargs.get("model_input_names", None)
+
+        if rbln_max_seq_len is None:
+            for tokenizer in preprocessors:
+                if hasattr(tokenizer, "model_max_length"):
+                    rbln_max_seq_len = tokenizer.model_max_length
+                    break
+            if rbln_max_seq_len is None:
+                raise ValueError("`rbln_max_seq_len` should be specified!")
+
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
+
+        if rbln_model_input_names is None:
+            for tokenizer in preprocessors:
+                if hasattr(tokenizer, "model_input_names"):
+                    rbln_model_input_names = tokenizer.model_input_names
+                    break
+            if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
+                rbln_model_input_names = cls.rbln_model_input_names
+            elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
+                input_names_order = inspect.signature(cls.hf_class.forward).parameters.keys()
+                raise ValueError(
+                    "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
+                    f"and be sure to make the order of the inputs same as QuestionAnswering forward() arguments like ({list(input_names_order)})"
+                )
+
+        input_info = [
+            (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
+            for model_input_name in rbln_model_input_names
+        ]
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        return rbln_config
+
+
+class RBLNModelForImageClassification(RBLNModel):
+    """
+    This is a generic model class that will be instantiated as one of the model classes of the library (with a image classification head) when created with the from_pretrained() class method
+    """
+
+    auto_model_class = AutoModelForImageClassification
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_kwargs: Dict[str, Any] = {},
+    ) -> RBLNConfig:
+        rbln_image_size = rbln_kwargs.get("image_size", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+
+        if rbln_image_size is None:
+            for processor in preprocessors:
+                if hasattr(processor, "size"):
+                    if all(required_key in processor.size.keys() for required_key in ["height", "width"]):
+                        rbln_image_size = (processor.size["height"], processor.size["width"])
+                    elif "shortest_edge" in processor.size.keys():
+                        rbln_image_size = (processor.size["shortest_edge"], processor.size["shortest_edge"])
+                    elif "longest_edge" in processor.size.keys():
+                        rbln_image_size = (processor.size["longest_edge"], processor.size["longest_edge"])
+                    break
+
+            if rbln_image_size is None:
+                rbln_image_size = model_config.image_size
+
+            if rbln_image_size is None:
+                raise ValueError("`rbln_image_size` should be specified!")
+
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
+
+        if isinstance(rbln_image_size, int):
+            rbln_image_height, rbln_image_width = rbln_image_size, rbln_image_size
+        elif isinstance(rbln_image_size, (list, tuple)):
+            rbln_image_height, rbln_image_width = rbln_image_size[0], rbln_image_size[1]
+        elif isinstance(rbln_image_size, dict):
+            rbln_image_height, rbln_image_width = rbln_image_size["height"], rbln_image_size["width"]
+        else:
+            raise ValueError(
+                "`rbln_image_size` should be `int` (ex. 224), `tuple` (ex. 224, 224), `dict` (ex. {'height': 224, 'width': 224}) format"
+            )
+
+        input_info = [
+            (
+                "pixel_values",
+                [rbln_batch_size, 3, rbln_image_height, rbln_image_width],
+                "float32",
+            )
+        ]
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        return RBLNConfig(rbln_cls=cls.__name__, compile_cfgs=[rbln_compile_config], rbln_kwargs=rbln_kwargs)
+
+
+class RBLNModelForAudioClassification(RBLNModel):
+    """
+    This is a generic model class that will be instantiated as one of the model classes of the library (with a audio classification head) when created with the from_pretrained() class method
+    This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based AudioClassification models on RBLN devices.
+    It implements the methods to convert a pre-trained transformers AudioClassification model into a RBLN transformer model by:
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+
+    Currently, this model class only supports the 'AST' model from the transformers library. Future updates may include support for additional model types.
+    """
+
+    auto_model_class = AutoModelForAudioClassification
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: "AutoFeatureExtractor",
+        model_config: "PretrainedConfig",
+        rbln_kwargs: Dict[str, Any] = {},
+    ) -> RBLNConfig:
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+        rbln_max_length = rbln_kwargs.get("max_length", None)
+        rbln_num_mel_bins = rbln_kwargs.get("num_mel_bins", None)
+
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
+
+        if rbln_num_mel_bins is None:
+            rbln_num_mel_bins = getattr(model_config, "num_mel_bins", None)
+            if rbln_num_mel_bins is None:
+                for feature_extractor in preprocessors:
+                    if hasattr(feature_extractor, "num_mel_bins"):
+                        rbln_num_mel_bins = feature_extractor.num_mel_bins
+                        break
+
+        if rbln_num_mel_bins is None:
+            raise ValueError("`rbln_num_mel_bins` should be specified!")
+
+        if rbln_max_length is None:
+            rbln_max_length = getattr(model_config, "max_length", None)
+            for feature_extractor in preprocessors:
+                if hasattr(feature_extractor, "max_length"):
+                    rbln_max_length = feature_extractor.max_length
+                    break
+
+        if rbln_max_length is None:
+            raise ValueError("`rbln_max_length` should be specified!")
+
+        input_info = [
+            (
+                "input_values",
+                [rbln_batch_size, rbln_max_length, rbln_num_mel_bins],
+                "float32",
+            ),
+        ]
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+        rbln_config.model_cfg.update(
+            {
+                "batch_size": rbln_batch_size,
+                "max_length": rbln_max_length,
+                "num_mel_bins": rbln_num_mel_bins,
+            }
+        )
+        return rbln_config
+
+
+class RBLNModelForSequenceClassification(RBLNModel):
+    """
+    This is a generic model class that will be instantiated as one of the model classes of the library (with a sequence classification head) when created with the from_pretrained() class method
+    This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based SequenceClassification models on RBLN devices.
+    It implements the methods to convert a pre-trained transformers SequenceClassification model into a RBLN transformer model by:
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+
+    Currently, this model class supports the 'XLMRoberta' and 'Roberta' model from the transformers library. Future updates may include support for additional model types.
+    """
+
+    auto_model_class = AutoModelForSequenceClassification
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_kwargs: Dict[str, Any] = {},
+    ) -> RBLNConfig:
+        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
+        rbln_model_input_names = rbln_kwargs.get("model_input_names", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+
+        max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
+            model_config, "max_position_embeddings", None
+        )
+
+        if rbln_max_seq_len is None:
+            rbln_max_seq_len = max_position_embeddings
+            if rbln_max_seq_len is None:
+                for tokenizer in preprocessors:
+                    if hasattr(tokenizer, "model_max_length"):
+                        rbln_max_seq_len = tokenizer.model_max_length
+                        break
+                if rbln_max_seq_len is None:
+                    raise ValueError("`rbln_max_seq_len` should be specified!")
+
+        if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
+            raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
+
+        if rbln_model_input_names is None:
+            for tokenizer in preprocessors:
+                if hasattr(tokenizer, "model_input_names"):
+                    rbln_model_input_names = tokenizer.model_input_names
+                    break
+            if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
+                rbln_model_input_names = cls.rbln_model_input_names
+            elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
+                original_model_class = getattr(transformers, model_config.architectures[0])
+                input_names_order = inspect.signature(original_model_class.forward).parameters.keys()
+                raise ValueError(
+                    "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
+                    f"and be sure to make the order of the inputs same as SequenceClassification forward() arguments like ({list(input_names_order)})"
+                )
+
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
+
+        input_info = [
+            (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
+            for model_input_name in rbln_model_input_names
+        ]
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        return rbln_config
+
+
+class RBLNModelForMaskedLM(RBLNModel):
+    auto_model_class = AutoModelForMaskedLM
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_kwargs: Dict[str, Any] = {},
+    ) -> RBLNConfig:
+        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
+        rbln_model_input_names = rbln_kwargs.get("model_input_names", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+
+        max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
+            model_config, "max_position_embeddings", None
+        )
+
+        if rbln_max_seq_len is None:
+            rbln_max_seq_len = max_position_embeddings
+            if rbln_max_seq_len is None:
+                for tokenizer in preprocessors:
+                    if hasattr(tokenizer, "model_max_length"):
+                        rbln_max_seq_len = tokenizer.model_max_length
+                        break
+                if rbln_max_seq_len is None:
+                    raise ValueError("`rbln_max_seq_len` should be specified!")
+
+        if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
+            raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
+
+        if rbln_model_input_names is None:
+            for tokenizer in preprocessors:
+                if hasattr(tokenizer, "model_input_names"):
+                    rbln_model_input_names = tokenizer.model_input_names
+                    break
+            if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
+                rbln_model_input_names = cls.rbln_model_input_names
+            elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
+                input_names_order = inspect.signature(cls.hf_class.forward).parameters.keys()
+                raise ValueError(
+                    "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
+                    f"and be sure to make the order of the inputs same as MaskedLM forward() arguments like ({list(input_names_order)})"
+                )
+
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
+
+        input_info = [
+            (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
+            for model_input_name in rbln_model_input_names
+        ]
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        return rbln_config

optimum/rbln/transformers/models/auto/__init__.py

@@ -1,3 +1,26 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
 from .modeling_auto import (
     RBLNAutoModel,
     RBLNAutoModelForAudioClassification,

optimum/rbln/transformers/models/auto/auto_factory.py

@@ -22,8 +22,16 @@
 # from Rebellions Inc.
 
 import importlib
+import inspect
+import warnings
 
-from transformers import AutoConfig
+from transformers import AutoConfig, PretrainedConfig
+from transformers.dynamic_module_utils import get_class_from_dynamic_module
+from transformers.models.auto.auto_factory import _get_model_class
+
+from optimum.rbln.modeling_base import RBLNBaseModel
+from optimum.rbln.modeling_config import RBLNConfig
+from optimum.rbln.utils.model_utils import convert_hf_to_rbln_model_name, convert_rbln_to_hf_model_name
 
 
 class _BaseAutoModelClass:
@@ -33,46 +41,132 @@ class _BaseAutoModelClass:
     def __init__(self, *args, **kwargs):
         raise EnvironmentError(
             f"{self.__class__.__name__} is designed to be instantiated "
-            f"using the `{self.__class__.__name__}.from_pretrained(pretrained_model_name_or_path)` or "
-            f"`{self.__class__.__name__}.from_config(config)` methods."
+            f"using the `{self.__class__.__name__}.from_pretrained(pretrained_model_name_or_path)`"
         )
 
     @classmethod
     def get_rbln_cls(
         cls,
-        model_id,
+        pretrained_model_name_or_path,
         *args,
+        export=True,
         **kwargs,
     ):
-        # kwargs.update({"return_unused_kwargs": True})
-        config = AutoConfig.from_pretrained(model_id, return_unused_kwargs=True, **kwargs)[0]
-
-        if len(config.architectures) > 1:
-            raise ValueError(
-                f"Model with ID '{model_id}' has multiple architectures defined in the configuration: "
-                f"{config.architectures}. `_BaseAutoModelClass` require exactly one architecture. "
-            )
-
-        architecture_name = config.architectures[0]
-        if architecture_name not in cls._model_mapping.values():
-            raise ValueError(
-                f"The 'RBLN{architecture_name}' architecture is not supported by `{cls.__name__}.from_pretrained()`."
-                "Please use the appropriate class's `from_pretrained()` method to load this model."
-            )
-
-        rbln_class_name = "RBLN" + architecture_name
-        module = importlib.import_module("optimum.rbln")
+        """
+        Determine the appropriate RBLN model class based on the given model ID and configuration.
+
+        Args:
+            pretrained_model_name_or_path (str): Identifier or path to the pretrained model.
+            export (bool): Whether to infer the class based on Hugging Face (HF) architecture.
+            kwargs: Additional arguments for configuration and loading.
+
+        Returns:
+            RBLNBaseModel: The corresponding RBLN model class.
+        """
+        if export:
+            hf_model_class = cls.infer_hf_model_class(pretrained_model_name_or_path, **kwargs)
+            rbln_class_name = convert_hf_to_rbln_model_name(hf_model_class.__name__)
+        else:
+            rbln_class_name = cls.get_rbln_model_class_name(pretrained_model_name_or_path, **kwargs)
+
+            if convert_rbln_to_hf_model_name(rbln_class_name) not in cls._model_mapping_names.values():
+                raise ValueError(
+                    f"The architecture '{rbln_class_name}' is not supported by the `{cls.__name__}.from_pretrained()` method. "
+                    "Please use the `from_pretrained()` method of the appropriate class to load this model, "
+                    f"or directly use '{rbln_class_name}.from_pretrained()`."
+                )
 
         try:
+            module = importlib.import_module("optimum.rbln")
             rbln_cls = getattr(module, rbln_class_name)
         except AttributeError as e:
             raise AttributeError(
-                f"Class '{rbln_class_name}' not found in 'optimum.rbln' module for model ID '{model_id}'. "
+                f"Class '{rbln_class_name}' not found in 'optimum.rbln' module for model ID '{pretrained_model_name_or_path}'. "
                 "Ensure that the class name is correctly mapped and available in the 'optimum.rbln' module."
             ) from e
 
         return rbln_cls
 
+    @classmethod
+    def infer_hf_model_class(
+        cls,
+        pretrained_model_name_or_path,
+        *args,
+        **kwargs,
+    ):
+        """
+        Infer the Hugging Face model class based on the configuration or model name.
+
+        Args:
+            pretrained_model_name_or_path (str): Identifier or path to the pretrained model.
+            kwargs: Additional arguments for configuration and loading.
+
+        Returns:
+            PretrainedModel: The inferred Hugging Face model class.
+        """
+
+        # Try to load configuration if provided or retrieve it from the model ID
+        config = kwargs.pop("config", None)
+        kwargs.update({"trust_remote_code": True})
+        kwargs["_from_auto"] = True
+
+        # Load configuration if not already provided
+        if not isinstance(config, PretrainedConfig):
+            config, kwargs = AutoConfig.from_pretrained(
+                pretrained_model_name_or_path,
+                return_unused_kwargs=True,
+                **kwargs,
+            )
+
+        # Get hf_model_class from Config
+        has_remote_code = (
+            hasattr(config, "auto_map") and convert_rbln_to_hf_model_name(cls.__name__) in config.auto_map
+        )
+        if has_remote_code:
+            class_ref = config.auto_map[convert_rbln_to_hf_model_name(cls.__name__)]
+            model_class = get_class_from_dynamic_module(class_ref, pretrained_model_name_or_path, **kwargs)
+        elif type(config) in cls._model_mapping.keys():
+            model_class = _get_model_class(config, cls._model_mapping)
+        else:
+            raise ValueError(
+                f"Unrecognized configuration class {config.__class__} for this kind of AutoModel: {cls.__name__}.\n"
+                f"Model type should be one of {', '.join(c.__name__ for c in cls._model_mapping.keys())}."
+            )
+
+        if model_class.__name__ != config.architectures[0]:
+            warnings.warn(
+                f"`{cls.__name__}.from_pretrained()` is invoking `{convert_hf_to_rbln_model_name(model_class.__name__)}.from_pretrained()`, which does not match the "
+                f"expected architecture `RBLN{config.architectures[0]}` from config. This mismatch could cause some operations to not be properly loaded "
+                f"from the checkpoint, leading to potential unintended behavior. If this is not intentional, consider calling the "
+                f"`from_pretrained()` method directly from the `RBLN{config.architectures[0]}` class instead.",
+                UserWarning,
+            )
+
+        return model_class
+
+    @classmethod
+    def get_rbln_model_class_name(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Retrieve the path to the compiled model directory for a given RBLN model.
+
+        Args:
+            pretrained_model_name_or_path (str): Identifier of the model.
+            kwargs: Additional arguments that match the parameters of `_load_compiled_model_dir`.
+
+        Returns:
+            str: Path to the compiled model directory.
+        """
+        sig = inspect.signature(RBLNBaseModel._load_compiled_model_dir)
+        valid_params = sig.parameters.keys()
+        filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+
+        model_path_subfolder = RBLNBaseModel._load_compiled_model_dir(
+            model_id=pretrained_model_name_or_path, **filtered_kwargs
+        )
+        rbln_config = RBLNConfig.load(model_path_subfolder)
+
+        return rbln_config.meta["cls"]
+
     @classmethod
     def from_pretrained(
         cls,