optimum-rbln 0.8.0.post2__py3-none-any.whl → 0.8.1__py3-none-any.whl

optimum/rbln/modeling_base.py

@@ -15,7 +15,7 @@
 import importlib
 import os
 import shutil
-from abc import ABC, abstractmethod
+from abc import ABC
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
@@ -24,10 +24,10 @@ import rebel
 import torch
 from transformers import AutoConfig, AutoModel, GenerationConfig, PretrainedConfig
 
-from .configuration_utils import RBLNAutoConfig, RBLNCompileConfig, RBLNModelConfig
+from .configuration_utils import RBLNAutoConfig, RBLNCompileConfig, RBLNModelConfig, get_rbln_config_class
 from .utils.hub import PushToHubMixin, pull_compiled_model_from_hub, validate_files
 from .utils.logging import get_logger
-from .utils.runtime_utils import UnavailableRuntime
+from .utils.runtime_utils import UnavailableRuntime, tp_and_devices_are_ok
 from .utils.save_utils import maybe_load_preprocessors
 from .utils.submodule import SubModulesMixin
 
@@ -47,40 +47,6 @@ class RBLNBaseModelConfig(RBLNModelConfig):
 
 
 class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
-    """
-    An abstract base class for compiling, loading, and saving neural network models from the huggingface
-    transformers and diffusers libraries to run on RBLN NPU devices.
-
-    This class supports loading and saving models using the `from_pretrained` and `save_pretrained` methods,
-    similar to the huggingface libraries.
-
-    The `from_pretrained` method loads a model corresponding to the given `model_id` from a local repository
-    or the huggingface hub onto the NPU. If the model is a PyTorch model and `export=True` is passed as a
-    kwarg, it compiles the PyTorch model corresponding to the given `model_id` before loading. If `model_id`
-    is an already rbln-compiled model, it can be directly loaded onto the NPU with `export=False`.
-
-    `rbln_npu` is a kwarg required for compilation, specifying the name of the NPU to be used. If this
-    keyword is not specified, the NPU installed on the host machine is used. If no NPU is installed on the
-    host machine, an error occurs.
-
-    `rbln_device` specifies the device to be used at runtime. If not specified, device 0 is used.
-
-    `rbln_create_runtimes` indicates whether to create runtime objects. If False, the runtime does not load
-    the model onto the NPU. This option is particularly useful when you want to perform compilation only on a
-    host machine without an NPU.
-
-    `RBLNModel`, `RBLNModelFor*`, etc. are all child classes of RBLNBaseModel.
-
-    Models compiled in this way can be saved to a local repository using `save_pretrained` or uploaded to
-    the huggingface hub.
-
-    It also supports generation through `generate` (for transformers models that support generation).
-
-    RBLNBaseModel is a class for models consisting of an arbitrary number of `torch.nn.Module`s, and
-    therefore is an abstract class without explicit implementations of `forward` or `export` functions.
-    To inherit from this class, `forward`, `export`, etc. must be implemented.
-    """
-
     model_type = "rbln_model"
     auto_model_class = AutoModel
     config_class = AutoConfig
@@ -156,7 +122,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         subfolder: str = "",
         local_files_only: bool = False,
     ) -> str:
-        """Load the directory containing the compiled model files."""
+        # Load the directory containing the compiled model files.
         model_path = Path(model_id)
 
         if model_path.is_dir():
@@ -372,22 +338,59 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     def prepare_rbln_config(
         cls, rbln_config: Optional[Union[Dict[str, Any], RBLNModelConfig]] = None, **kwargs
     ) -> Tuple[RBLNModelConfig, Dict[str, Any]]:
-        """
-        Extract rbln-config from kwargs and convert it to RBLNModelConfig.
-        """
+        # Extract rbln-config from kwargs and convert it to RBLNModelConfig.
+
         config_cls = cls.get_rbln_config_class()
         rbln_config, kwargs = config_cls.initialize_from_kwargs(rbln_config, **kwargs)
         return rbln_config, kwargs
 
     @classmethod
