optimum-rbln 0.1.15__py3-none-any.whl → 0.2.1a0__py3-none-any.whl

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
 # Copyright 2024 Rebellions Inc.
 
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,7 +34,6 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionXLPipeline class for inference of diffusion models on rbln devices."""
 
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -32,8 +45,8 @@ from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffus
 from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module
 
-from ....modeling_diffusers import RBLNDiffusionMixin
 from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
@@ -45,6 +58,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableD
     original_class = StableDiffusionXLControlNetImg2ImgPipeline
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]
 
+    # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl_img2img.py
     def check_inputs(
         self,
         prompt,
@@ -246,6 +260,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableD
                     f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
                 )
 
+    # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl_img2img.py
     @torch.no_grad()
     @remove_compile_time_kwargs
     def __call__(

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

@@ -20,11 +20,10 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionPipeline class for inference of diffusion models on rbln devices."""
 
 from diffusers import StableDiffusionPipeline
 
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionPipeline(RBLNDiffusionMixin, StableDiffusionPipeline):

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

@@ -20,11 +20,10 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionPipeline class for inference of diffusion models on rbln devices."""
 
 from diffusers import StableDiffusionImg2ImgPipeline
 
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionImg2ImgPipeline):

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py

@@ -20,11 +20,10 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionInpaintPipeline class for inference of diffusion models on rbln devices."""
 
 from diffusers import StableDiffusionInpaintPipeline
 
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionInpaintPipeline(RBLNDiffusionMixin, StableDiffusionInpaintPipeline):

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py

@@ -20,11 +20,10 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusion3Pipeline class for inference of diffusion models on rbln devices."""
 
 from diffusers import StableDiffusion3Pipeline
 
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3Pipeline(RBLNDiffusionMixin, StableDiffusion3Pipeline):

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py

@@ -20,11 +20,10 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusion3Img2ImgPipeline class for inference of diffusion models on rbln devices."""
 
 from diffusers import StableDiffusion3Img2ImgPipeline
 
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3Img2ImgPipeline(RBLNDiffusionMixin, StableDiffusion3Img2ImgPipeline):

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py

@@ -20,11 +20,10 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusion3InpaintPipeline class for inference of diffusion models on rbln devices."""
 
 from diffusers import StableDiffusion3InpaintPipeline
 
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3InpaintPipeline(RBLNDiffusionMixin, StableDiffusion3InpaintPipeline):

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/__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 .pipeline_stable_diffusion_xl import RBLNStableDiffusionXLPipeline
 from .pipeline_stable_diffusion_xl_img2img import RBLNStableDiffusionXLImg2ImgPipeline
 from .pipeline_stable_diffusion_xl_inpaint import RBLNStableDiffusionXLInpaintPipeline

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

@@ -20,11 +20,10 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionXLPipeline class for inference of diffusion models on rbln devices."""
 
 from diffusers import StableDiffusionXLPipeline
 
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLPipeline(RBLNDiffusionMixin, StableDiffusionXLPipeline):

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py

@@ -20,11 +20,10 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionXLPipeline class for inference of diffusion models on rbln devices."""
 
 from diffusers import StableDiffusionXLImg2ImgPipeline
 
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionXLImg2ImgPipeline):

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py

@@ -20,11 +20,10 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionXLInpaintPipeline class for inference of diffusion models on rbln devices."""
 
 from diffusers import StableDiffusionXLInpaintPipeline
 
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLInpaintPipeline(RBLNDiffusionMixin, StableDiffusionXLInpaintPipeline):

optimum/rbln/modeling.py

@@ -20,7 +20,7 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-import inspect
+
 import logging
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -28,29 +28,16 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import rebel
 import torch
-import transformers
 from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
-from transformers import (
-    AutoConfig,
-    AutoModelForAudioClassification,
-    AutoModelForImageClassification,
-    AutoModelForMaskedLM,
-    AutoModelForQuestionAnswering,
-    AutoModelForSequenceClassification,
-    PretrainedConfig,
-)
+from transformers import AutoConfig, PretrainedConfig
 
 from .modeling_base import RBLNBaseModel
-from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig, use_rbln_config
+from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, use_rbln_config
 
 
 if TYPE_CHECKING:
-    from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        AutoTokenizer,
-        PreTrainedModel,
-    )
+    from transformers import PreTrainedModel
+
 
 logger = logging.getLogger(__name__)
 
@@ -235,338 +222,17 @@ class RBLNModel(RBLNBaseModel):
         cls,
         compiled_models: List[rebel.RBLNCompiledModel],
         rbln_device_map: Dict[str, int],
+        activate_profiler: Optional[bool] = None,
     ) -> List[rebel.Runtime]:
+        if DEFAULT_COMPILED_MODEL_NAME not in rbln_device_map:
+            cls._raise_missing_compiled_file_error([DEFAULT_COMPILED_MODEL_NAME])
+
         device = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
-        return [compiled_model.create_runtime(tensor_type="pt", device=device) for compiled_model in compiled_models]
+        return [
+            compiled_model.create_runtime(tensor_type="pt", device=device, activate_profiler=activate_profiler)
+            for compiled_model in compiled_models
+        ]
 
     def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
         output = self.model[0](*args, **kwargs)
         return output
-
-
-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/modeling_base.py

@@ -21,7 +21,6 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-
 import importlib
 import logging
 import os
@@ -211,6 +210,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         **kwargs,
     ) -> "RBLNBaseModel":
         from_export_method = isinstance(rbln_config, RBLNConfig) and rbln_compiled_models is not None
+
         if not from_export_method:
             # from compiled dir
             rbln_kwargs = rbln_config or {}
@@ -304,7 +304,6 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     ):
         if isinstance(model_save_dir, str):
             model_save_dir = Path(model_save_dir)
-
         # FIXME:: Should we convert it?
         compiled_model_names = [cfg.compiled_model_name for cfg in rbln_config.compile_cfgs]
         rbln_compiled_models = [rbln_compiled_models[cm_name] for cm_name in compiled_model_names]
@@ -312,7 +311,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         # create runtimes only if `rbln_create_runtimes` is enabled
         try:
             models = (
-                cls._create_runtimes(rbln_compiled_models, rbln_config.device_map)
+                cls._create_runtimes(rbln_compiled_models, rbln_config.device_map, rbln_config.activate_profiler)
                 if rbln_config.create_runtimes
                 else UnavailableRuntime()
             )
@@ -373,13 +372,14 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         return from_pretrained_method(model_id=model_id, **kwargs)
 
     @classmethod
-    def compile(cls, model, rbln_compile_config: Optional[RBLNCompileConfig] = None):
+    def compile(cls, model, rbln_compile_config: Optional[RBLNCompileConfig] = None, **kwargs):
         compiled_model = rebel.compile_from_torch(
             model,
             input_info=rbln_compile_config.input_info,
             fusion=rbln_compile_config.fusion,
             npu=rbln_compile_config.npu,
             tensor_parallel_size=rbln_compile_config.tensor_parallel_size,
+            **kwargs,
         )
         return compiled_model
 
@@ -475,6 +475,25 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         if push_to_hub:
             return super().push_to_hub(save_directory, **kwargs)
 
+    @staticmethod
+    def _raise_missing_compiled_file_error(missing_files: List[str]):
+        """Raises a KeyError with a message indicating missing compiled model files."""
+
+        if len(missing_files) == 1:
+            message = f"The rbln model folder is missing the required '{missing_files[0]}.rbln' file. "
+        else:
+            files_str = ", ".join([f"'{f}.rbln'" for f in missing_files])
+            message = (
+                "The rbln model folder is missing required files. "
+                f"Ensure that {files_str} files are present in the folder. "
+            )
+        message += (
+            "These files are necessary for loading the rbln model. "
+            "If these files are missing, please recompile the model using the latest optimum-rbln "
+            "and ensure the compilation completes successfully."
+        )
+        raise KeyError(message)
+
     @classmethod
     @abstractmethod
     def _get_rbln_config(cls, **rbln_config_kwargs) -> RBLNConfig:
@@ -486,6 +505,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         cls,
         compiled_models: List[rebel.RBLNCompiledModel],
         rbln_device_map: Dict[str, int],
+        activate_profiler: Optional[bool] = None,
     ) -> List[rebel.Runtime]:
         # compiled_models -> runtimes
         pass