-    def from_pretrained(cls, model_id: Union[str, Path], export: bool = False, **kwargs) -> "RBLNBaseModel":
+    def from_pretrained(
+        cls: Type["RBLNBaseModel"],
+        model_id: Union[str, Path],
+        export: bool = False,
+        rbln_config: Optional[Union[Dict, RBLNModelConfig]] = None,
+        **kwargs: Dict[str, Any],
+    ) -> "RBLNBaseModel":
+        """
+        The `from_pretrained()` function is utilized in its standard form as in the HuggingFace transformers library.
+        User can use this function to load a pre-trained model from the HuggingFace library and convert it to a RBLN model to be run on RBLN NPUs.
+
+        Args:
+            model_id: The model id of the pre-trained model to be loaded. It can be downloaded from the HuggingFace model hub or a local path, or a model id of a compiled model using the RBLN Compiler.
+            export: A boolean flag to indicate whether the model should be compiled.
+            rbln_config: Configuration for RBLN model compilation and runtime. This can be provided as a dictionary or an instance of the model's configuration class (e.g., `RBLNLlamaForCausalLMConfig` for Llama models).
+                For detailed configuration options, see the specific model's configuration class documentation.
+
+            kwargs: Additional keyword arguments. Arguments with the prefix 'rbln_' are passed to rbln_config, while the remaining arguments are passed to the HuggingFace library.
+
+        Returns:
+            A RBLN model instance ready for inference on RBLN NPU devices.
+        """
+
         if isinstance(model_id, Path):
             model_id = model_id.as_posix()
         from_pretrained_method = cls._export if export else cls._from_pretrained
-        return from_pretrained_method(model_id=model_id, **kwargs)
+        return from_pretrained_method(model_id=model_id, **kwargs, rbln_config=rbln_config)
 
     @classmethod
-    def compile(cls, model, rbln_compile_config: Optional[RBLNCompileConfig] = None, **kwargs):
+    def compile(
+        cls,
+        model,
+        rbln_compile_config: RBLNCompileConfig,
+        create_runtimes: bool,
+        device: Union[int, List[int]],
+        **kwargs,
+    ):
+        if create_runtimes:
+            runtime_cannot_be_created = tp_and_devices_are_ok(
+                tensor_parallel_size=rbln_compile_config.tensor_parallel_size,
+                device=device,
+                npu=rbln_compile_config.npu,
+            )
+            if runtime_cannot_be_created:
+                raise ValueError(runtime_cannot_be_created)
+
         compiled_model = rebel.compile_from_torch(
             model,
             input_info=rbln_compile_config.input_info,
@@ -411,15 +414,13 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
     @classmethod
     def get_hf_class(cls):
-        """
-        Lazily loads and caches the corresponding HuggingFace model class.
-        Removes 'RBLN' prefix from the class name to get the original class name
-        (e.g., RBLNLlamaForCausalLM -> LlamaForCausalLM) and imports it from
-        the transformers/diffusers module.
+        # Lazily loads and caches the corresponding HuggingFace model class.
+        # Removes 'RBLN' prefix from the class name to get the original class name
+        # (e.g., RBLNLlamaForCausalLM -> LlamaForCausalLM) and imports it from
+        # the transformers/diffusers module.
 
-        Returns:
-            type: The original HuggingFace model class
-        """
+        # Returns:
+        #     type: The original HuggingFace model class
         if cls._hf_class is None:
             hf_cls_name = cls.__name__[4:]
             library = importlib.import_module(cls.hf_library_name)
@@ -428,18 +429,10 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
     @classmethod
     def get_rbln_config_class(cls) -> Type[RBLNModelConfig]:
-        """
-        Lazily loads and caches the corresponding RBLN model config class.
-        """
+        # Lazily loads and caches the corresponding RBLN model config class.
         if cls._rbln_config_class is None:
             rbln_config_class_name = cls.__name__ + "Config"
-            library = importlib.import_module("optimum.rbln")
-            cls._rbln_config_class = getattr(library, rbln_config_class_name, None)
-            if cls._rbln_config_class is None:
-                raise ValueError(
-                    f"RBLN config class {rbln_config_class_name} not found. This is an internal error. "
-                    "Please report it to the developers."
-                )
+            cls._rbln_config_class = get_rbln_config_class(rbln_config_class_name)
         return cls._rbln_config_class
 
     def can_generate(self):
@@ -449,17 +442,15 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         return self
 
     def parameters(self):
-        """
-        Provides a dummy parameter generator for compatibility.
+        # A dummy parameter generator for compatibility.
 
-        This method mimics the interface of torch.nn.Module.parameters()
-        specifically for code that uses `next(model.parameters())` to infer
-        the device or dtype. It yields a single dummy tensor on CPU with float32 dtype.
+        # This method mimics the interface of torch.nn.Module.parameters()
+        # specifically for code that uses `next(model.parameters())` to infer
+        # the device or dtype. It yields a single dummy tensor on CPU with float32 dtype.
 
-        Warning:
-            This does NOT yield the actual model parameters used by the RBLN runtime.
-            Code relying on iterating through all model parameters will not work as expected.
-        """
+        # Warning:
+        #     This does NOT yield the actual model parameters used by the RBLN runtime.
+        #     Code relying on iterating through all model parameters will not work as expected.
         yield torch.tensor([1.0], dtype=torch.float32, device=torch.device("cpu"))
 
     def __call__(self, *args, **kwargs):
@@ -547,7 +538,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
     @staticmethod
     def _raise_missing_compiled_file_error(missing_files: List[str]):
-        """Raises a KeyError with a message indicating missing compiled model files."""
+        # 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. "
@@ -563,40 +554,3 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             "and ensure the compilation completes successfully."
         )
         raise KeyError(message)
-
-    @classmethod
-    @abstractmethod
-    def _update_rbln_config(cls, **rbln_config_kwargs) -> RBLNModelConfig:
-        pass
-
-    @classmethod
-    @abstractmethod
-    def _create_runtimes(
-        cls,
-        compiled_models: List[rebel.RBLNCompiledModel],
-        rbln_config: RBLNModelConfig,
-    ) -> List[rebel.Runtime]:
-        # compiled_models -> runtimes
-        pass
-
-    @classmethod
-    @abstractmethod
-    def get_pytorch_model(cls, *args, **kwargs):
-        pass
-
-    @classmethod
-    @abstractmethod
-    def from_model(
-        cls,
-        model: "PreTrainedModel",
-        config: Optional[PretrainedConfig] = None,
-        rbln_config: Optional[RBLNModelConfig] = None,
-        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
-        subfolder: str = "",
-        **kwargs,
-    ):
-        pass
-
-    @abstractmethod
-    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
-        pass

optimum/rbln/transformers/__init__.py

@@ -18,16 +18,9 @@ from transformers.utils import _LazyModule
 
 
 _import_structure = {
-    "configuration_alias": [
-        "RBLNASTForAudioClassificationConfig",
-        "RBLNDistilBertForQuestionAnsweringConfig",
-        "RBLNResNetForImageClassificationConfig",
-        "RBLNXLMRobertaForSequenceClassificationConfig",
-        "RBLNRobertaForSequenceClassificationConfig",
-        "RBLNRobertaForMaskedLMConfig",
-        "RBLNViTForImageClassificationConfig",
-    ],
     "models": [
+        "RBLNASTForAudioClassification",
+        "RBLNASTForAudioClassificationConfig",
         "RBLNAutoModel",
         "RBLNAutoModelForAudioClassification",
         "RBLNAutoModelForCausalLM",
@@ -51,12 +44,14 @@ _import_structure = {
         "RBLNBertForQuestionAnsweringConfig",
         "RBLNBertModel",
         "RBLNBertModelConfig",
-        "RBLNBlip2VisionModelConfig",
-        "RBLNBlip2VisionModel",
-        "RBLNBlip2QFormerModel",
-        "RBLNBlip2QFormerModelConfig",
         "RBLNBlip2ForConditionalGeneration",
         "RBLNBlip2ForConditionalGenerationConfig",
+        "RBLNBlip2QFormerModel",
+        "RBLNBlip2QFormerModelConfig",
+        "RBLNBlip2VisionModel",
+        "RBLNBlip2VisionModelConfig",
+        "RBLNColPaliForRetrieval",
+        "RBLNColPaliForRetrievalConfig",
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelConfig",
         "RBLNCLIPTextModelWithProjection",
@@ -67,40 +62,48 @@ _import_structure = {
         "RBLNCLIPVisionModelWithProjectionConfig",
         "RBLNDecoderOnlyModelForCausalLM",
         "RBLNDecoderOnlyModelForCausalLMConfig",
+        "RBLNDistilBertForQuestionAnswering",
+        "RBLNDistilBertForQuestionAnsweringConfig",
         "RBLNDPTForDepthEstimation",
         "RBLNDPTForDepthEstimationConfig",
         "RBLNExaoneForCausalLM",
         "RBLNExaoneForCausalLMConfig",
-        "RBLNGemmaForCausalLM",
-        "RBLNGemmaForCausalLMConfig",
         "RBLNGemma3ForCausalLM",
         "RBLNGemma3ForCausalLMConfig",
         "RBLNGemma3ForConditionalGeneration",
         "RBLNGemma3ForConditionalGenerationConfig",
+        "RBLNGemmaForCausalLM",
+        "RBLNGemmaForCausalLMConfig",
         "RBLNGPT2LMHeadModel",
         "RBLNGPT2LMHeadModelConfig",
-        "RBLNIdefics3VisionTransformer",
         "RBLNIdefics3ForConditionalGeneration",
         "RBLNIdefics3ForConditionalGenerationConfig",
+        "RBLNIdefics3VisionTransformer",
         "RBLNIdefics3VisionTransformerConfig",
         "RBLNLlamaForCausalLM",
         "RBLNLlamaForCausalLMConfig",
-        "RBLNOPTForCausalLM",
-        "RBLNOPTForCausalLMConfig",
         "RBLNLlavaNextForConditionalGeneration",
         "RBLNLlavaNextForConditionalGenerationConfig",
         "RBLNMidmLMHeadModel",
         "RBLNMidmLMHeadModelConfig",
         "RBLNMistralForCausalLM",
         "RBLNMistralForCausalLMConfig",
+        "RBLNOPTForCausalLM",
+        "RBLNOPTForCausalLMConfig",
         "RBLNPhiForCausalLM",
         "RBLNPhiForCausalLMConfig",
-        "RBLNQwen2ForCausalLM",
-        "RBLNQwen2ForCausalLMConfig",
         "RBLNQwen2_5_VisionTransformerPretrainedModel",
         "RBLNQwen2_5_VisionTransformerPretrainedModelConfig",
         "RBLNQwen2_5_VLForConditionalGeneration",
         "RBLNQwen2_5_VLForConditionalGenerationConfig",
+        "RBLNQwen2ForCausalLM",
+        "RBLNQwen2ForCausalLMConfig",
+        "RBLNResNetForImageClassification",
+        "RBLNResNetForImageClassificationConfig",
+        "RBLNRobertaForMaskedLM",
+        "RBLNRobertaForMaskedLMConfig",
+        "RBLNRobertaForSequenceClassification",
+        "RBLNRobertaForSequenceClassificationConfig",
         "RBLNSiglipVisionModel",
         "RBLNSiglipVisionModelConfig",
         "RBLNT5EncoderModel",
@@ -109,44 +112,23 @@ _import_structure = {
         "RBLNT5ForConditionalGenerationConfig",
         "RBLNTimeSeriesTransformerForPrediction",
         "RBLNTimeSeriesTransformerForPredictionConfig",
+        "RBLNViTForImageClassification",
+        "RBLNViTForImageClassificationConfig",
         "RBLNWav2Vec2ForCTC",
         "RBLNWav2Vec2ForCTCConfig",
         "RBLNWhisperForConditionalGeneration",
         "RBLNWhisperForConditionalGenerationConfig",
+        "RBLNXLMRobertaForSequenceClassification",
+        "RBLNXLMRobertaForSequenceClassificationConfig",
         "RBLNXLMRobertaModel",
         "RBLNXLMRobertaModelConfig",
     ],
-    "modeling_alias": [
-        "RBLNASTForAudioClassification",
-        "RBLNDistilBertForQuestionAnswering",
-        "RBLNResNetForImageClassification",
-        "RBLNXLMRobertaForSequenceClassification",
-        "RBLNRobertaForSequenceClassification",
-        "RBLNRobertaForMaskedLM",
-        "RBLNViTForImageClassification",
-    ],
 }
 
 if TYPE_CHECKING:
-    from .configuration_alias import (
-        RBLNASTForAudioClassificationConfig,
-        RBLNDistilBertForQuestionAnsweringConfig,
-        RBLNResNetForImageClassificationConfig,
-        RBLNRobertaForMaskedLMConfig,
-        RBLNRobertaForSequenceClassificationConfig,
-        RBLNViTForImageClassificationConfig,
-        RBLNXLMRobertaForSequenceClassificationConfig,
-    )
-    from .modeling_alias import (
-        RBLNASTForAudioClassification,
-        RBLNDistilBertForQuestionAnswering,
-        RBLNResNetForImageClassification,
-        RBLNRobertaForMaskedLM,
-        RBLNRobertaForSequenceClassification,
-        RBLNViTForImageClassification,
-        RBLNXLMRobertaForSequenceClassification,
-    )
     from .models import (
+        RBLNASTForAudioClassification,
+        RBLNASTForAudioClassificationConfig,
         RBLNAutoModel,
         RBLNAutoModelForAudioClassification,
         RBLNAutoModelForCausalLM,
@@ -186,6 +168,8 @@ if TYPE_CHECKING:
         RBLNCLIPVisionModelWithProjectionConfig,
         RBLNDecoderOnlyModelForCausalLM,
         RBLNDecoderOnlyModelForCausalLMConfig,
+        RBLNDistilBertForQuestionAnswering,
+        RBLNDistilBertForQuestionAnsweringConfig,
         RBLNDPTForDepthEstimation,
         RBLNDPTForDepthEstimationConfig,
         RBLNExaoneForCausalLM,
@@ -220,6 +204,12 @@ if TYPE_CHECKING:
         RBLNQwen2_5_VLForConditionalGenerationConfig,
         RBLNQwen2ForCausalLM,
         RBLNQwen2ForCausalLMConfig,
+        RBLNResNetForImageClassification,
+        RBLNResNetForImageClassificationConfig,
+        RBLNRobertaForMaskedLM,
+        RBLNRobertaForMaskedLMConfig,
+        RBLNRobertaForSequenceClassification,
+        RBLNRobertaForSequenceClassificationConfig,
         RBLNSiglipVisionModel,
         RBLNSiglipVisionModelConfig,
         RBLNT5EncoderModel,
@@ -228,10 +218,14 @@ if TYPE_CHECKING:
         RBLNT5ForConditionalGenerationConfig,
         RBLNTimeSeriesTransformerForPrediction,
         RBLNTimeSeriesTransformerForPredictionConfig,
+        RBLNViTForImageClassification,
+        RBLNViTForImageClassificationConfig,
         RBLNWav2Vec2ForCTC,
         RBLNWav2Vec2ForCTCConfig,
         RBLNWhisperForConditionalGeneration,
         RBLNWhisperForConditionalGenerationConfig,
+        RBLNXLMRobertaForSequenceClassification,
+        RBLNXLMRobertaForSequenceClassificationConfig,
         RBLNXLMRobertaModel,
         RBLNXLMRobertaModelConfig,
     )

optimum/rbln/transformers/configuration_generic.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from ..configuration_utils import RBLNModelConfig
 
 
-class _RBLNTransformerEncoderConfig(RBLNModelConfig):
+class RBLNTransformerEncoderConfig(RBLNModelConfig):
     rbln_model_input_names: Optional[List[str]] = None
 
     def __init__(
@@ -25,7 +25,7 @@ class _RBLNTransformerEncoderConfig(RBLNModelConfig):
         max_seq_len: Optional[int] = None,
         batch_size: Optional[int] = None,
         model_input_names: Optional[List[str]] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -47,9 +47,12 @@ class _RBLNTransformerEncoderConfig(RBLNModelConfig):
         self.model_input_names = model_input_names or self.rbln_model_input_names
 
 
-class _RBLNImageModelConfig(RBLNModelConfig):
+class RBLNImageModelConfig(RBLNModelConfig):
     def __init__(
-        self, image_size: Optional[Union[int, Tuple[int, int]]] = None, batch_size: Optional[int] = None, **kwargs
+        self,
+        image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        batch_size: Optional[int] = None,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -86,32 +89,32 @@ class _RBLNImageModelConfig(RBLNModelConfig):
             return self.image_size["height"]
 
 
-class RBLNModelForQuestionAnsweringConfig(_RBLNTransformerEncoderConfig):
+class RBLNModelForQuestionAnsweringConfig(RBLNTransformerEncoderConfig):
     pass
 
 
-class RBLNModelForSequenceClassificationConfig(_RBLNTransformerEncoderConfig):
+class RBLNModelForSequenceClassificationConfig(RBLNTransformerEncoderConfig):
     pass
 
 
-class RBLNModelForMaskedLMConfig(_RBLNTransformerEncoderConfig):
+class RBLNModelForMaskedLMConfig(RBLNTransformerEncoderConfig):
     pass
 
 
-class RBLNModelForTextEncodingConfig(_RBLNTransformerEncoderConfig):
+class RBLNModelForTextEncodingConfig(RBLNTransformerEncoderConfig):
     pass
 
 
 # FIXME : Appropriate name ?
-class RBLNTransformerEncoderForFeatureExtractionConfig(_RBLNTransformerEncoderConfig):
+class RBLNTransformerEncoderForFeatureExtractionConfig(RBLNTransformerEncoderConfig):
     pass
 
 
-class RBLNModelForImageClassificationConfig(_RBLNImageModelConfig):
+class RBLNModelForImageClassificationConfig(RBLNImageModelConfig):
     pass
 
 
-class RBLNModelForDepthEstimationConfig(_RBLNImageModelConfig):
+class RBLNModelForDepthEstimationConfig(RBLNImageModelConfig):
     pass
 
 
@@ -121,7 +124,7 @@ class RBLNModelForAudioClassificationConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         max_length: Optional[int] = None,
         num_mel_bins: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

optimum/rbln/transformers/modeling_generic.py

@@ -43,9 +43,9 @@ from ..configuration_utils import RBLNCompileConfig
 from ..modeling import RBLNModel
 from ..utils.logging import get_logger
 from .configuration_generic import (
+    RBLNImageModelConfig,
     RBLNModelForAudioClassificationConfig,
-    _RBLNImageModelConfig,
-    _RBLNTransformerEncoderConfig,
+    RBLNTransformerEncoderConfig,
 )
 
 
@@ -55,7 +55,7 @@ if TYPE_CHECKING:
 logger = get_logger()
 
 
-class _RBLNTransformerEncoder(RBLNModel):
+class RBLNTransformerEncoder(RBLNModel):
     auto_model_class = AutoModel
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
     rbln_dtype = "int64"
@@ -66,8 +66,8 @@ class _RBLNTransformerEncoder(RBLNModel):
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
         model: Optional["PreTrainedModel"] = None,
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_config: Optional[_RBLNTransformerEncoderConfig] = None,
-    ) -> _RBLNTransformerEncoderConfig:
+        rbln_config: Optional[RBLNTransformerEncoderConfig] = None,
+    ) -> RBLNTransformerEncoderConfig:
         return cls.update_rbln_config_for_transformers_encoder(
             preprocessors=preprocessors,
             model=model,
@@ -81,8 +81,8 @@ class _RBLNTransformerEncoder(RBLNModel):
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
         model: Optional["PreTrainedModel"] = None,
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_config: Optional[_RBLNTransformerEncoderConfig] = None,
-    ) -> _RBLNTransformerEncoderConfig:
+        rbln_config: Optional[RBLNTransformerEncoderConfig] = None,
+    ) -> RBLNTransformerEncoderConfig:
         max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
             model_config, "max_position_embeddings", None
         )
@@ -139,7 +139,7 @@ class _RBLNTransformerEncoder(RBLNModel):
         return rbln_config
 
 
-class _RBLNImageModel(RBLNModel):
+class RBLNImageModel(RBLNModel):
     auto_model_class = AutoModel
     main_input_name = "pixel_values"
     output_class = BaseModelOutput
@@ -150,8 +150,8 @@ class _RBLNImageModel(RBLNModel):
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
         model: Optional["PreTrainedModel"] = None,
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_config: Optional[_RBLNImageModelConfig] = None,
-    ) -> _RBLNImageModelConfig:
+        rbln_config: Optional[RBLNImageModelConfig] = None,
+    ) -> RBLNImageModelConfig:
         return cls.update_rbln_config_for_image_model(
             preprocessors=preprocessors,
             model=model,
@@ -165,8 +165,8 @@ class _RBLNImageModel(RBLNModel):
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
         model: Optional["PreTrainedModel"] = None,
         model_config: Optional["PretrainedConfig"] = None,
-        rbln_config: Optional[_RBLNImageModelConfig] = None,
-    ) -> _RBLNImageModelConfig:
+        rbln_config: Optional[RBLNImageModelConfig] = None,
+    ) -> RBLNImageModelConfig:
         if rbln_config.image_size is None:
             for processor in preprocessors:
                 if hasattr(processor, "size"):
@@ -196,15 +196,14 @@ class _RBLNImageModel(RBLNModel):
         return rbln_config
 
 
-class RBLNModelForQuestionAnswering(_RBLNTransformerEncoder):
+class RBLNModelForQuestionAnswering(RBLNTransformerEncoder):
     auto_model_class = AutoModelForQuestionAnswering
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
     output_class = QuestionAnsweringModelOutput
 
     def _prepare_output(self, output, return_dict):
-        """
-        Prepare QuestionAnswering specific output format.
-        """
+        # Prepare QuestionAnswering specific output format.
+
         start_logits, end_logits = output
 
         if not return_dict:
@@ -213,32 +212,32 @@ class RBLNModelForQuestionAnswering(_RBLNTransformerEncoder):
             return QuestionAnsweringModelOutput(start_logits=start_logits, end_logits=end_logits)
 
 
-class RBLNModelForSequenceClassification(_RBLNTransformerEncoder):
+class RBLNModelForSequenceClassification(RBLNTransformerEncoder):
     auto_model_class = AutoModelForSequenceClassification
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
-class RBLNModelForMaskedLM(_RBLNTransformerEncoder):
+class RBLNModelForMaskedLM(RBLNTransformerEncoder):
     auto_model_class = AutoModelForMaskedLM
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
-class RBLNModelForTextEncoding(_RBLNTransformerEncoder):
+class RBLNModelForTextEncoding(RBLNTransformerEncoder):
     auto_model_class = AutoModelForTextEncoding
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
-class RBLNTransformerEncoderForFeatureExtraction(_RBLNTransformerEncoder):
+class RBLNTransformerEncoderForFeatureExtraction(RBLNTransformerEncoder):
     # TODO: RBLNModel is also for feature extraction.
     auto_model_class = AutoModel
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
-class RBLNModelForImageClassification(_RBLNImageModel):
+class RBLNModelForImageClassification(RBLNImageModel):
     auto_model_class = AutoModelForImageClassification
 
 
-class RBLNModelForDepthEstimation(_RBLNImageModel):
+class RBLNModelForDepthEstimation(RBLNImageModel):
     auto_model_class = AutoModelForDepthEstimation
 
 

optimum/rbln/transformers/modeling_rope_utils.py

@@ -48,10 +48,13 @@ def _compute_default_rope_parameters(
         Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
         post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
     """
-
     base = config.rope_theta
     partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
-    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    head_dim = (
+        config.head_dim
+        if hasattr(config, "head_dim") and config.head_dim is not None
+        else config.hidden_size // config.num_attention_heads
+    )
     dim = int(head_dim * partial_rotary_factor)
 
     attention_factor = 1.0  # Unused in this type of RoPE